diff --git a/django_spanner/_opentelemetry_tracing.py b/django_spanner/_opentelemetry_tracing.py new file mode 100644 index 0000000000..b8c3c1e15a --- /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 index b2d42a5174..a5c05e7a02 100644 --- 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 index f37e04b0c1..c310fda167 100644 --- 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") diff --git a/testing/constraints-3.6.txt b/testing/constraints-3.6.txt index 0136bc8c30..7573802344 100644 --- 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 index 0000000000..4faff45f2c --- /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 index 1fcb92bd29..0c66d04dfe 100644 --- 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 index 0000000000..076b9a83d9 --- /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)