generated sync code

Author: daniel-sanche

google/cloud/bigtable/data/_sync_autogen/_replaceable_channel.py

@@ -0,0 +1,133 @@
+# Copyright 2023 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+#
+
+# This file is automatically generated by CrossSync. Do not edit manually.
+
+from __future__ import annotations
+from typing import Callable
+from grpc import ChannelConnectivity
+from grpc import Call
+from grpc import Channel
+from grpc import UnaryUnaryMultiCallable
+from grpc import UnaryStreamMultiCallable
+from grpc import StreamUnaryMultiCallable
+from grpc import StreamStreamMultiCallable
+
+
+class _WrappedMultiCallable:
+    """
+    Wrapper class that implements the grpc MultiCallable interface.
+    Allows generic functions that return calls to pass checks for
+    MultiCallable objects.
+    """
+
+    def __init__(self, call_factory: Callable[..., Call]):
+        self._call_factory = call_factory
+
+    def __call__(self, *args, **kwargs) -> Call:
+        return self._call_factory(*args, **kwargs)
+
+
+class WrappedUnaryUnaryMultiCallable(_WrappedMultiCallable, UnaryUnaryMultiCallable):
+    pass
+
+
+class WrappedUnaryStreamMultiCallable(_WrappedMultiCallable, UnaryStreamMultiCallable):
+    pass
+
+
+class WrappedStreamUnaryMultiCallable(_WrappedMultiCallable, StreamUnaryMultiCallable):
+    pass
+
+
+class WrappedStreamStreamMultiCallable(
+    _WrappedMultiCallable, StreamStreamMultiCallable
+):
+    pass
+
+
+class _WrappedChannel(Channel):
+    """
+    A wrapper around a gRPC channel. All methods are passed
+    through to the underlying channel.
+    """
+
+    def __init__(self, channel: Channel):
+        self._channel = channel
+
+    def unary_unary(self, *args, **kwargs) -> UnaryUnaryMultiCallable:
+        return WrappedUnaryUnaryMultiCallable(
+            lambda *call_args, **call_kwargs: self._channel.unary_unary(
+                *args, **kwargs
+            )(*call_args, **call_kwargs)
+        )
+
+    def unary_stream(self, *args, **kwargs) -> UnaryStreamMultiCallable:
+        return WrappedUnaryStreamMultiCallable(
+            lambda *call_args, **call_kwargs: self._channel.unary_stream(
+                *args, **kwargs
+            )(*call_args, **call_kwargs)
+        )
+
+    def stream_unary(self, *args, **kwargs) -> StreamUnaryMultiCallable:
+        return WrappedStreamUnaryMultiCallable(
+            lambda *call_args, **call_kwargs: self._channel.stream_unary(
+                *args, **kwargs
+            )(*call_args, **call_kwargs)
+        )
+
+    def stream_stream(self, *args, **kwargs) -> StreamStreamMultiCallable:
+        return WrappedStreamStreamMultiCallable(
+            lambda *call_args, **call_kwargs: self._channel.stream_stream(
+                *args, **kwargs
+            )(*call_args, **call_kwargs)
+        )
+
+    def close(self, grace=None):
+        return self._channel.close(grace=grace)
+
+    def channel_ready(self):
+        return self._channel.channel_ready()
+
+    def __enter__(self):
+        self._channel.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self._channel.__exit__(exc_type, exc_val, exc_tb)
+
+    def get_state(self, try_to_connect: bool = False) -> ChannelConnectivity:
+        return self._channel.get_state(try_to_connect=try_to_connect)
+
+    def wait_for_state_change(self, last_observed_state):
+        return self._channel.wait_for_state_change(last_observed_state)
+
+    def __getattr__(self, name):
+        return getattr(self._channel, name)
+
+
+class _ReplaceableChannel(_WrappedChannel):
+    def __init__(self, channel_fn: Callable[[], Channel]):
+        self._channel_fn = channel_fn
+        self._channel = channel_fn()
+
+    def create_channel(self) -> Channel:
+        new_channel = self._channel_fn()
+        return new_channel
+
+    def replace_wrapped_channel(self, new_channel: Channel) -> Channel:
+        old_channel = self._channel
+        self._channel = new_channel
+        return old_channel

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

@@ -75,11 +75,11 @@
 from google.cloud.bigtable.data._cross_sync import CrossSync
 from typing import Iterable
 from grpc import insecure_channel
-from grpc import intercept_channel
 from google.cloud.bigtable_v2.services.bigtable.transports import (
     BigtableGrpcTransport as TransportType,
 )
 from google.cloud.bigtable.data._sync_autogen.mutations_batcher import _MB_SIZE
