From a51ed5911005ecea465b1b1aae22491c44ad28b0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aar=C3=B3n=20Ayerdis=20Espinoza?= Date: Mon, 1 Dec 2025 21:37:42 -0600 Subject: [PATCH 1/2] feat: Implement Cloud Run Logging trace correlation with contextvars --- .../core/middleware/logging_trace_context.py | 35 ++++++ testbed/core/tests/test_logging_trace.py | 100 ++++++++++++++++++ testbed/core/utils/logging_filters.py | 80 ++++++++------ testbed/settings/production.py | 9 ++ 4 files changed, 191 insertions(+), 33 deletions(-) create mode 100644 testbed/core/middleware/logging_trace_context.py create mode 100644 testbed/core/tests/test_logging_trace.py diff --git a/testbed/core/middleware/logging_trace_context.py b/testbed/core/middleware/logging_trace_context.py new file mode 100644 index 0000000..15c96f4 --- /dev/null +++ b/testbed/core/middleware/logging_trace_context.py @@ -0,0 +1,35 @@ +""" +Middleware for extracting Cloud Run trace context and request metadata. + +Reads and extract X-Cloud-Trace-Context header from Cloud Run at request start +and stores it in contextvars so CloudRunTraceFilter can add trace correlation to log records. +Lastly, clears the context at request end (response or exception). + +GCP DOC: https://docs.cloud.google.com/trace/docs/trace-context +Format: TRACE_ID/SPAN_ID;o=TRACE_ENABLED +""" +from django.utils.deprecation import MiddlewareMixin +from testbed.core.utils.logging_filters import set_trace_context, clear_trace_context + + +class LoggingTraceContextMiddleware(MiddlewareMixin): + + X_CLOUD_HEADER_NAME = "HTTP_X_CLOUD_TRACE_CONTEXT" + + def process_request(self, request): + + trace_header = request.META.get(self.X_CLOUD_HEADER_NAME) + + set_trace_context( + trace_header, + request_path=request.path, + request_method=request.method, + ) + + def process_response(self, request, response): + clear_trace_context() + return response + + def process_exception(self, request, exception): + clear_trace_context() + return None diff --git a/testbed/core/tests/test_logging_trace.py b/testbed/core/tests/test_logging_trace.py new file mode 100644 index 0000000..c66d396 --- /dev/null +++ b/testbed/core/tests/test_logging_trace.py @@ -0,0 +1,100 @@ +import os +import logging +from unittest.mock import patch, Mock +import pytest +from django.test import RequestFactory +from testbed.core.utils.logging_filters import ( + CloudRunTraceFilter, + set_trace_context, + clear_trace_context, +) +from testbed.core.middleware.logging_trace_context import LoggingTraceContextMiddleware + + +@pytest.fixture +def mock_environ(): + return { + 'GOOGLE_CLOUD_PROJECT': 'testbed-project', + 'ENVIRONMENT': 'test-env', + 'K_SERVICE': 'testbed-service', + 'K_REVISION': 'testbed-rev-001' + } + +@pytest.fixture +def log_record(): + return logging.LogRecord('test', logging.INFO, 'test.py', 1, 'test', (), None) + +@pytest.fixture +def request_factory(): + return RequestFactory() + +@pytest.fixture(autouse=True) +def cleanup_context(): + yield + clear_trace_context() + +# Tests CloudRunTraceFilter and LoggingTraceContextMiddleware + +# Filter should parse X-Cloud-Trace-Context and add trace/spanId fields +def test_filter_adds_trace_correlation(mock_environ, log_record): + with patch.dict(os.environ, mock_environ): + filter_instance = CloudRunTraceFilter() + set_trace_context('abc123/7890;o=1', '/callback', 'GET') + + filter_instance.filter(log_record) + + assert log_record.trace == 'projects/testbed-project/traces/abc123' + assert log_record.spanId == '7890' + assert log_record.request_path == '/callback' + assert log_record.request_method == 'GET' + +# Filter should add environment, service, and revision as searchable labels +def test_filter_adds_environment_labels(mock_environ, log_record): + with patch.dict(os.environ, mock_environ): + filter_instance = CloudRunTraceFilter() + filter_instance.filter(log_record) + + assert log_record.environment == 'test-env' + assert log_record.service == 'testbed-service' + assert log_record.revision == 'testbed-rev-001' + +# Filter should work when no trace context is set (e.g., local development) +def test_filter_handles_missing_trace(log_record): + with patch.dict(os.environ, {'ENVIRONMENT': 'test-env'}): + filter_instance = CloudRunTraceFilter() + clear_trace_context() + + result = filter_instance.filter(log_record) + + assert result is True + assert not hasattr(log_record, 'trace') + +# Middleware should extract X-Cloud-Trace-Context from request headers +def test_middleware_extracts_trace_header(request_factory): + middleware = LoggingTraceContextMiddleware(lambda r: None) + request = request_factory.post( + '/api/actors/1', + HTTP_X_CLOUD_TRACE_CONTEXT='trace123/span456;o=1' + ) + + middleware.process_request(request) + +# Middleware should clear context after response to prevent leakage +def test_middleware_clears_context_on_response(): + middleware = LoggingTraceContextMiddleware(lambda r: None) + request = Mock() + response = Mock() + + result = middleware.process_response(request, response) + + assert result == response + +# Middleware should clear context on exception to prevent leakage +def test_middleware_clears_context_on_exception(): + middleware = LoggingTraceContextMiddleware(lambda r: None) + request = Mock() + exception = Exception("Test") + + result = middleware.process_exception(request, exception) + + assert result is None diff --git a/testbed/core/utils/logging_filters.py b/testbed/core/utils/logging_filters.py index 30b2755..df9c97b 100644 --- a/testbed/core/utils/logging_filters.py +++ b/testbed/core/utils/logging_filters.py @@ -8,70 +8,84 @@ """ import logging import os -import threading +import contextvars -# Thread-local storage for trace context -_trace_local = threading.local() +_trace_context_var = contextvars.ContextVar("trace_context", default=None) +_request_path_var = contextvars.ContextVar("request_path", default=None) +_request_method_var = contextvars.ContextVar("request_method", default=None) -def set_trace_context(trace_context): - # Store trace context for the current thread/request - _trace_local.trace_context = trace_context +def set_trace_context(trace_context, request_path=None, request_method=None): + _trace_context_var.set(trace_context) + _request_path_var.set(request_path) + _request_method_var.set(request_method) def get_trace_context(): - # Retrieve trace context for the current thread/request - return getattr(_trace_local, 'trace_context', None) + return _trace_context_var.get() def clear_trace_context(): - # Clear trace context for the current thread/request - _trace_local.trace_context = None + _trace_context_var.set(None) + _request_path_var.set(None) + _request_method_var.set(None) class CloudRunTraceFilter(logging.Filter): """ - Enriches log records with Cloud Run trace context and environment metadata. + Enrich log records with Cloud Run trace context and environment metadata. Automatically adds: - trace: projects/PROJECT_ID/traces/TRACE_ID (for Cloud Logging correlation) - - span_id: Span ID from X-Cloud-Trace-Context - - environment: ENV_NAME (staging/production) + - spanId: Span ID from X-Cloud-Trace-Context + - environment: ENVIRONMENT (staging/production) - service: K_SERVICE (Cloud Run service name) - revision: K_REVISION (Cloud Run revision name) + - request_path: HTTP request path + - request_method: HTTP request method """ def __init__(self, name=""): + super().__init__(name) self.project_id = os.environ.get('GOOGLE_CLOUD_PROJECT', 'unknown') - self.environment = os.environ.get('ENV_NAME', 'unknown') + self.environment = os.environ.get('ENVIRONMENT', 'unknown') self.service = os.environ.get('K_SERVICE', 'unknown') self.revision = os.environ.get('K_REVISION', 'unknown') def filter(self, record): - # Add trace and environment context to log record - # Add environment metadata - record.environment = self.environment - record.service = self.service - record.revision = self.revision + if self.environment: + record.environment = self.environment + if self.service: + record.service = self.service + if self.revision: + record.revision = self.revision - # Add trace context if available - trace_context = get_trace_context() - if trace_context: + # Add trace correlation if trace header present + trace_context = _trace_context_var.get() + if trace_context and self.project_id: try: - # Parse X-Cloud-Trace-Context: TRACE_ID/SPAN_ID;o=TRACE_TRUE - # Convert to: projects/PROJECT_ID/traces/TRACE_ID - parts = trace_context.split('/') - if len(parts) >= 1: - trace_id = parts[0] - record.trace = f"projects/{self.project_id}/traces/{trace_id}" + # Parse X-Cloud-Trace-Context: TRACE_ID/SPAN_ID;o=1 + # https://docs.cloud.google.com/trace/docs/trace-context + x_cloud_trace_context_parts = trace_context.split(";")[0].split("/") + if len(x_cloud_trace_context_parts) >= 1: + trace_id = x_cloud_trace_context_parts[0] + if trace_id: + record.trace = f"projects/{self.project_id}/traces/{trace_id}" - if len(parts) >= 2: - span_parts = parts[1].split(';') - if span_parts: - record.span_id = span_parts[0] - except (IndexError, AttributeError): + if len(x_cloud_trace_context_parts) >= 2: + span_id = x_cloud_trace_context_parts[1] + if span_id: + record.spanId = span_id + except Exception: pass + request_path = _request_path_var.get() + request_method = _request_method_var.get() + if request_path: + record.request_path = request_path + if request_method: + record.request_method = request_method + return True diff --git a/testbed/settings/production.py b/testbed/settings/production.py index 79b19ca..cf3b607 100644 --- a/testbed/settings/production.py +++ b/testbed/settings/production.py @@ -11,6 +11,15 @@ SITE_URL = "https://ap-testbed.dtinit.org" BASE_URL = "https://ap-testbed.dtinit.org" +# Add logging trace context middleware for Cloud Run trace correlation +# Insert after SecurityMiddleware to capture trace context early in request processing +MIDDLEWARE = MIDDLEWARE.copy() +security_index = MIDDLEWARE.index("django.middleware.security.SecurityMiddleware") +MIDDLEWARE.insert( + security_index + 1, + "testbed.core.middleware.logging_trace_context.LoggingTraceContextMiddleware" +) + # Cloud Run uses X-Forwarded-Proto header for HTTPS detection SECURE_PROXY_SSL_HEADER = ('HTTP_X_FORWARDED_PROTO', 'https') From 448b312bfb7f331e497d19b525fd7a8c2249687f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aar=C3=B3n=20Ayerdis=20Espinoza?= Date: Tue, 6 Jan 2026 08:13:17 -0600 Subject: [PATCH 2/2] feat: Implement Cloud Run logging with setup_logging() --- .../core/middleware/logging_trace_context.py | 35 ------ testbed/core/tests/test_logging_trace.py | 100 ------------------ testbed/core/utils/logging_filters.py | 91 ---------------- testbed/core/utils/logging_utils.py | 67 ++++++------ testbed/settings/production.py | 52 +-------- 5 files changed, 38 insertions(+), 307 deletions(-) delete mode 100644 testbed/core/middleware/logging_trace_context.py delete mode 100644 testbed/core/tests/test_logging_trace.py delete mode 100644 testbed/core/utils/logging_filters.py diff --git a/testbed/core/middleware/logging_trace_context.py b/testbed/core/middleware/logging_trace_context.py deleted file mode 100644 index 15c96f4..0000000 --- a/testbed/core/middleware/logging_trace_context.py +++ /dev/null @@ -1,35 +0,0 @@ -""" -Middleware for extracting Cloud Run trace context and request metadata. - -Reads and extract X-Cloud-Trace-Context header from Cloud Run at request start -and stores it in contextvars so CloudRunTraceFilter can add trace correlation to log records. -Lastly, clears the context at request end (response or exception). - -GCP DOC: https://docs.cloud.google.com/trace/docs/trace-context -Format: TRACE_ID/SPAN_ID;o=TRACE_ENABLED -""" -from django.utils.deprecation import MiddlewareMixin -from testbed.core.utils.logging_filters import set_trace_context, clear_trace_context - - -class LoggingTraceContextMiddleware(MiddlewareMixin): - - X_CLOUD_HEADER_NAME = "HTTP_X_CLOUD_TRACE_CONTEXT" - - def process_request(self, request): - - trace_header = request.META.get(self.X_CLOUD_HEADER_NAME) - - set_trace_context( - trace_header, - request_path=request.path, - request_method=request.method, - ) - - def process_response(self, request, response): - clear_trace_context() - return response - - def process_exception(self, request, exception): - clear_trace_context() - return None diff --git a/testbed/core/tests/test_logging_trace.py b/testbed/core/tests/test_logging_trace.py deleted file mode 100644 index c66d396..0000000 --- a/testbed/core/tests/test_logging_trace.py +++ /dev/null @@ -1,100 +0,0 @@ -import os -import logging -from unittest.mock import patch, Mock -import pytest -from django.test import RequestFactory -from testbed.core.utils.logging_filters import ( - CloudRunTraceFilter, - set_trace_context, - clear_trace_context, -) -from testbed.core.middleware.logging_trace_context import LoggingTraceContextMiddleware - - -@pytest.fixture -def mock_environ(): - return { - 'GOOGLE_CLOUD_PROJECT': 'testbed-project', - 'ENVIRONMENT': 'test-env', - 'K_SERVICE': 'testbed-service', - 'K_REVISION': 'testbed-rev-001' - } - -@pytest.fixture -def log_record(): - return logging.LogRecord('test', logging.INFO, 'test.py', 1, 'test', (), None) - -@pytest.fixture -def request_factory(): - return RequestFactory() - -@pytest.fixture(autouse=True) -def cleanup_context(): - yield - clear_trace_context() - -# Tests CloudRunTraceFilter and LoggingTraceContextMiddleware - -# Filter should parse X-Cloud-Trace-Context and add trace/spanId fields -def test_filter_adds_trace_correlation(mock_environ, log_record): - with patch.dict(os.environ, mock_environ): - filter_instance = CloudRunTraceFilter() - set_trace_context('abc123/7890;o=1', '/callback', 'GET') - - filter_instance.filter(log_record) - - assert log_record.trace == 'projects/testbed-project/traces/abc123' - assert log_record.spanId == '7890' - assert log_record.request_path == '/callback' - assert log_record.request_method == 'GET' - -# Filter should add environment, service, and revision as searchable labels -def test_filter_adds_environment_labels(mock_environ, log_record): - with patch.dict(os.environ, mock_environ): - filter_instance = CloudRunTraceFilter() - filter_instance.filter(log_record) - - assert log_record.environment == 'test-env' - assert log_record.service == 'testbed-service' - assert log_record.revision == 'testbed-rev-001' - -# Filter should work when no trace context is set (e.g., local development) -def test_filter_handles_missing_trace(log_record): - with patch.dict(os.environ, {'ENVIRONMENT': 'test-env'}): - filter_instance = CloudRunTraceFilter() - clear_trace_context() - - result = filter_instance.filter(log_record) - - assert result is True - assert not hasattr(log_record, 'trace') - -# Middleware should extract X-Cloud-Trace-Context from request headers -def test_middleware_extracts_trace_header(request_factory): - middleware = LoggingTraceContextMiddleware(lambda r: None) - request = request_factory.post( - '/api/actors/1', - HTTP_X_CLOUD_TRACE_CONTEXT='trace123/span456;o=1' - ) - - middleware.process_request(request) - -# Middleware should clear context after response to prevent leakage -def test_middleware_clears_context_on_response(): - middleware = LoggingTraceContextMiddleware(lambda r: None) - request = Mock() - response = Mock() - - result = middleware.process_response(request, response) - - assert result == response - -# Middleware should clear context on exception to prevent leakage -def test_middleware_clears_context_on_exception(): - middleware = LoggingTraceContextMiddleware(lambda r: None) - request = Mock() - exception = Exception("Test") - - result = middleware.process_exception(request, exception) - - assert result is None diff --git a/testbed/core/utils/logging_filters.py b/testbed/core/utils/logging_filters.py deleted file mode 100644 index df9c97b..0000000 --- a/testbed/core/utils/logging_filters.py +++ /dev/null @@ -1,91 +0,0 @@ -""" -Logging filters for enriching log records with Cloud Run trace context and environment metadata. - -Automatically extracts trace information from Cloud Run's X-Cloud-Trace-Context header -and adds it to all log records for correlation in Google Cloud Logging. - -Reference: https://cloud.google.com/logging/docs/reference/v2/rest/v2/LogEntry -""" -import logging -import os -import contextvars - - -_trace_context_var = contextvars.ContextVar("trace_context", default=None) -_request_path_var = contextvars.ContextVar("request_path", default=None) -_request_method_var = contextvars.ContextVar("request_method", default=None) - - -def set_trace_context(trace_context, request_path=None, request_method=None): - _trace_context_var.set(trace_context) - _request_path_var.set(request_path) - _request_method_var.set(request_method) - - -def get_trace_context(): - return _trace_context_var.get() - - -def clear_trace_context(): - _trace_context_var.set(None) - _request_path_var.set(None) - _request_method_var.set(None) - - -class CloudRunTraceFilter(logging.Filter): - """ - Enrich log records with Cloud Run trace context and environment metadata. - - Automatically adds: - - trace: projects/PROJECT_ID/traces/TRACE_ID (for Cloud Logging correlation) - - spanId: Span ID from X-Cloud-Trace-Context - - environment: ENVIRONMENT (staging/production) - - service: K_SERVICE (Cloud Run service name) - - revision: K_REVISION (Cloud Run revision name) - - request_path: HTTP request path - - request_method: HTTP request method - """ - - def __init__(self, name=""): - - super().__init__(name) - self.project_id = os.environ.get('GOOGLE_CLOUD_PROJECT', 'unknown') - self.environment = os.environ.get('ENVIRONMENT', 'unknown') - self.service = os.environ.get('K_SERVICE', 'unknown') - self.revision = os.environ.get('K_REVISION', 'unknown') - - def filter(self, record): - if self.environment: - record.environment = self.environment - if self.service: - record.service = self.service - if self.revision: - record.revision = self.revision - - # Add trace correlation if trace header present - trace_context = _trace_context_var.get() - if trace_context and self.project_id: - try: - # Parse X-Cloud-Trace-Context: TRACE_ID/SPAN_ID;o=1 - # https://docs.cloud.google.com/trace/docs/trace-context - x_cloud_trace_context_parts = trace_context.split(";")[0].split("/") - if len(x_cloud_trace_context_parts) >= 1: - trace_id = x_cloud_trace_context_parts[0] - if trace_id: - record.trace = f"projects/{self.project_id}/traces/{trace_id}" - - if len(x_cloud_trace_context_parts) >= 2: - span_id = x_cloud_trace_context_parts[1] - if span_id: - record.spanId = span_id - except Exception: - pass - - request_path = _request_path_var.get() - request_method = _request_method_var.get() - if request_path: - record.request_path = request_path - if request_method: - record.request_method = request_method - - return True diff --git a/testbed/core/utils/logging_utils.py b/testbed/core/utils/logging_utils.py index b6e0dd9..18f7714 100644 --- a/testbed/core/utils/logging_utils.py +++ b/testbed/core/utils/logging_utils.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 diff --git a/testbed/settings/production.py b/testbed/settings/production.py index cf3b607..7b6a26a 100644 --- a/testbed/settings/production.py +++ b/testbed/settings/production.py @@ -1,5 +1,4 @@ # ruff: noqa: F405, F403 -import sys from google.oauth2 import service_account from .base import * @@ -11,15 +10,6 @@ SITE_URL = "https://ap-testbed.dtinit.org" BASE_URL = "https://ap-testbed.dtinit.org" -# Add logging trace context middleware for Cloud Run trace correlation -# Insert after SecurityMiddleware to capture trace context early in request processing -MIDDLEWARE = MIDDLEWARE.copy() -security_index = MIDDLEWARE.index("django.middleware.security.SecurityMiddleware") -MIDDLEWARE.insert( - security_index + 1, - "testbed.core.middleware.logging_trace_context.LoggingTraceContextMiddleware" -) - # Cloud Run uses X-Forwarded-Proto header for HTTPS detection SECURE_PROXY_SSL_HEADER = ('HTTP_X_FORWARDED_PROTO', 'https') @@ -67,43 +57,7 @@ EMAIL_HOST_USER = "noreply@dtinit.org" 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()