feat: OpenTelemetry trace/spanID integration for Python handlers (#889)

* feat: OpenTelemetry trace/spanID integration for Python handlers * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Added more tests for OTel Python integration * linting * more linting * renamed _parse_current_open_telemetry_span and fixed otel testcases * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * linting + removed print statements * added opentelemetry sdk module cleanup to system test * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md * Refactored get_request_and_trace_data back into get_request_data * 🦉 Updates from OwlBot post-processor See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
googleapis · May 22, 2024 · 78168a3 · 78168a3
1 parent e46bbf8
commit 78168a3
Show file tree

Hide file tree

Showing 9 changed files with 538 additions and 6 deletions.
diff --git a/google/cloud/logging_v2/handlers/_helpers.py b/google/cloud/logging_v2/handlers/_helpers.py
@@ -24,6 +24,8 @@
 except ImportError:  # pragma: NO COVER
     flask = None
 
+import opentelemetry.trace
+
 from google.cloud.logging_v2.handlers.middleware.request import _get_django_request
 
 _DJANGO_CONTENT_LENGTH = "CONTENT_LENGTH"
@@ -191,23 +193,71 @@ def _parse_xcloud_trace(header):
     return trace_id, span_id, trace_sampled
 
 
+def _retrieve_current_open_telemetry_span():
+    """Helper to retrieve trace, span ID, and trace sampled information from the current
+    OpenTelemetry span.
+
+    Returns:
+        Tuple[Optional[str], Optional[str], bool]:
+            Data related to the current trace_id, span_id, and trace_sampled for the
+            current OpenTelemetry span. If a span is not found, return None/False for all
+            fields.
+    """
+    span = opentelemetry.trace.get_current_span()
+    if span != opentelemetry.trace.span.INVALID_SPAN:
+        context = span.get_span_context()
+        trace_id = opentelemetry.trace.format_trace_id(context.trace_id)
+        span_id = opentelemetry.trace.format_span_id(context.span_id)
+        trace_sampled = context.trace_flags.sampled
+
+        return trace_id, span_id, trace_sampled
+
+    return None, None, False
+
+
 def get_request_data():
     """Helper to get http_request and trace data from supported web
-    frameworks (currently supported: Flask and Django).
+    frameworks (currently supported: Flask and Django), as well as OpenTelemetry. Attempts
+    to retrieve trace/spanID from OpenTelemetry first, before going to Traceparent then XCTC.
+    HTTP request data is taken from a supporting web framework (currently Flask or Django).
+    Because HTTP request data is decoupled from OpenTelemetry, it is possible to get as a
+    return value the HTTP request from the web framework of choice, and trace/span data from
+    OpenTelemetry, even if trace data is present in the HTTP request headers.
 
     Returns:
         Tuple[Optional[dict], Optional[str], Optional[str], bool]:
             Data related to the current http request, trace_id, span_id, and trace_sampled
             for the request. All fields will be None if a http request isn't found.
     """
+
+    (
+        otel_trace_id,
+        otel_span_id,
+        otel_trace_sampled,
+    ) = _retrieve_current_open_telemetry_span()
+
+    # Get HTTP request data
     checkers = (
         get_request_data_from_django,
         get_request_data_from_flask,
     )
 
-    for checker in checkers:
-        http_request, trace_id, span_id, trace_sampled = checker()
-        if http_request is not None:
-            return http_request, trace_id, span_id, trace_sampled
+    http_request, http_trace_id, http_span_id, http_trace_sampled = (
+        None,
+        None,
+        None,
+        False,
+    )
 
-    return None, None, None, False
+    for checker in checkers:
+        http_request, http_trace_id, http_span_id, http_trace_sampled = checker()
+        if http_request is None:
+            http_trace_id, http_span_id, http_trace_sampled = None, None, False
+        else:
+            break
+
+    # otel_trace_id existing means the other return values are non-null
+    if otel_trace_id:
+        return http_request, otel_trace_id, otel_span_id, otel_trace_sampled
+    else:
+        return http_request, http_trace_id, http_span_id, http_trace_sampled
diff --git a/noxfile.py b/noxfile.py
@@ -63,6 +63,7 @@
     "google-cloud-pubsub",
     "google-cloud-storage",
     "google-cloud-testutils",
+    "opentelemetry-sdk",
 ]
 SYSTEM_TEST_LOCAL_DEPENDENCIES: List[str] = []
 SYSTEM_TEST_DEPENDENCIES: List[str] = []

