feat: support span inference (#267)

Author: daniel-sanche

packages/google-cloud-logging/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

packages/google-cloud-logging/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

packages/google-cloud-logging/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", "")

packages/google-cloud-logging/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"} }'
 )

packages/google-cloud-logging/tests/environment/tests/common/common.py

@@ -137,6 +137,21 @@ def test_receive_log(self):
                 found_log = log
         self.assertIsNotNone(found_log, "expected log text not found")
 
+    def test_receive_unicode_log(self):
+        log_text = f"{inspect.currentframe().f_code.co_name} 嗨 世界 😀"
+        log_list = self.trigger_and_retrieve(log_text, "simplelog")
+
+        found_log = None
+        for log in log_list:
+            message = (
+                log.payload.get("message", None)
+                if isinstance(log.payload, dict)
+                else str(log.payload)
+            )
+            if message and log_text in message:
+                found_log = log
+        self.assertIsNotNone(found_log, "expected unicode log not found")
+
     # add back after v3.0.0
     # def test_monitored_resource(self):
     #     if self.language != "python":

packages/google-cloud-logging/tests/environment/tests/common/python.py

@@ -22,7 +22,7 @@
 
 
 class CommonPython:
-    def pylogging_test_receive_log(self):
+    def test_pylogging_receive_log(self):
         log_text = f"{inspect.currentframe().f_code.co_name}"
         log_list = self.trigger_and_retrieve(log_text, "pylogging")
 
@@ -37,6 +37,21 @@ def pylogging_test_receive_log(self):
                 found_log = log
         self.assertIsNotNone(found_log, "expected log text not found")
 
+    def test_pylogging_receive_unicode_log(self):
+        log_text = f"{inspect.currentframe().f_code.co_name} 嗨 世界 😀"
+        log_list = self.trigger_and_retrieve(log_text, "pylogging")
+
+        found_log = None
+        for log in log_list:
+            message = (
+                log.payload.get("message", None)
+                if isinstance(log.payload, dict)
+                else str(log.payload)
+            )
+            if message and log_text in message:
+                found_log = log
+        self.assertIsNotNone(found_log, "expected unicode log not found")
+
     def test_monitored_resource_pylogging(self):
         log_text = f"{inspect.currentframe().f_code.co_name}"
         log_list = self.trigger_and_retrieve(log_text, "pylogging")
@@ -90,12 +105,14 @@ def test_flask_http_request_pylogging(self):
         expected_base_url = "http://test"
         expected_path = "/pylogging"
         expected_trace = "123"
+        expected_span = "456"
+        trace_header = f"{expected_trace}/{expected_span};o=1"
 
         log_list = self.trigger_and_retrieve(
             log_text,
             "pylogging_flask",
             path=expected_path,
-            trace=expected_trace,
+            trace=trace_header,
             base_url=expected_base_url,
             agent=expected_agent,
         )
@@ -112,8 +129,13 @@ def test_flask_http_request_pylogging(self):
         self.assertEqual(found_request["protocol"], "HTTP/1.1")
 
         found_trace = log_list[-1].trace
+        found_span = log_list[-1].span_id
         self.assertIsNotNone(found_trace)
         self.assertIn("projects/", found_trace)
+        if self.environment != "functions":
+            # functions seems to override the user's trace value
+            self.assertIn(expected_trace, found_trace)
+            self.assertEqual(expected_span, found_span)
 
     def test_pylogging_extras(self):
         if self.environment == "kubernetes" or "appengine" in self.environment:
@@ -123,6 +145,7 @@ def test_pylogging_extras(self):
         log_text = f"{inspect.currentframe().f_code.co_name}"
         kwargs = {
             "trace": "123",
+            "span_id": "456",
             "requestMethod": "POST",
             "requestUrl": "http://test",
             "userAgent": "agent",
@@ -138,6 +161,7 @@ def test_pylogging_extras(self):
         if self.environment != "functions":
             # functions seems to override the user's trace value
             self.assertEqual(found_log.trace, kwargs["trace"])
+            self.assertEqual(found_log.span_id, kwargs["span_id"])
 
         # check that custom http request fields were set
         self.assertIsNotNone(found_log.http_request)