(trying to) reduce sql connections

Author: bigabig

backend/src/core/auth/auth_endpoint.py

@@ -192,7 +192,7 @@ async def oidc_callback(
         )
 
     try:
-        user = await oauth_service.authenticate_oidc(request)
+        user = await oauth_service.authenticate_oidc(db=db, request=request)
     except Exception as e:
         raise HTTPException(status_code=401, detail=str(e))
 

backend/src/core/auth/oauth_service.py

@@ -4,13 +4,13 @@
 from authlib.integrations.starlette_client import OAuth, OAuthError
 from fastapi import Request
 from loguru import logger
+from sqlalchemy.orm import Session
 
 from common.singleton_meta import SingletonMeta
 from config import conf
 from core.user.user_crud import crud_user
 from core.user.user_dto import UserCreate
 from core.user.user_orm import UserORM
-from repos.db.sql_repo import SQLRepo
 from repos.mail_repo import MailRepo
 
 
@@ -41,7 +41,7 @@ def __new__(cls, *args, **kwargs):
 
         return super(OAuthService, cls).__new__(cls)
 
-    async def authenticate_oidc(self, request: Request) -> UserORM:
+    async def authenticate_oidc(self, db: Session, request: Request) -> UserORM:
         try:
             token = await self.authentik.authorize_access_token(request)
         except OAuthError as error:
@@ -52,35 +52,34 @@ async def authenticate_oidc(self, request: Request) -> UserORM:
             userinfo = token.get("userinfo")
             print(f"Userinfo: {userinfo}")
 
-            with SQLRepo().db_session() as db:
-                try:
-                    # Warning: Security concern
-                    user = crud_user.read_by_email(db=db, email=userinfo["email"])
-                    return user
-                except Exception as e:
-                    logger.info(f"User not found, creating new user: {e}")
-                    # Create user if not exists
-                    user = crud_user.create(
-                        db=db,
-                        create_dto=UserCreate(
-                            email=userinfo["email"],
-                            first_name=userinfo.get("given_name", "Unknown"),
-                            last_name=userinfo.get("family_name", "Unknown"),
-                            # Set a random password since we'll only use OIDC
-                            password="".join(
-                                random.choices(
-                                    string.ascii_letters + string.digits,
-                                    k=32,
-                                )
-                            ),
+            try:
+                # Warning: Security concern
+                user = crud_user.read_by_email(db=db, email=userinfo["email"])
+                return user
+            except Exception as e:
+                logger.info(f"User not found, creating new user: {e}")
+                # Create user if not exists
+                user = crud_user.create(
+                    db=db,
+                    create_dto=UserCreate(
+                        email=userinfo["email"],
+                        first_name=userinfo.get("given_name", "Unknown"),
+                        last_name=userinfo.get("family_name", "Unknown"),
+                        # Set a random password since we'll only use OIDC
+                        password="".join(
+                            random.choices(
+                                string.ascii_letters + string.digits,
+                                k=32,
+                            )
                         ),
-                    )
-                    await MailRepo().send_welcome_mail(
-                        email=user.email,
-                        first_name=user.first_name,
-                        last_name=user.last_name,
-                    )
-                    return user
+                    ),
+                )
+                await MailRepo().send_welcome_mail(
+                    email=user.email,
+                    first_name=user.first_name,
+                    last_name=user.last_name,
+                )
+                return user
         except Exception as e:
             logger.error(f"Error processing OIDC authentication: {e}")
             raise Exception("Authentication failed")

backend/src/core/project/project_crud.py

@@ -26,7 +26,7 @@ class CRUDProject(CRUDBase[ProjectORM, ProjectCreate, ProjectUpdate]):
     ### CREATE OPERATIONS ###
 
     def create(
-        self, db: Session, *, create_dto: ProjectCreate, creating_user: UserORM
+        self, db: Session, *, create_dto: ProjectCreate, creating_user_id: int
     ) -> ProjectORM:
         # 1) create the project
         dto_obj_data = jsonable_encoder(create_dto)
@@ -44,13 +44,13 @@ def create(
         self.associate_user(db=db, proj_id=project_id, user_id=ASSISTANT_TRAINED_ID)
 
         # 3) associate the user that created the project
-        if creating_user.id not in [
+        if creating_user_id not in [
             SYSTEM_USER_ID,
             ASSISTANT_ZEROSHOT_ID,
             ASSISTANT_FEWSHOT_ID,
             ASSISTANT_TRAINED_ID,
         ]:
-            self.associate_user(db=db, proj_id=project_id, user_id=creating_user.id)
+            self.associate_user(db=db, proj_id=project_id, user_id=creating_user_id)
 
         # 4) create system codes
         crud_code.create_system_codes_for_project(db=db, proj_id=project_id)

backend/src/core/project/project_endpoint.py

@@ -34,7 +34,9 @@ def create_new_project(
     proj: ProjectCreate,
     current_user: UserORM = Depends(get_current_user),
 ) -> ProjectRead:
-    db_obj = crud_project.create(db=db, create_dto=proj, creating_user=current_user)
+    db_obj = crud_project.create(
+        db=db, create_dto=proj, creating_user_id=current_user.id
+    )
     return ProjectRead.model_validate(db_obj)
 
 

backend/src/modules/analysis/analysis_endpoint.py

@@ -30,6 +30,7 @@
 )
 def code_frequencies(
     *,
+    db: Session = Depends(get_db_session),
     project_id: int,
     code_ids: list[int],
     user_ids: list[int],
@@ -39,7 +40,11 @@ def code_frequencies(
     authz_user.assert_in_project(project_id)
 
     return find_code_frequencies(
-        project_id=project_id, code_ids=code_ids, user_ids=user_ids, doctypes=doctypes
+        db=db,
+        project_id=project_id,
+        code_ids=code_ids,
+        user_ids=user_ids,
+        doctypes=doctypes,
     )
 
 
@@ -50,6 +55,7 @@ def code_frequencies(
 )
 def code_occurrences(
     *,
+    db: Session = Depends(get_db_session),
     project_id: int,
     user_ids: list[int],
     code_id: int,
@@ -58,7 +64,7 @@ def code_occurrences(
     authz_user.assert_in_project(project_id)
 
     return find_code_occurrences(
-        project_id=project_id, user_ids=user_ids, code_id=code_id
+        db=db, project_id=project_id, user_ids=user_ids, code_id=code_id
     )
 
 
@@ -69,13 +75,15 @@ def code_occurrences(
 )
 def count_sdocs_with_date_metadata(
     *,
+    db: Session = Depends(get_db_session),
     project_id: int,
     date_metadata_id: int,
     authz_user: AuthzUser = Depends(),
 ) -> tuple[int, int]:
     authz_user.assert_in_project(project_id)
 
     return compute_num_sdocs_with_date_metadata(
+        db=db,
         project_id=project_id,
         date_metadata_id=date_metadata_id,
     )
@@ -97,5 +105,5 @@ def sample_sdocs_by_tags(
 ) -> list[SampledSdocsResults]:
     authz_user.assert_in_project(project_id)
     return document_sampler_by_tags(
-        project_id=project_id, tag_ids=tag_groups, n=n, frac=frac
+        db=db, project_id=project_id, tag_ids=tag_groups, n=n, frac=frac
     )