all: update review comments + show type for BeginTransaction + remove prints

Author: odeke-em

google/cloud/spanner_v1/_opentelemetry_tracing.py

@@ -62,9 +62,6 @@ def trace_call(name, session=None, extra_attributes=None, observability_options=
         yield None
         return
 
-    if session:
-        session._last_use_time = datetime.now()
-
     tracer_provider = None
 
     # By default enable_extended_tracing=True because in a bid to minimize
@@ -104,16 +101,6 @@ def trace_call(name, session=None, extra_attributes=None, observability_options=
     if not enable_extended_tracing:
         attributes.pop("db.statement", False)
         attributes.pop("sql", False)
-    else:
-        # Otherwise there are places where the annotated sql was inserted
-        # directly from the arguments as "sql", and transform those into "db.statement".
-        db_statement = attributes.get("db.statement", None)
-        if not db_statement:
-            sql = attributes.get("sql", None)
-            if sql:
-                attributes = attributes.copy()
-                attributes.pop("sql", False)
-                attributes["db.statement"] = sql
 
     with tracer.start_as_current_span(
         name, kind=trace.SpanKind.CLIENT, attributes=attributes

google/cloud/spanner_v1/session.py

@@ -471,7 +471,6 @@ def run_in_transaction(self, func, *args, **kw):
         ) as span:
             while True:
                 if self._transaction is None:
-                    add_span_event(span, "Creating Transaction")
                     txn = self.transaction()
                     txn.transaction_tag = transaction_tag
                     txn.exclude_txn_from_change_streams = (
@@ -499,7 +498,7 @@ def run_in_transaction(self, func, *args, **kw):
                         record_span_exception_and_status(span, exc)
                         add_span_event(
                             span,
-                            "Transaction was aborted, retrying afresh",
+                            "Transaction was aborted in user operation, retrying",
                             attributes,
                         )
 
@@ -516,7 +515,7 @@ def run_in_transaction(self, func, *args, **kw):
                 except Exception:
                     add_span_event(
                         span,
-                        "Invoking Transaction.rollback(), not retrying",
+                        "User operation failed. Invoking Transaction.rollback(), not retrying",
                         span_attributes,
                     )
                     txn.rollback()
@@ -537,7 +536,7 @@ def run_in_transaction(self, func, *args, **kw):
                         record_span_exception_and_status(span, exc)
                         add_span_event(
                             span,
-                            "Transaction was aborted, retrying afresh",
+                            "Transaction got aborted during commit, retrying afresh",
                             attributes,
                         )
 

google/cloud/spanner_v1/snapshot.py

@@ -926,7 +926,7 @@ def begin(self):
             )
         txn_selector = self._make_txn_selector()
         with trace_call(
-            "CloudSpanner.BeginTransaction",
+            f"CloudSpanner.{type(self).__name__}.begin",
             self._session,
             observability_options=getattr(database, "observability_options", None),
         ):

google/cloud/spanner_v1/transaction.py

@@ -157,7 +157,7 @@ def begin(self):
         )
         observability_options = getattr(database, "observability_options", None)
         with trace_call(
-            "CloudSpanner.BeginTransaction",
+            f"CloudSpanner.{type(self).__name__}.begin",
             self._session,
             observability_options=observability_options,
         ) as span:

tests/_helpers.py

@@ -79,7 +79,6 @@ def tearDown(self):
     def assertNoSpans(self):
         if HAS_OPENTELEMETRY_INSTALLED:
             span_list = self.get_finished_spans()
-            print("got_spans", [span.name for span in span_list])
             self.assertEqual(len(span_list), 0)
 
     def assertSpanAttributes(

tests/system/_helpers.py

@@ -137,3 +137,21 @@ def cleanup_old_instances(spanner_client):
 
 def unique_id(prefix, separator="-"):
     return f"{prefix}{system.unique_resource_id(separator)}"
+
+
+class FauxCall:
+    def __init__(self, code, details="FauxCall"):
+        self._code = code
+        self._details = details
+
+    def initial_metadata(self):
+        return {}
+
+    def trailing_metadata(self):
+        return {}
+
+    def code(self):
+        return self._code
+
+    def details(self):
+        return self._details

tests/system/test_observability_options.py

@@ -131,6 +131,161 @@ def test_propagation(enable_extended_tracing):
     test_propagation(False)
 
 
+@pytest.mark.skipif(
+    not _helpers.USE_EMULATOR,
+    reason="Emulator needed to run this tests",
+)
+@pytest.mark.skipif(
+    not HAS_OTEL_INSTALLED,
+    reason="Tracing requires OpenTelemetry",
+)
+def test_transaction_abort_then_retry_spans():
+    from google.auth.credentials import AnonymousCredentials
+    from google.api_core.exceptions import Aborted
+    from google.rpc import code_pb2
+    from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+    from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
+        InMemorySpanExporter,
+    )
+    from opentelemetry.trace.status import StatusCode
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.sampling import ALWAYS_ON
+    from opentelemetry import trace
+
+    PROJECT = _helpers.EMULATOR_PROJECT
+    CONFIGURATION_NAME = "config-name"
+    INSTANCE_ID = _helpers.INSTANCE_ID
+    DISPLAY_NAME = "display-name"
+    DATABASE_ID = _helpers.unique_id("temp_db")
+    NODE_COUNT = 5
+    LABELS = {"test": "true"}
+
+    counters = dict(aborted=0)
+    already_aborted = False
+
+    def select_in_txn(txn):
+        from google.rpc import error_details_pb2
+
+        results = txn.execute_sql("SELECT 1")
+        for row in results:
+            _ = row
+
+        if counters["aborted"] == 0:
+            counters["aborted"] = 1
+            raise Aborted(
+                "Thrown from ClientInterceptor for testing",
+                errors=[_helpers.FauxCall(code_pb2.ABORTED)],
+            )
+
+    tracer_provider = TracerProvider(sampler=ALWAYS_ON)
+    trace_exporter = InMemorySpanExporter()
+    tracer_provider.add_span_processor(SimpleSpanProcessor(trace_exporter))
+    observability_options = dict(
+        tracer_provider=tracer_provider,
+        enable_extended_tracing=True,
+    )
+
+    client = Client(
+        project=PROJECT,
+        observability_options=observability_options,
+        credentials=AnonymousCredentials(),
+    )
+
+    instance = client.instance(
+        INSTANCE_ID,
+        CONFIGURATION_NAME,
+        display_name=DISPLAY_NAME,
+        node_count=NODE_COUNT,
+        labels=LABELS,
+    )
+
+    try:
+        instance.create()
+    except Exception:
+        pass
+
+    db = instance.database(DATABASE_ID)
+    try:
+        db.create()
+    except Exception:
+        pass
+
+    db.run_in_transaction(select_in_txn)
+
+    span_list = trace_exporter.get_finished_spans()
+    got_span_names = [span.name for span in span_list]
+    want_span_names = [
+        "CloudSpanner.CreateSession",
+        "CloudSpanner.Transaction.execute_streaming_sql",
+        "CloudSpanner.Transaction.execute_streaming_sql",
+        "CloudSpanner.Transaction.commit",
+        "CloudSpanner.Session.run_in_transaction",
+        "CloudSpanner.Database.run_in_transaction",
+    ]
+
+    assert got_span_names == want_span_names
+
+    got_events = []
+    got_statuses = []
+
+    # Some event attributes are noisy/highly ephemeral
+    # and can't be directly compared against.
+    imprecise_event_attributes = ["exception.stacktrace", "delay_seconds"]
+    for span in span_list:
+        got_statuses.append(
+            (span.name, span.status.status_code, span.status.description)
+        )
+        for event in span.events:
+            evt_attributes = event.attributes.copy()
+            for attr_name in imprecise_event_attributes:
+                if attr_name in evt_attributes:
+                    evt_attributes[attr_name] = "EPHEMERAL"
+
+            got_events.append((event.name, evt_attributes))
+
+    # Check for the series of events
+    want_events = [
+        ("Starting Commit", {}),
+        ("Commit Done", {}),
+        ("Using Transaction", {"attempt": 1}),
+        (
+            "exception",
+            {
+                "exception.type": "google.api_core.exceptions.Aborted",
+                "exception.message": "409 Thrown from ClientInterceptor for testing",
+                "exception.stacktrace": "EPHEMERAL",
+                "exception.escaped": "False",
+            },
+        ),
+        (
+            "Transaction was aborted in user operation, retrying",
+            {"delay_seconds": "EPHEMERAL", "attempt": 1},
+        ),
+        ("Using Transaction", {"attempt": 2}),
+        ("Acquiring session", {"kind": "BurstyPool"}),
+        ("Waiting for a session to become available", {"kind": "BurstyPool"}),
+        ("No sessions available in pool. Creating session", {"kind": "BurstyPool"}),
+        ("Creating Session", {}),
+    ]
+    assert got_events == want_events
+
+    # Check for the statues.
+    codes = StatusCode
+    want_statuses = [
+        ("CloudSpanner.CreateSession", codes.OK, None),
+        ("CloudSpanner.Transaction.execute_streaming_sql", codes.OK, None),
+        ("CloudSpanner.Transaction.execute_streaming_sql", codes.OK, None),
+        ("CloudSpanner.Transaction.commit", codes.OK, None),
+        (
+            "CloudSpanner.Session.run_in_transaction",
+            codes.ERROR,
+            "409 Thrown from ClientInterceptor for testing",
+        ),
+        ("CloudSpanner.Database.run_in_transaction", codes.OK, None),
+    ]
+    assert got_statuses == want_statuses
+
+
 def _make_credentials():
     from google.auth.credentials import AnonymousCredentials