diff --git a/owlbot.py b/owlbot.py
@@ -93,6 +93,7 @@ def place_before(path, text, *before_text, escape=None):
         "google-cloud-pubsub",
         "google-cloud-storage",
         "google-cloud-testutils",
+        "opentelemetry-sdk"
     ],
     unit_test_external_dependencies=["flask", "webob", "django"],
     samples=True,

diff --git a/setup.py b/setup.py
@@ -44,6 +44,7 @@
     "google-cloud-audit-log >= 0.1.0, < 1.0.0dev",
     "google-cloud-core >= 2.0.0, <3.0.0dev",
     "grpc-google-iam-v1 >=0.12.4, <1.0.0dev",
+    "opentelemetry-api >= 1.0.0",
     "proto-plus >= 1.22.0, <2.0.0dev",
     "proto-plus >= 1.22.2, <2.0.0dev; python_version>='3.11'",
     "protobuf>=3.19.5,<5.0.0dev,!=3.20.0,!=3.20.1,!=4.21.0,!=4.21.1,!=4.21.2,!=4.21.3,!=4.21.4,!=4.21.5",

diff --git a/tests/system/test_system.py b/tests/system/test_system.py
@@ -19,6 +19,7 @@
 import numbers
 import os
 import pytest
+import sys
 import unittest
 import uuid
 
@@ -117,6 +118,25 @@ def setUpModule():
 )
 
 