+from google.cloud.bigtable.data._async._replaceable_channel import _ReplaceableChannel
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._helpers import RowKeySamples
@@ -131,15 +131,13 @@ def __init__(
         client_options = cast(
             Optional[client_options_lib.ClientOptions], client_options
         )
-        custom_channel = None
         self._emulator_host = os.getenv(BIGTABLE_EMULATOR)
         if self._emulator_host is not None:
             warnings.warn(
                 "Connecting to Bigtable emulator at {}".format(self._emulator_host),
                 RuntimeWarning,
                 stacklevel=2,
             )
-            custom_channel = insecure_channel(self._emulator_host)
             if credentials is None:
                 credentials = google.auth.credentials.AnonymousCredentials()
             if project is None:
@@ -155,7 +153,7 @@ def __init__(
             client_options=client_options,
             client_info=self.client_info,
             transport=lambda *args, **kwargs: TransportType(
-                *args, **kwargs, channel=custom_channel
+                *args, **kwargs, channel=self._build_grpc_channel
             ),
         )
         self._is_closed = CrossSync._Sync_Impl.Event()
@@ -179,6 +177,13 @@ def __init__(
                     stacklevel=2,
                 )
 
+    def _build_grpc_channel(self, *args, **kwargs) -> _ReplaceableChannel:
+        if self._emulator_host is not None:
+            create_channel_fn = partial(insecure_channel, self._emulator_host)
+        else:
+            create_channel_fn = partial(TransportType.create_channel, *args, **kwargs)
+        return _ReplaceableChannel(create_channel_fn)
+
     @staticmethod
     def _client_version() -> str:
         """Helper function to return the client version string for this client"""
@@ -277,27 +282,28 @@ def _manage_channel(
                 between `refresh_interval_min` and `refresh_interval_max`
             grace_period: time to allow previous channel to serve existing
                 requests before closing, in seconds"""
+        if not isinstance(self.transport.grpc_channel, _AsyncReplaceableChannel):
+            warnings.warn("Channel does not support auto-refresh.")
+            return
+        super_channel: _AsyncReplaceableChannel = self.transport.grpc_channel
         first_refresh = self._channel_init_time + random.uniform(
             refresh_interval_min, refresh_interval_max
         )
         next_sleep = max(first_refresh - time.monotonic(), 0)
         if next_sleep > 0:
-            self._ping_and_warm_instances(channel=self.transport.grpc_channel)
+            self._ping_and_warm_instances(channel=super_channel)
         while not self._is_closed.is_set():
             CrossSync._Sync_Impl.event_wait(
                 self._is_closed, next_sleep, async_break_early=False
             )
             if self._is_closed.is_set():
                 break
             start_timestamp = time.monotonic()
-            old_channel = self.transport.grpc_channel
-            new_channel = self.transport.create_channel()
-            new_channel = intercept_channel(new_channel, self.transport._interceptor)
+            new_channel = super_channel.create_channel()
             self._ping_and_warm_instances(channel=new_channel)
-            self.transport._grpc_channel = new_channel
-            self.transport._logged_channel = new_channel
-            self.transport._stubs = {}
-            self.transport._prep_wrapped_messages(self.client_info)
+            old_channel = super_channel.replace_wrapped_channel(
+                new_channel, grace_period
+            )
             if grace_period:
                 self._is_closed.wait(grace_period)
             old_channel.close()

tests/system/data/test_system_autogen.py

@@ -201,13 +201,14 @@ def test_channel_refresh(self, table_id, instance_id, temp_rows):
             CrossSync._Sync_Impl.yield_to_event_loop()
             with client.get_table(instance_id, table_id) as table:
                 rows = table.read_rows({})
-                first_channel = client.transport.grpc_channel
+                channel_wrapper = client.transport.grpc_channel
+                first_channel = client.transport.grpc_channel._channel
                 assert len(rows) == 2
                 CrossSync._Sync_Impl.sleep(2)
                 rows_after_refresh = table.read_rows({})
                 assert len(rows_after_refresh) == 2
-                assert client.transport.grpc_channel is not first_channel
-                print(table)
+                assert client.transport.grpc_channel is channel_wrapper
+                assert client.transport.grpc_channel._channel is not first_channel
         finally:
             client.close()
 

tests/unit/data/_sync_autogen/test_client.py

@@ -176,7 +176,7 @@ def test__start_background_channel_refresh_task_exists(self):
         client.close()
 
     def test__start_background_channel_refresh(self):
-        client = self._make_client(project="project-id")
+        client = self._make_client(project="project-id", use_emulator=False)
         with mock.patch.object(
             client, "_ping_and_warm_instances", CrossSync._Sync_Impl.Mock()
         ) as ping_and_warm:
@@ -282,7 +282,7 @@ def test__manage_channel_first_sleep(
             with mock.patch.object(CrossSync._Sync_Impl, "event_wait") as sleep:
                 sleep.side_effect = asyncio.CancelledError
                 try:
-                    client = self._make_client(project="project-id")
+                    client = self._make_client(project="project-id", use_emulator=False)
                     client._channel_init_time = -wait_time
                     client._manage_channel(refresh_interval, refresh_interval)
                 except asyncio.CancelledError:
@@ -296,36 +296,29 @@ def test__manage_channel_first_sleep(
 
     def test__manage_channel_ping_and_warm(self):
         """_manage channel should call ping and warm internally"""
-        import time
         import threading
-        from google.cloud.bigtable_v2.services.bigtable.transports.grpc import (
-            _LoggingClientInterceptor as Interceptor,
-        )
 
-        client_mock = mock.Mock()
-        client_mock.transport._interceptor = Interceptor()
-        client_mock._is_closed.is_set.return_value = False
-        client_mock._channel_init_time = time.monotonic()
-        orig_channel = client_mock.transport.grpc_channel
+        client = self._make_client(project="project-id", use_emulator=True)
+        orig_channel = client.transport.grpc_channel
         sleep_tuple = (
             (asyncio, "sleep")
             if CrossSync._Sync_Impl.is_async
             else (threading.Event, "wait")
         )
-        with mock.patch.object(*sleep_tuple):
-            orig_channel.close.side_effect = asyncio.CancelledError
+        with mock.patch.object(*sleep_tuple) as sleep_mock:
+            sleep_mock.side_effect = [None, asyncio.CancelledError]
             ping_and_warm = (
-                client_mock._ping_and_warm_instances
+                client._ping_and_warm_instances
             ) = CrossSync._Sync_Impl.Mock()
             try:
-                self._get_target_class()._manage_channel(client_mock, 10)
+                client._manage_channel(10)
             except asyncio.CancelledError:
                 pass
             assert ping_and_warm.call_count == 2
-            assert client_mock.transport._grpc_channel != orig_channel
+            assert client.transport.grpc_channel._channel != orig_channel
             called_with = [call[1]["channel"] for call in ping_and_warm.call_args_list]
             assert orig_channel in called_with
-            assert client_mock.transport.grpc_channel in called_with
+            assert client.transport.grpc_channel._channel in called_with
 
     @pytest.mark.parametrize(
         "refresh_interval, num_cycles, expected_sleep",
@@ -335,8 +328,6 @@ def test__manage_channel_sleeps(self, refresh_interval, num_cycles, expected_sle
         import time
         import random
 
-        channel = mock.Mock()
-        channel.close = CrossSync._Sync_Impl.Mock()
         with mock.patch.object(random, "uniform") as uniform:
             uniform.side_effect = lambda min_, max_: min_
             with mock.patch.object(time, "time") as time_mock:
@@ -345,8 +336,7 @@ def test__manage_channel_sleeps(self, refresh_interval, num_cycles, expected_sle
                     sleep.side_effect = [None for i in range(num_cycles - 1)] + [
                         asyncio.CancelledError
                     ]
-                    client = self._make_client(project="project-id")
-                    client.transport._grpc_channel = channel
+                    client = self._make_client(project="project-id", use_emulator=True)
                     with mock.patch.object(
                         client.transport, "create_channel", CrossSync._Sync_Impl.Mock
                     ):
@@ -374,7 +364,7 @@ def test__manage_channel_random(self):
                 uniform.return_value = 0
                 try:
                     uniform.side_effect = asyncio.CancelledError
-                    client = self._make_client(project="project-id")
+                    client = self._make_client(project="project-id", use_emulator=False)
                 except asyncio.CancelledError:
                     uniform.side_effect = None
                     uniform.reset_mock()
@@ -405,7 +395,7 @@ def test__manage_channel_refresh(self, num_cycles):
                 CrossSync._Sync_Impl.grpc_helpers, "create_channel"
             ) as create_channel:
                 create_channel.return_value = new_channel
-                client = self._make_client(project="project-id")
+                client = self._make_client(project="project-id", use_emulator=False)
                 create_channel.reset_mock()
                 try:
                     client._manage_channel(