Merge pull request #235 from dtinit/feat/Cloud-Run-trace-correlation-contextvars

feat: Implement Cloud Run Logging trace correlation with contextvars

Author: aaronjae22

testbed/core/utils/logging_filters.py

@@ -1,52 +1,55 @@
 """
-This module provides functions for Google Cloud Logging handlers,
-properly isolates Cloud-specific dependencies to production/staging
+Uses Google's setup_logging() which automatically:
+- Detects Django framework
+- Extracts X-Cloud-Trace-Context headers from requests
+- Adds trace/spanId to all log entries
+- Groups logs by request in Cloud Logging
+
+References:
+- https://cloud.google.com/python/docs/reference/logging/latest/auto-trace-span-extraction
+- https://cloud.google.com/trace/docs/trace-log-integration
+- https://docs.cloud.google.com/python/docs/reference/logging/latest/client
+- https://github.com/googleapis/python-logging/blob/main/google/cloud/logging_v2/handlers/handlers.py
 """
 
 import os
 import logging
 
+logger = logging.getLogger(__name__)
+
 
-def get_cloud_logging_handler():
+def setup_cloud_logging():
     """
-    1. Checks if Cloud Logging is enabled (USE_GCLOUD_LOGGING env var)
-    2. Imports google-cloud-logging packages ONLY if enabled
-    3. Returns configured handler with trace correlation filter
-    
+    Initialize Google Cloud Logging with automatic trace correlation.
+
     Returns:
-        logging.Handler: CloudLoggingHandler configured with custom logName and
-                        trace filter, or NullHandler if Cloud Logging is disabled
-    
+        bool: True if Cloud Logging was successfully initialized, False if not
+
     Environment Variables:
         USE_GCLOUD_LOGGING: Set to "1" to enable Cloud Logging
         GOOGLE_CLOUD_PROJECT: GCP project ID (auto-detected on Cloud Run)
     """
-    
     if os.environ.get('USE_GCLOUD_LOGGING', '0') != '1':
-        return logging.NullHandler()
-    
-    # Import Google Cloud packages only when Cloud Logging is enabled
-    # This ensures dev/CI/test environments never import these packages
+        logger.info(
+            "Cloud Logging disabled (USE_GCLOUD_LOGGING != 1)"
+        )
+        return False
+
     try:
-        from google.cloud.logging import Client as CloudLoggingClient
-        from google.cloud.logging.handlers import CloudLoggingHandler
-        from testbed.core.utils.logging_filters import CloudRunTraceFilter
-        
-        client = CloudLoggingClient()
-        
-        cloud_logging_handler = CloudLoggingHandler(
-            client,
-            name="testbed"
+        import google.cloud.logging
+
+        client = google.cloud.logging.Client()
+
+        client.setup_logging(log_level=logging.INFO)
+
+        logger.info(
+            "Cloud Logging initialized with automatic trace correlation"
         )
-        
-        cloud_logging_handler.addFilter(CloudRunTraceFilter())
-        
-        return cloud_logging_handler
-        
+        return True
+
     except Exception as e:
-        logger = logging.getLogger(__name__)
         logger.warning(
             f"Failed to initialize Cloud Logging: {e}. "
-            "Falling back to NullHandler. Logs will not appear in Cloud Logging."
+            "Falling back to console logging."
         )
-        return logging.NullHandler()
+        return False

testbed/core/utils/logging_utils.py

@@ -1,5 +1,4 @@
 # ruff: noqa: F405, F403
-import sys
 from google.oauth2 import service_account
 from .base import *
 
@@ -58,43 +57,7 @@
 EMAIL_HOST_USER = "[email protected]"
 EMAIL_HOST_PASSWORD = env.str('EMAIL_HOST_PASSWORD')
 
-# This section configures Google Cloud Logging for Cloud Run environments.
+# Google Cloud Logging with automatic trace correlation.
 # Enabled via USE_GCLOUD_LOGGING=1 environment variable.
-if os.environ.get('USE_GCLOUD_LOGGING', '0') == '1':
-
-    LOGGING["handlers"]["cloud_logging"] = {
-        "()": "testbed.core.utils.logging_utils.get_cloud_logging_handler",
-    }
-    
-    LOGGING["root"] = {
-        "handlers": ["cloud_logging"],
-        "level": "INFO",
-    }
-
-    LOGGING["loggers"]["django"]["handlers"] = ["cloud_logging"]
-    LOGGING["loggers"]["django"]["propagate"] = False
-    
-    LOGGING["loggers"]["testbed"]["handlers"] = ["cloud_logging"]
-    LOGGING["loggers"]["testbed"]["propagate"] = False
-    
-    LOGGING["loggers"]["django.request"] = {
-        "handlers": ["cloud_logging"],
-        "level": "INFO",
-        "propagate": False,
-    }
-    
-    LOGGING["loggers"]["gunicorn"] = {
-        "handlers": ["cloud_logging"],
-        "level": "INFO",
-        "propagate": False,
-    }
-    LOGGING["loggers"]["gunicorn.error"] = {
-        "handlers": ["cloud_logging"],
-        "level": "INFO",
-        "propagate": False,
-    }
-    LOGGING["loggers"]["gunicorn.access"] = {
-        "handlers": ["cloud_logging"],
-        "level": "INFO",
-        "propagate": False,
-    }
+from testbed.core.utils.logging_utils import setup_cloud_logging
+setup_cloud_logging()