Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"""
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

from __future__ import annotations

import logging
from typing import (
TYPE_CHECKING,
Any,
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment thread
daniel-sanche marked this conversation as resolved.
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):
Expand Down Expand Up @@ -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
Comment thread
daniel-sanche marked this conversation as resolved.
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(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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",
[
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Comment thread
daniel-sanche marked this conversation as resolved.
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()
Loading
Loading