tests

Author: surbhigarg92

examples/trace.py

@@ -22,6 +22,8 @@
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.sdk.trace.sampling import ALWAYS_ON
 from opentelemetry import trace
+from opentelemetry.propagate import set_global_textmap
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
 
 
 def main():
@@ -36,10 +38,13 @@ def main():
     # Setup the Cloud Spanner Client.
     spanner_client = spanner.Client(
         project_id,
-        observability_options=dict(tracer_provider=tracer_provider, enable_extended_tracing=True),
+        observability_options=dict(tracer_provider=tracer_provider, enable_extended_tracing=True, enable_end_to_end_tracing=True),
     )
     instance = spanner_client.instance('test-instance')
     database = instance.database('test-db')
+    
+    # Set W3C Trace Context as the global propagator for end to end tracing.
+    set_global_textmap(TraceContextTextMapPropagator())
 
     # Retrieve a tracer from our custom tracer provider.
     tracer = tracer_provider.get_tracer('MyApp')

google/cloud/spanner_v1/_helpers.py

@@ -34,8 +34,12 @@
 from google.cloud.spanner_v1 import JsonObject
 from google.cloud.spanner_v1.request_id_header import with_request_id
 from google.rpc.error_details_pb2 import RetryInfo
-from opentelemetry.propagate import inject
-from opentelemetry.propagators.textmap import Setter, Getter
+try:
+    from opentelemetry.propagate import inject
+    from opentelemetry.propagators.textmap import Setter
+    HAS_OPENTELEMETRY_INSTALLED = True
+except ImportError:
+    HAS_OPENTELEMETRY_INSTALLED = False
 from typing import List, Tuple
 import random
 
@@ -583,8 +587,9 @@ def _metadata_with_span_context(metadata: List[Tuple[str, str]], **kw) -> None:
     Returns:
         None
     """
-    metadata.append(("x-goog-spanner-end-to-end-tracing", "true"))
-    inject(setter=OpenTelemetryContextSetter(), carrier=metadata)
+    if HAS_OPENTELEMETRY_INSTALLED:
+        metadata.append(("x-goog-spanner-end-to-end-tracing", "true"))
+        inject(setter=OpenTelemetryContextSetter(), carrier=metadata)
 
 
 def _delay_until_retry(exc, deadline, attempts):

tests/unit/test_batch.py

@@ -563,7 +563,10 @@ def test_batch_write_grpc_error(self):
         )
 
     def _test_batch_write_with_request_options(
-        self, request_options=None, exclude_txn_from_change_streams=False
+        self,
+        request_options=None,
+        exclude_txn_from_change_streams=False,
+        enable_end_to_end_tracing=False,
     ):
         import datetime
         from google.cloud.spanner_v1 import BatchWriteResponse
@@ -577,7 +580,7 @@ def _test_batch_write_with_request_options(
         response = BatchWriteResponse(
             commit_timestamp=now_pb, indexes=[0], status=status_pb
         )
-        database = _Database()
+        database = _Database(enable_end_to_end_tracing=enable_end_to_end_tracing)
         api = database.spanner_api = _FauxSpannerAPI(_batch_write_response=[response])
         session = _Session(database)
         groups = self._make_one(session)
@@ -600,13 +603,22 @@ def _test_batch_write_with_request_options(
         ) = api._batch_request
         self.assertEqual(session, self.SESSION_NAME)
         self.assertEqual(mutation_groups, groups._mutation_groups)
-        self.assertEqual(
-            metadata,
-            [
-                ("google-cloud-resource-prefix", database.name),
-                ("x-goog-spanner-route-to-leader", "true"),
-            ],
-        )
+        expected_metadata = [
+            ("google-cloud-resource-prefix", database.name),
+            ("x-goog-spanner-route-to-leader", "true"),
+        ]
+
+        if enable_end_to_end_tracing:
+            expected_metadata.append(("x-goog-spanner-end-to-end-tracing", "true"))
+            self.assertTrue(
+                any(key == "traceparent" for key, _ in metadata),
+                "traceparent is missing in metadata",
+            )
+
+        # Remove traceparent from actual metadata for comparison
+        filtered_metadata = [item for item in metadata if item[0] != "traceparent"]
+
+        self.assertEqual(filtered_metadata, expected_metadata)
         if request_options is None:
             expected_request_options = RequestOptions()
         elif type(request_options) is dict:
@@ -626,6 +638,9 @@ def _test_batch_write_with_request_options(
 
     def test_batch_write_no_request_options(self):
         self._test_batch_write_with_request_options()
+        
+    def test_batch_write_end_to_end_tracing_enabled(self):
+        self._test_batch_write_with_request_options(enable_end_to_end_tracing=True)
 
     def test_batch_write_w_transaction_tag_success(self):
         self._test_batch_write_with_request_options(
@@ -656,8 +671,11 @@ def session_id(self):
 
 
 class _Database(object):
-    name = "testing"
-    _route_to_leader_enabled = True
+    def __init__(self, enable_end_to_end_tracing=False):
+        self.name = "testing"
+        self._route_to_leader_enabled = True
+        if enable_end_to_end_tracing:
+            self.observability_options = dict(enable_end_to_end_tracing=True)
 
 
 class _FauxSpannerAPI: