Address review feedback

Author: odeke-em

google/cloud/spanner_v1/_opentelemetry_tracing.py

@@ -103,7 +103,10 @@ def trace_call(name, session, extra_attributes=None, observability_options=None)
             yield span
         except Exception as error:
             span.set_status(Status(StatusCode.ERROR, str(error)))
-            span.record_exception(error)
+            # OpenTelemetry-Python imposes invoking span.record_exception on __exit__
+            # on any exception. We should file a bug later on with them to only
+            # invoke .record_exception if not already invoked, hence we should not
+            # invoke .record_exception on our own else we shall have 2 exceptions.
             raise
         else:
             span.set_status(Status(StatusCode.OK))
@@ -115,6 +118,6 @@ def get_current_span():
     return trace.get_current_span()
 
 
-def add_span_event(span, commentary, event_attributes=None):
+def add_span_event(span, event_name, event_attributes=None):
     if span:
-        span.add_event(commentary, event_attributes)
+        span.add_event(event_name, event_attributes)

google/cloud/spanner_v1/pool.py

@@ -197,23 +197,25 @@ def bind(self, database):
                          when needed.
         """
         self._database = database
+        requested_session_count = self.size - self._sessions.qsize()
+        span = get_current_span()
+        span_event_attributes = {"kind": type(self).__name__}
+
+        if requested_session_count <= 0:
+            add_span_event(
+                span,
+                f"Invalid session pool size({requested_session_count}) <= 0",
+                span_event_attributes,
+            )
+            return
+
         api = database.spanner_api
         metadata = _metadata_with_prefix(database.name)
         if database._route_to_leader_enabled:
             metadata.append(
                 _metadata_with_leader_aware_routing(database._route_to_leader_enabled)
             )
         self._database_role = self._database_role or self._database.database_role
-
-        requested_session_count = self.size - self._sessions.qsize()
-        request = BatchCreateSessionsRequest(
-            database=database.name,
-            session_count=requested_session_count,
-            session_template=Session(creator_role=self.database_role),
-        )
-
-        span = get_current_span()
-        span_event_attributes = {"kind": type(self).__name__}
         if requested_session_count > 0:
             add_span_event(
                 span,
@@ -225,6 +227,12 @@ def bind(self, database):
             add_span_event(span, "Session pool is already full", span_event_attributes)
             return
 
+        request = BatchCreateSessionsRequest(
+            database=database.name,
+            session_count=requested_session_count,
+            session_template=Session(creator_role=self.database_role),
+        )
+
         returned_session_count = 0
         while not self._sessions.full():
             request.session_count = requested_session_count - self._sessions.qsize()
@@ -277,7 +285,9 @@ def get(self, timeout=None):
             )
             session = self._sessions.get(block=True, timeout=timeout)
         except queue.Empty as e:
-            add_span_event(current_span, "No session available", span_event_attributes)
+            add_span_event(
+                current_span, "No sessions available in the pool", span_event_attributes
+            )
             raise e
 
         span_event_attributes["session.id"] = session._session_id
@@ -290,6 +300,7 @@ def get(self, timeout=None):
             )
             session = self._database.session()
             session.create()
+            # Replacing with the updated session.id.
             span_event_attributes["session.id"] = session._session_id
 
         span_event_attributes["time.elapsed"] = time.time() - start_time
@@ -377,10 +388,10 @@ def get(self):
                 span_event_attributes,
             )
             session = self._sessions.get_nowait()
-        except queue.Empty:
+        except queue.Empty as e:
             add_span_event(
                 current_span,
-                "No session available. Creating session",
+                "No sessions available in pool. Creating session",
                 span_event_attributes,
             )
             session = self._new_session()
@@ -502,6 +513,14 @@ def bind(self, database):
         span_event_attributes = {"kind": type(self).__name__}
 
         requested_session_count = request.session_count
+        if requested_session_count <= 0:
+            add_span_event(
+                current_span,
+                f"Invalid session pool size({requested_session_count}) <= 0",
+                span_event_attributes,
+            )
+            return
+
         current_span = get_current_span()
         add_span_event(
             current_span,
@@ -571,7 +590,11 @@ def get(self, timeout=None):
         try:
             ping_after, session = self._sessions.get(block=True, timeout=timeout)
         except queue.Empty as e:
-            add_span_event(current_span, "No session available", span_event_attributes)
+            add_span_event(
+                current_span,
+                "No sessions available in the pool within the specified timeout",
+                span_event_attributes,
+            )
             raise e
 
         if _NOW() > ping_after:

google/cloud/spanner_v1/session.py

@@ -174,7 +174,8 @@ def exists(self):
         current_span = get_current_span()
         if self._session_id is None:
             add_span_event(
-                current_span, "Checking if Session failed due to unset session_id"
+                current_span,
+                "Checking session existence: Session does not exist as it has not been created yet",
             )
             return False
 

google/cloud/spanner_v1/transaction.py

@@ -283,9 +283,6 @@ def commit(
             trace_attributes,
             observability_options,
         ) as span:
-            if self._transaction_id is None and len(self._mutations) > 0:
-                self.begin()
-
             add_span_event(span, "Starting Commit")
 
             method = functools.partial(

tests/_helpers.py

@@ -108,3 +108,22 @@ def assertSpanEvents(self, name, wantEventNames=[], span=None):
         for event in span.events:
             actualEventNames.append(event.name)
         self.assertEqual(actualEventNames, wantEventNames)
+
+    def assertSpanNames(self, want_span_names):
+        if not HAS_OPENTELEMETRY_INSTALLED:
+            return
+
+        span_list = self.get_finished_spans()
+        got_span_names = [span.name for span in span_list]
+        self.assertEqual(got_span_names, want_span_names)
+
+    def get_finished_spans(self):
+        if HAS_OPENTELEMETRY_INSTALLED:
+            return list(
+                filter(
+                    lambda span: span and span.name,
+                    self.ot_exporter.get_finished_spans(),
+                )
+            )
+        else:
+            return []

tests/unit/test_pool.py

@@ -242,6 +242,33 @@ def test_spans_bind_get(self):
         ]
         self.assertSpanEvents("pool.Get", wantEventNames)
 
+    def test_spans_bind_get_empty_pool(self):
+        pool = self._make_one(size=0)
+        database = _Database("name")
+        session1 = _Session(database)
+        with trace_call("pool.Get", session1) as span:
+            try:
+                pool.bind(database)
+                database._sessions = database._sessions[:0]
+                pool.get()
+            except:
+                pass
+
+        wantEventNames = [
+            "Invalid session pool size(0) <= 0",
+            "Acquiring session",
+            "Waiting for a session to become available",
+            "No sessions available in the pool",
+        ]
+        self.assertSpanEvents("pool.Get", wantEventNames)
+
+        # Check for the overall spans too.
+        self.assertSpanNames(["pool.Get"])
+        self.assertSpanAttributes(
+            "pool.Get",
+            attributes=TestFixedSizePool.BASE_ATTRIBUTES,
+        )
+
     def test_spans_pool_bind(self):
         pool = self._make_one(size=1)
         database = _Database("name")
@@ -259,7 +286,6 @@ def test_spans_pool_bind(self):
             "Requesting 1 sessions",
             "Creating 1 sessions",
             "exception",
-            "exception",
         ]
         self.assertSpanEvents("testBind", wantEventNames)
 
@@ -429,7 +455,7 @@ def test_spans_get_empty_pool(self):
         wantEventNames = [
             "Acquiring session",
             "Waiting for a session to become available",
-            "No session available. Creating session",
+            "No sessions available. Creating session",
         ]
         self.assertSpanEvents("pool.Get", wantEventNames)
 

tests/unit/test_session.py

@@ -212,43 +212,6 @@ def test_create_session_span_annotations(self):
             wantEventNames = ["Creating Session"]
             self.assertSpanEvents("TestSessionSpan", wantEventNames, span)
 
-    def test_get_session_span_annotations(self):
-        from google.cloud.spanner_v1 import CreateSessionRequest
-        from google.cloud.spanner_v1 import Session as SessionRequestProto
-
-        session_pb = self._make_session_pb(
-            self.SESSION_NAME, database_role=self.DATABASE_ROLE
-        )
-
-        gax_api = self._make_spanner_api()
-        gax_api.create_session.return_value = session_pb
-        database = self._make_database(database_role=self.DATABASE_ROLE)
-        database.spanner_api = gax_api
-        session = self._make_one(database, database_role=self.DATABASE_ROLE)
-
-        with trace_call("TestSessionSpan", session) as span:
-            session.create()
-
-            self.assertEqual(session.session_id, self.SESSION_ID)
-            self.assertEqual(session.database_role, self.DATABASE_ROLE)
-            session_template = SessionRequestProto(creator_role=self.DATABASE_ROLE)
-
-            request = CreateSessionRequest(
-                database=database.name,
-                session=session_template,
-            )
-
-            gax_api.create_session.assert_called_once_with(
-                request=request,
-                metadata=[
-                    ("google-cloud-resource-prefix", database.name),
-                    ("x-goog-spanner-route-to-leader", "true"),
-                ],
-            )
-
-            wantEventNames = ["Creating Session"]
-            self.assertSpanEvents("TestSessionSpan", wantEventNames, span)
-
     def test_create_wo_database_role(self):
         from google.cloud.spanner_v1 import CreateSessionRequest