Skip to content

Commit

Permalink
Browse files Browse the repository at this point in the history
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
vi3k6i5 committed Jun 17, 2021
1 parent ecf241a commit 2ba879a
Show file tree
Hide file tree
Showing 7 changed files with 283 additions and 4 deletions.
53 changes: 53 additions & 0 deletions 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
5 changes: 5 additions & 0 deletions noxfile.py
Expand Up @@ -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", ".")

Expand Down
8 changes: 7 additions & 1 deletion setup.py
Expand Up @@ -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")
Expand Down
5 changes: 4 additions & 1 deletion testing/constraints-3.6.txt
Expand Up @@ -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
google-cloud-spanner==3.0.0
opentelemetry-api==1.1.0
opentelemetry-sdk==1.1.0
opentelemetry-instrumentation==0.20b0
79 changes: 79 additions & 0 deletions 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)
7 changes: 5 additions & 2 deletions tests/unit/django_spanner/simple_test.py
Expand Up @@ -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"
Expand Down
130 changes: 130 additions & 0 deletions 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)

0 comments on commit 2ba879a

Please sign in to comment.