googleapis · HemangChothani · May 25, 2021 · Apr 13, 2021 · May 7, 2021 · May 7, 2021
diff --git a/google/cloud/sqlalchemy_spanner/_opentelemetry_tracing.py 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
@@ -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
@@ -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,