diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py index 048171474648..d9aba8c16c01 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_async/client.py @@ -416,11 +416,18 @@ 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._manage_channel, - sync_executor=self._executor, - task_name=f"{self.__class__.__name__} channel refresh", - ) + try: + self._channel_refresh_task = CrossSync.create_task( + self._manage_channel, + sync_executor=self._executor, + task_name=f"{self.__class__.__name__} channel refresh", + ) + except Exception as e: + _LOGGER.warning( + f"Failed to start background channel refresh task: {e}. " + "Channel refresh will be disabled." + ) + self._channel_refresh_task = None @CrossSync.convert async def close(self, timeout: float | None = 2.0): @@ -1151,18 +1158,29 @@ def __init__( default_retryable_errors or () ) + if CrossSync.is_async: + try: + CrossSync.verify_async_event_loop() + except RuntimeError as e: + raise RuntimeError( + f"{self.__class__.__name__} must be created within an async event loop context." + ) from e try: - self._register_instance_future = CrossSync.create_task( - self.client._register_instance, - self.instance_id, - self.app_profile_id, - id(self), - sync_executor=self.client._executor, + self._register_instance_future: CrossSync.Future[None] | None = ( + CrossSync.create_task( + self.client._register_instance, + self.instance_id, + self.app_profile_id, + id(self), + sync_executor=self.client._executor, + ) + ) + except Exception as e: + _LOGGER.warning( + f"Failed to start background instance registration: {e}. " + "Requests will proceed without proactive channel warming." ) - except RuntimeError as e: - raise RuntimeError( - f"{self.__class__.__name__} must be created within an async event loop context." - ) from e + self._register_instance_future = None def _create_operation( self, op_type: OperationType, **kwargs diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py index 24ece63767d4..dd03f05c0dec 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/_sync_autogen/client.py @@ -306,11 +306,17 @@ 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._manage_channel, - sync_executor=self._executor, - task_name=f"{self.__class__.__name__} channel refresh", - ) + try: + self._channel_refresh_task = CrossSync._Sync_Impl.create_task( + self._manage_channel, + sync_executor=self._executor, + task_name=f"{self.__class__.__name__} channel refresh", + ) + except Exception as e: + _LOGGER.warning( + f"Failed to start background channel refresh task: {e}. Channel refresh will be disabled." + ) + self._channel_refresh_task = None def close(self, timeout: float | None = 2.0): """Cancel all background tasks""" @@ -908,17 +914,20 @@ def __init__( default_retryable_errors or () ) try: - self._register_instance_future = CrossSync._Sync_Impl.create_task( - self.client._register_instance, - self.instance_id, - self.app_profile_id, - id(self), - sync_executor=self.client._executor, + self._register_instance_future: CrossSync._Sync_Impl.Future[None] | None = ( + CrossSync._Sync_Impl.create_task( + self.client._register_instance, + self.instance_id, + self.app_profile_id, + id(self), + sync_executor=self.client._executor, + ) + ) + except Exception as e: + _LOGGER.warning( + f"Failed to start background instance registration: {e}. Requests will proceed without proactive channel warming." ) - except RuntimeError as e: - raise RuntimeError( - f"{self.__class__.__name__} must be created within an async event loop context." - ) from e + self._register_instance_future = None def _create_operation( self, op_type: OperationType, **kwargs diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py index 4a207b057bdd..dbf00e23209d 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/execute_query/_async/execute_query_iterator.py @@ -14,6 +14,7 @@ from __future__ import annotations +import logging from typing import ( TYPE_CHECKING, Any, @@ -60,6 +61,8 @@ "google.cloud.bigtable.data.execute_query._sync_autogen.execute_query_iterator" ) +_LOGGER = logging.getLogger(__name__) + def _has_resume_token(response: ExecuteQueryResponse) -> bool: response_pb = response._pb # proto-plus attribute retrieval is slow. @@ -144,6 +147,13 @@ def __init__( ) self._req_metadata = req_metadata self._column_info = column_info + if CrossSync.is_async: + try: + CrossSync.verify_async_event_loop() + except RuntimeError as e: + raise RuntimeError( + f"{self.__class__.__name__} must be created within an async event loop context." + ) from e try: self._register_instance_task = CrossSync.create_task( self._client._register_instance, @@ -152,10 +162,12 @@ def __init__( id(self), sync_executor=self._client._executor, ) - except RuntimeError as e: - raise RuntimeError( - f"{self.__class__.__name__} must be created within an async event loop context." - ) from e + except Exception as e: + _LOGGER.warning( + f"Failed to start background instance registration: {e}. " + "Requests will proceed without proactive channel warming." + ) + self._register_instance_task = None @property def is_closed(self) -> bool: diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py b/packages/google-cloud-bigtable/google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py index 3f2febdaa557..263b25b874e5 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/data/execute_query/_sync_autogen/execute_query_iterator.py @@ -17,6 +17,7 @@ from __future__ import annotations +import logging from typing import TYPE_CHECKING, Any, Dict, Optional, Sequence, Tuple from google.api_core import retry as retries @@ -46,6 +47,7 @@ if TYPE_CHECKING: from google.cloud.bigtable.data import BigtableDataClient as DataClientType +_LOGGER = logging.getLogger(__name__) def _has_resume_token(response: ExecuteQueryResponse) -> bool: @@ -127,10 +129,11 @@ def __init__( id(self), sync_executor=self._client._executor, ) - except RuntimeError as e: - raise RuntimeError( - f"{self.__class__.__name__} must be created within an async event loop context." - ) from e + except Exception as e: + _LOGGER.warning( + f"Failed to start background instance registration: {e}. Requests will proceed without proactive channel warming." + ) + self._register_instance_task = None @property def is_closed(self) -> bool: diff --git a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py index 815f26a367e3..f5b1516c3879 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_async/test_client.py @@ -404,6 +404,33 @@ async def test__start_background_channel_refresh_disable_background_refresh(self assert client._channel_refresh_task is None ping_and_warm.assert_not_called() + @CrossSync.pytest + async def test__start_background_channel_refresh_task_creation_failure( + self, caplog + ): + import logging + + client = self._make_client( + project="project-id", + _disable_background_refresh=True, + ) + client._disable_background_refresh = False + client._emulator_host = None + with mock.patch.object( + CrossSync, "create_task", side_effect=RuntimeError("can't start new thread") + ) as mock_create_task: + with caplog.at_level(logging.WARNING): + client._start_background_channel_refresh() + assert client._channel_refresh_task is None + assert "Failed to start background channel refresh task" in caplog.text + assert "can't start new thread" in caplog.text + mock_create_task.assert_called_once_with( + client._manage_channel, + sync_executor=client._executor, + task_name=f"{client.__class__.__name__} channel refresh", + ) + await client.close() + @CrossSync.drop @CrossSync.pytest async def test__start_background_channel_refresh_task_names(self): @@ -1507,6 +1534,31 @@ def test_table_ctor_sync(self): TableAsync(client, "instance-id", "table-id") assert e.match("TableAsync must be created within an async event loop context.") + @CrossSync.pytest + async def test_table_ctor_task_creation_failure_warns_and_continues(self, caplog): + import logging + + client = self._make_client() + with mock.patch.object( + CrossSync, "create_task", side_effect=RuntimeError("can't start new thread") + ) as mock_create_task: + with caplog.at_level(logging.WARNING): + table = self._make_one(client) + assert table._register_instance_future is None + assert "Failed to start background instance registration" in caplog.text + assert "can't start new thread" in caplog.text + mock_create_task.assert_called_once_with( + client._register_instance, + table.instance_id, + table.app_profile_id, + id(table), + sync_executor=client._executor, + ) + async with table: + pass + await table.close() + await client.close() + @CrossSync.pytest # iterate over all retryable rpcs @pytest.mark.parametrize( diff --git a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py index 08d628f4aef7..ff1a39490afd 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py +++ b/packages/google-cloud-bigtable/tests/unit/data/_sync_autogen/test_client.py @@ -346,6 +346,31 @@ def test__start_background_channel_refresh_disable_background_refresh(self): assert client._channel_refresh_task is None ping_and_warm.assert_not_called() + def test__start_background_channel_refresh_task_creation_failure(self, caplog): + import logging + + client = self._make_client( + project="project-id", _disable_background_refresh=True + ) + client._disable_background_refresh = False + client._emulator_host = None + with mock.patch.object( + CrossSync._Sync_Impl, + "create_task", + side_effect=RuntimeError("can't start new thread"), + ) as mock_create_task: + with caplog.at_level(logging.WARNING): + client._start_background_channel_refresh() + assert client._channel_refresh_task is None + assert "Failed to start background channel refresh task" in caplog.text + assert "can't start new thread" in caplog.text + mock_create_task.assert_called_once_with( + client._manage_channel, + sync_executor=client._executor, + task_name=f"{client.__class__.__name__} channel refresh", + ) + client.close() + def test__ping_and_warm_instances(self): """test ping and warm with mocked asyncio.gather""" client_mock = mock.Mock() @@ -1262,6 +1287,32 @@ def test_ctor_invalid_timeout_values(self): assert "operation_timeout must be greater than 0" in str(e.value) client.close() + def test_table_ctor_task_creation_failure_warns_and_continues(self, caplog): + import logging + + client = self._make_client() + with mock.patch.object( + CrossSync._Sync_Impl, + "create_task", + side_effect=RuntimeError("can't start new thread"), + ) as mock_create_task: + with caplog.at_level(logging.WARNING): + table = self._make_one(client) + assert table._register_instance_future is None + assert "Failed to start background instance registration" in caplog.text + assert "can't start new thread" in caplog.text + mock_create_task.assert_called_once_with( + client._register_instance, + table.instance_id, + table.app_profile_id, + id(table), + sync_executor=client._executor, + ) + with table: + pass + table.close() + client.close() + @pytest.mark.parametrize( "fn_name,fn_args,is_stream,extra_retryables", [ diff --git a/packages/google-cloud-bigtable/tests/unit/data/execute_query/_async/test_query_iterator.py b/packages/google-cloud-bigtable/tests/unit/data/execute_query/_async/test_query_iterator.py index f6298f6bfd0d..07625189887d 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/execute_query/_async/test_query_iterator.py +++ b/packages/google-cloud-bigtable/tests/unit/data/execute_query/_async/test_query_iterator.py @@ -401,3 +401,61 @@ async def __anext__(self): # The close method should be called by the finally block on error client_mock._remove_instance_registration.assert_called_once() + + @CrossSync.drop + def test_iterator_ctor_sync(self): + # initializing iterator in a sync context should raise RuntimeError + client_mock = mock.Mock() + client_mock._register_instance = CrossSync.Mock() + with pytest.raises(RuntimeError) as e: + self._make_one( + client=client_mock, + instance_id="test-instance", + app_profile_id="test_profile", + request_body={}, + prepare_metadata=_pb_metadata_to_metadata_types( + metadata(column("test1", int64_type())) + ), + attempt_timeout=10, + operation_timeout=10, + ) + assert e.match( + f"{self._target_class().__name__} must be created within an async event loop context." + ) + + @CrossSync.pytest + async def test_iterator_ctor_task_creation_failure_warns_and_continues( + self, caplog + ): + import logging + + client_mock = mock.Mock() + client_mock._register_instance = CrossSync.Mock() + client_mock._remove_instance_registration = CrossSync.Mock() + client_mock._executor = concurrent.futures.ThreadPoolExecutor() + with mock.patch.object( + CrossSync, "create_task", side_effect=RuntimeError("can't start new thread") + ) as mock_create_task: + with caplog.at_level(logging.WARNING): + iterator = self._make_one( + client=client_mock, + instance_id="test-instance", + app_profile_id="test_profile", + request_body={}, + prepare_metadata=_pb_metadata_to_metadata_types( + metadata(column("test1", int64_type())) + ), + attempt_timeout=10, + operation_timeout=10, + ) + assert iterator._register_instance_task is None + assert "Failed to start background instance registration" in caplog.text + assert "can't start new thread" in caplog.text + mock_create_task.assert_called_once_with( + client_mock._register_instance, + "test-instance", + "test_profile", + id(iterator), + sync_executor=client_mock._executor, + ) + await iterator.close() diff --git a/packages/google-cloud-bigtable/tests/unit/data/execute_query/_sync_autogen/test_query_iterator.py b/packages/google-cloud-bigtable/tests/unit/data/execute_query/_sync_autogen/test_query_iterator.py index 01d73d8a2292..3f2c217f60ea 100644 --- a/packages/google-cloud-bigtable/tests/unit/data/execute_query/_sync_autogen/test_query_iterator.py +++ b/packages/google-cloud-bigtable/tests/unit/data/execute_query/_sync_autogen/test_query_iterator.py @@ -350,3 +350,39 @@ def __next__(self): for _ in iterator: pass client_mock._remove_instance_registration.assert_called_once() + + def test_iterator_ctor_task_creation_failure_warns_and_continues(self, caplog): + import logging + + client_mock = mock.Mock() + client_mock._register_instance = CrossSync._Sync_Impl.Mock() + client_mock._remove_instance_registration = CrossSync._Sync_Impl.Mock() + client_mock._executor = concurrent.futures.ThreadPoolExecutor() + with mock.patch.object( + CrossSync._Sync_Impl, + "create_task", + side_effect=RuntimeError("can't start new thread"), + ) as mock_create_task: + with caplog.at_level(logging.WARNING): + iterator = self._make_one( + client=client_mock, + instance_id="test-instance", + app_profile_id="test_profile", + request_body={}, + prepare_metadata=_pb_metadata_to_metadata_types( + metadata(column("test1", int64_type())) + ), + attempt_timeout=10, + operation_timeout=10, + ) + assert iterator._register_instance_task is None + assert "Failed to start background instance registration" in caplog.text + assert "can't start new thread" in caplog.text + mock_create_task.assert_called_once_with( + client_mock._register_instance, + "test-instance", + "test_profile", + id(iterator), + sync_executor=client_mock._executor, + ) + iterator.close()