From beb8bb91ccb2bbffe6c5a80a698f10578e57abe2 Mon Sep 17 00:00:00 2001
From: Emmanuel T Odeke <emmanuel@orijtech.com>
Date: Wed, 27 Nov 2024 22:54:07 -0800
Subject: [PATCH] Add docs for trace_call and trace_call_end_lazily

---
 .../spanner_v1/_opentelemetry_tracing.py      | 24 ++++++++++++++-----
 google/cloud/spanner_v1/batch.py              |  6 ++---
 google/cloud/spanner_v1/database.py           |  8 +++----
 3 files changed, 25 insertions(+), 13 deletions(-)

diff --git a/google/cloud/spanner_v1/_opentelemetry_tracing.py b/google/cloud/spanner_v1/_opentelemetry_tracing.py
index 79a501ddc3..b4c68a2172 100644
--- a/google/cloud/spanner_v1/_opentelemetry_tracing.py
+++ b/google/cloud/spanner_v1/_opentelemetry_tracing.py
@@ -54,8 +54,8 @@ def get_tracer(tracer_provider=None):
     return tracer_provider.get_tracer(TRACER_NAME, TRACER_VERSION)
 
 
-def _prepare_span_attributes(
-    name, session=None, extra_attributes=None, observability_options=None
+def _make_tracer_and_span_attributes(
+    session=None, extra_attributes=None, observability_options=None
 ):
     if not HAS_OPENTELEMETRY_INSTALLED:
         return None, None
@@ -104,8 +104,13 @@ def _prepare_span_attributes(
 
 @contextmanager
 def trace_call(name, session=None, extra_attributes=None, observability_options=None):
-    tracer, span_attributes = _prepare_span_attributes(
-        name, session, extra_attributes, observability_options
+    """
+    trace_call is used in situations where you need to end a span with a context manager
+    or after a scope is exited. If you need to keep a span alive and lazily end it, please
+    invoke `trace_call_end_lazily`.
+    """
+    tracer, span_attributes = _make_tracer_and_span_attributes(
+        session, extra_attributes, observability_options
     )
     if not tracer:
         yield None
@@ -124,10 +129,17 @@ def trace_call(name, session=None, extra_attributes=None, observability_options=
             span.set_status(Status(StatusCode.OK))
 
 
-def trace_end_explicitly(
+def trace_call_end_lazily(
     name, session=None, extra_attributes=None, observability_options=None
 ):
-    tracer, span_attributes = _prepare_span_attributes(
+    """
+    trace_call_end_lazily is used in situations where you won't have a context manager
+    and need to end a span explicitly when a specific condition happens. If you need a
+    context manager, please invoke `trace_call` with which you can invoke
+        `with trace_call(...) as span:`
+    It is the caller's responsibility to explicitly invoke span.end()
+    """
+    tracer, span_attributes = _make_tracer_and_span_attributes(
         session, extra_attributes, observability_options
     )
     if not tracer:
diff --git a/google/cloud/spanner_v1/batch.py b/google/cloud/spanner_v1/batch.py
index 1f177dd0c0..30a11e3713 100644
--- a/google/cloud/spanner_v1/batch.py
+++ b/google/cloud/spanner_v1/batch.py
@@ -28,7 +28,7 @@
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import (
     trace_call,
-    trace_end_explicitly,
+    trace_call_end_lazily,
 )
 from google.cloud.spanner_v1 import RequestOptions
 from google.cloud.spanner_v1._helpers import _retry
@@ -52,7 +52,7 @@ def __init__(self, session):
         observability_options = getattr(
             self._session._database, "observability_options", None
         )
-        self.__span = trace_end_explicitly(
+        self.__span = trace_call_end_lazily(
             "CloudSpanner." + type(self).__name__,
             self._session,
             observability_options=observability_options,
@@ -263,7 +263,7 @@ def __enter__(self):
         observability_options = getattr(
             self._session._database, "observability_options", None
         )
-        self.__span = trace_end_explicitly(
+        self.__span = trace_call_end_lazily(
             "CloudSpanner.Batch",
             self._session,
             observability_options=observability_options,
diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py
index 737b6464fd..0691ca0c2f 100644
--- a/google/cloud/spanner_v1/database.py
+++ b/google/cloud/spanner_v1/database.py
@@ -55,7 +55,7 @@
 )
 from google.cloud.spanner_v1._opentelemetry_tracing import (
     trace_call,
-    trace_end_explicitly,
+    trace_call_end_lazily,
 )
 from google.cloud.spanner_v1.batch import Batch
 from google.cloud.spanner_v1.batch import MutationGroups
@@ -1195,7 +1195,7 @@ def __init__(
     def __enter__(self):
         """Begin ``with`` block."""
         observability_options = self._database.observability_options
-        self.__span = trace_end_explicitly(
+        self.__span = trace_call_end_lazily(
             "CloudSpanner.Database.batch",
             None,
             observability_options=observability_options,
@@ -1291,7 +1291,7 @@ def __enter__(self):
         attributes = dict()
         if self._kw:
             attributes["multi_use"] = self._kw["multi_use"]
-        self.__span = trace_end_explicitly(
+        self.__span = trace_call_end_lazily(
             "CloudSpanner.Database.snapshot",
             None,
             attributes,
@@ -1346,7 +1346,7 @@ def __init__(
         self._exact_staleness = exact_staleness
         observability_options = getattr(self._database, "observability_options", {})
         self.__observability_options = observability_options
-        self.__span = trace_end_explicitly(
+        self.__span = trace_call_end_lazily(
             "CloudSpanner.BatchSnapshot",
             self._session,
             observability_options=observability_options,