hook up permission syncs to celery task

Author: trial-danswer

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

@@ -56,6 +56,18 @@
 from onyx.db.enums import SyncStatus
 from onyx.db.enums import SyncType
 from onyx.db.models import ConnectorCredentialPair
+from onyx.db.permission_sync_attempt import create_doc_permission_sync_attempt
+from onyx.db.permission_sync_attempt import (
+    mark_doc_permission_sync_attempt_completed_with_errors,
+)
+from onyx.db.permission_sync_attempt import mark_doc_permission_sync_attempt_failed
+from onyx.db.permission_sync_attempt import (
+    mark_doc_permission_sync_attempt_in_progress,
+)
+from onyx.db.permission_sync_attempt import (
+    mark_doc_permission_sync_attempt_succeeded,
+)
+from onyx.db.permission_sync_attempt import update_doc_permission_sync_progress
 from onyx.db.sync_record import insert_sync_record
 from onyx.db.sync_record import update_sync_record_status
 from onyx.db.users import batch_add_ext_perm_user_if_not_exists
@@ -379,6 +391,15 @@ def connector_permission_sync_generator_task(
     doc_permission_sync_ctx_dict["request_id"] = self.request.id
     doc_permission_sync_ctx.set(doc_permission_sync_ctx_dict)
 
+    with get_session_with_current_tenant() as db_session:
+        attempt_id = create_doc_permission_sync_attempt(
+            connector_credential_pair_id=cc_pair_id,
+            db_session=db_session,
+        )
+        task_logger.info(
+            f"Created doc permission sync attempt: {attempt_id} for cc_pair={cc_pair_id}"
+        )
+
     redis_connector = RedisConnector(tenant_id, cc_pair_id)
 
     r = get_redis_client()
@@ -483,6 +504,8 @@ def connector_permission_sync_generator_task(
 
             logger.info(f"Syncing docs for {source_type} with cc_pair={cc_pair_id}")
 
+            mark_doc_permission_sync_attempt_in_progress(attempt_id, db_session)
+
             payload = redis_connector.permissions.payload
             if not payload:
                 raise ValueError(f"No fence payload found: cc_pair={cc_pair_id}")
@@ -533,22 +556,53 @@ def fetch_all_existing_docs_ids_fn() -> list[str]:
             )
 
             tasks_generated = 0
+            docs_with_errors = 0
             for doc_external_access in document_external_accesses:
-                redis_connector.permissions.update_db(
-                    lock=lock,
-                    new_permissions=[doc_external_access],
-                    source_string=source_type,
-                    connector_id=cc_pair.connector.id,
-                    credential_id=cc_pair.credential.id,
-                    task_logger=task_logger,
-                )
-                tasks_generated += 1
+                try:
+                    redis_connector.permissions.update_db(
+                        lock=lock,
+                        new_permissions=[doc_external_access],
+                        source_string=source_type,
+                        connector_id=cc_pair.connector.id,
+                        credential_id=cc_pair.credential.id,
+                        task_logger=task_logger,
+                    )
+                    tasks_generated += 1
+                except Exception as e:
+                    docs_with_errors += 1
+                    task_logger.exception(
+                        f"Error updating permissions for doc {doc_external_access.doc_id}: {e}"
+                    )
+                    # Continue processing other documents
 
             task_logger.info(
                 f"RedisConnector.permissions.generate_tasks finished. "
-                f"cc_pair={cc_pair_id} tasks_generated={tasks_generated}"
+                f"cc_pair={cc_pair_id} tasks_generated={tasks_generated} docs_with_errors={docs_with_errors}"
             )
 
+            update_doc_permission_sync_progress(
+                db_session=db_session,
+                attempt_id=attempt_id,
+                total_docs_synced=tasks_generated,
+                docs_with_permission_errors=docs_with_errors,
+            )
+            task_logger.info(
+                f"Updated progress for attempt {attempt_id}: {tasks_generated} docs, {docs_with_errors} errors"
+            )
+
+            if docs_with_errors > 0:
+                mark_doc_permission_sync_attempt_completed_with_errors(
+                    attempt_id, db_session
+                )
+                task_logger.info(
+                    f"Marked doc permission sync attempt {attempt_id} as completed with errors"
+                )
+            else:
+                mark_doc_permission_sync_attempt_succeeded(attempt_id, db_session)
+                task_logger.info(
+                    f"Marked doc permission sync attempt {attempt_id} as succeeded"
+                )
+
             redis_connector.permissions.generator_complete = tasks_generated
 
     except Exception as e:
@@ -561,6 +615,12 @@ def fetch_all_existing_docs_ids_fn() -> list[str]:
             f"Permission sync exceptioned: cc_pair={cc_pair_id} payload_id={payload_id}"
         )
 
