feat: support span inference (#267)

Author: daniel-sanche

google/cloud/logging_v2/handlers/_helpers.py

@@ -16,6 +16,7 @@
 
 import math
 import json
+import re
 
 try:
     import flask
@@ -55,12 +56,13 @@ def get_request_data_from_flask():
     """Get http_request and trace data from flask request headers.
 
     Returns:
-        Tuple[Optional[dict], Optional[str]]:
-            Data related to the current http request and the trace_id for the
-            request. Both fields will be None if a flask request isn't found.
+        Tuple[Optional[dict], Optional[str], Optional[str]]:
+            Data related to the current http request, trace_id, and span_id for
+            the request. All fields will be None if a django request isn't
+            found.
     """
     if flask is None or not flask.request:
-        return None, None
+        return None, None, None
 
     # build http_request
     http_request = {
@@ -73,27 +75,26 @@ def get_request_data_from_flask():
         "protocol": flask.request.environ.get(_PROTOCOL_HEADER),
     }
 
-    # find trace id
-    trace_id = None
+    # find trace id and span id
     header = flask.request.headers.get(_FLASK_TRACE_HEADER)
-    if header:
-        trace_id = header.split("/", 1)[0]
+    trace_id, span_id = _parse_trace_span(header)
 
-    return http_request, trace_id
+    return http_request, trace_id, span_id
 
 
 def get_request_data_from_django():
     """Get http_request and trace data from django request headers.
 
     Returns:
-        Tuple[Optional[dict], Optional[str]]:
-            Data related to the current http request and the trace_id for the
-            request. Both fields will be None if a django request isn't found.
+        Tuple[Optional[dict], Optional[str], Optional[str]]:
+            Data related to the current http request, trace_id, and span_id for
+            the request. All fields will be None if a django request isn't
+            found.
     """
     request = _get_django_request()
 
     if request is None:
-        return None, None
+        return None, None, None
 
     # convert content_length to int if it exists
     content_length = None
@@ -112,32 +113,55 @@ def get_request_data_from_django():
         "protocol": request.META.get(_PROTOCOL_HEADER),
     }
 
-    # find trace id
-    trace_id = None
+    # find trace id and span id
     header = request.META.get(_DJANGO_TRACE_HEADER)
-    if header:
-        trace_id = header.split("/", 1)[0]
+    trace_id, span_id = _parse_trace_span(header)
 
-    return http_request, trace_id
+    return http_request, trace_id, span_id
+
+
+def _parse_trace_span(header):
+    """Given an X_CLOUD_TRACE header, extract the trace and span ids.
+
+    Args:
+        header (str): the string extracted from the X_CLOUD_TRACE header
+    Returns:
+        Tuple[Optional[dict], Optional[str]]:
+            The trace_id and span_id extracted from the header
+            Each field will be None if not found.
+    """
+    trace_id = None
+    span_id = None
+    if header:
+        try:
+            split_header = header.split("/", 1)
+            trace_id = split_header[0]
+            header_suffix = split_header[1]
+            # the span is the set of alphanumeric characters after the /
+            span_id = re.findall(r"^\w+", header_suffix)[0]
+        except IndexError:
+            pass
+    return trace_id, span_id
 
 
 def get_request_data():
     """Helper to get http_request and trace data from supported web
     frameworks (currently supported: Flask and Django).
 
     Returns:
-        Tuple[Optional[dict], Optional[str]]:
-            Data related to the current http request and the trace_id for the
-            request. Both fields will be None if a supported web request isn't found.
+        Tuple[Optional[dict], Optional[str], Optional[str]]:
+            Data related to the current http request, trace_id, and span_id for
+            the request. All fields will be None if a django request isn't
+            found.
     """
     checkers = (
         get_request_data_from_django,
         get_request_data_from_flask,
     )
 
     for checker in checkers:
-        http_request, trace_id = checker()
+        http_request, trace_id, span_id = checker()
         if http_request is not None:
-            return http_request, trace_id
+            return http_request, trace_id, span_id
 
-    return None, None
+    return None, None, None

google/cloud/logging_v2/handlers/app_engine.py

@@ -90,7 +90,7 @@ def get_gae_labels(self):
         """
         gae_labels = {}
 
-        _, trace_id = get_request_data()
+        _, trace_id, _ = get_request_data()
         if trace_id is not None:
             gae_labels[_TRACE_ID_LABEL] = trace_id
 
@@ -107,7 +107,7 @@ def emit(self, record):
             record (logging.LogRecord): The record to be logged.
         """
         message = super(AppEngineHandler, self).format(record)
-        inferred_http, inferred_trace = get_request_data()
+        inferred_http, inferred_trace, _ = get_request_data()
         if inferred_trace is not None:
             inferred_trace = f"projects/{self.project_id}/traces/{inferred_trace}"
         # allow user overrides

google/cloud/logging_v2/handlers/handlers.py

@@ -59,7 +59,7 @@ def filter(self, record):
                 }
         record.msg = "" if record.msg is None else record.msg
         # find http request data
-        inferred_http, inferred_trace = get_request_data()
+        inferred_http, inferred_trace, inferred_span = get_request_data()
         if inferred_trace is not None and self.project is not None:
             inferred_trace = f"projects/{self.project}/traces/{inferred_trace}"
         # set labels
@@ -70,6 +70,7 @@ def filter(self, record):
         )
 
         record.trace = getattr(record, "trace", inferred_trace) or ""
+        record.span_id = getattr(record, "span_id", inferred_span) or ""
         record.http_request = getattr(record, "http_request", inferred_http) or {}
         record.request_method = record.http_request.get("requestMethod", "")
         record.request_url = record.http_request.get("requestUrl", "")

google/cloud/logging_v2/handlers/structured_log.py

@@ -24,6 +24,7 @@
     '"severity": "%(levelname)s", '
     '"logging.googleapis.com/labels": { %(total_labels_str)s }, '
     '"logging.googleapis.com/trace": "%(trace)s", '
+    '"logging.googleapis.com/spanId": "%(span_id)s", '
     '"logging.googleapis.com/sourceLocation": { "file": "%(file)s", "line": "%(line)d", "function": "%(function)s"}, '
     '"httpRequest": {"requestMethod": "%(request_method)s", "requestUrl": "%(request_url)s", "userAgent": "%(user_agent)s", "protocol": "%(protocol)s"} }'
 )