More upates for test span results

google/cloud/spanner_v1/batch.py

@@ -51,7 +51,7 @@ def __init__(self, session):
         super(_BatchBase, self).__init__(session)
         self._mutations = []
         self.__base_discard_span = trace_call_end_lazily(
-            f"CloudSpanner.{type(self).__name__}",
+            f"CloudSpanner.a{type(self).__name__}",
             self._session,
             None,
             getattr(self._session._database, "observability_options", None),

google/cloud/spanner_v1/database.py

@@ -910,7 +910,6 @@ def run_in_transaction(self, func, *args, **kw):
             observability_options = getattr(self, "observability_options", None)
             with trace_call(
                 "CloudSpanner.Database.run_in_transaction",
-                None,
                 observability_options=observability_options,
             ):
                 with SessionCheckout(self._pool) as session:
@@ -1566,7 +1565,6 @@ def process_read_batch(
         observability_options = self.observability_options or {}
         with trace_call(
             f"CloudSpanner.{type(self).__name__}.process_read_batch",
-            session,
             observability_options=observability_options,
         ):
             kwargs = copy.deepcopy(batch["read"])

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,
         ):

tests/_helpers.py

@@ -78,7 +78,6 @@ def tearDown(self):
     def assertNoSpans(self):
         if HAS_OPENTELEMETRY_INSTALLED:
             span_list = self.get_finished_spans()
-            print("got_span_list", [span.name for span in span_list])
             self.assertEqual(len(span_list), 0)
 
     def assertSpanAttributes(
@@ -92,8 +91,6 @@ def assertSpanAttributes(
 
             self.assertEqual(span.name, name)
             self.assertEqual(span.status.status_code, status)
-            print("got_span_attributes ", dict(span.attributes))
-            print("want_span_attributes", attributes)
             self.assertEqual(dict(span.attributes), attributes)
 
     def get_finished_spans(self):
@@ -104,3 +101,6 @@ def get_finished_spans(self):
         # A span with name=None is the result from invoking trace_call without
         # intention to trace, hence these have to be filtered out.
         return list(filter(lambda span: span.name, spans))
+
+    def reset(self):
+        self.tearDown()

tests/system/test_session_api.py

@@ -658,7 +658,7 @@ def test_transaction_read_and_insert_then_rollback(
         )
         assert_span_attributes(
             ot_exporter,
-            "CloudSpanner.BeginTransaction",
+            "CloudSpanner.Transaction.begin",
             attributes=_make_attributes(db_name),
             span=span_list[5],
         )
@@ -1370,38 +1370,67 @@ def unit_of_work(transaction):
     with tracer.start_as_current_span("Test Span"):
         session.run_in_transaction(unit_of_work)
 
-    span_list = ot_exporter.get_finished_spans()
+    span_list = []
+    for span in ot_exporter.get_finished_spans():
+        if span and span.name:
+            span_list.append(span)
+
+    span_list = sorted(span_list, key=lambda v1: v1.start_time)
+
     expected_span_names = [
         "CloudSpanner.CreateSession",
-        "CloudSpanner.Batch.commit",
         "CloudSpanner.Batch",
         "CloudSpanner.Batch",
+        "CloudSpanner.Batch.commit",
+        "Test Span",
+        "CloudSpanner.ReadWriteTransaction",
+        "CloudSpanner.Transaction",
         "CloudSpanner.DMLTransaction",
         "CloudSpanner.Transaction.commit",
-        "CloudSpanner.Transaction",
-        "CloudSpanner.ReadWriteTransaction",
-        "Test Span",
     ]
+
     got_span_names = [span.name for span in span_list]
     assert got_span_names == expected_span_names
 
+    # We expect:
+    # |------CloudSpanner.CreateSession--------
+    #
+    # |---Test Span----------------------------|
+    #  |>--ReadWriteTransaction-----------------
+    #    |>-Transaction-------------------------
+    #      |--------------DMLTransaction--------
+    #
+    #  |>---Batch-------------------------------
+    #
+    # |>----------Batch-------------------------
+    #  |>------------Batch.commit---------------
+
+    # CreateSession should have a trace of its own, with no children
+    # nor being a child of any other span.
+    session_span = span_list[0]
+    test_span = span_list[4]
+    # assert session_span.context.trace_id != test_span.context.trace_id
+    for span in span_list[1:]:
+        if span.parent:
+            assert span.parent.span_id != session_span.context.span_id
+
+    def assert_parent_and_children(parent_span, children):
+        for span in children:
+            assert span.context.trace_id == parent_span.context.trace_id
+            assert span.parent.span_id == parent_span.context.span_id
+
     # [CreateSession --> Batch] should have their own trace.
-    session_parent_span = span_list[0]
-    session_creation_child_spans = span_list[1:1]
-    for span in session_creation_child_spans:
-        assert span.context.trace_id == session_parent_span.context.trace_id
-        assert span.parent.span_id == session_parent_span.context.span_id
-
-    # [Test Span --> DMLTransaction] should have their own trace.
-    overall_userland_span = span_list[-1]
-    current_context_spans = span_list[3:-1]
-    for span in current_context_spans:
-        assert span.context.trace_id == overall_userland_span.context.trace_id
-        assert span.parent.span_id == overall_userland_span.context.span_id
-
-    assert (
-        session_parent_span.context.trace_id != overall_userland_span.context.trace_id
-    )
+    rw_txn_span = span_list[5]
+    children_of_test_span = [rw_txn_span]
+    assert_parent_and_children(test_span, children_of_test_span)
+
+    children_of_rw_txn_span = [span_list[6]]
+    assert_parent_and_children(rw_txn_span, children_of_rw_txn_span)
+
+    # Batch_first should have no parent, should be in its own trace.
+    batch_0_span = span_list[2]
+    children_of_batch_0 = [span_list[1]]
+    assert_parent_and_children(rw_txn_span, children_of_rw_txn_span)
 
 
 def test_execute_partitioned_dml(

tests/unit/test_snapshot.py

@@ -1701,8 +1701,13 @@ def test_begin_w_other_error(self):
         with self.assertRaises(RuntimeError):
             snapshot.begin()
 
+        span_list = self.get_finished_spans()
+        got_span_names = [span.name for span in span_list]
+        want_span_names = ["CloudSpanner.Snapshot.begin"]
+        assert got_span_names == want_span_names
+
         self.assertSpanAttributes(
-            "CloudSpanner.BeginTransaction",
+            "CloudSpanner.Snapshot.begin",
             status=StatusCode.ERROR,
             attributes=BASE_ATTRIBUTES,
         )
@@ -1760,7 +1765,7 @@ def test_begin_ok_exact_staleness(self):
         )
 
         self.assertSpanAttributes(
-            "CloudSpanner.BeginTransaction",
+            "CloudSpanner.Snapshot.begin",
             status=StatusCode.OK,
             attributes=BASE_ATTRIBUTES,
         )
@@ -1796,7 +1801,7 @@ def test_begin_ok_exact_strong(self):
         )
 
         self.assertSpanAttributes(
-            "CloudSpanner.BeginTransaction",
+            "CloudSpanner.Snapshot.begin",
             status=StatusCode.OK,
             attributes=BASE_ATTRIBUTES,
         )

tests/unit/test_transaction.py

@@ -125,6 +125,7 @@ def test__make_txn_selector(self):
         self.assertEqual(selector.id, self.TRANSACTION_ID)
 
     def test_begin_already_begun(self):
+        self.reset()
         session = _Session()
         transaction = self._make_one(session)
         transaction._transaction_id = self.TRANSACTION_ID
@@ -134,6 +135,7 @@ def test_begin_already_begun(self):
         self.assertNoSpans()
 
     def test_begin_already_rolled_back(self):
+        self.reset()
         session = _Session()
         transaction = self._make_one(session)
         transaction.rolled_back = True
@@ -143,6 +145,7 @@ def test_begin_already_rolled_back(self):
         self.assertNoSpans()
 
     def test_begin_already_committed(self):
+        self.reset()
         session = _Session()
         transaction = self._make_one(session)
         transaction.committed = object()
@@ -152,6 +155,7 @@ def test_begin_already_committed(self):
         self.assertNoSpans()
 
     def test_begin_w_other_error(self):
+        self.reset()
         database = _Database()
         database.spanner_api = self._make_spanner_api()
         database.spanner_api.begin_transaction.side_effect = RuntimeError()
@@ -161,13 +165,20 @@ def test_begin_w_other_error(self):
         with self.assertRaises(RuntimeError):
             transaction.begin()
 
+        span_list = self.get_finished_spans()
+        got_span_names = [span.name for span in span_list]
+        want_span_names = ["CloudSpanner.Transaction.begin"]
+        print(got_span_names)
+        assert got_span_names == want_span_names
+
         self.assertSpanAttributes(
-            "CloudSpanner.BeginTransaction",
+            "CloudSpanner.Transaction.begin",
             status=StatusCode.ERROR,
             attributes=TestTransaction.BASE_ATTRIBUTES,
         )
 
     def test_begin_ok(self):
+        self.reset()
         from google.cloud.spanner_v1 import Transaction as TransactionPB
 
         transaction_pb = TransactionPB(id=self.TRANSACTION_ID)
@@ -195,10 +206,11 @@ def test_begin_ok(self):
         )
 
         self.assertSpanAttributes(
-            "CloudSpanner.BeginTransaction", attributes=TestTransaction.BASE_ATTRIBUTES
+            "CloudSpanner.Transaction.begin", attributes=TestTransaction.BASE_ATTRIBUTES
         )
 
     def test_begin_w_retry(self):
+        self.reset()
         from google.cloud.spanner_v1 import (
             Transaction as TransactionPB,
         )