Merge branch 'main' into samples_1_run_tests

Author: daniel-sanche

packages/google-cloud-bigquery-storage/samples/to_dataframe/requirements.txt

@@ -1,19 +1,10 @@
 google-auth==2.52.0
 google-cloud-bigquery-storage==2.38.0
-google-cloud-bigquery==3.30.0; python_version <= '3.8'
-google-cloud-bigquery==3.41.0; python_version >= '3.9'
-pyarrow==12.0.1; python_version == '3.7'
-pyarrow==17.0.0; python_version == '3.8'
-pyarrow==24.0.0; python_version >= '3.9'
-ipython==7.31.1; python_version == '3.7'
-ipython==8.10.0; python_version == '3.8'
-ipython==8.18.1; python_version == '3.9'
-ipython==8.33.0; python_version == '3.10'
+google-cloud-bigquery==3.41.0
+pyarrow==24.0.0
+ipython===8.33.0; python_version == '3.10'
 ipython==9.13.0; python_version >= '3.11'
 ipywidgets==8.1.8
-pandas==1.3.5; python_version == '3.7'
-pandas==2.0.3; python_version == '3.8'
-pandas==2.3.3; python_version >= '3.9'
+pandas==2.3.3
 tqdm==4.67.3
-db-dtypes==1.4.2; python_version <= '3.8'
-db-dtypes==1.6.0; python_version >= '3.9'
+db-dtypes==1.6.0

packages/google-cloud-bigtable/google/cloud/bigtable/batcher.py

@@ -204,7 +204,10 @@ def flush(self):
             * :exc:`~batcher.MutationsBatchError` if there's any error in the mutations.
         """
         self._close_batcher()
+        # Set atexit registration after internal batcher.
+        atexit.unregister(self.close)
         self._init_batcher()
+        atexit.register(self.close)
 
     def __exit__(self, exc_type, exc_value, exc_traceback):
         """Clean up resources. Flush and shutdown the ThreadPoolExecutor."""

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

@@ -416,9 +416,8 @@ def _start_background_channel_refresh(self) -> None:
         ):
             # raise error if not in an event loop in async client
             CrossSync.verify_async_event_loop()
-            self._channel_refresh_task = CrossSync.create_task(
+            self._channel_refresh_task = CrossSync.create_daemon_task(
                 self._manage_channel,
-                sync_executor=self._executor,
                 task_name=f"{self.__class__.__name__} channel refresh",
             )
 

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

@@ -287,8 +287,8 @@ def __init__(
             if not CrossSync.is_async
             else None
         )
-        self._flush_timer = CrossSync.create_task(
-            self._timer_routine, flush_interval, sync_executor=self._sync_flush_executor
+        self._flush_timer = CrossSync.create_daemon_task(
+            self._timer_routine, flush_interval
         )
         self._flush_jobs: set[CrossSync.Future[None]] = set()
         # MutationExceptionGroup reports number of successful entries along with failures
@@ -392,18 +392,25 @@ async def _flush_internal(self, new_entries: list[RowMutationEntry]):
                 CrossSync.Future[list[FailedMutationEntryError]], list[RowMutationEntry]
             ]
         ] = []
+        found_exceptions: list[FailedMutationEntryError] = []
         async for batch, metric in self._flow_control.add_to_flow_with_metrics(
             new_entries, self._target.client._metrics
         ):
-            batch_task = CrossSync.create_task(
-                self._execute_mutate_rows,
-                batch,
-                metric,
-                sync_executor=self._sync_rpc_executor,
-            )
-            in_process_requests.append((batch_task, batch))
+            try:
+                batch_task = CrossSync.create_task(
+                    self._execute_mutate_rows,
+                    batch,
+                    metric,
+                    sync_executor=self._sync_rpc_executor,
+                )
+                in_process_requests.append((batch_task, batch))
+            except RuntimeError:
+                # Executor is already shut down (e.g. during atexit); flush synchronously.
+                found_exceptions.extend(await self._execute_mutate_rows(batch, metric))
         # wait for all inflight requests to complete
-        found_exceptions = await self._wait_for_batch_results(*in_process_requests)
+        found_exceptions.extend(
+            await self._wait_for_batch_results(*in_process_requests)
+        )
         # update exception data to reflect any new errors
         self._entries_processed_since_last_raise += len(new_entries)
         self._add_exceptions(found_exceptions)
@@ -545,7 +552,12 @@ async def close(self):
         """
         self._closed.set()
         self._flush_timer.cancel()
-        self._schedule_flush()
+        entries = self._staged_entries
+        try:
+            self._schedule_flush()
+        except RuntimeError:
+            # Executor is already shut down (e.g. during atexit); flush synchronously.
+            await self._flush_internal(entries)
         # shut down executors
         if self._sync_flush_executor:
             with self._sync_flush_executor:

packages/google-cloud-bigtable/google/cloud/bigtable/data/_cross_sync/cross_sync.py

@@ -205,6 +205,19 @@ def create_task(
             task.set_name(task_name)
         return task
 
+    @staticmethod
+    def create_daemon_task(
+        fn: Callable[..., Coroutine[Any, Any, T]],
+        *fn_args,
+        task_name: str | None = None,
+        **fn_kwargs,
+    ) -> CrossSync.Task[T]:
+        """
+        Abstraction over asyncio.create_task for long-lived background tasks.
+        Sync version runs on a dedicated daemon thread instead of a ThreadPoolExecutor.
+        """
+        return CrossSync.create_task(fn, *fn_args, task_name=task_name, **fn_kwargs)
+
     @staticmethod
     async def yield_to_event_loop() -> None:
         """
@@ -319,6 +332,30 @@ def create_task(
                 raise ValueError("sync_executor is required for sync version")
             return sync_executor.submit(fn, *fn_args, **fn_kwargs)
 
+        @staticmethod
+        def create_daemon_task(
+            fn: Callable[..., T],
+            *fn_args,
+            task_name: str | None = None,
+            **fn_kwargs,
+        ) -> CrossSync._Sync_Impl.Task[T]:
+            """
+            Run a long-lived background task on a dedicated daemon thread and
+            return a Future representing its completion.
+            """
+            future: concurrent.futures.Future[T] = concurrent.futures.Future()
+
+            def _run() -> None:
+                if not future.set_running_or_notify_cancel():  # pragma: NO COVER
+                    return
+                try:
+                    future.set_result(fn(*fn_args, **fn_kwargs))
+                except BaseException as exc:
+                    future.set_exception(exc)
+
+            threading.Thread(target=_run, name=task_name, daemon=True).start()
+            return future
+
         @staticmethod
         def yield_to_event_loop() -> None:
             """

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

@@ -306,9 +306,8 @@ def _start_background_channel_refresh(self) -> None:
             and (not self._disable_background_refresh)
         ):
             CrossSync._Sync_Impl.verify_async_event_loop()
-            self._channel_refresh_task = CrossSync._Sync_Impl.create_task(
+            self._channel_refresh_task = CrossSync._Sync_Impl.create_daemon_task(
                 self._manage_channel,
-                sync_executor=self._executor,
                 task_name=f"{self.__class__.__name__} channel refresh",
             )
 

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

@@ -254,8 +254,8 @@ def __init__(
             if not CrossSync._Sync_Impl.is_async
             else None
         )
-        self._flush_timer = CrossSync._Sync_Impl.create_task(
-            self._timer_routine, flush_interval, sync_executor=self._sync_flush_executor
+        self._flush_timer = CrossSync._Sync_Impl.create_daemon_task(
+            self._timer_routine, flush_interval
         )
         self._flush_jobs: set[CrossSync._Sync_Impl.Future[None]] = set()
         self._entries_processed_since_last_raise: int = 0
@@ -340,17 +340,21 @@ def _flush_internal(self, new_entries: list[RowMutationEntry]):
                 list[RowMutationEntry],
             ]
         ] = []
+        found_exceptions: list[FailedMutationEntryError] = []
         for batch, metric in self._flow_control.add_to_flow_with_metrics(
             new_entries, self._target.client._metrics
         ):
