feat: end to end tracing

Author: surbhigarg92

google/cloud/spanner_v1/_helpers.py

@@ -34,7 +34,9 @@
 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
+from typing import List, Tuple
 import random
 
 # Validation error messages
@@ -47,6 +49,27 @@
 )
 
 
+class OpenTelemetryContextSetter(Setter):
+    """
+    Used by Open Telemetry for context propagation.
+    """
+
+    def set(self, carrier: List[Tuple[str, str]], key: str, value: str) -> None:
+        """
+        Injects trace context into Spanner metadata
+
+        Args:
+            carrier(PubsubMessage): The Pub/Sub message which is the carrier of Open Telemetry
+            data.
+            key(str): The key for which the Open Telemetry context data needs to be set.
+            value(str): The Open Telemetry context value to be set.
+
+        Returns:
+            None
+        """
+        carrier.append((key, value))
+
+
 def _try_to_coerce_bytes(bytestring):
     """Try to coerce a byte string into the right thing based on Python
     version and whether or not it is base64 encoded.
@@ -550,6 +573,20 @@ def _metadata_with_leader_aware_routing(value, **kw):
     return ("x-goog-spanner-route-to-leader", str(value).lower())
 
 
+def _metadata_with_span_context(metadata: List[Tuple[str, str]], **kw) -> None:
+    """
+    Appends metadata with end to end tracing header and OpenTelemetry span context .
+
+    Args:
+        metadata (list[tuple[str, str]]): The metadata carrier where the OpenTelemetry context
+                                          should be injected.
+    Returns:
+        None
+    """
+    metadata.append(("x-goog-spanner-end-to-end-tracing", "true"))
+    inject(setter=OpenTelemetryContextSetter(), carrier=metadata)
+
+
 def _delay_until_retry(exc, deadline, attempts):
     """Helper for :meth:`Session.run_in_transaction`.
 

google/cloud/spanner_v1/_opentelemetry_tracing.py

@@ -20,6 +20,9 @@
 
 from google.cloud.spanner_v1 import SpannerClient
 from google.cloud.spanner_v1 import gapic_version
+from google.cloud.spanner_v1._helpers import (
+    _metadata_with_span_context,
+)
 
 try:
     from opentelemetry import trace
@@ -38,6 +41,9 @@
 extended_tracing_globally_disabled = (
     os.getenv("SPANNER_ENABLE_EXTENDED_TRACING", "").lower() == "false"
 )
+end_to_end_tracing_globally_enabled = (
+    os.getenv("SPANNER_ENABLE_END_TO_END_TRACING", "").lower() == "true"
+)
 
 
 def get_tracer(tracer_provider=None):
@@ -56,7 +62,9 @@ def get_tracer(tracer_provider=None):
 
 
 @contextmanager
-def trace_call(name, session=None, extra_attributes=None, observability_options=None):
+def trace_call(
+    name, session=None, extra_attributes=None, observability_options=None, metadata=None
+):
     if session:
         session._last_use_time = datetime.now()
 
