fix: fix instance registration cleanup on early iterator termination

Author: jackdingilian

google/cloud/bigtable/data/_async/client.py

@@ -476,7 +476,8 @@ async def _manage_channel(
     async def _register_instance(
         self,
         instance_id: str,
-        owner: _DataApiTargetAsync | ExecuteQueryIteratorAsync,
+        app_profile_id: Optional[str],
+        owner_id: int,
     ) -> None:
         """
         Registers an instance with the client, and warms the channel for the instance
@@ -486,13 +487,15 @@ async def _register_instance(
 
         Args:
           instance_id: id of the instance to register.
-          owner: table that owns the instance. Owners will be tracked in
+          app_profile_id: id of the app profile calling the instance.
+          owner_id: integer id of the object owning the instance. Owners will be tracked in
               _instance_owners, and instances will only be unregistered when all
-              owners call _remove_instance_registration
+              owners call _remove_instance_registration. Can be obtained by calling
+              id(owner)
         """
         instance_name = self._gapic_client.instance_path(self.project, instance_id)
-        instance_key = _WarmedInstanceKey(instance_name, owner.app_profile_id)
-        self._instance_owners.setdefault(instance_key, set()).add(id(owner))
+        instance_key = _WarmedInstanceKey(instance_name, app_profile_id)
+        self._instance_owners.setdefault(instance_key, set()).add(owner_id)
         if instance_key not in self._active_instances:
             self._active_instances.add(instance_key)
             if self._channel_refresh_task:
@@ -510,10 +513,11 @@ async def _register_instance(
             "_DataApiTargetAsync": "_DataApiTarget",
         }
     )
