got tests running

Author: daniel-sanche

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

@@ -55,7 +55,15 @@ def __call__(self, *args, **kwargs) -> Call:
 class WrappedUnaryUnaryMultiCallable(
     _WrappedMultiCallable, UnaryUnaryMultiCallable
 ):
-    pass
+    if not CrossSync.is_async:
+        # add missing functions for sync unary callable
+        
+        def with_call(self, *args, **kwargs):
+            call = self.__call__(self, *args, **kwargs)
+            return call(), call
+        
+        def future(self, *args, **kwargs):
+            raise NotImplementedError
 
 
 class WrappedUnaryStreamMultiCallable(
@@ -114,6 +122,8 @@ def stream_stream(self, *args, **kwargs) -> StreamStreamMultiCallable:
             )(*call_args, **call_kwargs)
         )
 
+    # grace not supported by sync version
+    @CrossSync.drop
     async def close(self, grace=None):
         return await self._channel.close(grace=grace)
 
@@ -138,6 +148,18 @@ async def wait_for_state_change(self, last_observed_state):
     def __getattr__(self, name):
         return getattr(self._channel, name)
 
+    if not CrossSync.is_async:
+
+        def close(self):
+            return self._channel.close()
+
+        def subscribe(self, callback, try_to_connect=False):
+            return self._channel.subscribe(callback, try_to_connect)
+
+        def unsubscribe(self, callback):
+            return self._channel.unsubscribe(callback)
+
+
 @CrossSync.convert_class(sync_name="_ReplaceableChannel", replace_symbols={"_AsyncWrappedChannel": "_WrappedChannel"})
 class _AsyncReplaceableChannel(_AsyncWrappedChannel):
 

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

@@ -103,7 +103,7 @@
     from grpc import intercept_channel
     from google.cloud.bigtable_v2.services.bigtable.transports import BigtableGrpcTransport as TransportType  # type: ignore
     from google.cloud.bigtable.data._sync_autogen.mutations_batcher import _MB_SIZE
-    from google.cloud.bigtable.data._async._replaceable_channel import _ReplaceableChannel
+    from google.cloud.bigtable.data._sync_autogen._replaceable_channel import _ReplaceableChannel
 
 
 if TYPE_CHECKING:
@@ -346,7 +346,7 @@ async def _ping_and_warm_instances(
         )
         return [r or None for r in result_list]
 
-    @CrossSync.convert
+    @CrossSync.convert(replace_symbols={"_AsyncReplaceableChannel": "_ReplaceableChannel"})
     async def _manage_channel(
         self,
         refresh_interval_min: float = 60 * 35,
@@ -397,7 +397,7 @@ async def _manage_channel(
             new_channel = super_channel.create_channel()
             await self._ping_and_warm_instances(channel=new_channel)
             # cycle channel out of use, with long grace window before closure
-            old_channel = super_channel.replace_wrapped_channel(new_channel, grace_period)
+            old_channel = super_channel.replace_wrapped_channel(new_channel)
             # give old_channel a chance to complete existing rpcs
             if CrossSync.is_async:
                 await old_channel.close(grace_period)

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

@@ -41,7 +41,12 @@ def __call__(self, *args, **kwargs) -> Call:
 
 
 class WrappedUnaryUnaryMultiCallable(_WrappedMultiCallable, UnaryUnaryMultiCallable):
-    pass
+    def with_call(self, *args, **kwargs):
+        call = self.__call__(self, *args, **kwargs)
+        return (call(), call)
+
+    def future(self, *args, **kwargs):
+        raise NotImplementedError
 
 
 class WrappedUnaryStreamMultiCallable(_WrappedMultiCallable, UnaryStreamMultiCallable):
@@ -95,9 +100,6 @@ def stream_stream(self, *args, **kwargs) -> StreamStreamMultiCallable:
             )(*call_args, **call_kwargs)
         )
 
-    def close(self, grace=None):
-        return self._channel.close(grace=grace)
-
     def channel_ready(self):
         return self._channel.channel_ready()
 
@@ -117,6 +119,15 @@ def wait_for_state_change(self, last_observed_state):
     def __getattr__(self, name):
         return getattr(self._channel, name)
 
+    def close(self):
+        return self._channel.close()
+
+    def subscribe(self, callback, try_to_connect=False):
+        return self._channel.subscribe(callback, try_to_connect)
+
+    def unsubscribe(self, callback):
+        return self._channel.unsubscribe(callback)
+
 
 class _ReplaceableChannel(_WrappedChannel):
     def __init__(self, channel_fn: Callable[[], Channel]):

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

@@ -79,7 +79,9 @@
     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
+from google.cloud.bigtable.data._sync_autogen._replaceable_channel import (
+    _ReplaceableChannel,
+)
 
 if TYPE_CHECKING:
     from google.cloud.bigtable.data._helpers import RowKeySamples
@@ -282,10 +284,10 @@ 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):
+        if not isinstance(self.transport.grpc_channel, _ReplaceableChannel):
             warnings.warn("Channel does not support auto-refresh.")
             return
-        super_channel: _AsyncReplaceableChannel = self.transport.grpc_channel
+        super_channel: _ReplaceableChannel = self.transport.grpc_channel
         first_refresh = self._channel_init_time + random.uniform(
             refresh_interval_min, refresh_interval_max
         )
@@ -301,9 +303,7 @@ def _manage_channel(
             start_timestamp = time.monotonic()
             new_channel = super_channel.create_channel()
             self._ping_and_warm_instances(channel=new_channel)
-            old_channel = super_channel.replace_wrapped_channel(
-                new_channel, grace_period
-            )
+            old_channel = super_channel.replace_wrapped_channel(new_channel)
             if grace_period:
                 self._is_closed.wait(grace_period)
             old_channel.close()