update

Author: pablonyx

backend/onyx/background/celery/tasks/doc_permission_syncing/tasks.py

@@ -15,7 +15,6 @@
 from redis import Redis
 from redis.exceptions import LockError
 from redis.lock import Lock as RedisLock
-from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 
 from ee.onyx.db.connector_credential_pair import get_all_auto_sync_cc_pairs
@@ -463,20 +462,11 @@ def update_external_document_permissions_task(
 
     try:
         with get_session_with_tenant(tenant_id) as db_session:
-            try:
-                # Add the users to the DB if they don't exist
-                batch_add_ext_perm_user_if_not_exists(
-                    db_session=db_session,
-                    emails=list(external_access.external_user_emails),
-                )
-            except IntegrityError:
-                db_session.rollback()
-                task_logger.error(
-                    "IntegrityError (likely 'UniqueViolation') occurred while adding users. "
-                    f"Emails: {list(external_access.external_user_emails)}. Retrying once."
-                )
-                return False
-
+            batch_add_ext_perm_user_if_not_exists(
+                db_session=db_session,
+                emails=list(external_access.external_user_emails),
+                continue_on_error=True,
+            )
             # Then upsert the document's external permissions
             created_new_doc = upsert_document_external_perms(
                 db_session=db_session,

backend/onyx/db/users.py

@@ -6,6 +6,7 @@
 from fastapi_users.password import PasswordHelper
 from sqlalchemy import func
 from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 from sqlalchemy.sql import expression
 from sqlalchemy.sql.elements import ColumnElement
@@ -20,8 +21,11 @@
 from onyx.db.models import SamlAccount
 from onyx.db.models import User
 from onyx.db.models import User__UserGroup
+from onyx.setup import setup_logger
 from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop
 
+logger = setup_logger()
+
 
 def validate_user_role_update(requested_role: UserRole, current_role: UserRole) -> None:
     """
@@ -274,7 +278,7 @@ def _generate_ext_permissioned_user(email: str) -> User:
 
 
 def batch_add_ext_perm_user_if_not_exists(
-    db_session: Session, emails: list[str]
+    db_session: Session, emails: list[str], continue_on_error: bool = False
 ) -> list[User]:
     lower_emails = [email.lower() for email in emails]
     found_users, missing_lower_emails = _get_users_by_emails(db_session, lower_emails)
@@ -283,10 +287,27 @@ def batch_add_ext_perm_user_if_not_exists(
     for email in missing_lower_emails:
         new_users.append(_generate_ext_permissioned_user(email=email))
 
-    db_session.add_all(new_users)
-    db_session.commit()
-
-    return found_users + new_users
+    try:
+        db_session.add_all(new_users)
+        db_session.commit()
+    except IntegrityError:
+        db_session.rollback()
+        logger.warning(
+            "IntegrityError occurred during batch insert. Falling back to individual inserts."
+        )
+        if not continue_on_error:
+            raise
+        for user in new_users:
+            try:
+                db_session.add(user)
+                db_session.commit()
+            except IntegrityError:
+                db_session.rollback()
+                logger.warning(f"Skipping duplicate user: {user.email}")
+                continue
+    # Fetch all users again to ensure we have the most up-to-date list
+    all_users, _ = _get_users_by_emails(db_session, lower_emails)
+    return all_users
 
 
 def delete_user_from_db(