all: implement custom tracer_provider injection

odeke-em · odeke-em · commit c169c4f6db94 · 2024-11-11T21:35:36.000-08:00
An important feature for observability is to allow the injection
of a custom tracer_provider instead of always using the global
tracer_provider by sending in

    observability_options=dict(
        tracer_provider=tracer_provider,
        enable_extended_tracing=True,
    )
diff --git a/docs/opentelemetry-tracing.rst b/docs/opentelemetry-tracing.rst
@@ -25,12 +25,16 @@ We also need to tell OpenTelemetry which exporter to use. To export Spanner trac
 
     # Create and export one trace every 1000 requests
     sampler = TraceIdRatioBased(1/1000)
-    # Use the default tracer provider
-    trace.set_tracer_provider(TracerProvider(sampler=sampler))
-    trace.get_tracer_provider().add_span_processor(
+    tracer_provider = TracerProvider(sampler=sampler)
+    tracer_provider.add_span_processor(
         # Initialize the cloud tracing exporter
         BatchSpanProcessor(CloudTraceSpanExporter())
     )
+    observability_options = dict(
+        tracer_provider=tracer_provider,
+        enable_extended_tracing=True,
+    )
+    spanner = spanner.NewClient(project_id, observability_options=observability_options)
 
 
 To get more fine-grained traces from gRPC, you can enable the gRPC instrumentation by the following
diff --git a/examples/trace.py b/examples/trace.py
@@ -32,15 +32,18 @@ def main():
     tracer_provider = TracerProvider(sampler=ALWAYS_ON)
     trace_exporter = CloudTraceSpanExporter(project_id=project_id)
     tracer_provider.add_span_processor(BatchSpanProcessor(trace_exporter))
-    trace.set_tracer_provider(tracer_provider)
-    # Retrieve a tracer from the global tracer provider.
-    tracer = tracer_provider.get_tracer('MyApp')
 
     # Setup the Cloud Spanner Client.
-    spanner_client = spanner.Client(project_id)
+    spanner_client = spanner.Client(
+        project_id,
+        observability_options=dict(tracer_provider=tracer_provider, enable_extended_tracing=True),
+    )
     instance = spanner_client.instance('test-instance')
     database = instance.database('test-db')
 
+    # Retrieve a tracer from our custom tracer provider.
+    tracer = tracer_provider.get_tracer('MyApp')
+
     # Now run our queries
     with tracer.start_as_current_span('QueryInformationSchema'):
         with database.snapshot() as snapshot:
diff --git a/google/cloud/spanner_v1/_opentelemetry_tracing.py b/google/cloud/spanner_v1/_opentelemetry_tracing.py
@@ -35,6 +35,20 @@
 TRACER_VERSION = gapic_version.__version__
 
 
+class ObservabilityOptions:
+    def __init__(self, tracer_provider=None, enable_extended_tracing=False):
+        self.__tracer_provider = tracer_provider
+        self.__enable_extended_tracing = enable_extended_tracing
+
+    @property
+    def tracer_provider(self):
+        return self.__tracer_provider
+
+    @property
+    def enable_extended_tracing(self):
+        return self.__enable_extended_tracing
+
+
 def get_tracer(tracer_provider=None):
     """
     get_tracer is a utility to unify and simplify retrieval of the tracer, without
@@ -51,13 +65,21 @@ def get_tracer(tracer_provider=None):
 
 
 @contextmanager
-def trace_call(name, session, extra_attributes=None):
+def trace_call(name, session, extra_attributes=None, observability_options=None):
     if not HAS_OPENTELEMETRY_INSTALLED or not session:
         # Empty context manager. Users will have to check if the generated value is None or a span
         yield None
         return
 
-    tracer = get_tracer()
+    tracer_provider = None
+    enable_extended_tracing = False
+    if getattr(session, "_observability_options", None):
+        opts = session._observability_options
+        if opts:
+            tracer_provider = opts.tracer_provider
+            enable_extended_tracing = opts.enable_extended_tracing
+
+    tracer = get_tracer(tracer_provider)
 
     # Set base attributes that we know for every trace created
     attributes = {
@@ -72,6 +94,13 @@ def trace_call(name, session, extra_attributes=None):
     if extra_attributes:
         attributes.update(extra_attributes)
 
+    # TODO(@odeke-em) enable after discussion with team and agreement
+    # over extended tracing changes as the legacy default is always to
+    # record SQL statements on spans.
+    if False and not enable_extended_tracing:
+        attributes.pop("db.statement", False)
+        attributes.pop("sql", False)
+
     with tracer.start_as_current_span(
         name, kind=trace.SpanKind.CLIENT, attributes=attributes
     ) as span:
diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py
@@ -126,6 +126,13 @@ class Client(ClientWithProject):
             for all ReadRequests and ExecuteSqlRequests that indicates which replicas
             or regions should be used for non-transactional reads or queries.
 
+    :type labels: dict (str -> any) or None
+    :param observability_options: (Optional) the configuration to control
+           the tracer's behavior.
+           tracer_provider is the injected tracer provider
+           enable_extended_tracing: :type:boolean when set to true will allow for
+           spans that issue SQL statements to be annotated with SQL.
+
     :raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
              and ``admin`` are :data:`True`
     """
@@ -146,6 +153,7 @@ def __init__(
         query_options=None,
         route_to_leader_enabled=True,
         directed_read_options=None,
+        observability_options=None,
     ):
         self._emulator_host = _get_spanner_emulator_host()
 
@@ -187,6 +195,7 @@ def __init__(
 
         self._route_to_leader_enabled = route_to_leader_enabled
         self._directed_read_options = directed_read_options
+        self._observability_options = observability_options
 
     @property
     def credentials(self):
@@ -371,6 +380,7 @@ def instance(
             self._emulator_host,
             labels,
             processing_units,
+            observability_options=self._observability_options,
         )
 
     def list_instances(self, filter_="", page_size=None):
diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py
@@ -156,6 +156,7 @@ def __init__(
         database_role=None,
         enable_drop_protection=False,
         proto_descriptors=None,
+        observability_options=None,
     ):
         self.database_id = database_id
         self._instance = instance
@@ -178,11 +179,16 @@ def __init__(
         self._reconciling = False
         self._directed_read_options = self._instance._client.directed_read_options
         self._proto_descriptors = proto_descriptors
+        self._observability_options = observability_options
 
         if pool is None:
-            pool = BurstyPool(database_role=database_role)
+            pool = BurstyPool(
+                database_role=database_role,
+                observability_options=self._observability_options,
+            )
 
         self._pool = pool
+        self._pool._observability_options = observability_options
         pool.bind(self)
 
     @classmethod
@@ -742,7 +748,12 @@ def session(self, labels=None, database_role=None):
         # If role is specified in param, then that role is used
         # instead.
         role = database_role or self._database_role
-        return Session(self, labels=labels, database_role=role)
+        return Session(
+            self,
+            labels=labels,
+            database_role=role,
+            observability_options=self._observability_options,
+        )
 
     def snapshot(self, **kw):
         """Return an object which wraps a snapshot.
diff --git a/google/cloud/spanner_v1/instance.py b/google/cloud/spanner_v1/instance.py
@@ -110,6 +110,13 @@ class Instance(object):
 
     :type labels: dict (str -> str) or None
     :param labels: (Optional) User-assigned labels for this instance.
+
+    :type labels: dict (str -> any) or None
+    :param observability_options: (Optional) the configuration to control
+           the tracer's behavior.
+           tracer_provider is the injected tracer provider
+           enable_extended_tracing: :type:boolean when set to true will allow for
+           spans that issue SQL statements to be annotated with SQL.
     """
 
     def __init__(
@@ -122,6 +129,7 @@ def __init__(
         emulator_host=None,
         labels=None,
         processing_units=None,
+        observability_options=None,
     ):
         self.instance_id = instance_id
         self._client = client
@@ -145,6 +153,7 @@ def __init__(
         if labels is None:
             labels = {}
         self.labels = labels
+        self._observability_options = observability_options
 
     def _update_from_pb(self, instance_pb):
         """Refresh self from the server-provided protobuf.
@@ -499,6 +508,7 @@ def database(
                 database_role=database_role,
                 enable_drop_protection=enable_drop_protection,
                 proto_descriptors=proto_descriptors,
+                observability_options=self._observability_options,
             )
         else:
             return TestDatabase(
diff --git a/google/cloud/spanner_v1/pool.py b/google/cloud/spanner_v1/pool.py
@@ -42,11 +42,12 @@ class AbstractSessionPool(object):
 
     _database = None
 
-    def __init__(self, labels=None, database_role=None):
+    def __init__(self, labels=None, database_role=None, observability_options=None):
         if labels is None:
             labels = {}
         self._labels = labels
         self._database_role = database_role
+        self._observability_options = observability_options
 
     @property
     def labels(self):
@@ -178,8 +179,13 @@ def __init__(
         default_timeout=DEFAULT_TIMEOUT,
         labels=None,
         database_role=None,
+        observability_options=None,
     ):
-        super(FixedSizePool, self).__init__(labels=labels, database_role=database_role)
+        super(FixedSizePool, self).__init__(
+            labels=labels,
+            database_role=database_role,
+            observability_options=observability_options,
+        )
         self.size = size
         self.default_timeout = default_timeout
         self._sessions = queue.LifoQueue(size)
@@ -284,8 +290,18 @@ class BurstyPool(AbstractSessionPool):
     :param database_role: (Optional) user-assigned database_role for the session.
     """
 
-    def __init__(self, target_size=10, labels=None, database_role=None):
-        super(BurstyPool, self).__init__(labels=labels, database_role=database_role)
+    def __init__(
+        self,
+        target_size=10,
+        labels=None,
+        database_role=None,
+        observability_options=None,
+    ):
+        super(BurstyPool, self).__init__(
+            labels=labels,
+            database_role=database_role,
+            observability_options=observability_options,
+        )
         self.target_size = target_size
         self._database = None
         self._sessions = queue.LifoQueue(target_size)
@@ -392,8 +408,13 @@ def __init__(
         ping_interval=3000,
         labels=None,
         database_role=None,
+        observability_options=None,
     ):
-        super(PingingPool, self).__init__(labels=labels, database_role=database_role)
+        super(PingingPool, self).__init__(
+            labels=labels,
+            database_role=database_role,
+            observability_options=observability_options,
+        )
         self.size = size
         self.default_timeout = default_timeout
         self._delta = datetime.timedelta(seconds=ping_interval)
@@ -546,6 +567,7 @@ def __init__(
         ping_interval=3000,
         labels=None,
         database_role=None,
+        observability_options=None,
     ):
         """This throws a deprecation warning on initialization."""
         warn(
@@ -561,6 +583,7 @@ def __init__(
             ping_interval,
             labels=labels,
             database_role=database_role,
+            observability_options=observability_options,
         )
 
         self.begin_pending_transactions()
diff --git a/google/cloud/spanner_v1/session.py b/google/cloud/spanner_v1/session.py
@@ -63,12 +63,15 @@ class Session(object):
     _session_id = None
     _transaction = None
 
-    def __init__(self, database, labels=None, database_role=None):
+    def __init__(
+        self, database, labels=None, database_role=None, observability_options=None
+    ):
         self._database = database
         if labels is None:
             labels = {}
         self._labels = labels
         self._database_role = database_role
+        self._observability_options = observability_options
 
     def __lt__(self, other):
         return self._session_id < other._session_id
diff --git a/tests/unit/test__opentelemetry_tracing.py b/tests/unit/test__opentelemetry_tracing.py
@@ -30,7 +30,13 @@ def _make_rpc_error(error_cls, trailing_metadata=None):
 def _make_session():
     from google.cloud.spanner_v1.session import Session
 
-    return mock.Mock(autospec=Session, instance=True)
+    session = mock.Mock(autospec=Session, instance=True)
+    # Setting _observability_options to None is to avoid the nasty spill-over
+    # of mock._tracer_provider spuriously failing tests, because per
+    # unittest.mock.Mock's definition invoking any attribute or method
+    # returns another mock.
+    setattr(session, "_observability_options", None)
+    return session
 
 
 # Skip all of these tests if we don't have OpenTelemetry
