feat(spanner): implement custom tracer_provider injection for opentelemetry traces (#1229)

* all: implement custom tracer_provider injection 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, ) * Address review feedback by attaching observability_options to Client only * Attach observability_options directly before trace_call * More reverts for formatting * Plumb observability_options into _restart_on_unavailable * completely decouple observability_options from session * apply SPANNER_ENABLE_EXTENDED_TRACING but in inverse due to compatibility * Document SPANNER_ENABLE_EXTENDED_TRACING in environment * Revert a vestige of mock * tests: add unit test for propagating TracerProvider * Add preliminary end-to-end test to check for injection of observability_options * Document default enable_extended_tracing value * Carve out observability_options test * Ensure that observability_options test sets up and deletes database * Ensure instance.create() is invoked in system tests * Use getattr for mock _Client * Update with code review suggestions * Deal with mock.Mock false positives failing tests * Address review feedback

Author: odeke-em

docs/opentelemetry-tracing.rst

@@ -25,12 +25,21 @@ 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,
+
+ # By default extended_tracing is set to True due
+ # to legacy reasons to avoid breaking changes, you
+ # can modify it though using the environment variable
+ # SPANNER_ENABLE_EXTENDED_TRACING=false.
+ enable_extended_tracing=False,
+ )
+ 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
@@ -52,3 +61,13 @@ Generated spanner traces should now be available on `Cloud Trace <https://consol
 
 Tracing is most effective when many libraries are instrumented to provide insight over the entire lifespan of a request.
 For a list of libraries that can be instrumented, see the `OpenTelemetry Integrations` section of the `OpenTelemetry Python docs <https://opentelemetry-python.readthedocs.io/en/stable/>`_
+
+Annotating spans with SQL
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+By default your spans will be annotated with SQL statements where appropriate, but that can be a PII (Personally Identifiable Information)
+leak. Sadly due to legacy behavior, we cannot simply turn off this behavior by default. However you can control this behavior by setting
+
+ SPANNER_ENABLE_EXTENDED_TRACING=false
+
+to turn it off globally or when creating each SpannerClient, please set `observability_options.enable_extended_tracing=false`

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:

google/cloud/spanner_v1/_opentelemetry_tracing.py

@@ -15,6 +15,7 @@
 """Manages OpenTelemetry trace creation and handling"""
 
 from contextlib import contextmanager
+import os
 
 from google.cloud.spanner_v1 import SpannerClient
 from google.cloud.spanner_v1 import gapic_version
@@ -33,6 +34,9 @@
 
 TRACER_NAME = "cloud.google.com/python/spanner"
 TRACER_VERSION = gapic_version.__version__
+extended_tracing_globally_disabled = (
+ os.getenv("SPANNER_ENABLE_EXTENDED_TRACING", "").lower() == "false"
+)
 
 
 def get_tracer(tracer_provider=None):
@@ -51,13 +55,26 @@ 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
+
+ # By default enable_extended_tracing=True because in a bid to minimize
+ # breaking changes and preserve legacy behavior, we are keeping it turned
+ # on by default.
+ enable_extended_tracing = True
+
+ if isinstance(observability_options, dict): # Avoid false positives with mock.Mock
+ tracer_provider = observability_options.get("tracer_provider", None)
+ enable_extended_tracing = observability_options.get(
+ "enable_extended_tracing", enable_extended_tracing
+ )
+
+ tracer = get_tracer(tracer_provider)
 
  # Set base attributes that we know for every trace created
  attributes = {
@@ -72,6 +89,12 @@ def trace_call(name, session, extra_attributes=None):
  if extra_attributes:
  attributes.update(extra_attributes)
 
+ if extended_tracing_globally_disabled:
+ enable_extended_tracing = False
+
+ if not enable_extended_tracing:
+ attributes.pop("db.statement", False)
+
  with tracer.start_as_current_span(
  name, kind=trace.SpanKind.CLIENT, attributes=attributes
  ) as span:

google/cloud/spanner_v1/batch.py

@@ -205,7 +205,13 @@ def commit(
  max_commit_delay=max_commit_delay,
  request_options=request_options,
  )
- with trace_call("CloudSpanner.Commit", self._session, trace_attributes):
+ observability_options = getattr(database, "observability_options", None)
+ with trace_call(
+ "CloudSpanner.Commit",
+ self._session,
+ trace_attributes,
+ observability_options=observability_options,
+ ):
  method = functools.partial(
  api.commit,
  request=request,
@@ -318,7 +324,13 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
  request_options=request_options,
  exclude_txn_from_change_streams=exclude_txn_from_change_streams,
  )
- with trace_call("CloudSpanner.BatchWrite", self._session, trace_attributes):
+ observability_options = getattr(database, "observability_options", None)
+ with trace_call(
+ "CloudSpanner.BatchWrite",
+ self._session,
+ trace_attributes,
+ observability_options=observability_options,
+ ):
  method = functools.partial(
  api.batch_write,
  request=request,

google/cloud/spanner_v1/client.py

@@ -126,6 +126,16 @@ class Client(ClientWithProject):
  for all ReadRequests and ExecuteSqlRequests that indicates which replicas
  or regions should be used for non-transactional reads or queries.
 
+ :type observability_options: 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.
+ 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.
+
  :raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
  and ``admin`` are :data:`True`
  """
@@ -146,6 +156,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 +198,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):
@@ -268,6 +280,15 @@ def route_to_leader_enabled(self):
  """
  return self._route_to_leader_enabled
 
+ @property
+ def observability_options(self):
+ """Getter for observability_options.
+
+ :rtype: dict
+ :returns: The configured observability_options if set.
+ """
+ return self._observability_options
+
  @property
  def directed_read_options(self):
  """Getter for directed_read_options.

google/cloud/spanner_v1/database.py

@@ -718,6 +718,7 @@ def execute_pdml():
  method=method,
  request=request,
  transaction_selector=txn_selector,
+ observability_options=self.observability_options,
  )
 
  result_set = StreamedResultSet(iterator)
@@ -1106,6 +1107,17 @@ def set_iam_policy(self, policy):
  response = api.set_iam_policy(request=request, metadata=metadata)
  return response
 
+ @property
+ def observability_options(self):
+ """
+ Returns the observability options that you set when creating
+ the SpannerClient.
+ """
+ if not (self._instance and self._instance._client):
+ return None
+
+ return getattr(self._instance._client, "observability_options", None)
+
 
 class BatchCheckout(object):
  """Context manager for using a batch from a database.

google/cloud/spanner_v1/session.py

@@ -142,7 +142,13 @@ def create(self):
  if self._labels:
  request.session.labels = self._labels
 
- with trace_call("CloudSpanner.CreateSession", self, self._labels):
+ observability_options = getattr(self._database, "observability_options", None)
+ with trace_call(
+ "CloudSpanner.CreateSession",
+ self,
+ self._labels,
+ observability_options=observability_options,
+ ):
  session_pb = api.create_session(
  request=request,
  metadata=metadata,
@@ -169,7 +175,10 @@ def exists(self):
  )
  )
 
- with trace_call("CloudSpanner.GetSession", self) as span:
+ observability_options = getattr(self._database, "observability_options", None)
+ with trace_call(
+ "CloudSpanner.GetSession", self, observability_options=observability_options
+ ) as span:
  try:
  api.get_session(name=self.name, metadata=metadata)
  if span:
@@ -194,7 +203,12 @@ def delete(self):
  raise ValueError("Session ID not set by back-end")
  api = self._database.spanner_api
  metadata = _metadata_with_prefix(self._database.name)
- with trace_call("CloudSpanner.DeleteSession", self):
+ observability_options = getattr(self._database, "observability_options", None)
+ with trace_call(
+ "CloudSpanner.DeleteSession",
+ self,
+ observability_options=observability_options,
+ ):
  api.delete_session(name=self.name, metadata=metadata)
 
  def ping(self):