-    async def _remove_instance_registration(
+    def _remove_instance_registration(
         self,
         instance_id: str,
-        owner: _DataApiTargetAsync | ExecuteQueryIteratorAsync,
+        app_profile_id: Optional[str],
+        owner_id: int,
     ) -> bool:
         """
         Removes an instance from the client's registered instances, to prevent
@@ -523,17 +527,17 @@ async def _remove_instance_registration(
 
         Args:
             instance_id: id of the instance to remove
-            owner: table that owns the instance. Owners will be tracked in
-              _instance_owners, and instances will only be unregistered when all
-              owners call _remove_instance_registration
+            app_profile_id: id of the app profile calling the instance.
+            owner_id: integer id of the object owning the instance. Can be
+                obtained by calling id(owner)
         Returns:
             bool: True if instance was removed, else False
         """
         instance_name = self._gapic_client.instance_path(self.project, instance_id)
-        instance_key = _WarmedInstanceKey(instance_name, owner.app_profile_id)
+        instance_key = _WarmedInstanceKey(instance_name, app_profile_id)
         owner_list = self._instance_owners.get(instance_key, set())
         try:
-            owner_list.remove(id(owner))
+            owner_list.remove(owner_id)
             if len(owner_list) == 0:
                 self._active_instances.remove(instance_key)
             return True
@@ -1014,7 +1018,8 @@ def __init__(
             self._register_instance_future = CrossSync.create_task(
                 self.client._register_instance,
                 self.instance_id,
-                self,
+                self.app_profile_id,
+                id(self),
                 sync_executor=self.client._executor,
             )
         except RuntimeError as e:
@@ -1725,7 +1730,9 @@ async def close(self):
         """
         if self._register_instance_future:
             self._register_instance_future.cancel()
-        await self.client._remove_instance_registration(self.instance_id, self)
+        self.client._remove_instance_registration(
+            self.instance_id, self.app_profile_id, id(self)
+        )
 
     @CrossSync.convert(sync_name="__enter__")
     async def __aenter__(self):

google/cloud/bigtable/data/_sync_autogen/client.py

@@ -354,7 +354,7 @@ def _manage_channel(
             next_sleep = max(next_refresh - (time.monotonic() - start_timestamp), 0)
 
     def _register_instance(
-        self, instance_id: str, owner: _DataApiTarget | ExecuteQueryIterator
+        self, instance_id: str, app_profile_id: Optional[str], owner_id: int
     ) -> None:
         """Registers an instance with the client, and warms the channel for the instance
         The client will periodically refresh grpc channel used to make
@@ -363,12 +363,14 @@ def _register_instance(
 
         Args:
           instance_id: id of the instance to register.
-          owner: table that owns the instance. Owners will be tracked in
+          app_profile_id: id of the app profile calling the instance.
+          owner_id: integer id of the object owning the instance. Owners will be tracked in
               _instance_owners, and instances will only be unregistered when all
-              owners call _remove_instance_registration"""
+              owners call _remove_instance_registration. Can be obtained by calling
+              id(owner)"""
         instance_name = self._gapic_client.instance_path(self.project, instance_id)
-        instance_key = _WarmedInstanceKey(instance_name, owner.app_profile_id)
-        self._instance_owners.setdefault(instance_key, set()).add(id(owner))
+        instance_key = _WarmedInstanceKey(instance_name, app_profile_id)
+        self._instance_owners.setdefault(instance_key, set()).add(owner_id)
         if instance_key not in self._active_instances:
             self._active_instances.add(instance_key)
             if self._channel_refresh_task:
@@ -377,7 +379,7 @@ def _register_instance(
                 self._start_background_channel_refresh()
 
     def _remove_instance_registration(
-        self, instance_id: str, owner: _DataApiTarget | ExecuteQueryIterator
+        self, instance_id: str, app_profile_id: Optional[str], owner_id: int
     ) -> bool:
         """Removes an instance from the client's registered instances, to prevent
         warming new channels for the instance
@@ -386,16 +388,16 @@ def _remove_instance_registration(
 
         Args:
             instance_id: id of the instance to remove
-            owner: table that owns the instance. Owners will be tracked in
-              _instance_owners, and instances will only be unregistered when all
-              owners call _remove_instance_registration
+            app_profile_id: id of the app profile calling the instance.
+            owner_id: integer id of the object owning the instance. Can be
+                obtained by calling id(owner)
         Returns:
             bool: True if instance was removed, else False"""
         instance_name = self._gapic_client.instance_path(self.project, instance_id)
-        instance_key = _WarmedInstanceKey(instance_name, owner.app_profile_id)
+        instance_key = _WarmedInstanceKey(instance_name, app_profile_id)
         owner_list = self._instance_owners.get(instance_key, set())
         try:
-            owner_list.remove(id(owner))
+            owner_list.remove(owner_id)
             if len(owner_list) == 0:
                 self._active_instances.remove(instance_key)
             return True
@@ -806,7 +808,8 @@ def __init__(
             self._register_instance_future = CrossSync._Sync_Impl.create_task(
                 self.client._register_instance,
                 self.instance_id,
-                self,
+                self.app_profile_id,
+                id(self),
                 sync_executor=self.client._executor,
             )
         except RuntimeError as e:
@@ -1460,7 +1463,9 @@ def close(self):
         """Called to close the Table instance and release any resources held by it."""
         if self._register_instance_future:
             self._register_instance_future.cancel()
-        self.client._remove_instance_registration(self.instance_id, self)
+        self.client._remove_instance_registration(
+            self.instance_id, self.app_profile_id, id(self)
+        )
 
     def __enter__(self):
         """Implement async context manager protocol

google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py

@@ -127,6 +127,7 @@ def __init__(
         self.has_received_token = False
         self._result_generator = self._next_impl()
         self._register_instance_task = None
+        self._fully_consumed = False
         self._is_closed = False
         self._request_body = request_body
         self._attempt_timeout_gen = _attempt_timeout_generator(
@@ -145,7 +146,8 @@ def __init__(
             self._register_instance_task = CrossSync.create_task(
                 self._client._register_instance,
                 self._instance_id,
-                self,
+                self.app_profile_id,
+                id(self),
                 sync_executor=self._client._executor,
             )
         except RuntimeError as e:
@@ -193,39 +195,42 @@ async def _next_impl(self) -> CrossSync.Iterator[QueryResultRow]:
         Generator wrapping the response stream which parses the stream results
         and returns full `QueryResultRow`s.
         """
-        async for response in self._stream:
-            try:
-                # we've received a resume token, so we can finalize the metadata
-                if self._final_metadata is None and _has_resume_token(response):
-                    self._finalize_metadata()
-
-                batches_to_parse = self._byte_cursor.consume(response)
-                if not batches_to_parse:
-                    continue
-                # metadata must be set at this point since there must be a resume_token
-                # for byte_cursor to yield data
-                if not self.metadata:
-                    raise ValueError(
-                        "Error parsing response before finalizing metadata"
+        try:
+            async for response in self._stream:
+                try:
+                    # we've received a resume token, so we can finalize the metadata
+                    if self._final_metadata is None and _has_resume_token(response):
+                        self._finalize_metadata()
+
+                    batches_to_parse = self._byte_cursor.consume(response)
+                    if not batches_to_parse:
+                        continue
+                    # metadata must be set at this point since there must be a resume_token
+                    # for byte_cursor to yield data
+                    if not self.metadata:
+                        raise ValueError(
+                            "Error parsing response before finalizing metadata"
+                        )
+                    results = self._reader.consume(
+                        batches_to_parse, self.metadata, self._column_info
                     )
-                results = self._reader.consume(
-                    batches_to_parse, self.metadata, self._column_info
-                )
-                if results is None:
-                    continue
-
-            except ValueError as e:
-                raise InvalidExecuteQueryResponse(
-                    "Invalid ExecuteQuery response received"
-                ) from e
-
-            for result in results:
-                yield result
-        # this means the stream has finished with no responses. In that case we know the
-        # latest_prepare_reponses was used successfully so we can finalize the metadata
-        if self._final_metadata is None:
-            self._finalize_metadata()
-        await self.close()
+                    if results is None:
+                        continue
+
+                except ValueError as e:
+                    raise InvalidExecuteQueryResponse(
+                        "Invalid ExecuteQuery response received"
+                    ) from e
+
+                for result in results:
+                    yield result
+            # this means the stream has finished with no responses. In that case we know the
+            # latest_prepare_reponses was used successfully so we can finalize the metadata
+            if self._final_metadata is None:
+                self._finalize_metadata()
+            self._fully_consumed = True
+        finally:
+            self._close_internal()
 
     @CrossSync.convert(sync_name="__next__", replace_symbols={"__anext__": "__next__"})
     async def __anext__(self) -> QueryResultRow:
@@ -285,15 +290,26 @@ def metadata(self) -> Metadata:
     @CrossSync.convert
     async def close(self) -> None:
         """
-        Cancel all background tasks. Should be called all rows were processed.
+        Cancel all background tasks. Should be called after all rows were processed.
+
+        Called automatically by iterator
 
         :raises: :class:`ValueError <exceptions.ValueError>` if called in an invalid state
         """
+        # this doesn't need to be async anymore but we wrap the sync api to avoid a breaking
+        # change
+        self._close_internal()
+
+    def _close_internal(self) -> None:
         if self._is_closed:
             return
-        if not self._byte_cursor.empty():
+        # Throw an error if the iterator has been successfully consumed but there is
+        # still buffered data
+        if self._fully_consumed and not self._byte_cursor.empty():
             raise ValueError("Unexpected buffered data at end of executeQuery reqest")
         self._is_closed = True
         if self._register_instance_task is not None:
             self._register_instance_task.cancel()
-        await self._client._remove_instance_registration(self._instance_id, self)
+        self._client._remove_instance_registration(
+            self._instance_id, self.app_profile_id, id(self)
+        )

google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py

@@ -102,6 +102,7 @@ def __init__(
         self.has_received_token = False
         self._result_generator = self._next_impl()
         self._register_instance_task = None
+        self._fully_consumed = False
         self._is_closed = False
         self._request_body = request_body
         self._attempt_timeout_gen = _attempt_timeout_generator(
@@ -120,7 +121,8 @@ def __init__(
             self._register_instance_task = CrossSync._Sync_Impl.create_task(
                 self._client._register_instance,
                 self._instance_id,
-                self,
+                self.app_profile_id,
+                id(self),
                 sync_executor=self._client._executor,
             )
         except RuntimeError as e:
@@ -159,31 +161,34 @@ def _make_request_with_resume_token(self):
     def _next_impl(self) -> CrossSync._Sync_Impl.Iterator[QueryResultRow]:
         """Generator wrapping the response stream which parses the stream results
         and returns full `QueryResultRow`s."""
-        for response in self._stream:
-            try:
-                if self._final_metadata is None and _has_resume_token(response):
-                    self._finalize_metadata()
-                batches_to_parse = self._byte_cursor.consume(response)
-                if not batches_to_parse:
-                    continue
-                if not self.metadata:
-                    raise ValueError(
-                        "Error parsing response before finalizing metadata"
+        try:
+            for response in self._stream:
+                try:
+                    if self._final_metadata is None and _has_resume_token(response):
+                        self._finalize_metadata()
+                    batches_to_parse = self._byte_cursor.consume(response)
+                    if not batches_to_parse:
+                        continue
+                    if not self.metadata:
+                        raise ValueError(
+                            "Error parsing response before finalizing metadata"
+                        )
+                    results = self._reader.consume(
+                        batches_to_parse, self.metadata, self._column_info
                     )
-                results = self._reader.consume(
-                    batches_to_parse, self.metadata, self._column_info
-                )
-                if results is None:
-                    continue
-            except ValueError as e:
-                raise InvalidExecuteQueryResponse(
-                    "Invalid ExecuteQuery response received"
-                ) from e
-            for result in results:
-                yield result
-        if self._final_metadata is None:
-            self._finalize_metadata()
-        self.close()
+                    if results is None:
+                        continue
+                except ValueError as e:
+                    raise InvalidExecuteQueryResponse(
+                        "Invalid ExecuteQuery response received"
+                    ) from e
+                for result in results:
+                    yield result
+            if self._final_metadata is None:
+                self._finalize_metadata()
+            self._fully_consumed = True
+        finally:
+            self._close_internal()
 
     def __next__(self) -> QueryResultRow:
         """Yields QueryResultRows representing the results of the query.
@@ -233,15 +238,22 @@ def metadata(self) -> Metadata:
         return self._final_metadata
 
     def close(self) -> None:
-        """Cancel all background tasks. Should be called all rows were processed.
+        """Cancel all background tasks. Should be called after all rows were processed.
+
+        Called automatically by iterator
 
         :raises: :class:`ValueError <exceptions.ValueError>` if called in an invalid state
         """
+        self._close_internal()
+
+    def _close_internal(self) -> None:
         if self._is_closed:
             return
-        if not self._byte_cursor.empty():
+        if self._fully_consumed and (not self._byte_cursor.empty()):
             raise ValueError("Unexpected buffered data at end of executeQuery reqest")
         self._is_closed = True
         if self._register_instance_task is not None:
             self._register_instance_task.cancel()
-        self._client._remove_instance_registration(self._instance_id, self)
+        self._client._remove_instance_registration(
+            self._instance_id, self.app_profile_id, id(self)
+        )