@@ -72,6 +80,8 @@ def trace_call(name, session=None, extra_attributes=None, observability_options=
     # on by default.
     enable_extended_tracing = True
 
+    enable_end_to_end_tracing = False
+
     db_name = ""
     if session and getattr(session, "_database", None):
         db_name = session._database.name
@@ -81,6 +91,9 @@ def trace_call(name, session=None, extra_attributes=None, observability_options=
         enable_extended_tracing = observability_options.get(
             "enable_extended_tracing", enable_extended_tracing
         )
+        enable_end_to_end_tracing = observability_options.get(
+            "enable_end_to_end_tracing", enable_end_to_end_tracing
+        )
         db_name = observability_options.get("db_name", db_name)
 
     tracer = get_tracer(tracer_provider)
@@ -104,10 +117,15 @@ def trace_call(name, session=None, extra_attributes=None, observability_options=
     if not enable_extended_tracing:
         attributes.pop("db.statement", False)
 
+    if end_to_end_tracing_globally_enabled:
+        enable_end_to_end_tracing = True
+
     with tracer.start_as_current_span(
         name, kind=trace.SpanKind.CLIENT, attributes=attributes
     ) as span:
         try:
+            if enable_end_to_end_tracing:
+                _metadata_with_span_context(metadata)
             yield span
         except Exception as error:
             span.set_status(Status(StatusCode.ERROR, str(error)))

google/cloud/spanner_v1/batch.py

@@ -226,6 +226,7 @@ def commit(
             self._session,
             trace_attributes,
             observability_options=observability_options,
+            metadata=metadata,
         ):
             method = functools.partial(
                 api.commit,
@@ -348,6 +349,7 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
             self._session,
             trace_attributes,
             observability_options=observability_options,
+            metadata=metadata,
         ):
             method = functools.partial(
                 api.batch_write,

google/cloud/spanner_v1/client.py

@@ -135,6 +135,10 @@ class Client(ClientWithProject):
            Default `True`, please set it to `False` to turn it off
            or you can use the environment variable `SPANNER_ENABLE_EXTENDED_TRACING=<boolean>`
            to control it.
+           enable_end_to_end_tracing: :type:boolean when set to true will allow for spans from Spanner server side.
+           Default `False`, please set it to `True` to turn it on
+           or you can use the environment variable `SPANNER_ENABLE_END_TO_END_TRACING=<boolean>`
+           to control it.
 
     :raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
              and ``admin`` are :data:`True`

google/cloud/spanner_v1/database.py

@@ -728,6 +728,7 @@ def execute_pdml():
                         method=method,
                         trace_name="CloudSpanner.ExecuteStreamingSql",
                         request=request,
+                        metadata=metadata,
                         transaction_selector=txn_selector,
                         observability_options=self.observability_options,
                     )

google/cloud/spanner_v1/pool.py

@@ -242,6 +242,7 @@ def bind(self, database):
         with trace_call(
             "CloudSpanner.FixedPool.BatchCreateSessions",
             observability_options=observability_options,
+            metadata=metadata,
         ) as span:
             returned_session_count = 0
             while not self._sessions.full():
@@ -552,6 +553,7 @@ def bind(self, database):
         with trace_call(
             "CloudSpanner.PingingPool.BatchCreateSessions",
             observability_options=observability_options,
+            metadata=metadata,
         ) as span:
             returned_session_count = 0
             while returned_session_count < self.size:

google/cloud/spanner_v1/session.py

@@ -165,6 +165,7 @@ def create(self):
             self,
             self._labels,
             observability_options=observability_options,
+            metadata=metadata,
         ):
             session_pb = api.create_session(
                 request=request,
@@ -204,7 +205,10 @@ def exists(self):
 
         observability_options = getattr(self._database, "observability_options", None)
         with trace_call(
-            "CloudSpanner.GetSession", self, observability_options=observability_options
+            "CloudSpanner.GetSession",
+            self,
+            observability_options=observability_options,
+            metadata=metadata,
         ) as span:
             try:
                 api.get_session(name=self.name, metadata=metadata)
@@ -248,6 +252,7 @@ def delete(self):
                 "session.name": self.name,
             },
             observability_options=observability_options,
+            metadata=metadata,
         ):
             api.delete_session(name=self.name, metadata=metadata)
 

google/cloud/spanner_v1/snapshot.py

@@ -52,6 +52,7 @@
 def _restart_on_unavailable(
     method,
     request,
+    metadata=None,
     trace_name=None,
     session=None,
     attributes=None,
@@ -96,8 +97,9 @@ def _restart_on_unavailable(
                     session,
                     attributes,
                     observability_options=observability_options,
+                    metadata=metadata,
                 ):
-                    iterator = method(request=request)
+                    iterator = method(request=request, metadata=metadata)
             for item in iterator:
                 item_buffer.append(item)
                 # Setting the transaction id because the transaction begin was inlined for first rpc.
@@ -119,6 +121,7 @@ def _restart_on_unavailable(
                 session,
                 attributes,
                 observability_options=observability_options,
+                metadata=metadata,
             ):
                 request.resume_token = resume_token
                 if transaction is not None:
@@ -139,6 +142,7 @@ def _restart_on_unavailable(
                 session,
                 attributes,
                 observability_options=observability_options,
+                metadata=metadata,
             ):
                 request.resume_token = resume_token
                 if transaction is not None:
@@ -340,6 +344,7 @@ def read(
                 iterator = _restart_on_unavailable(
                     restart,
                     request,
+                    metadata,
                     f"CloudSpanner.{type(self).__name__}.read",
                     self._session,
                     trace_attributes,
@@ -362,6 +367,7 @@ def read(
             iterator = _restart_on_unavailable(
                 restart,
                 request,
+                metadata,
                 f"CloudSpanner.{type(self).__name__}.read",
                 self._session,
                 trace_attributes,
@@ -556,6 +562,7 @@ def execute_sql(
                 return self._get_streamed_result_set(
                     restart,
                     request,
+                    metadata,
                     trace_attributes,
                     column_info,
                     observability_options,
@@ -565,6 +572,7 @@ def execute_sql(
             return self._get_streamed_result_set(
                 restart,
                 request,
+                metadata,
                 trace_attributes,
                 column_info,
                 observability_options,
@@ -575,6 +583,7 @@ def _get_streamed_result_set(
         self,
         restart,
         request,
+        metadata,
         trace_attributes,
         column_info,
         observability_options=None,
@@ -583,6 +592,7 @@ def _get_streamed_result_set(
         iterator = _restart_on_unavailable(
             restart,
             request,
+            metadata,
             f"CloudSpanner.{type(self).__name__}.execute_sql",
             self._session,
             trace_attributes,
@@ -689,6 +699,7 @@ def partition_read(
             self._session,
             extra_attributes=trace_attributes,
             observability_options=getattr(database, "observability_options", None),
+            metadata=metadata,
         ):
             method = functools.partial(
                 api.partition_read,
@@ -792,6 +803,7 @@ def partition_query(
             self._session,
             trace_attributes,
             observability_options=getattr(database, "observability_options", None),
+            metadata=metadata,
         ):
             method = functools.partial(
                 api.partition_query,
@@ -938,6 +950,7 @@ def begin(self):
             f"CloudSpanner.{type(self).__name__}.begin",
             self._session,
             observability_options=getattr(database, "observability_options", None),
+            metadata=metadata,
         ):
             method = functools.partial(
                 api.begin_transaction,