Skip to content
This repository was archived by the owner on Mar 20, 2026. It is now read-only.

Commit 2ba879a

Browse files
vi3k6i5vi3k6i5
authored
feat: Add support for open telemetry (#633)
* fix: lint_setup_py was failing in Kokoro is not fixed * feat: adding opentelemetry tracing * feat: added opentelemetry support * feat: added open telemetry tracing support and tests * refactor: lint fixes * refactor: lint fixes * refactor: added license text * ci: corrrected version for google-cloud-spanner * refactor: removed schema changes and tests related to ot, will send PR for that separately * refactor: removed commented lines of code * refactor: lint corrections
1 parent ecf241a commit 2ba879a

7 files changed

Lines changed: 283 additions & 4 deletions

File tree

‎django_spanner/_opentelemetry_tracing.py‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff 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

‎noxfile.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,11 @@ def default(session):
7575
"pytest",
7676
"pytest-cov",
7777
"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",
7883
)
7984
session.install("-e", ".")
8085

‎setup.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,13 @@
1818
# 'Development Status :: 5 - Production/Stable'
1919
release_status = "Development Status :: 3 - Alpha"
2020
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+
}
2228

2329
BASE_DIR = os.path.dirname(__file__)
2430
VERSION_FILENAME = os.path.join(BASE_DIR, "version.py")

‎testing/constraints-3.6.txt‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,7 @@
66
# e.g., if setup.py has "foo >= 1.14.0, < 2.0.0dev",
77
# Then this file should have foo==1.14.0
88
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

‎tests/_helpers.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff 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)

‎tests/unit/django_spanner/simple_test.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,16 @@
77
from django_spanner.client import DatabaseClient
88
from django_spanner.base import DatabaseWrapper
99
from django_spanner.operations import DatabaseOperations
10-
from unittest import TestCase
10+
11+
# from unittest import TestCase
12+
from tests._helpers import OpenTelemetryBase
1113
import os
1214

1315

14-
class SpannerSimpleTestClass(TestCase):
16+
class SpannerSimpleTestClass(OpenTelemetryBase):
1517
@classmethod
1618
def setUpClass(cls):
19+
super(SpannerSimpleTestClass, cls).setUpClass()
1720
cls.PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
1821

1922
cls.INSTANCE_ID = "instance_id"

‎tests/unit/django_spanner/test__opentelemetry_tracing.py‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff 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)

0 commit comments

Comments
 (0)