Merge branch 'csm_1_data_model' into csm_2_instrumentation

Author: daniel-sanche

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

@@ -1715,6 +1715,7 @@ async def close(self):
         """
         Called to close the Table instance and release any resources held by it.
         """
+        self._metrics.close()
         if self._register_instance_future:
             self._register_instance_future.cancel()
         await self.client._remove_instance_registration(self.instance_id, self)

google/cloud/bigtable/data/_metrics/handlers/_base.py

@@ -33,3 +33,6 @@ def on_attempt_complete(
         self, attempt: CompletedAttemptMetric, op: ActiveOperationMetric
     ) -> None:
         pass
+
+    def close(self):
+        pass

google/cloud/bigtable/data/_metrics/metrics_controller.py

@@ -54,3 +54,10 @@ def create_operation(
         Creates a new operation and registers it with the subscribed handlers.
         """
         return ActiveOperationMetric(op_type, **kwargs, handlers=self.handlers)
+
+    def close(self):
+        """
+        Close all handlers.
+        """
+        for handler in self.handlers:
+            handler.close()

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

@@ -1442,6 +1442,7 @@ def read_modify_write_row(
 
     def close(self):
         """Called to close the Table instance and release any resources held by it."""
+        self._metrics.close()
         if self._register_instance_future:
             self._register_instance_future.cancel()
         self.client._remove_instance_registration(self.instance_id, self)

tests/unit/data/_async/test_client.py

@@ -1464,6 +1464,20 @@ async def test_call_metadata(self, include_app_profile, fn_name, fn_args, gapic_
             # empty app_profile_id should send empty string
             assert "app_profile_id=" in routing_str
 
+    @CrossSync.pytest
+    async def test_close(self):
+        client = self._make_client()
+        table = self._make_one(client)
+        with mock.patch.object(
+            table._metrics, "close", mock.Mock()
+        ) as metric_close_mock:
+            with mock.patch.object(
+                client, "_remove_instance_registration"
+            ) as remove_mock:
+                await table.close()
+                remove_mock.assert_called_once_with(table.instance_id, table)
+                metric_close_mock.assert_called_once()
+
 
 @CrossSync.convert_class(
     "TestAuthorizedView", add_mapping_for_name="TestAuthorizedView"

tests/unit/data/_metrics/test_metrics_controller.py

@@ -87,3 +87,10 @@ def test_create_operation(self):
         assert op.zone is expected_zone
         assert len(op.handlers) == 1
         assert op.handlers[0] is handler
+
+    def test_close(self):
+        handlers = [mock.Mock() for _ in range(3)]
+        controller = self._make_one(handlers=handlers)
+        controller.close()
+        for handler in handlers:
+            handler.close.assert_called_once()

tests/unit/data/_sync_autogen/test_client.py

@@ -1173,6 +1173,19 @@ def test_call_metadata(self, include_app_profile, fn_name, fn_args, gapic_fn):
         else:
             assert "app_profile_id=" in routing_str
 
+    def test_close(self):
+        client = self._make_client()
+        table = self._make_one(client)
+        with mock.patch.object(
+            table._metrics, "close", mock.Mock()
+        ) as metric_close_mock:
+            with mock.patch.object(
+                client, "_remove_instance_registration"
+            ) as remove_mock:
+                table.close()
+                remove_mock.assert_called_once_with(table.instance_id, table)
+                metric_close_mock.assert_called_once()
+
 
 @CrossSync._Sync_Impl.add_mapping_decorator("TestAuthorizedView")
 class TestAuthorizedView(CrossSync._Sync_Impl.TestTable):