feat: Add support for open telemetry

django_spanner/_opentelemetry_tracing.py

@@ -0,0 +1,53 @@
+# Copyright 2021 Google LLC
+#
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file or at
+# https://developers.google.com/open-source/licenses/bsd
+
+"""Manages OpenTelemetry trace creation and handling"""
+
+from contextlib import contextmanager
+
+from google.api_core.exceptions import GoogleAPICallError
+
+try:
+    from opentelemetry import trace
+    from opentelemetry.trace.status import Status, StatusCode
+
+    HAS_OPENTELEMETRY_INSTALLED = True
+except ImportError:
+    HAS_OPENTELEMETRY_INSTALLED = False
+
+
+@contextmanager
+def trace_call(name, connection, extra_attributes=None):
+    if not HAS_OPENTELEMETRY_INSTALLED or not connection:
+        # Empty context manager. Users will have to check if the generated value
+        # is None or a span.
+        yield None
+        return
+
+    tracer = trace.get_tracer(__name__)
+
+    # Set base attributes that we know for every trace created
+    attributes = {
+        "db.type": "spanner",
+        "db.engine": "django_spanner",
+        "db.project": connection.settings_dict["PROJECT"],
+        "db.instance": connection.settings_dict["INSTANCE"],
+        "db.name": connection.settings_dict["NAME"],
+    }
+
+    if extra_attributes:
+        attributes.update(extra_attributes)
+
+    with tracer.start_as_current_span(
+        name, kind=trace.SpanKind.CLIENT, attributes=attributes
+    ) as span:
+        try:
+            span.set_status(Status(StatusCode.OK))
+            yield span
+        except GoogleAPICallError as error:
+            span.set_status(Status(StatusCode.ERROR))
+            span.record_exception(error)
+            raise

noxfile.py

@@ -75,6 +75,11 @@ def default(session):
         "pytest",
         "pytest-cov",
         "coverage",
+        "sqlparse==0.3.0",
+        "google-cloud-spanner==3.0.0",
+        "opentelemetry-api==1.1.0",
+        "opentelemetry-sdk==1.1.0",
+        "opentelemetry-instrumentation==0.20b0",
     )
     session.install("-e", ".")
 

setup.py

@@ -18,7 +18,13 @@
 # 'Development Status :: 5 - Production/Stable'
 release_status = "Development Status :: 3 - Alpha"
 dependencies = ["sqlparse >= 0.3.0", "google-cloud-spanner >= 3.0.0"]
-extras = {}
+extras = {
+    "tracing": [
+        "opentelemetry-api >= 1.1.0",
+        "opentelemetry-sdk >= 1.1.0",
+        "opentelemetry-instrumentation >= 0.20b0",
+    ]
+}
 
 BASE_DIR = os.path.dirname(__file__)
 VERSION_FILENAME = os.path.join(BASE_DIR, "version.py")

testing/constraints-3.6.txt

@@ -6,4 +6,7 @@
 # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev",
 # Then this file should have foo==1.14.0
 sqlparse==0.3.0
-google-cloud-spanner==2.1.0
+google-cloud-spanner==3.0.0
+opentelemetry-api==1.1.0
+opentelemetry-sdk==1.1.0
+opentelemetry-instrumentation==0.20b0

tests/_helpers.py

