-
Notifications
You must be signed in to change notification settings - Fork 28
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
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
- Loading branch information
Showing
7 changed files
with
283 additions
and
4 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
130 changes: 130 additions & 0 deletions
130
tests/unit/django_spanner/test__opentelemetry_tracing.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -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) |