minor clean up

onyx-dot-app · Nov 4, 2024 · 6efe24f · 6efe24f
1 parent dd40b12
commit 6efe24f
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 42 deletions.
diff --git a/backend/danswer/auth/users.py b/backend/danswer/auth/users.py
@@ -239,29 +239,29 @@ async def create(
         safe: bool = False,
         request: Optional[Request] = None,
     ) -> User:
-        try:
-            tenant_id = (
-                get_tenant_id_for_email(user_create.email)
-                if MULTI_TENANT
-                else POSTGRES_DEFAULT_SCHEMA
-            )
-        except exceptions.UserNotExists:
-            # Tenant does not exist; provision a new tenant
-            tenant_provisioning_service = TenantProvisioningService()
+        if MULTI_TENANT:
             try:
-                tenant_id = await tenant_provisioning_service.provision_tenant(
-                    user_create.email
-                )
-            except Exception as e:
-                logger.error(f"Tenant provisioning failed: {e}")
+                tenant_id = get_tenant_id_for_email(user_create.email)
+
+            except exceptions.UserNotExists:
+                # If tenant does not exist and in Multi tenant mode, provision a new tenant
+                tenant_provisioning_service = TenantProvisioningService()
+                try:
+                    tenant_id = await tenant_provisioning_service.provision_tenant(
+                        user_create.email
+                    )
+                except Exception as e:
+                    logger.error(f"Tenant provisioning failed: {e}")
+                    raise HTTPException(
+                        status_code=500, detail="Failed to provision tenant."
+                    )
+
+            if not tenant_id:
                 raise HTTPException(
-                    status_code=500, detail="Failed to provision tenant."
+                    status_code=401, detail="User does not belong to an organization"
                 )
-
-        if not tenant_id:
-            raise HTTPException(
-                status_code=401, detail="User does not belong to an organization"
-            )
+        else:
+            tenant_id = POSTGRES_DEFAULT_SCHEMA
 
         async with get_async_session_with_tenant(tenant_id) as db_session:
             token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
@@ -321,24 +321,23 @@ async def oauth_callback(
         is_verified_by_default: bool = False,
     ) -> models.UOAP:
         # Get tenant_id from mapping table
-        try:
-            tenant_id = (
-                get_tenant_id_for_email(account_email)
-                if MULTI_TENANT
-                else POSTGRES_DEFAULT_SCHEMA
-            )
-        except exceptions.UserNotExists:
-            # Tenant does not exist; provision a new tenant
-            tenant_provisioning_service = TenantProvisioningService()
+        if MULTI_TENANT:
             try:
-                tenant_id = await tenant_provisioning_service.provision_tenant(
-                    account_email
-                )
-            except Exception as e:
-                logger.error(f"Tenant provisioning failed: {e}")
-                raise HTTPException(
-                    status_code=500, detail="Failed to provision tenant."
-                )
+                tenant_id = get_tenant_id_for_email(account_email)
+            except exceptions.UserNotExists:
+                # Tenant does not exist; provision a new tenant
+                tenant_provisioning_service = TenantProvisioningService()
+                try:
+                    tenant_id = await tenant_provisioning_service.provision_tenant(
+                        account_email
+                    )
+                except Exception as e:
+                    logger.error(f"Tenant provisioning failed: {e}")
+                    raise HTTPException(
+                        status_code=500, detail="Failed to provision tenant."
+                    )
+        else:
+            tenant_id = POSTGRES_DEFAULT_SCHEMA
 
         if not tenant_id:
             raise HTTPException(status_code=401, detail="User not found")

diff --git a/backend/ee/danswer/server/tenants/provisioning.py b/backend/ee/danswer/server/tenants/provisioning.py
@@ -35,11 +35,6 @@
 logger = logging.getLogger(__name__)
 
 
-def drop_schema(tenant_id: str) -> None:
-    with get_sqlalchemy_engine().connect() as connection:
-        connection.execute(text(f"DROP SCHEMA IF EXISTS {tenant_id} CASCADE"))
-
-
 class TenantProvisioningService:
     async def provision_tenant(self, email: str) -> str:
         tenant_id = TENANT_ID_PREFIX + str(uuid.uuid4())  # Generate new tenant ID
@@ -244,3 +239,8 @@ def ensure_schema_exists(tenant_id: str) -> bool:
                 db_session.execute(stmt)
                 return True
             return False
+
+
+def drop_schema(tenant_id: str) -> None:
+    with get_sqlalchemy_engine().connect() as connection:
+        connection.execute(text(f"DROP SCHEMA IF EXISTS {tenant_id} CASCADE"))