Repository: googleapis/python-spanner-django

Issue:
feat: Add support for open telemetry
Add support for open telemetry tracing if open telemetry is installed.
Added the same as extra dependency.

Also added unit tests for the same.

Follow up PR to add tracing in schema with related checks in unit tests. https://github.com/googleapis/python-spanner-django/pull/648

Fixes #633 🦕


Discussion:

Resolving patch:
diff --git a/django_spanner/_opentelemetry_tracing.py b/django_spanner/_opentelemetry_tracing.py
new file mode 100644
--- /dev/null
+++ b/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
diff --git a/noxfile.py b/noxfile.py
--- a/noxfile.py
+++ b/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", ".")
 
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/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")


Test patch:
diff --git a/testing/constraints-3.6.txt b/testing/constraints-3.6.txt
--- a/testing/constraints-3.6.txt
+++ b/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
\ No newline at end of file
+google-cloud-spanner==3.0.0
+opentelemetry-api==1.1.0
+opentelemetry-sdk==1.1.0
+opentelemetry-instrumentation==0.20b0
diff --git a/tests/_helpers.py b/tests/_helpers.py
new file mode 100644
--- /dev/null
+++ b/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)
diff --git a/tests/unit/django_spanner/simple_test.py b/tests/unit/django_spanner/simple_test.py
--- a/tests/unit/django_spanner/simple_test.py
+++ b/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"
diff --git a/tests/unit/django_spanner/test__opentelemetry_tracing.py b/tests/unit/django_spanner/test__opentelemetry_tracing.py
new file mode 100644
--- /dev/null
+++ b/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)