feat: add opentelemetry tracing

google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py

@@ -0,0 +1,75 @@
+# Copyright 2021 Google LLC All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Manages OpenTelemetry trace creation and handling"""
+
+from contextlib import contextmanager
+
+from google.api_core.exceptions import GoogleAPICallError
+from google.cloud.spanner_v1 import SpannerClient
+from google.cloud.spanner_dbapi.exceptions import IntegrityError
+from google.cloud.spanner_dbapi.exceptions import InterfaceError
+from google.cloud.spanner_dbapi.exceptions import OperationalError
+from google.cloud.spanner_dbapi.exceptions import ProgrammingError
+
+try:
+ from opentelemetry import trace
+ from opentelemetry.trace.status import Status, StatusCanonicalCode
+ from opentelemetry.instrumentation.utils import http_status_to_canonical_code
+
+ HAS_OPENTELEMETRY_INSTALLED = True
+except ImportError:
+ HAS_OPENTELEMETRY_INSTALLED = False
+
+
+@contextmanager
+def trace_call(name, extra_attributes=None):
+ if not HAS_OPENTELEMETRY_INSTALLED:
+ # 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.url": SpannerClient.DEFAULT_ENDPOINT,
+ "net.host.name": SpannerClient.DEFAULT_ENDPOINT,
+ }
+
+ if extra_attributes:
+ attributes.update(extra_attributes)
+
+ with tracer.start_as_current_span(
+ name, kind=trace.SpanKind.CLIENT, attributes=attributes
+ ) as span:
+ try:
+ yield span
+ except (ValueError, InterfaceError) as e:
+ span.set_status(Status(StatusCanonicalCode.UNKNOWN, e.args[0]))
+ except GoogleAPICallError as error:
+ if error.code is not None:
+ span.set_status(Status(http_status_to_canonical_code(error.code)))
+ elif error.grpc_status_code is not None:
+ span.set_status(
+ # OpenTelemetry's StatusCanonicalCode maps 1-1 with grpc status
+ # codes
+ Status(StatusCanonicalCode(error.grpc_status_code.value[0]))
+ )
+ raise
+ except (IntegrityError, ProgrammingError, OperationalError) as e:
+ span.set_status(
+ Status(http_status_to_canonical_code(e.args[0].code), e.args[0].message)
+ )

google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py

@@ -27,6 +27,7 @@
  RESERVED_WORDS,
 )
 from google.cloud import spanner_dbapi
+from google.cloud.sqlalchemy_spanner._opentelemetry_tracing import trace_call
 
 # Spanner-to-SQLAlchemy types map
 _type_map = {
@@ -761,4 +762,31 @@ def do_rollback(self, dbapi_connection):
  ):
  pass
  else:
- dbapi_connection.rollback()
+ trace_attributes = {"db.connection": dbapi_connection}
+ with trace_call("SpannerSqlAlchemy.Rollback", trace_attributes):
+ dbapi_connection.rollback()
+
+ def do_commit(self, dbapi_connection):
+ trace_attributes = {"db.connection": dbapi_connection}
+ with trace_call("SpannerSqlAlchemy.Commit", trace_attributes):
+ dbapi_connection.commit()
+
+ def do_close(self, dbapi_connection):
+ trace_attributes = {"db.connection": dbapi_connection}
+ with trace_call("SpannerSqlAlchemy.Close", trace_attributes):
+ dbapi_connection.close()
+
+ def do_executemany(self, cursor, statement, parameters, context=None):
+ trace_attributes = {"db.statement": statement, "db.params": parameters}
+ with trace_call("SpannerSqlAlchemy.Executemany", trace_attributes):
+ cursor.executemany(statement, parameters)
+
+ def do_execute(self, cursor, statement, parameters, context=None):
+ trace_attributes = {"db.statement": statement, "db.params": parameters}
+ with trace_call("SpannerSqlAlchemy.Execute", trace_attributes):
+ cursor.execute(statement, parameters)
+
+ def do_execute_no_params(self, cursor, statement, context=None):
+ trace_attributes = {"db.statement": statement}
+ with trace_call("SpannerSqlAlchemy.ExecuteNoParams", trace_attributes):
+ cursor.execute(statement)

setup.py

@@ -20,6 +20,13 @@
 name = "sqlalchemy-spanner"
 description = "SQLAlchemy dialect integrated into Cloud Spanner database"
 dependencies = ["sqlalchemy>=1.1.13, <=1.3.23", "google-cloud-spanner>=3.3.0"]
+extras = {
+ "tracing": [
+ "opentelemetry-api==0.11b0",
+ "opentelemetry-sdk==0.11b0",
+ "opentelemetry-instrumentation==0.11b0",
+ ]
+}
 
 # Only include packages under the 'google' namespace. Do not include tests,
 # benchmarks, etc.
@@ -45,6 +52,7 @@
  ]
  },
  install_requires=dependencies,
+ extras_require=extras,
  name=name,
  namespace_packages=namespaces,
  packages=packages,