From 23333d249ded64328f1ceaa7b2f51415e5edb251 Mon Sep 17 00:00:00 2001
From: "Richard Kuo (Danswer)" <rkuo@onyx.app>
Date: Wed, 26 Feb 2025 23:04:55 -0800
Subject: [PATCH] fix dupe ids and session handling

---
 backend/onyx/db/document_set.py               | 20 +++++++++++++++----
 backend/onyx/db/session.py                    | 13 +++---------
 backend/onyx/server/documents/cc_pair.py      |  2 ++
 .../onyx/server/features/document_set/api.py  |  3 +++
 4 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/backend/onyx/db/document_set.py b/backend/onyx/db/document_set.py
index 0f91cbc71b89..03c6c21bff8b 100644
--- a/backend/onyx/db/document_set.py
+++ b/backend/onyx/db/document_set.py
@@ -323,15 +323,27 @@ def update_document_set(
         _mark_document_set_cc_pairs_as_outdated__no_commit(
             db_session=db_session, document_set_id=document_set_row.id
         )
+
+        # commit this before performing more updates on the row id or else
+        # we'll have conflicting updates in the same commit
+        db_session.commit()
+
         # add in rows for the new CC pairs
-        ds_cc_pairs = [
-            DocumentSet__ConnectorCredentialPair(
+        existing_cc_pair_ids: set[int] = set()  # use to avoid duplicates
+
+        ds_cc_pairs: list[DocumentSet__ConnectorCredentialPair] = []
+        for cc_pair_id in document_set_update_request.cc_pair_ids:
+            if cc_pair_id in existing_cc_pair_ids:
+                continue
+
+            item = DocumentSet__ConnectorCredentialPair(
                 document_set_id=document_set_update_request.id,
                 connector_credential_pair_id=cc_pair_id,
                 is_current=True,
             )
-            for cc_pair_id in document_set_update_request.cc_pair_ids
-        ]
+            ds_cc_pairs.append(item)
+            existing_cc_pair_ids.add(cc_pair_id)
+
         db_session.add_all(ds_cc_pairs)
         db_session.commit()
     except:
diff --git a/backend/onyx/db/session.py b/backend/onyx/db/session.py
index 61882aad6435..e70e30cf5d56 100644
--- a/backend/onyx/db/session.py
+++ b/backend/onyx/db/session.py
@@ -71,15 +71,15 @@ def get_session_with_tenant(*, tenant_id: str | None) -> Generator[Session, None
 
 
 def get_session() -> Generator[Session, None, None]:
-    return OnyxSession.get_session()
+    yield from OnyxSession.get_session()
 
 
 def get_multi_tenant_session(tenant_id: str) -> Generator[Session, None, None]:
-    return OnyxSession.get_multi_tenant_session(tenant_id)
+    yield from OnyxSession.get_multi_tenant_session(tenant_id)
 
 
 def get_single_tenant_session() -> Generator[Session, None, None]:
-    return OnyxSession.get_single_tenant_session()
+    yield from OnyxSession.get_single_tenant_session()
 
 
 async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
@@ -88,13 +88,6 @@ async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
         yield session
 
 
-# AsyncSessionLocal = sessionmaker(  # type: ignore
-#     bind=get_sqlalchemy_async_engine(),
-#     class_=AsyncSession,
-#     expire_on_commit=False,
-# )
-
-
 @asynccontextmanager
 async def get_async_session_with_tenant(
     tenant_id: str | None = None,
diff --git a/backend/onyx/server/documents/cc_pair.py b/backend/onyx/server/documents/cc_pair.py
index 6eeb5c973a5c..138738304729 100644
--- a/backend/onyx/server/documents/cc_pair.py
+++ b/backend/onyx/server/documents/cc_pair.py
@@ -652,6 +652,8 @@ def associate_credential_to_connector(
         # Ensures we don't leave invalid connectors in the database
         # NOTE: consensus is that it makes sense to unify connector and ccpair creation flows
         # which would rid us of needing to handle cases like these
+        logger.exception("ValidationError")
+
         delete_connector(db_session, connector_id)
         db_session.commit()
 
diff --git a/backend/onyx/server/features/document_set/api.py b/backend/onyx/server/features/document_set/api.py
index df73477cd766..bea2ffcf38be 100644
--- a/backend/onyx/server/features/document_set/api.py
+++ b/backend/onyx/server/features/document_set/api.py
@@ -21,9 +21,11 @@ from onyx.server.features.document_set.models import CheckDocSetPublicResponse
 from onyx.server.features.document_set.models import DocumentSet
 from onyx.server.features.document_set.models import DocumentSetCreationRequest
 from onyx.server.features.document_set.models import DocumentSetUpdateRequest
+from onyx.utils.logger import setup_logger
 from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop
 from shared_configs.contextvars import get_current_tenant_id
 
+logger = setup_logger()
 
 router = APIRouter(prefix="/manage")
 
@@ -83,6 +85,7 @@ def patch_document_set(
             user=user,
         )
     except Exception as e:
+        logger.exception("patch_document_set exceptioned")
         raise HTTPException(status_code=400, detail=str(e))
 
     primary_app.send_task(