refactor: improve code organization based on feedback

- Move imports to top of file instead of inside functions
- Extract common header parsing logic into _extract_context_from_headers helper
- Reduce code duplication between sync and async decorators
- Add comment explaining why crash handler needs to extract context from headers
- This addresses the context reset issue where decorators clean up before exception handlers run

Author: taeold

src/functions_framework/aio/__init__.py

@@ -15,9 +15,12 @@
 import asyncio
 import functools
 import inspect
+import json
+import logging
 import os
 import re
 import sys
+import traceback
 
 from typing import Any, Awaitable, Callable, Dict, Tuple, Union
 
@@ -55,41 +58,27 @@
 
 
 async def _crash_handler(request, exc):
-    import logging
-    import traceback
-    import json
-    
     # Log the exception
     logger = logging.getLogger()
     tb_lines = traceback.format_exception(type(exc), exc, exc.__traceback__)
     tb_text = ''.join(tb_lines)
     error_msg = f"Exception on {request.url.path} [{request.method}]\n{tb_text}".rstrip()
 
-    # Check if we need to set execution context for logging
+    # When execution ID logging is enabled, we need to extract execution_id from headers
+    # because the decorator resets the context before exception handlers run
     if _enable_execution_id_logging():
-        # Get execution_id from request headers
-        exec_id = request.headers.get(execution_id.EXECUTION_ID_REQUEST_HEADER)
-        span_id = None
-        
-        # Try to get span ID from trace context header
-        trace_context = request.headers.get(execution_id.TRACE_CONTEXT_REQUEST_HEADER, "")
-        trace_match = re.match(execution_id._TRACE_CONTEXT_REGEX_PATTERN, trace_context)
-        if trace_match:
-            span_id = trace_match.group("span_id")  # pragma: no cover
-        
-        if exec_id:
+        context = execution_id._extract_context_from_headers(request.headers)
+        if context.execution_id:
             # Temporarily set context for logging
-            token = execution_id.execution_context_var.set(
-                execution_id.ExecutionContext(exec_id, span_id)
-            )
+            token = execution_id.execution_context_var.set(context)
             try:
                 # Output as JSON so LoggingHandlerAddExecutionId can process it
                 log_entry = {"message": error_msg, "levelname": "ERROR"}
                 logger.error(json.dumps(log_entry))
             finally:
                 execution_id.execution_context_var.reset(token)
         else:  # pragma: no cover
-            # No execution ID, just log normally
+            # No execution ID in headers
             log_entry = {"message": error_msg, "levelname": "ERROR"}
             logger.error(json.dumps(log_entry))
     else:

src/functions_framework/execution_id.py

@@ -15,6 +15,7 @@
 import contextlib
 import contextvars
 import functools
+import inspect
 import io
 import json
 import logging
@@ -78,6 +79,20 @@ def _generate_execution_id():
     )
 
 
+def _extract_context_from_headers(headers):
+    """Extract execution context from request headers."""
+    execution_id = headers.get(EXECUTION_ID_REQUEST_HEADER)
+    
+    # Try to get span ID from trace context header
+    trace_context = re.match(
+        _TRACE_CONTEXT_REGEX_PATTERN,
+        headers.get(TRACE_CONTEXT_REQUEST_HEADER, ""),
+    )
+    span_id = trace_context.group("span_id") if trace_context else None
+    
+    return ExecutionContext(execution_id, span_id)
+
+
 # Middleware to add execution id to request header if one does not already exist
 class WsgiMiddleware:
     def __init__(self, wsgi_app):
@@ -147,13 +162,8 @@ def set_execution_context(request, enable_id_logging=False):
     def decorator(view_function):
         @functools.wraps(view_function)
         def wrapper(*args, **kwargs):
-            trace_context = re.match(
-                _TRACE_CONTEXT_REGEX_PATTERN,
-                request.headers.get(TRACE_CONTEXT_REQUEST_HEADER, ""),
-            )
-            execution_id = request.headers.get(EXECUTION_ID_REQUEST_HEADER)
-            span_id = trace_context.group("span_id") if trace_context else None
-            _set_current_context(ExecutionContext(execution_id, span_id))
+            context = _extract_context_from_headers(request.headers)
+            _set_current_context(context)
 
             with stderr_redirect, stdout_redirect:
                 return view_function(*args, **kwargs)
@@ -179,21 +189,15 @@ def set_execution_context_async(enable_id_logging=False):
     def decorator(view_function):
         @functools.wraps(view_function)
         async def async_wrapper(request, *args, **kwargs):
-            # Extract execution ID and span ID from Starlette request
-            trace_context = re.match(
-                _TRACE_CONTEXT_REGEX_PATTERN,
-                request.headers.get(TRACE_CONTEXT_REQUEST_HEADER, ""),
-            )
-            execution_id = request.headers.get(EXECUTION_ID_REQUEST_HEADER)
-            span_id = trace_context.group("span_id") if trace_context else None
+            # Extract execution context from headers
+            context = _extract_context_from_headers(request.headers)
 
             # Set context using contextvars
-            token = execution_context_var.set(ExecutionContext(execution_id, span_id))
+            token = execution_context_var.set(context)
 
             try:
                 with stderr_redirect, stdout_redirect:
                     # Handle both sync and async functions
-                    import inspect
                     if inspect.iscoroutinefunction(view_function):
                         return await view_function(request, *args, **kwargs)
                     else:
@@ -205,15 +209,10 @@ async def async_wrapper(request, *args, **kwargs):
         @functools.wraps(view_function)
         def sync_wrapper(request, *args, **kwargs):  # pragma: no cover
             # For sync functions, we still need to set up the context
-            trace_context = re.match(
-                _TRACE_CONTEXT_REGEX_PATTERN,
-                request.headers.get(TRACE_CONTEXT_REQUEST_HEADER, ""),
-            )
-            execution_id = request.headers.get(EXECUTION_ID_REQUEST_HEADER)
-            span_id = trace_context.group("span_id") if trace_context else None
+            context = _extract_context_from_headers(request.headers)
 
             # Set context using contextvars
-            token = execution_context_var.set(ExecutionContext(execution_id, span_id))
+            token = execution_context_var.set(context)
 
             try:
                 with stderr_redirect, stdout_redirect:
@@ -223,7 +222,6 @@ def sync_wrapper(request, *args, **kwargs):  # pragma: no cover
                 execution_context_var.reset(token)
 
         # Return appropriate wrapper based on whether the function is async
-        import inspect
         if inspect.iscoroutinefunction(view_function):
             return async_wrapper
         else: