diff --git a/google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py b/google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py new file mode 100644 index 00000000..faaae11e --- /dev/null +++ b/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) + ) diff --git a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index 52886740..36e147c1 100644 --- a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/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,42 @@ def do_rollback(self, dbapi_connection): ): pass else: - dbapi_connection.rollback() + trace_attributes = {"db.instance": dbapi_connection.database.name} + with trace_call("SpannerSqlAlchemy.Rollback", trace_attributes): + dbapi_connection.rollback() + + def do_commit(self, dbapi_connection): + trace_attributes = {"db.instance": dbapi_connection.database.name} + with trace_call("SpannerSqlAlchemy.Commit", trace_attributes): + dbapi_connection.commit() + + def do_close(self, dbapi_connection): + trace_attributes = {"db.instance": dbapi_connection.database.name} + 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, + "db.instance": cursor.connection.database.name, + } + 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, + "db.instance": cursor.connection.database.name, + } + 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, + "db.instance": cursor.connection.database.name, + } + with trace_call("SpannerSqlAlchemy.ExecuteNoParams", trace_attributes): + cursor.execute(statement) diff --git a/setup.py b/setup.py index 17fd6162..9db3c595 100644 --- a/setup.py +++ b/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,