add extra tags to pruning logs (#2994)

Co-authored-by: Richard Kuo <rkuo@rkuo.com>

backend/danswer/background/celery/tasks/pruning/tasks.py

@@ -34,6 +34,7 @@ from danswer.db.engine import get_session_with_tenant
 from danswer.db.enums import ConnectorCredentialPairStatus
 from danswer.db.models import ConnectorCredentialPair
 from danswer.redis.redis_pool import get_redis_client
+from danswer.utils.logger import pruning_ctx
 from danswer.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -229,10 +230,15 @@ def connector_pruning_generator_task(
     and compares those IDs to locally stored documents and deletes all locally stored IDs missing
     from the most recently pulled document ID list"""
 
-    r = get_redis_client(tenant_id=tenant_id)
+    pruning_ctx_dict = pruning_ctx.get()
+    pruning_ctx_dict["cc_pair_id"] = cc_pair_id
+    pruning_ctx_dict["request_id"] = self.request.id
+    pruning_ctx.set(pruning_ctx_dict)
 
     rcp = RedisConnectorPruning(cc_pair_id)
 
+    r = get_redis_client(tenant_id=tenant_id)
+
     lock = r.lock(
         DanswerRedisLocks.PRUNING_LOCK_PREFIX + f"_{rcp._id}",
         timeout=CELERY_PRUNING_LOCK_TIMEOUT,

backend/danswer/utils/logger.py

@@ -1,3 +1,4 @@
+import contextvars
 import logging
 import os
 from collections.abc import MutableMapping
@@ -16,6 +17,10 @@ from shared_configs.configs import TENANT_ID_PREFIX
 
 logging.addLevelName(logging.INFO + 5, "NOTICE")
 
+pruning_ctx: contextvars.ContextVar[dict[str, Any]] = contextvars.ContextVar(
+    "pruning_ctx", default=dict()
+)
+
 
 class IndexAttemptSingleton:
     """Used to tell if this process is an indexing job, and if so what is the
@@ -64,6 +69,14 @@ class DanswerLoggingAdapter(logging.LoggerAdapter):
         index_attempt_id = IndexAttemptSingleton.get_index_attempt_id()
         cc_pair_id = IndexAttemptSingleton.get_connector_credential_pair_id()
 
+        pruning_ctx_dict = pruning_ctx.get()
+        if len(pruning_ctx_dict) > 0:
+            if "request_id" in pruning_ctx_dict:
+                msg = f"[Prune: {pruning_ctx_dict['request_id']}] {msg}"
+
+            if "cc_pair_id" in pruning_ctx_dict:
+                msg = f"[CC Pair: {pruning_ctx_dict['cc_pair_id']}] {msg}"
+        else:
             if index_attempt_id is not None:
                 msg = f"[Index Attempt: {index_attempt_id}] {msg}"