| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
1 parent ecf241a commit 2ba879a
7 files changed
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -0,0 +1,53 @@ | |||
| 1 | + # Copyright 2021 Google LLC | ||
| 2 | + # | ||
| 3 | + # Use of this source code is governed by a BSD-style | ||
| 4 | + # license that can be found in the LICENSE file or at | ||
| 5 | + # https://developers.google.com/open-source/licenses/bsd | ||
| 6 | + | ||
| 7 | + """Manages OpenTelemetry trace creation and handling""" | ||
| 8 | + | ||
| 9 | + from contextlib import contextmanager | ||
| 10 | + | ||
| 11 | + from google.api_core.exceptions import GoogleAPICallError | ||
| 12 | + | ||
| 13 | + try: | ||
| 14 | + from opentelemetry import trace | ||
| 15 | + from opentelemetry.trace.status import Status, StatusCode | ||
| 16 | + | ||
| 17 | + HAS_OPENTELEMETRY_INSTALLED = True | ||
| 18 | + except ImportError: | ||
| 19 | + HAS_OPENTELEMETRY_INSTALLED = False | ||
| 20 | + | ||
| 21 | + | ||
| 22 | + @contextmanager | ||
| 23 | + def trace_call(name, connection, extra_attributes=None): | ||
| 24 | + if not HAS_OPENTELEMETRY_INSTALLED or not connection: | ||
| 25 | + # Empty context manager. Users will have to check if the generated value | ||
| 26 | + # is None or a span. | ||
| 27 | + yield None | ||
| 28 | + return | ||
| 29 | + | ||
| 30 | + tracer = trace.get_tracer(__name__) | ||
| 31 | + | ||
| 32 | + # Set base attributes that we know for every trace created | ||
| 33 | + attributes = { | ||
| 34 | + "db.type": "spanner", | ||
| 35 | + "db.engine": "django_spanner", | ||
| 36 | + "db.project": connection.settings_dict["PROJECT"], | ||
| 37 | + "db.instance": connection.settings_dict["INSTANCE"], | ||
| 38 | + "db.name": connection.settings_dict["NAME"], | ||
| 39 | + } | ||
| 40 | + | ||
| 41 | + if extra_attributes: | ||
| 42 | + attributes.update(extra_attributes) | ||
| 43 | + | ||
| 44 | + with tracer.start_as_current_span( | ||
| 45 | + name, kind=trace.SpanKind.CLIENT, attributes=attributes | ||
| 46 | + ) as span: | ||
| 47 | + try: | ||
| 48 | + span.set_status(Status(StatusCode.OK)) | ||
| 49 | + yield span | ||
| 50 | + except GoogleAPICallError as error: | ||
| 51 | + span.set_status(Status(StatusCode.ERROR)) | ||
| 52 | + span.record_exception(error) | ||
| 53 | + raise | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -75,6 +75,11 @@ def default(session): | |||
| 75 | 75 | "pytest", | |
| 76 | 76 | "pytest-cov", | |
| 77 | 77 | "coverage", | |
| 78 | + "sqlparse==0.3.0", | ||
| 79 | + "google-cloud-spanner==3.0.0", | ||
| 80 | + "opentelemetry-api==1.1.0", | ||
| 81 | + "opentelemetry-sdk==1.1.0", | ||
| 82 | + "opentelemetry-instrumentation==0.20b0", | ||
| 78 | 83 | ) | |
| 79 | 84 | session.install("-e", ".") | |
| 80 | 85 | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -18,7 +18,13 @@ | |||
| 18 | 18 | # 'Development Status :: 5 - Production/Stable' | |
| 19 | 19 | release_status = "Development Status :: 3 - Alpha" | |
| 20 | 20 | dependencies = ["sqlparse >= 0.3.0", "google-cloud-spanner >= 3.0.0"] | |
| 21 | - extras = {} | ||
| 21 | + extras = { | ||
| 22 | + "tracing": [ | ||
| 23 | + "opentelemetry-api >= 1.1.0", | ||
| 24 | + "opentelemetry-sdk >= 1.1.0", | ||
| 25 | + "opentelemetry-instrumentation >= 0.20b0", | ||
| 26 | + ] | ||
| 27 | + } | ||
| 22 | 28 | ||
| 23 | 29 | BASE_DIR = os.path.dirname(__file__) | |
| 24 | 30 | VERSION_FILENAME = os.path.join(BASE_DIR, "version.py") | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -6,4 +6,7 @@ | |||
| 6 | 6 | # e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev", | |
| 7 | 7 | # Then this file should have foo==1.14.0 | |
| 8 | 8 | sqlparse==0.3.0 | |
| 9 | - google-cloud-spanner==2.1.0 | ||
| 9 | + google-cloud-spanner==3.0.0 | ||
| 10 | + opentelemetry-api==1.1.0 | ||
| 11 | + opentelemetry-sdk==1.1.0 | ||
| 12 | + opentelemetry-instrumentation==0.20b0 | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -0,0 +1,79 @@ | |||
| 1 | + # Copyright 2021 Google LLC | ||
| 2 | + # | ||
| 3 | + # Use of this source code is governed by a BSD-style | ||
| 4 | + # license that can be found in the LICENSE file or at | ||
| 5 | + # https://developers.google.com/open-source/licenses/bsd | ||
| 6 | + | ||
| 7 | + import unittest | ||
| 8 | + import mock | ||
| 9 | + | ||
| 10 | + try: | ||
| 11 | + from opentelemetry import trace | ||
| 12 | + from opentelemetry.sdk.trace import TracerProvider | ||
| 13 | + from opentelemetry.sdk.trace.export import SimpleSpanProcessor | ||
| 14 | + from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( | ||
| 15 | + InMemorySpanExporter, | ||
| 16 | + ) | ||
| 17 | + from opentelemetry.trace.status import StatusCode | ||
| 18 | + | ||
| 19 | + trace.set_tracer_provider(TracerProvider()) | ||
| 20 | + | ||
| 21 | + HAS_OPENTELEMETRY_INSTALLED = True | ||
| 22 | + except ImportError: | ||
| 23 | + HAS_OPENTELEMETRY_INSTALLED = False | ||
| 24 | + | ||
| 25 | + StatusCode = mock.Mock() | ||
| 26 | + | ||
| 27 | + _TEST_OT_EXPORTER = None | ||
| 28 | + _TEST_OT_PROVIDER_INITIALIZED = False | ||
| 29 | + | ||
| 30 | + | ||
| 31 | + def get_test_ot_exporter(): | ||
| 32 | + global _TEST_OT_EXPORTER | ||
| 33 | + | ||
| 34 | + if _TEST_OT_EXPORTER is None: | ||
| 35 | + _TEST_OT_EXPORTER = InMemorySpanExporter() | ||
| 36 | + return _TEST_OT_EXPORTER | ||
| 37 | + | ||
| 38 | + | ||
| 39 | + def use_test_ot_exporter(): | ||
| 40 | + global _TEST_OT_PROVIDER_INITIALIZED | ||
| 41 | + | ||
| 42 | + if _TEST_OT_PROVIDER_INITIALIZED: | ||
| 43 | + return | ||
| 44 | + | ||
| 45 | + provider = trace.get_tracer_provider() | ||
| 46 | + if not hasattr(provider, "add_span_processor"): | ||
| 47 | + return | ||
| 48 | + provider.add_span_processor(SimpleSpanProcessor(get_test_ot_exporter())) | ||
| 49 | + _TEST_OT_PROVIDER_INITIALIZED = True | ||
| 50 | + | ||
| 51 | + | ||
| 52 | + class OpenTelemetryBase(unittest.TestCase): | ||
| 53 | + @classmethod | ||
| 54 | + def setUpClass(cls): | ||
| 55 | + if HAS_OPENTELEMETRY_INSTALLED: | ||
| 56 | + use_test_ot_exporter() | ||
| 57 | + cls.ot_exporter = get_test_ot_exporter() | ||
| 58 | + | ||
| 59 | + def tearDown(self): | ||
| 60 | + if HAS_OPENTELEMETRY_INSTALLED: | ||
| 61 | + self.ot_exporter.clear() | ||
| 62 | + | ||
| 63 | + def assertNoSpans(self): | ||
| 64 | + if HAS_OPENTELEMETRY_INSTALLED: | ||
| 65 | + span_list = self.ot_exporter.get_finished_spans() | ||
| 66 | + self.assertEqual(len(span_list), 0) | ||
| 67 | + | ||
| 68 | + def assertSpanAttributes( | ||
| 69 | + self, name, status=StatusCode.OK, attributes=None, span=None | ||
| 70 | + ): | ||
| 71 | + if HAS_OPENTELEMETRY_INSTALLED: | ||
| 72 | + if not span: | ||
| 73 | + span_list = self.ot_exporter.get_finished_spans() | ||
| 74 | + self.assertEqual(len(span_list), 1) | ||
| 75 | + span = span_list[0] | ||
| 76 | + | ||
| 77 | + self.assertEqual(span.name, name) | ||
| 78 | + self.assertEqual(span.status.status_code, status) | ||
| 79 | + self.assertEqual(dict(span.attributes), attributes) | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -7,13 +7,16 @@ | |||
| 7 | 7 | from django_spanner.client import DatabaseClient | |
| 8 | 8 | from django_spanner.base import DatabaseWrapper | |
| 9 | 9 | from django_spanner.operations import DatabaseOperations | |
| 10 | - from unittest import TestCase | ||
| 10 | + | ||
| 11 | + # from unittest import TestCase | ||
| 12 | + from tests._helpers import OpenTelemetryBase | ||
| 11 | 13 | import os | |
| 12 | 14 | ||
| 13 | 15 | ||
| 14 | - class SpannerSimpleTestClass(TestCase): | ||
| 16 | + class SpannerSimpleTestClass(OpenTelemetryBase): | ||
| 15 | 17 | @classmethod | |
| 16 | 18 | def setUpClass(cls): | |
| 19 | + super(SpannerSimpleTestClass, cls).setUpClass() | ||
| 17 | 20 | cls.PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"] | |
| 18 | 21 | ||
| 19 | 22 | cls.INSTANCE_ID = "instance_id" | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -0,0 +1,130 @@ | |||
| 1 | + # Copyright 2021 Google LLC | ||
| 2 | + # | ||
| 3 | + # Use of this source code is governed by a BSD-style | ||
| 4 | + # license that can be found in the LICENSE file or at | ||
| 5 | + # https://developers.google.com/open-source/licenses/bsd | ||
| 6 | + | ||
| 7 | + import importlib | ||
| 8 | + import mock | ||
| 9 | + import unittest | ||
| 10 | + import sys | ||
| 11 | + import os | ||
| 12 | + | ||
| 13 | + try: | ||
| 14 | + from opentelemetry import trace as trace_api | ||
| 15 | + from opentelemetry.trace.status import StatusCode | ||
| 16 | + except ImportError: | ||
| 17 | + pass | ||
| 18 | + | ||
| 19 | + from google.api_core.exceptions import GoogleAPICallError | ||
| 20 | + from django_spanner import _opentelemetry_tracing | ||
| 21 | + | ||
| 22 | + from tests._helpers import OpenTelemetryBase, HAS_OPENTELEMETRY_INSTALLED | ||
| 23 | + | ||
| 24 | + PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"] | ||
| 25 | + INSTANCE_ID = "instance_id" | ||
| 26 | + DATABASE_ID = "database_id" | ||
| 27 | + OPTIONS = {"option": "dummy"} | ||
| 28 | + | ||
| 29 | + | ||
| 30 | + def _make_rpc_error(error_cls, trailing_metadata=None): | ||
| 31 | + import grpc | ||
| 32 | + | ||
| 33 | + grpc_error = mock.create_autospec(grpc.Call, instance=True) | ||
| 34 | + grpc_error.trailing_metadata.return_value = trailing_metadata | ||
| 35 | + return error_cls("error", errors=(grpc_error,)) | ||
| 36 | + | ||
| 37 | + | ||
| 38 | + def _make_connection(): | ||
| 39 | + from django_spanner.base import DatabaseWrapper | ||
| 40 | + | ||
| 41 | + settings_dict = { | ||
| 42 | + "PROJECT": PROJECT, | ||
| 43 | + "INSTANCE": INSTANCE_ID, | ||
| 44 | + "NAME": DATABASE_ID, | ||
| 45 | + "OPTIONS": OPTIONS, | ||
| 46 | + } | ||
| 47 | + return DatabaseWrapper(settings_dict) | ||
| 48 | + | ||
| 49 | + | ||
| 50 | + # Skip all of these tests if we don't have OpenTelemetry | ||
| 51 | + if HAS_OPENTELEMETRY_INSTALLED: | ||
| 52 | + | ||
| 53 | + class TestNoTracing(unittest.TestCase): | ||
| 54 | + def setUp(self): | ||
| 55 | + self._temp_opentelemetry = sys.modules["opentelemetry"] | ||
| 56 | + | ||
| 57 | + sys.modules["opentelemetry"] = None | ||
| 58 | + importlib.reload(_opentelemetry_tracing) | ||
| 59 | + | ||
| 60 | + def tearDown(self): | ||
| 61 | + sys.modules["opentelemetry"] = self._temp_opentelemetry | ||
| 62 | + importlib.reload(_opentelemetry_tracing) | ||
| 63 | + | ||
| 64 | + def test_no_trace_call(self): | ||
| 65 | + with _opentelemetry_tracing.trace_call( | ||
| 66 | + "Test", _make_connection() | ||
| 67 | + ) as no_span: | ||
| 68 | + self.assertIsNone(no_span) | ||
| 69 | + | ||
| 70 | + class TestTracing(OpenTelemetryBase): | ||
| 71 | + def test_trace_call(self): | ||
| 72 | + extra_attributes = { | ||
| 73 | + "attribute1": "value1", | ||
| 74 | + # Since our database is mocked, we have to override the db.instance parameter so it is a string | ||
| 75 | + "db.instance": "database_name", | ||
| 76 | + } | ||
| 77 | + | ||
| 78 | + expected_attributes = { | ||
| 79 | + "db.type": "spanner", | ||
| 80 | + "db.engine": "django_spanner", | ||
| 81 | + "db.project": PROJECT, | ||
| 82 | + "db.instance": INSTANCE_ID, | ||
| 83 | + "db.name": DATABASE_ID, | ||
| 84 | + } | ||
| 85 | + expected_attributes.update(extra_attributes) | ||
| 86 | + | ||
| 87 | + with _opentelemetry_tracing.trace_call( | ||
| 88 | + "CloudSpannerDjango.Test", _make_connection(), extra_attributes | ||
| 89 | + ) as span: | ||
| 90 | + span.set_attribute("after_setup_attribute", 1) | ||
| 91 | + | ||
| 92 | + expected_attributes["after_setup_attribute"] = 1 | ||
| 93 | + | ||
| 94 | + span_list = self.ot_exporter.get_finished_spans() | ||
| 95 | + self.assertEqual(len(span_list), 1) | ||
| 96 | + span = span_list[0] | ||
| 97 | + self.assertEqual(span.kind, trace_api.SpanKind.CLIENT) | ||
| 98 | + self.assertEqual(span.attributes, expected_attributes) | ||
| 99 | + self.assertEqual(span.name, "CloudSpannerDjango.Test") | ||
| 100 | + self.assertEqual(span.status.status_code, StatusCode.OK) | ||
| 101 | + | ||
| 102 | + def test_trace_error(self): | ||
| 103 | + extra_attributes = {"db.instance": "database_name"} | ||
| 104 | + | ||
| 105 | + expected_attributes = { | ||
| 106 | + "db.type": "spanner", | ||
| 107 | + "db.engine": "django_spanner", | ||
| 108 | + "db.project": os.environ["GOOGLE_CLOUD_PROJECT"], | ||
| 109 | + "db.instance": "instance_id", | ||
| 110 | + "db.name": "database_id", | ||
| 111 | + } | ||
| 112 | + expected_attributes.update(extra_attributes) | ||
| 113 | + | ||
| 114 | + with self.assertRaises(GoogleAPICallError): | ||
| 115 | + with _opentelemetry_tracing.trace_call( | ||
| 116 | + "CloudSpannerDjango.Test", | ||
| 117 | + _make_connection(), | ||
| 118 | + extra_attributes, | ||
| 119 | + ) as span: | ||
| 120 | + from google.api_core.exceptions import InvalidArgument | ||
| 121 | + | ||
| 122 | + raise _make_rpc_error(InvalidArgument) | ||
| 123 | + | ||
| 124 | + span_list = self.ot_exporter.get_finished_spans() | ||
| 125 | + self.assertEqual(len(span_list), 1) | ||
| 126 | + span = span_list[0] | ||
| 127 | + self.assertEqual(span.kind, trace_api.SpanKind.CLIENT) | ||
| 128 | + self.assertEqual(dict(span.attributes), expected_attributes) | ||
| 129 | + self.assertEqual(span.name, "CloudSpannerDjango.Test") | ||
| 130 | + self.assertEqual(span.status.status_code, StatusCode.ERROR) | ||
| Back | FazBrowse Home | New Git URL |
0 commit comments