mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-09-26 20:08:38 +02:00
Migrate user milestone logic (#3493)
This commit is contained in:
@@ -2,6 +2,9 @@ from posthog import Posthog
|
|||||||
|
|
||||||
from ee.onyx.configs.app_configs import POSTHOG_API_KEY
|
from ee.onyx.configs.app_configs import POSTHOG_API_KEY
|
||||||
from ee.onyx.configs.app_configs import POSTHOG_HOST
|
from ee.onyx.configs.app_configs import POSTHOG_HOST
|
||||||
|
from onyx.utils.logger import setup_logger
|
||||||
|
|
||||||
|
logger = setup_logger()
|
||||||
|
|
||||||
posthog = Posthog(project_api_key=POSTHOG_API_KEY, host=POSTHOG_HOST)
|
posthog = Posthog(project_api_key=POSTHOG_API_KEY, host=POSTHOG_HOST)
|
||||||
|
|
||||||
@@ -11,4 +14,5 @@ def event_telemetry(
|
|||||||
event: str,
|
event: str,
|
||||||
properties: dict | None = None,
|
properties: dict | None = None,
|
||||||
) -> None:
|
) -> None:
|
||||||
|
logger.info(f"Capturing Posthog event: {distinct_id} {event} {properties}")
|
||||||
posthog.capture(distinct_id, event, properties)
|
posthog.capture(distinct_id, event, properties)
|
||||||
|
@@ -285,25 +285,6 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
|
|||||||
finally:
|
finally:
|
||||||
CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
|
CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
|
||||||
|
|
||||||
# Blocking but this should be very quick
|
|
||||||
with get_session_with_tenant(tenant_id) as db_session:
|
|
||||||
if not user_count:
|
|
||||||
create_milestone_and_report(
|
|
||||||
user=user,
|
|
||||||
distinct_id=user.email,
|
|
||||||
event_type=MilestoneRecordType.USER_SIGNED_UP,
|
|
||||||
properties=None,
|
|
||||||
db_session=db_session,
|
|
||||||
)
|
|
||||||
else:
|
|
||||||
create_milestone_and_report(
|
|
||||||
user=user,
|
|
||||||
distinct_id=user.email,
|
|
||||||
event_type=MilestoneRecordType.MULTIPLE_USERS,
|
|
||||||
properties=None,
|
|
||||||
db_session=db_session,
|
|
||||||
)
|
|
||||||
|
|
||||||
return user
|
return user
|
||||||
|
|
||||||
async def validate_password(self, password: str, _: schemas.UC | models.UP) -> None:
|
async def validate_password(self, password: str, _: schemas.UC | models.UP) -> None:
|
||||||
@@ -422,6 +403,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
|
|||||||
|
|
||||||
# Add OAuth account
|
# Add OAuth account
|
||||||
await self.user_db.add_oauth_account(user, oauth_account_dict)
|
await self.user_db.add_oauth_account(user, oauth_account_dict)
|
||||||
|
|
||||||
await self.on_after_register(user, request)
|
await self.on_after_register(user, request)
|
||||||
|
|
||||||
else:
|
else:
|
||||||
@@ -475,6 +457,39 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
|
|||||||
async def on_after_register(
|
async def on_after_register(
|
||||||
self, user: User, request: Optional[Request] = None
|
self, user: User, request: Optional[Request] = None
|
||||||
) -> None:
|
) -> None:
|
||||||
|
tenant_id = await fetch_ee_implementation_or_noop(
|
||||||
|
"onyx.server.tenants.provisioning",
|
||||||
|
"get_or_provision_tenant",
|
||||||
|
async_return_default_schema,
|
||||||
|
)(
|
||||||
|
email=user.email,
|
||||||
|
request=request,
|
||||||
|
)
|
||||||
|
|
||||||
|
token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
|
||||||
|
try:
|
||||||
|
user_count = await get_user_count()
|
||||||
|
|
||||||
|
with get_session_with_tenant(tenant_id=tenant_id) as db_session:
|
||||||
|
if user_count == 1:
|
||||||
|
create_milestone_and_report(
|
||||||
|
user=user,
|
||||||
|
distinct_id=user.email,
|
||||||
|
event_type=MilestoneRecordType.USER_SIGNED_UP,
|
||||||
|
properties=None,
|
||||||
|
db_session=db_session,
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
create_milestone_and_report(
|
||||||
|
user=user,
|
||||||
|
distinct_id=user.email,
|
||||||
|
event_type=MilestoneRecordType.MULTIPLE_USERS,
|
||||||
|
properties=None,
|
||||||
|
db_session=db_session,
|
||||||
|
)
|
||||||
|
finally:
|
||||||
|
CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
|
||||||
|
|
||||||
logger.notice(f"User {user.id} has registered.")
|
logger.notice(f"User {user.id} has registered.")
|
||||||
optional_telemetry(
|
optional_telemetry(
|
||||||
record_type=RecordType.SIGN_UP,
|
record_type=RecordType.SIGN_UP,
|
||||||
|
Reference in New Issue
Block a user