diff --git a/google/cloud/bigtable/client.py b/google/cloud/bigtable/client.py index d5ba3169c..c06aaa6f6 100644 --- a/google/cloud/bigtable/client.py +++ b/google/cloud/bigtable/client.py @@ -27,6 +27,7 @@ * a :class:`~google.cloud.bigtable.table.Table` owns a :class:`~google.cloud.bigtable.row.Row` (and all the cells in the row) """ +import copy import os import warnings import grpc # type: ignore @@ -34,9 +35,7 @@ from google.api_core.gapic_v1 import client_info as client_info_lib from google.auth.credentials import AnonymousCredentials # type: ignore -from google.cloud import bigtable_v2 from google.cloud.bigtable import admin -from google.cloud.bigtable_v2.services.bigtable.transports import BigtableGrpcTransport from google.cloud.bigtable.admin.services.bigtable_instance_admin.transports import ( BigtableInstanceAdminGrpcTransport, ) @@ -54,6 +53,11 @@ from google.cloud.bigtable.cluster import _CLUSTER_NAME_RE from google.cloud.environment_vars import BIGTABLE_EMULATOR # type: ignore +from google.cloud.bigtable.data import BigtableDataClient +from google.cloud.bigtable.data._helpers import ( + _DEFAULT_BIGTABLE_EMULATOR_CLIENT, +) + INSTANCE_TYPE_PRODUCTION = instance.Instance.Type.PRODUCTION INSTANCE_TYPE_DEVELOPMENT = instance.Instance.Type.DEVELOPMENT @@ -66,7 +70,6 @@ READ_ONLY_SCOPE = "https://www.googleapis.com/auth/bigtable.data.readonly" """Scope for reading table data.""" -_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator" _GRPC_CHANNEL_OPTIONS = ( ("grpc.max_send_message_length", -1), ("grpc.max_receive_message_length", -1), @@ -290,18 +293,7 @@ def table_data_client(self): :rtype: :class:`.bigtable_v2.BigtableClient` :returns: A BigtableClient object. """ - if self._table_data_client is None: - transport = self._create_gapic_client_channel( - bigtable_v2.BigtableClient, - BigtableGrpcTransport, - ) - klass = _create_gapic_client( - bigtable_v2.BigtableClient, - client_options=self._client_options, - transport=transport, - ) - self._table_data_client = klass(self) - return self._table_data_client + return self._veneer_data_client._gapic_client @property def table_admin_client(self): @@ -369,6 +361,21 @@ def instance_admin_client(self): self._instance_admin_client = klass(self) return self._instance_admin_client + @property + def _veneer_data_client(self): + """Getter for the new Data Table API.""" + if self._table_data_client is None: + client_info = copy.copy(self._client_info) + client_info.client_library_version = f"{bigtable.__version__}-data-shim" + self._table_data_client = BigtableDataClient( + project=self.project, + credentials=self._credentials, + client_options=self._client_options, + _client_info=client_info, + _disable_background_refresh=True, + ) + return self._table_data_client + def instance(self, instance_id, display_name=None, instance_type=None, labels=None): """Factory to create a instance associated with this client. diff --git a/google/cloud/bigtable/data/_async/client.py b/google/cloud/bigtable/data/_async/client.py index 0af7154a6..35fe42814 100644 --- a/google/cloud/bigtable/data/_async/client.py +++ b/google/cloud/bigtable/data/_async/client.py @@ -64,7 +64,6 @@ import google.auth.credentials import google.auth._default from google.api_core import client_options as client_options_lib -from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT from google.cloud.bigtable.data.row import Row from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.data.exceptions import FailedQueryShardError @@ -73,6 +72,7 @@ from google.cloud.bigtable.data._helpers import TABLE_DEFAULT, _align_timeouts from google.cloud.bigtable.data._helpers import _WarmedInstanceKey from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT +from google.cloud.bigtable.data._helpers import _DEFAULT_BIGTABLE_EMULATOR_CLIENT from google.cloud.bigtable.data._helpers import _retry_exception_factory from google.cloud.bigtable.data._helpers import _validate_timeouts from google.cloud.bigtable.data._helpers import _get_error_type @@ -184,9 +184,15 @@ def __init__( """ if "pool_size" in kwargs: warnings.warn("pool_size no longer supported") - # set up client info headers for veneer library - self.client_info = DEFAULT_CLIENT_INFO - self.client_info.client_library_version = self._client_version() + + # set up client info headers for veneer library. _client_info is for internal use only, + # for the legacy client shim. + if kwargs.get("_client_info"): + self.client_info = kwargs["_client_info"] + else: + self.client_info = DEFAULT_CLIENT_INFO + self.client_info.client_library_version = self._client_version() + # parse client options if type(client_options) is dict: client_options = client_options_lib.from_dict(client_options) @@ -236,6 +242,10 @@ def __init__( "is the default." ) self._is_closed = CrossSync.Event() + # Private argument, for internal use only + self._disable_background_refresh = bool( + kwargs.get("_disable_background_refresh", False) + ) self.transport = cast(TransportType, self._gapic_client.transport) # keep track of active instances to for warmup on channel refresh self._active_instances: Set[_WarmedInstanceKey] = set() @@ -329,6 +339,7 @@ def _start_background_channel_refresh(self) -> None: not self._channel_refresh_task and not self._emulator_host and not self._is_closed.is_set() + and not self._disable_background_refresh ): # raise error if not in an event loop in async client CrossSync.verify_async_event_loop() diff --git a/google/cloud/bigtable/data/_helpers.py b/google/cloud/bigtable/data/_helpers.py index 424a34486..db17b0e42 100644 --- a/google/cloud/bigtable/data/_helpers.py +++ b/google/cloud/bigtable/data/_helpers.py @@ -44,6 +44,9 @@ # used by read_rows_sharded to limit how many requests are attempted in parallel _CONCURRENCY_LIMIT = 10 +# used by every data client as a default project name for testing on Bigtable emulator. +_DEFAULT_BIGTABLE_EMULATOR_CLIENT = "google-cloud-bigtable-emulator" + # used to identify an active bigtable resource that needs to be warmed through PingAndWarm # each instance/app_profile_id pair needs to be individually tracked _WarmedInstanceKey = namedtuple( diff --git a/google/cloud/bigtable/data/_sync_autogen/client.py b/google/cloud/bigtable/data/_sync_autogen/client.py index adc849649..88136ddad 100644 --- a/google/cloud/bigtable/data/_sync_autogen/client.py +++ b/google/cloud/bigtable/data/_sync_autogen/client.py @@ -54,7 +54,6 @@ import google.auth.credentials import google.auth._default from google.api_core import client_options as client_options_lib -from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT from google.cloud.bigtable.data.row import Row from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery from google.cloud.bigtable.data.exceptions import FailedQueryShardError @@ -62,6 +61,7 @@ from google.cloud.bigtable.data._helpers import TABLE_DEFAULT, _align_timeouts from google.cloud.bigtable.data._helpers import _WarmedInstanceKey from google.cloud.bigtable.data._helpers import _CONCURRENCY_LIMIT +from google.cloud.bigtable.data._helpers import _DEFAULT_BIGTABLE_EMULATOR_CLIENT from google.cloud.bigtable.data._helpers import _retry_exception_factory from google.cloud.bigtable.data._helpers import _validate_timeouts from google.cloud.bigtable.data._helpers import _get_error_type @@ -127,8 +127,11 @@ def __init__( """ if "pool_size" in kwargs: warnings.warn("pool_size no longer supported") - self.client_info = DEFAULT_CLIENT_INFO - self.client_info.client_library_version = self._client_version() + if kwargs.get("_client_info"): + self.client_info = kwargs["_client_info"] + else: + self.client_info = DEFAULT_CLIENT_INFO + self.client_info.client_library_version = self._client_version() if type(client_options) is dict: client_options = client_options_lib.from_dict(client_options) client_options = cast( @@ -168,6 +171,9 @@ def __init__( f"The configured universe domain ({self.universe_domain}) does not match the universe domain found in the credentials ({self._credentials.universe_domain}). If you haven't configured the universe domain explicitly, `googleapis.com` is the default." ) self._is_closed = CrossSync._Sync_Impl.Event() + self._disable_background_refresh = bool( + kwargs.get("_disable_background_refresh", False) + ) self.transport = cast(TransportType, self._gapic_client.transport) self._active_instances: Set[_WarmedInstanceKey] = set() self._instance_owners: dict[_WarmedInstanceKey, Set[int]] = {} @@ -238,6 +244,7 @@ def _start_background_channel_refresh(self) -> None: not self._channel_refresh_task and (not self._emulator_host) and (not self._is_closed.is_set()) + and (not self._disable_background_refresh) ): CrossSync._Sync_Impl.verify_async_event_loop() self._channel_refresh_task = CrossSync._Sync_Impl.create_task( diff --git a/google/cloud/bigtable/table.py b/google/cloud/bigtable/table.py index 97d7322e9..9ce7c312a 100644 --- a/google/cloud/bigtable/table.py +++ b/google/cloud/bigtable/table.py @@ -132,6 +132,12 @@ def __init__(self, table_id, instance, mutation_timeout=None, app_profile_id=Non self._app_profile_id = app_profile_id self.mutation_timeout = mutation_timeout + self._table_impl = self._instance._client._veneer_data_client.get_table( + self._instance.instance_id, + self.table_id, + app_profile_id=self._app_profile_id, + ) + @property def name(self): """Table name used in requests. diff --git a/tests/unit/data/_async/test_client.py b/tests/unit/data/_async/test_client.py index a5ec1d02d..0bc8b921d 100644 --- a/tests/unit/data/_async/test_client.py +++ b/tests/unit/data/_async/test_client.py @@ -119,7 +119,13 @@ async def test_ctor(self): @CrossSync.pytest async def test_ctor_super_inits(self): from google.cloud.client import ClientWithProject + from google.cloud.bigtable import __version__ as bigtable_version from google.api_core import client_options as client_options_lib + from google.cloud.bigtable_v2.services.bigtable.transports.base import ( + DEFAULT_CLIENT_INFO, + ) + + import copy project = "project-id" credentials = AnonymousCredentials() @@ -147,6 +153,22 @@ async def test_ctor_super_inits(self): kwargs = bigtable_client_init.call_args[1] assert kwargs["credentials"] == credentials assert kwargs["client_options"] == options_parsed + + expected_client_info = copy.copy(DEFAULT_CLIENT_INFO) + expected_client_info.client_library_version = ( + f"{bigtable_version}-data" + if not CrossSync.is_async + else f"{bigtable_version}-data-async" + ) + assert ( + kwargs["client_info"].to_user_agent() + == expected_client_info.to_user_agent() + ) + assert ( + kwargs["client_info"].to_grpc_metadata() + == expected_client_info.to_grpc_metadata() + ) + # test mixin superclass init was called assert client_project_init.call_count == 1 kwargs = client_project_init.call_args[1] @@ -154,6 +176,40 @@ async def test_ctor_super_inits(self): assert kwargs["credentials"] == credentials assert kwargs["client_options"] == options_parsed + @CrossSync.pytest + async def test_ctor_legacy_client(self): + from google.api_core import client_options as client_options_lib + from google.api_core.gapic_v1.client_info import ClientInfo + + project = "project-id" + credentials = AnonymousCredentials() + client_info = ClientInfo(gapic_version="1.2.3", user_agent="test-client-") + client_options = {"api_endpoint": "foo.bar:1234"} + options_parsed = client_options_lib.from_dict(client_options) + with mock.patch.object( + CrossSync.GapicClient, "__init__" + ) as bigtable_client_init: + try: + client = self._make_client( + project=project, + credentials=credentials, + client_options=options_parsed, + use_emulator=False, + _client_info=client_info, + _disable_background_refresh=True, + ) + + assert client._disable_background_refresh + assert client.client_info is client_info + except TypeError: + pass + + # test gapic superclass init was called with the right arguments + assert bigtable_client_init.call_count == 1 + kwargs = bigtable_client_init.call_args[1] + assert kwargs["credentials"] == credentials + assert kwargs["client_options"] == options_parsed + @CrossSync.pytest async def test_ctor_dict_options(self): from google.api_core.client_options import ClientOptions @@ -245,6 +301,22 @@ async def test__start_background_channel_refresh(self): assert ping_and_warm.call_count == 1 await client.close() + @CrossSync.pytest + async def test__start_background_channel_refresh_disable_background_refresh(self): + client = self._make_client( + project="project-id", + _disable_background_refresh=True, + ) + # should create background tasks for each channel + with mock.patch.object( + client, "_ping_and_warm_instances", CrossSync.Mock() + ) as ping_and_warm: + client._emulator_host = None + client.transport._grpc_channel = CrossSync.SwappableChannel(mock.Mock) + client._start_background_channel_refresh() + assert client._channel_refresh_task is None + ping_and_warm.assert_not_called() + @CrossSync.drop @CrossSync.pytest @pytest.mark.skipif( diff --git a/tests/unit/data/_sync_autogen/test_client.py b/tests/unit/data/_sync_autogen/test_client.py index 6ad6c1063..c2a15a991 100644 --- a/tests/unit/data/_sync_autogen/test_client.py +++ b/tests/unit/data/_sync_autogen/test_client.py @@ -89,7 +89,12 @@ def test_ctor(self): def test_ctor_super_inits(self): from google.cloud.client import ClientWithProject + from google.cloud.bigtable import __version__ as bigtable_version from google.api_core import client_options as client_options_lib + from google.cloud.bigtable_v2.services.bigtable.transports.base import ( + DEFAULT_CLIENT_INFO, + ) + import copy project = "project-id" credentials = AnonymousCredentials() @@ -116,12 +121,56 @@ def test_ctor_super_inits(self): kwargs = bigtable_client_init.call_args[1] assert kwargs["credentials"] == credentials assert kwargs["client_options"] == options_parsed + expected_client_info = copy.copy(DEFAULT_CLIENT_INFO) + expected_client_info.client_library_version = ( + f"{bigtable_version}-data" + if not CrossSync._Sync_Impl.is_async + else f"{bigtable_version}-data-async" + ) + assert ( + kwargs["client_info"].to_user_agent() + == expected_client_info.to_user_agent() + ) + assert ( + kwargs["client_info"].to_grpc_metadata() + == expected_client_info.to_grpc_metadata() + ) assert client_project_init.call_count == 1 kwargs = client_project_init.call_args[1] assert kwargs["project"] == project assert kwargs["credentials"] == credentials assert kwargs["client_options"] == options_parsed + def test_ctor_legacy_client(self): + from google.api_core import client_options as client_options_lib + from google.api_core.gapic_v1.client_info import ClientInfo + + project = "project-id" + credentials = AnonymousCredentials() + client_info = ClientInfo(gapic_version="1.2.3", user_agent="test-client-") + client_options = {"api_endpoint": "foo.bar:1234"} + options_parsed = client_options_lib.from_dict(client_options) + with mock.patch.object( + CrossSync._Sync_Impl.GapicClient, "__init__" + ) as bigtable_client_init: + try: + client = self._make_client( + project=project, + credentials=credentials, + client_options=options_parsed, + use_emulator=False, + _client_info=client_info, + _disable_background_refresh=True, + ) + assert client._disable_background_refresh + assert client.client_info is client_info + except TypeError: + pass + assert bigtable_client_init.call_count == 1 + kwargs = bigtable_client_init.call_args[1] + assert kwargs["credentials"] == credentials + assert kwargs["client_options"] == options_parsed + def test_ctor_dict_options(self): from google.api_core.client_options import ClientOptions @@ -194,6 +243,21 @@ def test__start_background_channel_refresh(self): assert ping_and_warm.call_count == 1 client.close() + def test__start_background_channel_refresh_disable_background_refresh(self): + client = self._make_client( + project="project-id", _disable_background_refresh=True + ) + with mock.patch.object( + client, "_ping_and_warm_instances", CrossSync._Sync_Impl.Mock() + ) as ping_and_warm: + client._emulator_host = None + client.transport._grpc_channel = CrossSync._Sync_Impl.SwappableChannel( + mock.Mock + ) + client._start_background_channel_refresh() + assert client._channel_refresh_task is None + ping_and_warm.assert_not_called() + def test__ping_and_warm_instances(self): """test ping and warm with mocked asyncio.gather""" client_mock = mock.Mock() diff --git a/tests/unit/v2_client/_testing.py b/tests/unit/v2_client/_testing.py index 302d33ac1..6b34afffc 100644 --- a/tests/unit/v2_client/_testing.py +++ b/tests/unit/v2_client/_testing.py @@ -27,11 +27,22 @@ def __init__(self, *results): def _make_credentials(): + from google.cloud.bigtable_v2 import BigtableClient import google.auth.credentials - class _CredentialsWithScopes( - google.auth.credentials.Credentials, google.auth.credentials.Scoped + class _CredentialsWithScopesAndQuotaProject( + google.auth.credentials.Scoped, + google.auth.credentials.CredentialsWithQuotaProject, + google.auth.credentials.Credentials, ): pass - return mock.Mock(spec=_CredentialsWithScopes) + credentials = mock.Mock(spec=_CredentialsWithScopesAndQuotaProject) + + # Needed to mock universe domain and quota project for new data client tests + credentials_with_scopes = mock.Mock(spec=_CredentialsWithScopesAndQuotaProject) + credentials_with_scopes.universe_domain = BigtableClient._DEFAULT_UNIVERSE + credentials_with_scopes.with_quota_project.return_value = credentials_with_scopes + credentials.with_scopes.return_value = credentials_with_scopes + + return credentials diff --git a/tests/unit/v2_client/test_client.py b/tests/unit/v2_client/test_client.py index 45d260267..db647b908 100644 --- a/tests/unit/v2_client/test_client.py +++ b/tests/unit/v2_client/test_client.py @@ -171,15 +171,18 @@ def test_client_constructor_w_both_admin_and_read_only(): def test_client_constructor_w_emulator_host(): from google.cloud.environment_vars import BIGTABLE_EMULATOR - from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT - from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS + from google.cloud.bigtable.data._sync_autogen.client import ( + _DEFAULT_BIGTABLE_EMULATOR_CLIENT, + ) import grpc emulator_host = "localhost:8081" with mock.patch("os.environ", {BIGTABLE_EMULATOR: emulator_host}): channel = grpc.insecure_channel("no-host") - with mock.patch("grpc.insecure_channel", return_value=channel) as factory: - factory.return_value = channel + with mock.patch( + "google.cloud.bigtable.data._sync_autogen.client.insecure_channel", + return_value=channel, + ) as factory: client = _make_client() # don't test local_composite_credentials # client._local_composite_credentials = lambda: credentials @@ -189,21 +192,23 @@ def test_client_constructor_w_emulator_host(): assert client._emulator_host == emulator_host assert client.project == _DEFAULT_BIGTABLE_EMULATOR_CLIENT - factory.assert_called_once_with( - emulator_host, - options=_GRPC_CHANNEL_OPTIONS, - ) + factory.assert_called_once_with(emulator_host) + + assert client._table_data_client._emulator_host == emulator_host + assert client._table_data_client.project == _DEFAULT_BIGTABLE_EMULATOR_CLIENT def test_client_constructor_w_emulator_host_w_project(): from google.cloud.environment_vars import BIGTABLE_EMULATOR - from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS import grpc emulator_host = "localhost:8081" with mock.patch("os.environ", {BIGTABLE_EMULATOR: emulator_host}): channel = grpc.insecure_channel("no-host") - with mock.patch("grpc.insecure_channel", return_value=channel) as factory: + with mock.patch( + "google.cloud.bigtable.data._sync_autogen.client.insecure_channel", + return_value=channel, + ) as factory: client = _make_client(project=PROJECT) # channels are formed when needed, so access a client # create a gapic channel @@ -211,23 +216,27 @@ def test_client_constructor_w_emulator_host_w_project(): assert client._emulator_host == emulator_host assert client.project == PROJECT - factory.assert_called_once_with( - emulator_host, - options=_GRPC_CHANNEL_OPTIONS, - ) + factory.assert_called_once_with(emulator_host) + + assert client._table_data_client._emulator_host == emulator_host + assert client._table_data_client.project == PROJECT def test_client_constructor_w_emulator_host_w_credentials(): from google.cloud.environment_vars import BIGTABLE_EMULATOR - from google.cloud.bigtable.client import _DEFAULT_BIGTABLE_EMULATOR_CLIENT - from google.cloud.bigtable.client import _GRPC_CHANNEL_OPTIONS + from google.cloud.bigtable.data._sync_autogen.client import ( + _DEFAULT_BIGTABLE_EMULATOR_CLIENT, + ) import grpc emulator_host = "localhost:8081" credentials = _make_credentials() with mock.patch("os.environ", {BIGTABLE_EMULATOR: emulator_host}): channel = grpc.insecure_channel("no-host") - with mock.patch("grpc.insecure_channel", return_value=channel) as factory: + with mock.patch( + "google.cloud.bigtable.data._sync_autogen.client.insecure_channel", + return_value=channel, + ) as factory: client = _make_client(credentials=credentials) # channels are formed when needed, so access a client # create a gapic channel @@ -235,10 +244,10 @@ def test_client_constructor_w_emulator_host_w_credentials(): assert client._emulator_host == emulator_host assert client.project == _DEFAULT_BIGTABLE_EMULATOR_CLIENT - factory.assert_called_once_with( - emulator_host, - options=_GRPC_CHANNEL_OPTIONS, - ) + factory.assert_called_once_with(emulator_host) + + assert client._table_data_client._emulator_host == emulator_host + assert client._table_data_client.project == _DEFAULT_BIGTABLE_EMULATOR_CLIENT def test_client__get_scopes_default(): @@ -382,7 +391,76 @@ def test_client_project_path(): assert client.project_path == project_name -def test_client_table_data_client_not_initialized(): +def test_client_veneer_data_client_not_initialized(): + from google.cloud.bigtable.data import BigtableDataClient + from google.cloud.bigtable import __version__ + + credentials = _make_credentials() + client = _make_client(project=PROJECT, credentials=credentials) + + with mock.patch("copy.copy") as copy_mock: + data_client = client._veneer_data_client + + assert isinstance(data_client, BigtableDataClient) + assert client._table_data_client is data_client + + assert client._table_data_client._disable_background_refresh + assert ( + client._table_data_client.client_info.client_library_version + == f"{__version__}-data-shim" + ) + copy_mock.assert_called_once_with(client._client_info) + + +def test_client_veneer_data_client_not_initialized_w_client_info(): + from google.api_core.gapic_v1.client_info import ClientInfo + from google.cloud.bigtable import __version__ + + credentials = _make_credentials() + client_info = ClientInfo(gapic_version="1.2.3", user_agent="test-client-") + client = _make_client( + project=PROJECT, credentials=credentials, client_info=client_info + ) + + with mock.patch("copy.copy") as copy_mock: + data_client = client._veneer_data_client + + assert client._table_data_client is data_client + + assert client._client_info is client_info + assert client._table_data_client.client_info is copy_mock.return_value + assert client._table_data_client._disable_background_refresh + assert ( + client._table_data_client.client_info.client_library_version + == f"{__version__}-data-shim" + ) + copy_mock.assert_called_once_with(client_info) + + +def test_client_veneer_data_client_not_initialized_w_client_options(): + from google.api_core.client_options import ClientOptions + + credentials = _make_credentials() + client_options = ClientOptions(quota_project_id="QUOTA-PROJECT", api_endpoint="xyz") + client = _make_client( + project=PROJECT, credentials=credentials, client_options=client_options + ) + + data_client = client._veneer_data_client + assert client._table_data_client is data_client + assert client._table_data_client._disable_background_refresh + assert client._table_data_client._gapic_client._client_options == client_options + + +def test_client_veneer_data_client_initialized(): + credentials = _make_credentials() + client = _make_client(project=PROJECT, credentials=credentials, admin=True) + + already = client._table_data_client = object() + assert client._veneer_data_client is already + + +def test_client_data_gapic_client_not_initialized(): from google.cloud.bigtable_v2 import BigtableClient credentials = _make_credentials() @@ -390,10 +468,10 @@ def test_client_table_data_client_not_initialized(): table_data_client = client.table_data_client assert isinstance(table_data_client, BigtableClient) - assert client._table_data_client is table_data_client + assert client._table_data_client._gapic_client is table_data_client -def test_client_table_data_client_not_initialized_w_client_info(): +def test_client_data_gapic_client_not_initialized_w_client_info(): from google.cloud.bigtable_v2 import BigtableClient credentials = _make_credentials() @@ -402,14 +480,32 @@ def test_client_table_data_client_not_initialized_w_client_info(): project=PROJECT, credentials=credentials, client_info=client_info ) - table_data_client = client.table_data_client + mock_gapic_client = mock.MagicMock(spec=BigtableClient) + mock_gapic_client.universe_domain = BigtableClient._DEFAULT_UNIVERSE + + with mock.patch( + "google.cloud.bigtable.data._sync_autogen.client.GapicClient", + return_value=mock_gapic_client, + ) as gapic_mock: + with mock.patch("copy.copy") as copy_mock: + table_data_client = client.table_data_client + assert isinstance(table_data_client, BigtableClient) assert client._client_info is client_info - assert client._table_data_client is table_data_client + assert client._table_data_client._gapic_client is table_data_client + + copy_mock.assert_called_once_with(client._client_info) + gapic_mock.assert_called_once_with( + client_info=copy_mock.return_value, + credentials=mock.ANY, + transport=mock.ANY, + client_options=mock.ANY, + ) -def test_client_table_data_client_not_initialized_w_client_options(): +def test_client_data_gapic_client_not_initialized_w_client_options(): from google.api_core.client_options import ClientOptions + from google.cloud.bigtable_v2 import BigtableClient credentials = _make_credentials() client_options = ClientOptions(quota_project_id="QUOTA-PROJECT", api_endpoint="xyz") @@ -417,27 +513,34 @@ def test_client_table_data_client_not_initialized_w_client_options(): project=PROJECT, credentials=credentials, client_options=client_options ) - patch = mock.patch("google.cloud.bigtable_v2.BigtableClient") - with patch as mocked: + mock_gapic_client = mock.MagicMock() + mock_gapic_client.universe_domain = BigtableClient._DEFAULT_UNIVERSE + + with mock.patch( + "google.cloud.bigtable.data._sync_autogen.client.GapicClient", + return_value=mock_gapic_client, + ) as gapic_mock: table_data_client = client.table_data_client - assert table_data_client is mocked.return_value - assert client._table_data_client is table_data_client + assert client._table_data_client._gapic_client is table_data_client - mocked.assert_called_once_with( - client_info=client._client_info, - credentials=None, + gapic_mock.assert_called_once_with( + client_info=mock.ANY, + credentials=mock.ANY, transport=mock.ANY, client_options=client_options, ) def test_client_table_data_client_initialized(): + from google.cloud.bigtable.data import BigtableDataClient + credentials = _make_credentials() client = _make_client(project=PROJECT, credentials=credentials, admin=True) - already = client._table_data_client = object() - assert client.table_data_client is already + already = client._table_data_client = mock.Mock(spec=BigtableDataClient) + already._gapic_client = mock.Mock() + assert client.table_data_client is already._gapic_client def test_client_table_admin_client_not_initialized_no_admin_flag(): diff --git a/tests/unit/v2_client/test_instance.py b/tests/unit/v2_client/test_instance.py index c06036438..2fb534d4f 100644 --- a/tests/unit/v2_client/test_instance.py +++ b/tests/unit/v2_client/test_instance.py @@ -791,7 +791,7 @@ def test_instance_table_factory(): from google.cloud.bigtable.table import Table app_profile_id = "appProfileId1262094415" - instance = _make_instance(INSTANCE_ID, None) + instance = _make_instance(INSTANCE_ID, mock.MagicMock()) table = instance.table(TABLE_ID, app_profile_id=app_profile_id) assert isinstance(table, Table) diff --git a/tests/unit/v2_client/test_row.py b/tests/unit/v2_client/test_row.py index f04802f5c..894b4d036 100644 --- a/tests/unit/v2_client/test_row.py +++ b/tests/unit/v2_client/test_row.py @@ -446,7 +446,8 @@ def test_conditional_row_commit(): # Patch the stub used by the API method. api.check_and_mutate_row.side_effect = [response_pb] - client._table_data_client = api + client.table_data_client + client._table_data_client._gapic_client = api # Create expected_result. expected_result = predicate_matched @@ -589,7 +590,8 @@ def test_append_row_commit(): expected_result = object() # Patch API calls - client._table_data_client = api + client.table_data_client + client._table_data_client._gapic_client = api def mock_parse_rmw_row_response(row_response): row_responses.append(row_response) @@ -597,7 +599,7 @@ def mock_parse_rmw_row_response(row_response): # Perform the method and check the result. with _Monkey(MUT, _parse_rmw_row_response=mock_parse_rmw_row_response): - row._table._instance._client._table_data_client = api + row._table._instance._client._table_data_client._gapic_client = api row.append_cell_value(column_family_id, column, value) result = row.commit() call_args = api.read_modify_write_row.call_args_list[0] diff --git a/tests/unit/v2_client/test_table.py b/tests/unit/v2_client/test_table.py index 921f3aea8..57eb707c4 100644 --- a/tests/unit/v2_client/test_table.py +++ b/tests/unit/v2_client/test_table.py @@ -162,7 +162,14 @@ def _make_table(*args, **kwargs): def test_table_constructor_defaults(): - instance = mock.Mock(spec=[]) + from google.cloud.bigtable.client import Client + + client = mock.create_autospec(Client) + instance = mock.Mock( + _client=client, + instance_id=INSTANCE_ID, + spec=["_client", "instance_id"], + ) table = _make_table(TABLE_ID, instance) @@ -170,10 +177,24 @@ def test_table_constructor_defaults(): assert table._instance is instance assert table.mutation_timeout is None assert table._app_profile_id is None + assert table._table_impl is client._veneer_data_client.get_table.return_value + client._veneer_data_client.get_table.assert_called_once_with( + INSTANCE_ID, + TABLE_ID, + app_profile_id=None, + ) def test_table_constructor_explicit(): - instance = mock.Mock(spec=[]) + from google.cloud.bigtable.client import Client + + client = mock.create_autospec(Client) + instance = mock.Mock( + _client=client, + instance_id=INSTANCE_ID, + spec=["_client", "instance_id"], + ) + mutation_timeout = 123 app_profile_id = "profile-123" @@ -188,14 +209,22 @@ def test_table_constructor_explicit(): assert table._instance is instance assert table.mutation_timeout == mutation_timeout assert table._app_profile_id == app_profile_id + assert table._table_impl is client._veneer_data_client.get_table.return_value + client._veneer_data_client.get_table.assert_called_once_with( + INSTANCE_ID, + TABLE_ID, + app_profile_id=app_profile_id, + ) def test_table_name(): table_data_client = mock.Mock(spec=["table_path"]) + _veneer_data_client = mock.Mock() client = mock.Mock( project=PROJECT_ID, table_data_client=table_data_client, - spec=["project", "table_data_client"], + _veneer_data_client=_veneer_data_client, + spec=["project", "table_data_client", "_veneer_data_client"], ) instance = mock.Mock( _client=client, @@ -339,8 +368,9 @@ def test_table___ne__same_value(): def test_table___ne__(): - table1 = _make_table("table_id1", None) - table2 = _make_table("table_id2", None) + mock_instance = mock.Mock() + table1 = _make_table("table_id1", mock_instance) + table2 = _make_table("table_id2", mock_instance) assert table1 != table2 @@ -632,10 +662,23 @@ def test_table_get_encryption_info(): table_api.get_table.assert_called_once_with(request=expected_request) -def _make_data_api(): +def _make_data_api(client): + from google.cloud.bigtable.data import BigtableDataClient + + data_client_mock = mock.create_autospec(BigtableDataClient) + client._table_data_client = data_client_mock + + return data_client_mock + + +def _make_gapic_api(client): from google.cloud.bigtable_v2.services.bigtable import BigtableClient - return mock.create_autospec(BigtableClient) + data_client_mock = _make_data_api(client) + gapic_client_mock = mock.create_autospec(BigtableClient) + data_client_mock._gapic_client = gapic_client_mock + + return gapic_client_mock def _table_read_row_helper(chunks, expected_result, app_profile_id=None): @@ -665,8 +708,8 @@ def mock_create_row_request(table_name, **kwargs): response_pb = _ReadRowsResponsePB(chunks=chunks) response_iterator = iter([response_pb]) - data_api = client._table_data_client = _make_data_api() - data_api.read_rows.return_value = response_iterator + gapic_api = _make_gapic_api(client) + gapic_api.read_rows.return_value = response_iterator filter_obj = RowSampleFilter(0.33) @@ -692,7 +735,7 @@ def mock_create_row_request(table_name, **kwargs): assert result == expected_result assert mock_created == expected_request - data_api.read_rows.assert_called_once_with( + gapic_api.read_rows.assert_called_once_with( request_pb, timeout=61.0, retry=DEFAULT_RETRY_READ_ROWS ) @@ -863,7 +906,7 @@ def test_table_read_rows(): credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) - data_api = client._table_data_client = _make_data_api() + gapic_api = _make_gapic_api(client) instance = client.instance(instance_id=INSTANCE_ID) app_profile_id = "app-profile-id" table = _make_table(TABLE_ID, instance, app_profile_id=app_profile_id) @@ -879,7 +922,7 @@ def mock_create_row_request(table_name, **kwargs): # Create expected_result. expected_result = PartialRowsData( - client._table_data_client.transport.read_rows, request_pb, retry + client._table_data_client._gapic_client.transport.read_rows, request_pb, retry ) # Perform the method and check the result. @@ -909,7 +952,7 @@ def mock_create_row_request(table_name, **kwargs): } assert mock_created == [(table.name, created_kwargs)] - data_api.read_rows.assert_called_once_with(request_pb, timeout=61.0, retry=retry) + gapic_api.read_rows.assert_called_once_with(request_pb, timeout=61.0, retry=retry) def test_table_read_retry_rows(): @@ -917,7 +960,7 @@ def test_table_read_retry_rows(): credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) - data_api = client._table_data_client = _make_data_api() + gapic_api = _make_gapic_api(client) instance = client.instance(instance_id=INSTANCE_ID) table = _make_table(TABLE_ID, instance) @@ -948,11 +991,11 @@ def test_table_read_retry_rows(): response_failure_iterator_2 = _MockFailureIterator_2([response_1]) response_iterator = _MockReadRowsIterator(response_2) - data_api.table_path.return_value = ( + gapic_api.table_path.return_value = ( f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" ) - data_api.read_rows.side_effect = [ + gapic_api.read_rows.side_effect = [ response_failure_iterator_1, response_failure_iterator_2, response_iterator, @@ -968,7 +1011,7 @@ def test_table_read_retry_rows(): result = rows[1] assert result.row_key == ROW_KEY_2 - assert len(data_api.read_rows.mock_calls) == 3 + assert len(gapic_api.read_rows.mock_calls) == 3 def test_table_read_retry_rows_no_full_table_scan(): @@ -976,7 +1019,7 @@ def test_table_read_retry_rows_no_full_table_scan(): credentials = _make_credentials() client = _make_client(project="project-id", credentials=credentials, admin=True) - data_api = client._table_data_client = _make_data_api() + gapic_api = _make_gapic_api(client) instance = client.instance(instance_id=INSTANCE_ID) table = _make_table(TABLE_ID, instance) @@ -995,11 +1038,11 @@ def test_table_read_retry_rows_no_full_table_scan(): response_1 = _ReadRowsResponseV2([chunk_1]) response_failure_iterator_2 = _MockFailureIterator_2([response_1]) - data_api.table_path.return_value = ( + gapic_api.table_path.return_value = ( f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" ) - data_api.read_rows.side_effect = [ + gapic_api.read_rows.side_effect = [ response_failure_iterator_2, ] @@ -1013,9 +1056,9 @@ def test_table_read_retry_rows_no_full_table_scan(): result = rows[0] assert result.row_key == ROW_KEY_2 - assert len(data_api.read_rows.mock_calls) == 1 + assert len(gapic_api.read_rows.mock_calls) == 1 assert ( - len(data_api.read_rows.mock_calls[0].args[0].rows.row_ranges) > 0 + len(gapic_api.read_rows.mock_calls[0].args[0].rows.row_ranges) > 0 ) # not empty row_ranges @@ -1052,11 +1095,11 @@ def test_table_yield_retry_rows(): response_failure_iterator_2 = _MockFailureIterator_2([response_1]) response_iterator = _MockReadRowsIterator(response_2) - data_api = client._table_data_client = _make_data_api() - data_api.table_path.return_value = ( + gapic_api = _make_gapic_api(client) + gapic_api.table_path.return_value = ( f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" ) - data_api.read_rows.side_effect = [ + gapic_api.read_rows.side_effect = [ response_failure_iterator_1, response_failure_iterator_2, response_iterator, @@ -1078,7 +1121,7 @@ def test_table_yield_retry_rows(): start_key=ROW_KEY_1, end_key=ROW_KEY_2, ) - data_api.read_rows.mock_calls = [expected_request] * 3 + gapic_api.read_rows.mock_calls = [expected_request] * 3 def test_table_yield_rows_with_row_set(): @@ -1125,11 +1168,11 @@ def test_table_yield_rows_with_row_set(): response_3 = _ReadRowsResponseV2([chunk_3]) response_iterator = _MockReadRowsIterator(response_1, response_2, response_3) - data_api = client._table_data_client = _make_data_api() - data_api.table_path.return_value = ( + gapic_api = _make_gapic_api(client) + gapic_api.table_path.return_value = ( f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" ) - data_api.read_rows.side_effect = [response_iterator] + gapic_api.read_rows.side_effect = [response_iterator] rows = [] row_set = RowSet() @@ -1153,7 +1196,7 @@ def test_table_yield_rows_with_row_set(): end_key=ROW_KEY_2, ) expected_request.rows.row_keys.append(ROW_KEY_3) - data_api.read_rows.assert_called_once_with( + gapic_api.read_rows.assert_called_once_with( expected_request, timeout=61.0, retry=DEFAULT_RETRY_READ_ROWS ) @@ -1165,8 +1208,8 @@ def test_table_sample_row_keys(): table = _make_table(TABLE_ID, instance) response_iterator = object() - data_api = client._table_data_client = _make_data_api() - data_api.sample_row_keys.return_value = [response_iterator] + gapic_api = _make_gapic_api(client) + gapic_api.sample_row_keys.return_value = [response_iterator] result = table.sample_row_keys() @@ -1249,7 +1292,7 @@ def test_table_drop_by_prefix_w_timeout(): def test_table_mutations_batcher_factory(): flush_count = 100 max_row_bytes = 1000 - table = _make_table(TABLE_ID, None) + table = _make_table(TABLE_ID, mock.Mock()) mutation_batcher = table.mutations_batcher( flush_count=flush_count, max_row_bytes=max_row_bytes ) @@ -1356,8 +1399,10 @@ def test_table_test_iam_permissions(): def test_table_backup_factory_defaults(): from google.cloud.bigtable.backup import Backup + from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.client import Client - instance = _make_table(INSTANCE_ID, None) + instance = Instance(INSTANCE_ID, mock.create_autospec(Client)) table = _make_table(TABLE_ID, instance) backup = table.backup(BACKUP_ID) @@ -1381,8 +1426,9 @@ def test_table_backup_factory_non_defaults(): from google.cloud._helpers import UTC from google.cloud.bigtable.backup import Backup from google.cloud.bigtable.instance import Instance + from google.cloud.bigtable.client import Client - instance = Instance(INSTANCE_ID, None) + instance = Instance(INSTANCE_ID, mock.create_autospec(Client)) table = _make_table(TABLE_ID, instance) timestamp = datetime.datetime.utcnow().replace(tzinfo=UTC) backup = table.backup( @@ -1536,9 +1582,9 @@ def test_rmrw_callable_empty_rows(): client = _make_client(project="project-id", credentials=credentials, admin=True) instance = client.instance(instance_id=INSTANCE_ID) table = _make_table(TABLE_ID, instance) - data_api = client._table_data_client = _make_data_api() - data_api.mutate_rows.return_value = [] - data_api.table_path.return_value = ( + gapic_api = _make_gapic_api(client) + gapic_api.mutate_rows.return_value = [] + gapic_api.table_path.return_value = ( f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" ) @@ -1575,9 +1621,9 @@ def test_rmrw_callable_no_retry_strategy(): response_codes = [SUCCESS, RETRYABLE_1, NON_RETRYABLE] response = _make_responses(response_codes) - data_api = client._table_data_client = _make_data_api() - data_api.mutate_rows.return_value = [response] - data_api.table_path.return_value = ( + gapic_api = _make_gapic_api(client) + gapic_api.mutate_rows.return_value = [response] + gapic_api.table_path.return_value = ( f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" ) worker = _make_worker(client, table.name, [row_1, row_2, row_3]) @@ -1587,7 +1633,7 @@ def test_rmrw_callable_no_retry_strategy(): result = [status.code for status in statuses] assert result == response_codes - data_api.mutate_rows.assert_called_once() + gapic_api.mutate_rows.assert_called_once() def test_rmrw_callable_retry(): @@ -1618,9 +1664,9 @@ def test_rmrw_callable_retry(): response_1 = _make_responses([SUCCESS, RETRYABLE_1, NON_RETRYABLE]) response_2 = _make_responses([SUCCESS]) - data_api = client._table_data_client = _make_data_api() - data_api.mutate_rows.side_effect = [[response_1], [response_2]] - data_api.table_path.return_value = ( + gapic_api = _make_gapic_api(client) + gapic_api.mutate_rows.side_effect = [[response_1], [response_2]] + gapic_api.table_path.return_value = ( f"projects/{PROJECT_ID}/instances/{INSTANCE_ID}/tables/{TABLE_ID}" ) worker = _make_worker(client, table.name, [row_1, row_2, row_3]) @@ -1632,7 +1678,7 @@ def test_rmrw_callable_retry(): assert result == [SUCCESS, SUCCESS, NON_RETRYABLE] - assert client._table_data_client.mutate_rows.call_count == 2 + assert client._table_data_client._gapic_client.mutate_rows.call_count == 2 def _do_mutate_retryable_rows_helper( @@ -1670,16 +1716,16 @@ def _do_mutate_retryable_rows_helper( response = _make_responses(responses) - data_api = client._table_data_client = _make_data_api() + gapic_api = _make_gapic_api(client) if retryable_error: if mutate_rows_side_effect is not None: - data_api.mutate_rows.side_effect = mutate_rows_side_effect + gapic_api.mutate_rows.side_effect = mutate_rows_side_effect else: - data_api.mutate_rows.side_effect = ServiceUnavailable("testing") + gapic_api.mutate_rows.side_effect = ServiceUnavailable("testing") else: if mutate_rows_side_effect is not None: - data_api.mutate_rows.side_effect = mutate_rows_side_effect - data_api.mutate_rows.return_value = [response] + gapic_api.mutate_rows.side_effect = mutate_rows_side_effect + gapic_api.mutate_rows.return_value = [response] worker = _make_worker(client, table.name, rows=rows) @@ -1718,9 +1764,9 @@ def _do_mutate_retryable_rows_helper( assert result == expected_result if len(responses) == 0 and not retryable_error: - data_api.mutate_rows.assert_not_called() + gapic_api.mutate_rows.assert_not_called() else: - data_api.mutate_rows.assert_called_once_with( + gapic_api.mutate_rows.assert_called_once_with( table_name=table.name, entries=expected_entries, app_profile_id=None, @@ -1728,7 +1774,7 @@ def _do_mutate_retryable_rows_helper( **expected_kwargs, ) if timeout is not None: - called = data_api.mutate_rows.mock_calls[0] + called = gapic_api.mutate_rows.mock_calls[0] assert called.kwargs["timeout"]._deadline == timeout