fix: for open-telemetry status code spec change (#88)

* fix: for open-telemetry status code spec change

* fix: updated open telemetry version

* refactor: lint fixes

* refactor: removed commented and unused line

* refactor: change OpenTelemetryBase to inheri unittest.TestCase

* refactor: lint corrections

* refactor: lint correction

* refactor: lint fixes

* test: add separate unit test session in noxfile

* refactor: move open telemetry test inside test/unit tests

Author: vi3k6i5

google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py

@@ -18,15 +18,10 @@
 
 from google.api_core.exceptions import GoogleAPICallError
 from google.cloud.spanner_v1 import SpannerClient
-from google.cloud.spanner_dbapi.exceptions import IntegrityError
-from google.cloud.spanner_dbapi.exceptions import InterfaceError
-from google.cloud.spanner_dbapi.exceptions import OperationalError
-from google.cloud.spanner_dbapi.exceptions import ProgrammingError
 
 try:
     from opentelemetry import trace
-    from opentelemetry.trace.status import Status, StatusCanonicalCode
-    from opentelemetry.instrumentation.utils import http_status_to_canonical_code
+    from opentelemetry.trace.status import Status, StatusCode
 
     HAS_OPENTELEMETRY_INSTALLED = True
 except ImportError:
@@ -45,6 +40,7 @@ def trace_call(name, extra_attributes=None):
     # Set base attributes that we know for every trace created
     attributes = {
         "db.type": "spanner",
+        "db.engine": "sqlalchemy_spanner",
         "db.url": SpannerClient.DEFAULT_ENDPOINT,
         "net.host.name": SpannerClient.DEFAULT_ENDPOINT,
     }
@@ -56,20 +52,9 @@ def trace_call(name, extra_attributes=None):
         name, kind=trace.SpanKind.CLIENT, attributes=attributes
     ) as span:
         try:
+            span.set_status(Status(StatusCode.OK))
             yield span
-        except (ValueError, InterfaceError) as e:
-            span.set_status(Status(StatusCanonicalCode.UNKNOWN, e.args[0]))
         except GoogleAPICallError as error:
-            if error.code is not None:
-                span.set_status(Status(http_status_to_canonical_code(error.code)))
-            elif error.grpc_status_code is not None:
-                span.set_status(
-                    # OpenTelemetry's StatusCanonicalCode maps 1-1 with grpc status
-                    # codes
-                    Status(StatusCanonicalCode(error.grpc_status_code.value[0]))
-                )
+            span.set_status(Status(StatusCode.ERROR))
+            span.record_exception(error)
             raise
-        except (IntegrityError, ProgrammingError, OperationalError) as e:
-            span.set_status(
-                Status(http_status_to_canonical_code(e.args[0].code), e.args[0].message)
-            )

noxfile.py

@@ -18,6 +18,7 @@
 
 import configparser
 import nox
+import os
 
 ALEMBIC_CONF = """
 [alembic]
@@ -110,11 +111,25 @@ def lint_setup_py(session):
 def compliance_test(session):
     """Run SQLAlchemy dialect compliance test suite."""
     session.install("pytest")
+    session.install("mock")
     session.install("-e", ".")
     session.run("python", "create_test_database.py")
     session.run("pytest", "-v")
 
 
+@nox.session(python=DEFAULT_PYTHON_VERSION)
+def unit(session):
+    """Run unit tests."""
+    # Run SQLAlchemy dialect compliance test suite with OpenTelemetry.
+    session.install("pytest")
+    session.install("mock")
+    session.install("-e", ".")
+    session.install("opentelemetry-api==1.1.0")
+    session.install("opentelemetry-sdk==1.1.0")
+    session.install("opentelemetry-instrumentation==0.20b0")
+    session.run("py.test", "--quiet", os.path.join("test/unit"), *session.posargs)
+
+
 @nox.session(python=DEFAULT_PYTHON_VERSION)
 def migration_test(session):
     """Migrate with SQLAlchemy and Alembic and check the result."""

setup.py

@@ -22,9 +22,9 @@
 dependencies = ["sqlalchemy>=1.1.13, <=1.3.23", "google-cloud-spanner>=3.3.0"]
 extras = {
     "tracing": [
-        "opentelemetry-api==0.11b0",
-        "opentelemetry-sdk==0.11b0",
-        "opentelemetry-instrumentation==0.11b0",
+        "opentelemetry-api >= 1.1.0",
+        "opentelemetry-sdk >= 1.1.0",
+        "opentelemetry-instrumentation >= 0.20b0",
     ]
 }
 

test/__init__.py

@@ -0,0 +1,80 @@
+# 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 mock
+from sqlalchemy.testing import fixtures
+
+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(fixtures.TestBase):
+    @classmethod
+    def setup_class(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)

test/_helpers.py

@@ -0,0 +1,114 @@
+# 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 pytest
+import sys
+
+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 google.cloud.spanner_v1 import SpannerClient
+from google.cloud.sqlalchemy_spanner import _opentelemetry_tracing
+from test._helpers import OpenTelemetryBase, HAS_OPENTELEMETRY_INSTALLED
+
+
+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,))
+
+
+# Skip all of these tests if we don't have OpenTelemetry
+if HAS_OPENTELEMETRY_INSTALLED:
+
+    class TestNoTracing(OpenTelemetryBase):
+        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") as no_span:
+                assert no_span is None
+
+    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": "sqlalchemy_spanner",
+                "db.url": SpannerClient.DEFAULT_ENDPOINT,
+                "net.host.name": SpannerClient.DEFAULT_ENDPOINT,
+            }
+            expected_attributes.update(extra_attributes)
+
+            with _opentelemetry_tracing.trace_call(
+                "CloudSpannerSqlAlchemy.Test", 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()
+            assert len(span_list) == 1
+
+            span = span_list[0]
+            assert span.kind == trace_api.SpanKind.CLIENT
+            span_attr = dict(span.attributes)
+            for key in expected_attributes:
+                assert key in span_attr
+                assert span_attr[key] == expected_attributes[key]
+            assert span.name == "CloudSpannerSqlAlchemy.Test"
+            assert span.status.status_code == StatusCode.OK
+
+        def test_trace_error(self):
+            extra_attributes = {"db.instance": "database_name"}
+
+            expected_attributes = {
+                "db.type": "spanner",
+                "db.engine": "sqlalchemy_spanner",
+                "db.url": SpannerClient.DEFAULT_ENDPOINT,
+                "net.host.name": SpannerClient.DEFAULT_ENDPOINT,
+            }
+            expected_attributes.update(extra_attributes)
+
+            with pytest.raises(GoogleAPICallError):
+                with _opentelemetry_tracing.trace_call(
+                    "CloudSpannerSqlAlchemy.Test", extra_attributes,
+                ) as span:
+                    from google.api_core.exceptions import InvalidArgument
+
+                    raise _make_rpc_error(InvalidArgument)
+
+            span_list = self.ot_exporter.get_finished_spans()
+            assert len(span_list) == 1
+            span = span_list[0]
+            assert span.kind == trace_api.SpanKind.CLIENT
+            span_attr = dict(span.attributes)
+            for key in expected_attributes:
+                assert key in span_attr
+                assert span_attr[key] == expected_attributes[key]
+            assert span.name == "CloudSpannerSqlAlchemy.Test"
+            assert span.status.status_code == StatusCode.ERROR