From 076895992953a9d80429a7db5eb76ffce771c556 Mon Sep 17 00:00:00 2001 From: HemangChothani Date: Tue, 13 Apr 2021 17:39:47 +0530 Subject: [PATCH 1/5] feat: add opentelemetry tracing --- .../_opentelemetry_tracing.py | 66 +++++++++++++++++++ setup.py | 8 +++ 2 files changed, 74 insertions(+) create mode 100644 google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py diff --git a/google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py b/google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py new file mode 100644 index 00000000..b69a3a7a --- /dev/null +++ b/google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py @@ -0,0 +1,66 @@ +# 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 + +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, engine, extra_attributes=None): + if not HAS_OPENTELEMETRY_INSTALLED or not engine: + # 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": engine.url, + "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 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 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, From 8cb5ab0436b9fdf5190d3c57f5cb1174b1374c58 Mon Sep 17 00:00:00 2001 From: HemangChothani Date: Fri, 7 May 2021 17:03:48 +0530 Subject: [PATCH 2/5] feat: call create span code in methods --- .../_opentelemetry_tracing.py | 17 ++++++-- .../sqlalchemy_spanner/sqlalchemy_spanner.py | 40 +++++++++++++++++++ 2 files changed, 53 insertions(+), 4 deletions(-) diff --git a/google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py b/google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py index b69a3a7a..faaae11e 100644 --- a/google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py +++ b/google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py @@ -18,6 +18,10 @@ 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 @@ -30,19 +34,18 @@ @contextmanager -def trace_call(name, engine, extra_attributes=None): - if not HAS_OPENTELEMETRY_INSTALLED or not engine: +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": engine.url, + "db.url": SpannerClient.DEFAULT_ENDPOINT, "net.host.name": SpannerClient.DEFAULT_ENDPOINT, } @@ -54,6 +57,8 @@ def trace_call(name, engine, extra_attributes=None): ) 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))) @@ -64,3 +69,7 @@ def trace_call(name, engine, extra_attributes=None): 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 92f225b7..454a05f5 100644 --- a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -24,6 +24,7 @@ SQLCompiler, ) from google.cloud import spanner_dbapi +from google.cloud.sqlalchemy_spanner._opentelemetry_tracing import trace_call # Spanner-to-SQLAlchemy types map _type_map = { @@ -664,3 +665,42 @@ def get_isolation_level(self, conn_proxy): conn = conn_proxy.connection return "AUTOCOMMIT" if conn.autocommit else "SERIALIZABLE" + + def do_rollback(self, dbapi_connection): + """To prevent transaction rollback error, rollback is ignored if + DBAPI rollback is already executed.""" + if ( + not isinstance(dbapi_connection, spanner_dbapi.Connection) + and dbapi_connection.connection._transaction + and dbapi_connection.connection._transaction.rolled_back + ): + pass + else: + 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) From 9fc3111b0b3594c56e1e03e537591b223d7adc1d Mon Sep 17 00:00:00 2001 From: HemangChothani Date: Fri, 21 May 2021 16:43:45 +0530 Subject: [PATCH 3/5] feat: nit --- google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index ca3f3d69..28e7f64e 100644 --- a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -762,17 +762,17 @@ def do_rollback(self, dbapi_connection): ): pass else: - trace_attributes = {"db.connection": dbapi_connection} + 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.connection": 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.connection": dbapi_connection} + trace_attributes = {"db.instance": dbapi_connection.database.name} with trace_call("SpannerSqlAlchemy.Close", trace_attributes): dbapi_connection.close() From 617d4a0e1cc9847eb3ac96d95afa64f586694470 Mon Sep 17 00:00:00 2001 From: HemangChothani Date: Tue, 25 May 2021 11:18:07 +0530 Subject: [PATCH 4/5] feat: add db.instance in trace attribute --- .../sqlalchemy_spanner/sqlalchemy_spanner.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index 28e7f64e..7ddc7645 100644 --- a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -777,16 +777,27 @@ def do_close(self, dbapi_connection): dbapi_connection.close() def do_executemany(self, cursor, statement, parameters, context=None): - trace_attributes = {"db.statement": statement, "db.params": parameters} + 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} + 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} + trace_attributes = { + "db.statement": statement, + "db.instance": cursor.connection.database.name, + } with trace_call("SpannerSqlAlchemy.ExecuteNoParams", trace_attributes): cursor.execute(statement) From fa94ff60d948ba7beee0c2df48c2f3ac563584e3 Mon Sep 17 00:00:00 2001 From: HemangChothani Date: Tue, 25 May 2021 13:00:57 +0530 Subject: [PATCH 5/5] feat: nit --- google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index 7ddc7645..36e147c1 100644 --- a/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -782,7 +782,7 @@ def do_executemany(self, cursor, statement, parameters, context=None): "db.params": parameters, "db.instance": cursor.connection.database.name, } - with trace_call("SpannerSqlAlchemy.Executemany", trace_attributes): + with trace_call("SpannerSqlAlchemy.ExecuteMany", trace_attributes): cursor.executemany(statement, parameters) def do_execute(self, cursor, statement, parameters, context=None):