From b3ecf0d59ffb72a86dbd5e2f27661bfb4dc493eb Mon Sep 17 00:00:00 2001
From: pablonyx <pablo@danswer.ai>
Date: Tue, 17 Dec 2024 15:59:56 -0800
Subject: [PATCH] Migrate user milestone logic (#3493)

---
 backend/ee/onyx/utils/telemetry.py |  4 +++
 backend/onyx/auth/users.py         | 53 +++++++++++++++++++-----------
 2 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/backend/ee/onyx/utils/telemetry.py b/backend/ee/onyx/utils/telemetry.py
index 960ae46939a..42ac89c61bc 100644
--- a/backend/ee/onyx/utils/telemetry.py
+++ b/backend/ee/onyx/utils/telemetry.py
@@ -2,6 +2,9 @@
 
 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)
diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py
index ee1f84f9da3..69f78f50ae3 100644
--- a/backend/onyx/auth/users.py
+++ b/backend/onyx/auth/users.py
@@ -285,25 +285,6 @@ async def create(
             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 @@ async def oauth_callback(
 
                     # 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 @@ async def oauth_callback(
     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,