highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-04-10 21:09:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

Migrate user milestone logic (#3493)

Browse Source
This commit is contained in:
pablonyx 2024-12-17 15:59:56 -08:00 committed by GitHub
parent 8db6d49fe5
commit b3ecf0d59f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 38 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
backend/ee/onyx/utils/telemetry.py
View File

@ -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_HOST
from onyx.utils.logger import setup_logger
logger = setup_logger()
posthog = Posthog(project_api_key=POSTHOG_API_KEY, host=POSTHOG_HOST)
@ -11,4 +14,5 @@ def event_telemetry(
event: str,
properties: dict | None = None,
) -> None:
logger.info(f"Capturing Posthog event: {distinct_id} {event} {properties}")
posthog.capture(distinct_id, event, properties)

53
backend/onyx/auth/users.py
View File

@ -285,25 +285,6 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
finally:
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
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
await self.user_db.add_oauth_account(user, oauth_account_dict)
await self.on_after_register(user, request)
else:
@ -475,6 +457,39 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
async def on_after_register(
self, user: User, request: Optional[Request] = 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.")
optional_telemetry(
record_type=RecordType.SIGN_UP,