-            batch_task = CrossSync._Sync_Impl.create_task(
-                self._execute_mutate_rows,
-                batch,
-                metric,
-                sync_executor=self._sync_rpc_executor,
-            )
-            in_process_requests.append((batch_task, batch))
-        found_exceptions = self._wait_for_batch_results(*in_process_requests)
+            try:
+                batch_task = CrossSync._Sync_Impl.create_task(
+                    self._execute_mutate_rows,
+                    batch,
+                    metric,
+                    sync_executor=self._sync_rpc_executor,
+                )
+                in_process_requests.append((batch_task, batch))
+            except RuntimeError:
+                found_exceptions.extend(self._execute_mutate_rows(batch, metric))
+        found_exceptions.extend(self._wait_for_batch_results(*in_process_requests))
         self._entries_processed_since_last_raise += len(new_entries)
         self._add_exceptions(found_exceptions)
 
@@ -462,7 +466,11 @@ def close(self):
         """Flush queue and clean up resources"""
         self._closed.set()
         self._flush_timer.cancel()
-        self._schedule_flush()
+        entries = self._staged_entries
+        try:
+            self._schedule_flush()
+        except RuntimeError:
+            self._flush_internal(entries)
         if self._sync_flush_executor:
             with self._sync_flush_executor:
                 self._sync_flush_executor.shutdown(wait=True)

packages/google-cloud-bigtable/samples/legacy_client/beam/requirements.txt

@@ -1,5 +1,3 @@
-apache-beam==2.60.0; python_version == '3.8'
-apache-beam==2.69.0; python_version == '3.9'
-apache-beam==2.71.0; python_version >= '3.10'
+apache-beam==2.71.0
 google-cloud-bigtable
 google-cloud-core==2.5.0

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

@@ -1159,6 +1159,22 @@ async def test_close(self):
                     assert flush_mock.call_count == 1
                     assert raise_mock.call_count == 1
 
+    @CrossSync.pytest
+    async def test_close_executor_shutdown_fallback(self):
+        """If create_task raises RuntimeError during close (e.g. atexit), flush synchronously"""
+        async with self._make_one() as instance:
+            mutation = self._make_mutation(count=1, size=1)
+            instance._staged_entries = [mutation]
+            with mock.patch.object(
+                CrossSync, "create_task", side_effect=RuntimeError("shutdown")
+            ):
+                with mock.patch.object(
+                    instance, "_execute_mutate_rows", return_value=[]
+                ) as exec_mock:
+                    await instance.close()
+                    assert exec_mock.call_count == 1
+                    assert instance._staged_entries == []
+
     @CrossSync.pytest
     async def test_close_w_exceptions(self):
         """Raise exceptions on close"""

packages/google-cloud-bigtable/tests/unit/data/_cross_sync/test_cross_sync.py

@@ -450,6 +450,45 @@ def test_create_task_no_executor(self, cs_sync):
             cs_sync.create_task(lambda: None)
         assert "sync_executor is required" in str(e.value)
 
+    def test_create_daemon_task(self, cs_sync):
+        """
+        create_daemon_task should run on a daemon thread and return a Future
+        """
+        result = cs_sync.create_daemon_task(
+            lambda x, y: (x + y, threading.current_thread().daemon),
+            2,
+            y=3,
+            task_name="test-daemon",
+        )
+        assert isinstance(result, cs_sync.Task)
+        assert result.result(timeout=5) == (5, True)
+
+    def test_create_daemon_task_exception(self, cs_sync):
+        """
+        create_daemon_task should propagate exceptions onto the returned Future
+        """
+
+        def fail():
+            raise ValueError("boom")
+
+        result = cs_sync.create_daemon_task(fail)
+        with pytest.raises(ValueError, match="boom"):
+            result.result(timeout=5)
+
+    @pytest.mark.asyncio
+    async def test_create_daemon_task_async(self, cs_async):
+        """
+        Async create_daemon_task should create an asyncio.Task
+        """
+
+        async def coro_fn(x, y):
+            return x + y
+
+        result = cs_async.create_daemon_task(coro_fn, 2, y=3, task_name="async-daemon")
+        assert isinstance(result, asyncio.Task)
+        assert result.get_name() == "async-daemon"
+        assert await result == 5
+
     @pytest.mark.asyncio
     async def test_create_task_async(self, cs_async):
         """