support async callbacks

Author: daniel-sanche

packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/mutations_batcher.py

@@ -16,6 +16,7 @@
 
 import atexit
 import concurrent.futures
+import inspect
 import logging
 import time
 import warnings
@@ -444,9 +445,21 @@ async def _execute_mutate_rows(
             await self._flow_control.remove_from_flow(batch)
 
             # Call batch done callback with list of statuses.
+            # This is an internal callback used by the legacy synchronous shim,
+            # though async callbacks are awaited when running in async mode.
             if self._user_batch_completed_callback:
                 try:
-                    self._user_batch_completed_callback(statuses)
+                    result = self._user_batch_completed_callback(statuses)
+                    if CrossSync.is_async:
+                        if inspect.isawaitable(result):
+                            await result
+                    else:
+                        if inspect.isawaitable(result):
+                            if inspect.iscoroutine(result):
+                                result.close()
+                            raise TypeError(
+                                "_user_batch_completed_callback must be a synchronous callable"
+                            )
                 except Exception as exc:
                     _LOGGER.warning(
                         f"Exception raised in user batch completion callback: {exc}"

packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/mutations_batcher.py

@@ -19,6 +19,7 @@
 
 import atexit
 import concurrent.futures
+import inspect
 import logging
 import time
 import warnings
@@ -386,7 +387,13 @@ def _execute_mutate_rows(
             self._flow_control.remove_from_flow(batch)
             if self._user_batch_completed_callback:
                 try:
-                    self._user_batch_completed_callback(statuses)
+                    result = self._user_batch_completed_callback(statuses)
+                    if inspect.isawaitable(result):
+                        if inspect.iscoroutine(result):
+                            result.close()
+                        raise TypeError(
+                            "_user_batch_completed_callback must be a synchronous callable"
+                        )
                 except Exception as exc:
                     _LOGGER.warning(
                         f"Exception raised in user batch completion callback: {exc}"

packages/google-cloud-bigtable/tests/unit/data/_async/test_mutations_batcher.py

@@ -1075,6 +1075,31 @@ async def test__execute_mutate_rows_batch_completed_callback_exception(self):
                 callback.assert_called_once()
                 assert result == []
 
+    @CrossSync.pytest
+    async def test__execute_mutate_rows_batch_completed_callback_coroutine(self):
+        from google.rpc import code_pb2, status_pb2
+
+        with mock.patch.object(CrossSync, "_MutateRowsOperation") as mutate_rows:
+            mutate_rows.return_value = CrossSync.Mock()
+            table = mock.Mock()
+            table.default_mutate_rows_operation_timeout = 17
+            table.default_mutate_rows_attempt_timeout = 13
+            table.default_mutate_rows_retryable_errors = ()
+            called_with = []
+            async_callback = mock.AsyncMock(
+                side_effect=lambda statuses: called_with.append(statuses)
+            )
+
+            async with self._make_one(table) as instance:
+                instance._user_batch_completed_callback = async_callback
+                batch = [self._make_mutation()]
+                result = await instance._execute_mutate_rows(batch, mock.Mock())
+                assert result == []
+                if CrossSync.is_async:
+                    assert called_with == [[status_pb2.Status(code=code_pb2.OK)]]
+                else:
+                    assert called_with == []
+
     @CrossSync.pytest
     async def test__raise_exceptions(self):
         """Raise exceptions and reset error state"""

packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_mutations_batcher.py

@@ -957,6 +957,29 @@ def test__execute_mutate_rows_batch_completed_callback_exception(self):
                 callback.assert_called_once()
                 assert result == []
 
+    def test__execute_mutate_rows_batch_completed_callback_coroutine(self):
+        from google.rpc import code_pb2, status_pb2
+
+        with mock.patch.object(
+            CrossSync._Sync_Impl, "_MutateRowsOperation"
+        ) as mutate_rows:
+            mutate_rows.return_value = CrossSync._Sync_Impl.Mock()
+            table = mock.Mock()
+            table.default_mutate_rows_operation_timeout = 17
+            table.default_mutate_rows_attempt_timeout = 13
+            table.default_mutate_rows_retryable_errors = ()
+            called_with = []
+            async_callback = mock.AsyncMock(
+                side_effect=lambda statuses: called_with.append(statuses)
+            )
+
+            with self._make_one(table) as instance:
+                instance._user_batch_completed_callback = async_callback
+                batch = [self._make_mutation()]
+                result = instance._execute_mutate_rows(batch, mock.Mock())
+                assert result == []
+                assert called_with == []
+
     def test__raise_exceptions(self):
         """Raise exceptions and reset error state"""
         from google.cloud.bigtable.data import exceptions