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"} }'
 )