+        with get_session_with_current_tenant() as db_session:
+            mark_doc_permission_sync_attempt_failed(attempt_id, db_session)
+            task_logger.info(
+                f"Marked doc permission sync attempt {attempt_id} as failed"
+            )
+
         redis_connector.permissions.generator_clear()
         redis_connector.permissions.taskset_clear()
         redis_connector.permissions.set_fence(None)

backend/ee/onyx/background/celery/tasks/external_group_syncing/tasks.py

@@ -49,6 +49,21 @@
 from onyx.db.enums import SyncStatus
 from onyx.db.enums import SyncType
 from onyx.db.models import ConnectorCredentialPair
+from onyx.db.permission_sync_attempt import (
+    create_external_group_sync_attempt,
+)
+from onyx.db.permission_sync_attempt import (
+    mark_external_group_sync_attempt_failed,
+)
+from onyx.db.permission_sync_attempt import (
+    mark_external_group_sync_attempt_in_progress,
+)
+from onyx.db.permission_sync_attempt import (
+    mark_external_group_sync_attempt_succeeded,
+)
+from onyx.db.permission_sync_attempt import (
+    update_external_group_sync_progress,
+)
 from onyx.db.sync_record import insert_sync_record
 from onyx.db.sync_record import update_sync_record_status
 from onyx.redis.redis_connector import RedisConnector
@@ -449,6 +464,16 @@ def _perform_external_group_sync(
     cc_pair_id: int,
     tenant_id: str,
 ) -> None:
+    # Create attempt record at the start
+    with get_session_with_current_tenant() as db_session:
+        attempt_id = create_external_group_sync_attempt(
+            connector_credential_pair_id=cc_pair_id,
+            db_session=db_session,
+        )
+        logger.info(
+            f"Created external group sync attempt: {attempt_id} for cc_pair={cc_pair_id}"
+        )
+
     with get_session_with_current_tenant() as db_session:
         cc_pair = get_connector_credential_pair_from_id(
             db_session=db_session,
@@ -477,14 +502,27 @@ def _perform_external_group_sync(
         )
         mark_old_external_groups_as_stale(db_session, cc_pair_id)
 
+        # Mark attempt as in progress
+        mark_external_group_sync_attempt_in_progress(attempt_id, db_session)
+        logger.info(f"Marked external group sync attempt {attempt_id} as in progress")
+
         logger.info(
             f"Syncing external groups for {source_type} for cc_pair: {cc_pair_id}"
         )
         external_user_group_batch: list[ExternalUserGroup] = []
+        total_users_processed = 0
+        total_groups_processed = 0
+        total_group_memberships_synced = 0
         try:
             external_user_group_generator = ext_group_sync_func(tenant_id, cc_pair)
             for external_user_group in external_user_group_generator:
                 external_user_group_batch.append(external_user_group)
+
+                # Track progress
+                total_groups_processed += 1
+                total_users_processed += len(external_user_group.user_emails)
+                total_group_memberships_synced += len(external_user_group.user_emails)
+
                 if len(external_user_group_batch) >= _EXTERNAL_GROUP_BATCH_SIZE:
                     logger.debug(
                         f"New external user groups: {external_user_group_batch}"
@@ -506,6 +544,10 @@ def _perform_external_group_sync(
                     source=cc_pair.connector.source,
                 )
         except Exception as e:
+            # Mark attempt as failed
+            mark_external_group_sync_attempt_failed(attempt_id, db_session)
+            logger.info(f"Marked external group sync attempt {attempt_id} as failed")
+
             # TODO: add some notification to the admins here
             logger.exception(
                 f"Error syncing external groups for {source_type} for cc_pair: {cc_pair_id} {e}"
@@ -517,6 +559,23 @@ def _perform_external_group_sync(
         )
         remove_stale_external_groups(db_session, cc_pair_id)
 
+        # Update progress and mark as succeeded
+        update_external_group_sync_progress(
+            db_session=db_session,
+            attempt_id=attempt_id,
+            total_users_processed=total_users_processed,
+            total_groups_processed=total_groups_processed,
+            total_group_memberships_synced=total_group_memberships_synced,
+        )
+        logger.info(
+            f"Updated progress for attempt {attempt_id}: "
+            f"{total_groups_processed} groups, {total_users_processed} users, "
+            f"{total_group_memberships_synced} memberships"
+        )
+
+        mark_external_group_sync_attempt_succeeded(attempt_id, db_session)
+        logger.info(f"Marked external group sync attempt {attempt_id} as succeeded")
+
         mark_all_relevant_cc_pairs_as_external_group_synced(db_session, cc_pair)