Infer database._channel_id only once along with spanner_api

googleapis · Dec 27, 2024 · 227b93a · 227b93a
1 parent c4745fc
commit 227b93a
Show file tree

Hide file tree

Showing 3 changed files with 152 additions and 32 deletions.
diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py
@@ -188,6 +188,7 @@ def __init__(
         self._reconciling = False
         self._directed_read_options = self._instance._client.directed_read_options
         self._proto_descriptors = proto_descriptors
+        self._channel_id = 0  # It'll be created when _spanner_api is created.
 
         if pool is None:
             pool = BurstyPool(database_role=database_role)
@@ -446,22 +447,16 @@ def spanner_api(self):
                 client_info=client_info,
                 client_options=client_options,
             )
-        return self._spanner_api
 
-    @property
-    def _channel_id(self):
-        """
-        Helper to retrieve the associated channelID for the spanner_api.
-        This property is paramount to x-goog-spanner-request-id.
-        """
-        with self.__transport_lock:
-            api = self.spanner_api
-            channel_id = self.__transports_to_channel_id.get(api._transport, None)
-            if channel_id is None:
-                channel_id = len(self.__transports_to_channel_id) + 1
-                self.__transports_to_channel_id[api._transport] = channel_id
-
-            return channel_id
+            with self.__transport_lock:
+                transport = self._spanner_api._transport
+                channel_id = self.__transports_to_channel_id.get(transport, None)
+                if channel_id is None:
+                    channel_id = len(self.__transports_to_channel_id) + 1
+                    self.__transports_to_channel_id[transport] = channel_id
+                self._channel_id = channel_id
+
+        return self._spanner_api
 
     def metadata_with_request_id(self, nth_request, nth_attempt, prior_metadata=[]):
         return _metadata_with_request_id(

diff --git a/tests/mockserver_tests/mock_server_test_base.py b/tests/mockserver_tests/mock_server_test_base.py
@@ -133,11 +133,8 @@ def setup_class(cls):
     def teardown_class(cls):
         if MockServerTestBase.server is not None:
             MockServerTestBase.server.stop(grace=None)
+            Client.NTH_CLIENT.reset()
             MockServerTestBase.server = None
-        self.reset()
-
-    def reset(self):
-        Client.NTH_CLIENT.reset()
 
     def setup_method(self, *args, **kwargs):
         self._client = None