@@ -0,0 +1,79 @@
+# Copyright 2021 Google LLC
+#
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file or at
+# https://developers.google.com/open-source/licenses/bsd
+
+import unittest
+import mock
+
+try:
+    from opentelemetry import trace
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+    from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+        InMemorySpanExporter,
+    )
+    from opentelemetry.trace.status import StatusCode
+
+    trace.set_tracer_provider(TracerProvider())
+
+    HAS_OPENTELEMETRY_INSTALLED = True
+except ImportError:
+    HAS_OPENTELEMETRY_INSTALLED = False
+
+    StatusCode = mock.Mock()
+
+_TEST_OT_EXPORTER = None
+_TEST_OT_PROVIDER_INITIALIZED = False
+
+
+def get_test_ot_exporter():
+    global _TEST_OT_EXPORTER
+
+    if _TEST_OT_EXPORTER is None:
+        _TEST_OT_EXPORTER = InMemorySpanExporter()
+    return _TEST_OT_EXPORTER
+
+
+def use_test_ot_exporter():
+    global _TEST_OT_PROVIDER_INITIALIZED
+
+    if _TEST_OT_PROVIDER_INITIALIZED:
+        return
+
+    provider = trace.get_tracer_provider()
+    if not hasattr(provider, "add_span_processor"):
+        return
+    provider.add_span_processor(SimpleSpanProcessor(get_test_ot_exporter()))
+    _TEST_OT_PROVIDER_INITIALIZED = True
+
+
+class OpenTelemetryBase(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        if HAS_OPENTELEMETRY_INSTALLED:
+            use_test_ot_exporter()
+            cls.ot_exporter = get_test_ot_exporter()
+
+    def tearDown(self):
+        if HAS_OPENTELEMETRY_INSTALLED:
+            self.ot_exporter.clear()
+
+    def assertNoSpans(self):
+        if HAS_OPENTELEMETRY_INSTALLED:
+            span_list = self.ot_exporter.get_finished_spans()
+            self.assertEqual(len(span_list), 0)
+
+    def assertSpanAttributes(
+        self, name, status=StatusCode.OK, attributes=None, span=None
+    ):
+        if HAS_OPENTELEMETRY_INSTALLED:
+            if not span:
+                span_list = self.ot_exporter.get_finished_spans()
+                self.assertEqual(len(span_list), 1)
+                span = span_list[0]
+
+            self.assertEqual(span.name, name)
+            self.assertEqual(span.status.status_code, status)
+            self.assertEqual(dict(span.attributes), attributes)

tests/unit/django_spanner/simple_test.py

@@ -7,13 +7,16 @@
 from django_spanner.client import DatabaseClient
 from django_spanner.base import DatabaseWrapper
 from django_spanner.operations import DatabaseOperations
-from unittest import TestCase
+
+# from unittest import TestCase
+from tests._helpers import OpenTelemetryBase
 import os
 
 
-class SpannerSimpleTestClass(TestCase):
+class SpannerSimpleTestClass(OpenTelemetryBase):
     @classmethod
     def setUpClass(cls):
+        super(SpannerSimpleTestClass, cls).setUpClass()
         cls.PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
 
         cls.INSTANCE_ID = "instance_id"

tests/unit/django_spanner/test__opentelemetry_tracing.py

@@ -0,0 +1,130 @@
+# Copyright 2021 Google LLC
+#
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file or at
+# https://developers.google.com/open-source/licenses/bsd
+
+import importlib
+import mock
+import unittest
+import sys
+import os
+
+try:
+    from opentelemetry import trace as trace_api
+    from opentelemetry.trace.status import StatusCode
+except ImportError:
+    pass
+
+from google.api_core.exceptions import GoogleAPICallError
+from django_spanner import _opentelemetry_tracing
+
+from tests._helpers import OpenTelemetryBase, HAS_OPENTELEMETRY_INSTALLED
+
+PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
+INSTANCE_ID = "instance_id"
+DATABASE_ID = "database_id"
+OPTIONS = {"option": "dummy"}
+
+
+def _make_rpc_error(error_cls, trailing_metadata=None):
+    import grpc
+
+    grpc_error = mock.create_autospec(grpc.Call, instance=True)
+    grpc_error.trailing_metadata.return_value = trailing_metadata
+    return error_cls("error", errors=(grpc_error,))
+
+
+def _make_connection():
+    from django_spanner.base import DatabaseWrapper
+
+    settings_dict = {
+        "PROJECT": PROJECT,
+        "INSTANCE": INSTANCE_ID,
+        "NAME": DATABASE_ID,
+        "OPTIONS": OPTIONS,
+    }
+    return DatabaseWrapper(settings_dict)
+
+
+# Skip all of these tests if we don't have OpenTelemetry
+if HAS_OPENTELEMETRY_INSTALLED:
+
+    class TestNoTracing(unittest.TestCase):
+        def setUp(self):
+            self._temp_opentelemetry = sys.modules["opentelemetry"]
+
+            sys.modules["opentelemetry"] = None
+            importlib.reload(_opentelemetry_tracing)
+
+        def tearDown(self):
+            sys.modules["opentelemetry"] = self._temp_opentelemetry
+            importlib.reload(_opentelemetry_tracing)
+
+        def test_no_trace_call(self):
+            with _opentelemetry_tracing.trace_call(
+                "Test", _make_connection()
+            ) as no_span:
+                self.assertIsNone(no_span)
+
+    class TestTracing(OpenTelemetryBase):
+        def test_trace_call(self):
+            extra_attributes = {
+                "attribute1": "value1",
+                # Since our database is mocked, we have to override the db.instance parameter so it is a string
+                "db.instance": "database_name",
+            }
+
+            expected_attributes = {
+                "db.type": "spanner",
+                "db.engine": "django_spanner",
+                "db.project": PROJECT,
+                "db.instance": INSTANCE_ID,
+                "db.name": DATABASE_ID,
+            }
+            expected_attributes.update(extra_attributes)
+
+            with _opentelemetry_tracing.trace_call(
+                "CloudSpannerDjango.Test", _make_connection(), extra_attributes
+            ) as span:
+                span.set_attribute("after_setup_attribute", 1)
+
+            expected_attributes["after_setup_attribute"] = 1
+
+            span_list = self.ot_exporter.get_finished_spans()
+            self.assertEqual(len(span_list), 1)
+            span = span_list[0]
+            self.assertEqual(span.kind, trace_api.SpanKind.CLIENT)
+            self.assertEqual(span.attributes, expected_attributes)
+            self.assertEqual(span.name, "CloudSpannerDjango.Test")
+            self.assertEqual(span.status.status_code, StatusCode.OK)
+
+        def test_trace_error(self):
+            extra_attributes = {"db.instance": "database_name"}
+
+            expected_attributes = {
+                "db.type": "spanner",
+                "db.engine": "django_spanner",
+                "db.project": os.environ["GOOGLE_CLOUD_PROJECT"],
+                "db.instance": "instance_id",
+                "db.name": "database_id",
+            }
+            expected_attributes.update(extra_attributes)
+
+            with self.assertRaises(GoogleAPICallError):
+                with _opentelemetry_tracing.trace_call(
+                    "CloudSpannerDjango.Test",
+                    _make_connection(),
+                    extra_attributes,
+                ) as span:
+                    from google.api_core.exceptions import InvalidArgument
+
+                    raise _make_rpc_error(InvalidArgument)
+
+            span_list = self.ot_exporter.get_finished_spans()
+            self.assertEqual(len(span_list), 1)
+            span = span_list[0]
+            self.assertEqual(span.kind, trace_api.SpanKind.CLIENT)
+            self.assertEqual(dict(span.attributes), expected_attributes)
+            self.assertEqual(span.name, "CloudSpannerDjango.Test")
+            self.assertEqual(span.status.status_code, StatusCode.ERROR)