+def _cleanup_otel_sdk_modules(f):
+    """
+    Decorator to delete all references to opentelemetry SDK modules after a
+    testcase is run. Test case should import opentelemetry SDK modules inside
+    the function. This is to test situations where the opentelemetry SDK
+    is not imported at all.
+    """
+
+    def wrapped(*args, **kwargs):
+        f(*args, **kwargs)
+
+        # Deleting from sys.modules should be good enough in this use case
+        for module_name in list(sys.modules.keys()):
+            if module_name.startswith("opentelemetry.sdk"):
+                sys.modules.pop(module_name)
+
+    return wrapped
+
+
 class TestLogging(unittest.TestCase):
     JSON_PAYLOAD = {
         "message": "System test: test_log_struct",
@@ -662,6 +682,43 @@ def test_log_root_handler(self):
         self.assertEqual(len(entries), 1)
         self.assertEqual(entries[0].payload, expected_payload)
 
+    @_cleanup_otel_sdk_modules
+    def test_log_handler_otel_integration(self):
+        # Doing OTel imports here to not taint the other tests with OTel SDK imports
+        from opentelemetry import trace
+        from opentelemetry.sdk.trace import TracerProvider
+
+        LOG_MESSAGE = "This is a test of OpenTelemetry"
+        LOGGER_NAME = "otel-integration"
+        handler_name = self._logger_name(LOGGER_NAME)
+
+        handler = CloudLoggingHandler(
+            Config.CLIENT, name=handler_name, transport=SyncTransport
+        )
+        # only create the logger to delete, hidden otherwise
+        logger = Config.CLIENT.logger(handler.name)
+        self.to_delete.append(logger)
+
+        # Set up OTel SDK
+        provider = TracerProvider()
+
+        tracer = provider.get_tracer("test_system")
+        with tracer.start_as_current_span("test-span") as span:
+            context = span.get_span_context()
+            expected_trace_id = f"projects/{Config.CLIENT.project}/traces/{trace.format_trace_id(context.trace_id)}"
+            expected_span_id = trace.format_span_id(context.span_id)
+            expected_tracesampled = context.trace_flags.sampled
+
+            cloud_logger = logging.getLogger(LOGGER_NAME)
+            cloud_logger.addHandler(handler)
+            cloud_logger.warning(LOG_MESSAGE)
+
+            entries = _list_entries(logger)
+            self.assertEqual(len(entries), 1)
+            self.assertEqual(entries[0].trace, expected_trace_id)
+            self.assertEqual(entries[0].span_id, expected_span_id)
+            self.assertTrue(entries[0].trace_sampled, expected_tracesampled)
+
     def test_create_metric(self):
         METRIC_NAME = "test-create-metric%s" % (_RESOURCE_ID,)
         metric = Config.CLIENT.metric(

diff --git a/tests/unit/handlers/__init__.py b/tests/unit/handlers/__init__.py
@@ -11,3 +11,44 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+
+
+# Utility functions to setup mock OpenTelemetry spans, needed by multiple test
+# suites.
+
+import contextlib
+
+import opentelemetry.context
+import opentelemetry.trace
+
+from opentelemetry.trace import NonRecordingSpan
+from opentelemetry.trace.span import TraceFlags
+
+_OTEL_SPAN_CONTEXT_TRACE_ID = 0x123456789123456789
+_OTEL_SPAN_CONTEXT_SPAN_ID = 0x123456789
+_OTEL_SPAN_CONTEXT_TRACEFLAGS = TraceFlags(TraceFlags.SAMPLED)
+
+_EXPECTED_OTEL_TRACE_ID = "00000000000000123456789123456789"
+_EXPECTED_OTEL_SPAN_ID = "0000000123456789"
+_EXPECTED_OTEL_TRACESAMPLED = True
+
+
+@contextlib.contextmanager
+def _setup_otel_span_context():
+    """Sets up a nonrecording OpenTelemetry span with a mock span context that gets returned
+    by opentelemetry.trace.get_current_span, and returns it as a contextmanager
+    """
+    span_context = opentelemetry.trace.SpanContext(
+        _OTEL_SPAN_CONTEXT_TRACE_ID,
+        _OTEL_SPAN_CONTEXT_SPAN_ID,
+        False,
+        trace_flags=_OTEL_SPAN_CONTEXT_TRACEFLAGS,
+    )
+    ctx = opentelemetry.trace.set_span_in_context(NonRecordingSpan(span_context))
+    tracer = opentelemetry.trace.NoOpTracer()
+    token = opentelemetry.context.attach(ctx)
+    try:
+        with tracer.start_as_current_span("test-span", context=ctx):
+            yield
+    finally:
+        opentelemetry.context.detach(token)
diff --git a/tests/unit/handlers/test__helpers.py b/tests/unit/handlers/test__helpers.py
@@ -16,6 +16,13 @@
 
 import mock
 
+from tests.unit.handlers import (
+    _setup_otel_span_context,
+    _EXPECTED_OTEL_TRACE_ID,
+    _EXPECTED_OTEL_SPAN_ID,
+    _EXPECTED_OTEL_TRACESAMPLED,
+)
+
 _FLASK_TRACE_ID = "flask0id"
 _FLASK_SPAN_ID = "span0flask"
 _FLASK_HTTP_REQUEST = {"requestUrl": "https://flask.palletsprojects.com/en/1.1.x/"}
@@ -356,6 +363,120 @@ def test_wo_libraries(self):
         output = self._call_fut()
         self.assertEqual(output, (None, None, None, False))
 
+    def test_otel_span_exists_no_request(self):
+        flask_expected = (None, None, None, False)
+        django_expected = (None, None, None, False)
+
+        with _setup_otel_span_context():
+            _, _, output = self._helper(django_expected, flask_expected)
+            self.assertEqual(
+                output,
+                (
+                    None,
+                    _EXPECTED_OTEL_TRACE_ID,
+                    _EXPECTED_OTEL_SPAN_ID,
+                    _EXPECTED_OTEL_TRACESAMPLED,
+                ),
+            )
+
+    def test_otel_span_exists_django_request(self):
+        django_expected = (
+            _DJANGO_HTTP_REQUEST,
+            _DJANGO_TRACE_ID,
+            _DJANGO_SPAN_ID,
+            False,
+        )
+        flask_expected = (None, None, None, False)
+
+        with _setup_otel_span_context():
+            _, _, output = self._helper(django_expected, flask_expected)
+            self.assertEqual(
+                output,
+                (
+                    _DJANGO_HTTP_REQUEST,
+                    _EXPECTED_OTEL_TRACE_ID,
+                    _EXPECTED_OTEL_SPAN_ID,
+                    _EXPECTED_OTEL_TRACESAMPLED,
+                ),
+            )
+
+    def test_otel_span_exists_flask_request(self):
+        django_expected = (None, None, None, False)
+        flask_expected = (_FLASK_HTTP_REQUEST, _FLASK_TRACE_ID, _FLASK_SPAN_ID, False)
+
+        with _setup_otel_span_context():
+            _, _, output = self._helper(django_expected, flask_expected)
+            self.assertEqual(
+                output,
+                (
+                    _FLASK_HTTP_REQUEST,
+                    _EXPECTED_OTEL_TRACE_ID,
+                    _EXPECTED_OTEL_SPAN_ID,
+                    _EXPECTED_OTEL_TRACESAMPLED,
+                ),
+            )
+
+    def test_otel_span_exists_both_django_and_flask(self):
+        django_expected = (
+            _DJANGO_HTTP_REQUEST,
+            _DJANGO_TRACE_ID,
+            _DJANGO_SPAN_ID,
+            False,
+        )
+        flask_expected = (_FLASK_HTTP_REQUEST, _FLASK_TRACE_ID, _FLASK_SPAN_ID, False)
+
+        with _setup_otel_span_context():
+            _, _, output = self._helper(django_expected, flask_expected)
+
+            # Django wins
+            self.assertEqual(
+                output,
+                (
+                    _DJANGO_HTTP_REQUEST,
+                    _EXPECTED_OTEL_TRACE_ID,
+                    _EXPECTED_OTEL_SPAN_ID,
+                    _EXPECTED_OTEL_TRACESAMPLED,
+                ),
+            )
+
+    def test_no_otel_span_no_requests(self):
+        flask_expected = (None, None, None, False)
+        django_expected = (None, None, None, False)
+        _, _, output = self._helper(django_expected, flask_expected)
+        self.assertEqual(output, (None, None, None, False))
+
+    def test_no_otel_span_django_request(self):
+        django_expected = (
+            _DJANGO_HTTP_REQUEST,
+            _DJANGO_TRACE_ID,
+            _DJANGO_SPAN_ID,
+            False,
+        )
+        flask_expected = (None, None, None, False)
+        _, _, output = self._helper(django_expected, flask_expected)
+        self.assertEqual(output, django_expected)
+
+    def test_no_otel_span_flask_request(self):
+        django_expected = (None, None, None, False)
+        flask_expected = (_FLASK_HTTP_REQUEST, _FLASK_TRACE_ID, _FLASK_SPAN_ID, False)
+        _, _, output = self._helper(django_expected, flask_expected)
+
+        # Django wins
+        self.assertEqual(output, flask_expected)
+
+    def test_no_otel_span_both_django_and_flask(self):
+        django_expected = (
+            _DJANGO_HTTP_REQUEST,
+            _DJANGO_TRACE_ID,
+            _DJANGO_SPAN_ID,
+            False,
+        )
+        flask_expected = (_FLASK_HTTP_REQUEST, _FLASK_TRACE_ID, _FLASK_SPAN_ID, False)
+        _, _, output = self._helper(django_expected, flask_expected)
+
+        # Django wins
+        self.assertEqual(output, django_expected)
+
 
 class Test__parse_xcloud_trace(unittest.TestCase):
     @staticmethod
@@ -477,3 +598,25 @@ def test_invalid_headers(self):
             self.assertIsNone(trace_id)
             self.assertIsNone(span_id)
             self.assertEqual(sampled, False)
+
+
+class Test__parse_open_telemetry_data(unittest.TestCase):
+    @staticmethod
+    def _call_fut():
+        from google.cloud.logging_v2.handlers import _helpers
+
+        trace, span, sampled = _helpers._retrieve_current_open_telemetry_span()
+        return trace, span, sampled
+
+    def test_no_op(self):
+        trace_id, span_id, sampled = self._call_fut()
+        self.assertIsNone(trace_id)
+        self.assertIsNone(span_id)
+        self.assertEqual(sampled, False)
+
+    def test_span_exists(self):
+        with _setup_otel_span_context():
+            trace_id, span_id, sampled = self._call_fut()
+            self.assertEqual(trace_id, _EXPECTED_OTEL_TRACE_ID)
+            self.assertEqual(span_id, _EXPECTED_OTEL_SPAN_ID)
+            self.assertEqual(sampled, _EXPECTED_OTEL_TRACESAMPLED)