google-cloud-bigtable 2.32.0__tar.gz → 2.33.0__tar.gz

{google_cloud_bigtable-2.32.0/google_cloud_bigtable.egg-info → google_cloud_bigtable-2.33.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: google-cloud-bigtable
-Version: 2.32.0
+Version: 2.33.0
 Summary: Google Cloud Bigtable API client library
 Home-page: https://github.com/googleapis/python-bigtable
 Author: Google LLC
@@ -23,7 +23,7 @@ Requires-Python: >=3.7
 License-File: LICENSE
 Requires-Dist: google-api-core[grpc]<3.0.0,>=2.17.0
 Requires-Dist: google-cloud-core<3.0.0,>=1.4.4
-Requires-Dist: google-auth!=2.24.0,!=2.25.0,<3.0.0,>=2.14.1
+Requires-Dist: google-auth!=2.24.0,!=2.25.0,<3.0.0,>=2.23.0
 Requires-Dist: grpc-google-iam-v1<1.0.0,>=0.12.4
 Requires-Dist: proto-plus<2.0.0,>=1.22.3
 Requires-Dist: proto-plus<2.0.0,>=1.25.0; python_version >= "3.13"

google_cloud_bigtable-2.33.0/google/cloud/bigtable/data/_async/_swappable_channel.py

@@ -0,0 +1,139 @@
+# 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.
+#
+from __future__ import annotations
+
+from typing import Callable
+
+from google.cloud.bigtable.data._cross_sync import CrossSync
+
+from grpc import ChannelConnectivity
+
+if CrossSync.is_async:
+    from grpc.aio import Channel
+else:
+    from grpc import Channel
+
+__CROSS_SYNC_OUTPUT__ = "google.cloud.bigtable.data._sync_autogen._swappable_channel"
+
+
+@CrossSync.convert_class(sync_name="_WrappedChannel", rm_aio=True)
+class _AsyncWrappedChannel(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):
+        return self._channel.unary_unary(*args, **kwargs)
+
+    def unary_stream(self, *args, **kwargs):
+        return self._channel.unary_stream(*args, **kwargs)
+
+    def stream_unary(self, *args, **kwargs):
+        return self._channel.stream_unary(*args, **kwargs)
+
+    def stream_stream(self, *args, **kwargs):
+        return self._channel.stream_stream(*args, **kwargs)
+
+    async def channel_ready(self):
+        return await self._channel.channel_ready()
+
+    @CrossSync.convert(
+        sync_name="__enter__", replace_symbols={"__aenter__": "__enter__"}
+    )
+    async def __aenter__(self):
+        await self._channel.__aenter__()
+        return self
+
+    @CrossSync.convert(sync_name="__exit__", replace_symbols={"__aexit__": "__exit__"})
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        return await self._channel.__aexit__(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)
+
+    async def wait_for_state_change(self, last_observed_state):
+        return await self._channel.wait_for_state_change(last_observed_state)
+
+    def __getattr__(self, name):
+        return getattr(self._channel, name)
+
+    async def close(self, grace=None):
+        if CrossSync.is_async:
+            return await self._channel.close(grace=grace)
+        else:
+            # grace not supported by sync version
+            return self._channel.close()
+
+    if not CrossSync.is_async:
+        # add required sync methods
+
+        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="SwappableChannel",
+    replace_symbols={"_AsyncWrappedChannel": "_WrappedChannel"},
+)
+class AsyncSwappableChannel(_AsyncWrappedChannel):
+    """
+    Provides a grpc channel wrapper, that allows the internal channel to be swapped out
+
+    Args:
+      - channel_fn: a nullary function that returns a new channel instance.
+            It should be a partial with all channel configuration arguments built-in
+    """
+
+    def __init__(self, channel_fn: Callable[[], Channel]):
+        self._channel_fn = channel_fn
+        self._channel = channel_fn()
+
+    def create_channel(self) -> Channel:
+        """
+        Create a fresh channel using the stored `channel_fn` partial
+        """
+        new_channel = self._channel_fn()
+        if CrossSync.is_async:
+            # copy over interceptors
+            # this is needed because of how gapic attaches the LoggingClientAIOInterceptor
+            # sync channels add interceptors by wrapping, so this step isn't needed
+            new_channel._unary_unary_interceptors = (
+                self._channel._unary_unary_interceptors
+            )
+            new_channel._unary_stream_interceptors = (
+                self._channel._unary_stream_interceptors
+            )
+            new_channel._stream_unary_interceptors = (
+                self._channel._stream_unary_interceptors
+            )
+            new_channel._stream_stream_interceptors = (
+                self._channel._stream_stream_interceptors
+            )
+        return new_channel
+
+    def swap_channel(self, new_channel: Channel) -> Channel:
+        """
+        Replace the wrapped channel with a new instance. Typically created using `create_channel`
+        """
+        old_channel = self._channel
+        self._channel = new_channel
+        return old_channel

{google_cloud_bigtable-2.32.0 → google_cloud_bigtable-2.33.0}/google/cloud/bigtable/data/_async/client.py

@@ -58,6 +58,8 @@ from google.api_core import retry as retries
 from google.api_core.exceptions import DeadlineExceeded
 from google.api_core.exceptions import ServiceUnavailable
 from google.api_core.exceptions import Aborted
+from google.protobuf.message import Message
+from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper
 
 import google.auth.credentials
 import google.auth._default
@@ -92,13 +94,22 @@ if CrossSync.is_async:
     from google.cloud.bigtable_v2.services.bigtable.transports import (
         BigtableGrpcAsyncIOTransport as TransportType,
     )
+    from google.cloud.bigtable_v2.services.bigtable import (
+        BigtableAsyncClient as GapicClient,
+    )
     from google.cloud.bigtable.data._async.mutations_batcher import _MB_SIZE
+    from google.cloud.bigtable.data._async._swappable_channel import (
+        AsyncSwappableChannel,
+    )
 else:
     from typing import Iterable  # noqa: F401
     from grpc import insecure_channel
-    from grpc import intercept_channel
     from google.cloud.bigtable_v2.services.bigtable.transports import BigtableGrpcTransport as TransportType  # type: ignore
+    from google.cloud.bigtable_v2.services.bigtable import BigtableClient as GapicClient  # type: ignore
     from google.cloud.bigtable.data._sync_autogen.mutations_batcher import _MB_SIZE
+    from google.cloud.bigtable.data._sync_autogen._swappable_channel import (  # noqa: F401
+        SwappableChannel,
+    )
 
 
 if TYPE_CHECKING:
@@ -182,7 +193,6 @@ class BigtableDataClientAsync(ClientWithProject):
         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(
@@ -191,11 +201,11 @@ class BigtableDataClientAsync(ClientWithProject):
                 stacklevel=2,
             )
             # use insecure channel if emulator is set
-            custom_channel = insecure_channel(self._emulator_host)
             if credentials is None:
                 credentials = google.auth.credentials.AnonymousCredentials()
             if project is None:
                 project = _DEFAULT_BIGTABLE_EMULATOR_CLIENT
+
         # initialize client
         ClientWithProject.__init__(
             self,
@@ -203,14 +213,28 @@ class BigtableDataClientAsync(ClientWithProject):
             project=project,
             client_options=client_options,
         )
-        self._gapic_client = CrossSync.GapicClient(
+        self._gapic_client = GapicClient(
             credentials=credentials,
             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
             ),
         )
+        if (
+            credentials
+            and credentials.universe_domain != self.universe_domain
+            and self._emulator_host is None
+        ):
+            # validate that the universe domain of the credentials matches the
+            # universe domain configured in client_options
+            raise ValueError(
+                f"The configured universe domain ({self.universe_domain}) does "
+                "not match the universe domain found in the credentials "
+                f"({self._credentials.universe_domain}). If you haven't "
+                "configured the universe domain explicitly, `googleapis.com` "
+                "is the default."
+            )
         self._is_closed = CrossSync.Event()
         self.transport = cast(TransportType, self._gapic_client.transport)
         # keep track of active instances to for warmup on channel refresh
@@ -220,7 +244,7 @@ class BigtableDataClientAsync(ClientWithProject):
         self._instance_owners: dict[_WarmedInstanceKey, Set[int]] = {}
         self._channel_init_time = time.monotonic()
         self._channel_refresh_task: CrossSync.Task[None] | None = None
-        self._executor = (
+        self._executor: concurrent.futures.ThreadPoolExecutor | None = (
             concurrent.futures.ThreadPoolExecutor() if not CrossSync.is_async else None
         )
         if self._emulator_host is None:
@@ -235,6 +259,47 @@ class BigtableDataClientAsync(ClientWithProject):
                     stacklevel=2,
                 )
 
+    @CrossSync.convert(replace_symbols={"AsyncSwappableChannel": "SwappableChannel"})
+    def _build_grpc_channel(self, *args, **kwargs) -> AsyncSwappableChannel:
+        """
+        This method is called by the gapic transport to create a grpc channel.
+
+        The init arguments passed down are captured in a partial used by AsyncSwappableChannel
+        to create new channel instances in the future, as part of the channel refresh logic
+
+        Emulators always use an inseucre channel
+
+        Args:
+          - *args: positional arguments passed by the gapic layer to create a new channel with
+          - **kwargs: keyword arguments passed by the gapic layer to create a new channel with
+        Returns:
+          a custom wrapped swappable channel
+        """
+        if self._emulator_host is not None:
+            # emulators use insecure channel
+            create_channel_fn = partial(insecure_channel, self._emulator_host)
+        else:
+            create_channel_fn = partial(TransportType.create_channel, *args, **kwargs)
+        return AsyncSwappableChannel(create_channel_fn)
+
+    @property
+    def universe_domain(self) -> str:
+        """Return the universe domain used by the client instance.
+
+        Returns:
+            str: The universe domain used by the client instance.
+        """
+        return self._gapic_client.universe_domain
+
+    @property
+    def api_endpoint(self) -> str:
+        """Return the API endpoint used by the client instance.
+
+        Returns:
+            str: The API endpoint used by the client instance.
+        """
+        return self._gapic_client.api_endpoint
+
     @staticmethod
     def _client_version() -> str:
         """
@@ -332,7 +397,12 @@ class BigtableDataClientAsync(ClientWithProject):
         )
         return [r or None for r in result_list]
 
-    @CrossSync.convert
+    def _invalidate_channel_stubs(self):
+        """Helper to reset the cached stubs. Needed when changing out the grpc channel"""
+        self.transport._stubs = {}
+        self.transport._prep_wrapped_messages(self.client_info)
+
+    @CrossSync.convert(replace_symbols={"AsyncSwappableChannel": "SwappableChannel"})
     async def _manage_channel(
         self,
         refresh_interval_min: float = 60 * 35,
@@ -357,13 +427,17 @@ class BigtableDataClientAsync(ClientWithProject):
             grace_period: time to allow previous channel to serve existing
                 requests before closing, in seconds
         """
+        if not isinstance(self.transport.grpc_channel, AsyncSwappableChannel):
+            warnings.warn("Channel does not support auto-refresh.")
+            return
+        super_channel: AsyncSwappableChannel = 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:
             # warm the current channel immediately
-            await self._ping_and_warm_instances(channel=self.transport.grpc_channel)
+            await self._ping_and_warm_instances(channel=super_channel)
         # continuously refresh the channel every `refresh_interval` seconds
         while not self._is_closed.is_set():
             await CrossSync.event_wait(
@@ -376,24 +450,11 @@ class BigtableDataClientAsync(ClientWithProject):
                 break
             start_timestamp = time.monotonic()
             # prepare new channel for use
-            # TODO: refactor to avoid using internal references: https://github.com/googleapis/python-bigtable/issues/1094
-            old_channel = self.transport.grpc_channel
-            new_channel = self.transport.create_channel()
-            if CrossSync.is_async:
-                new_channel._unary_unary_interceptors.append(
-                    self.transport._interceptor
-                )
-            else:
-                new_channel = intercept_channel(
-                    new_channel, self.transport._interceptor
-                )
+            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
-            self.transport._grpc_channel = new_channel
-            self.transport._logged_channel = new_channel
-            # invalidate caches
-            self.transport._stubs = {}
-            self.transport._prep_wrapped_messages(self.client_info)
+            old_channel = super_channel.swap_channel(new_channel)
+            self._invalidate_channel_stubs()
             # give old_channel a chance to complete existing rpcs
             if CrossSync.is_async:
                 await old_channel.close(grace_period)
@@ -401,7 +462,7 @@ class BigtableDataClientAsync(ClientWithProject):
                 if grace_period:
                     self._is_closed.wait(grace_period)  # type: ignore
                 old_channel.close()  # type: ignore
-            # subtract thed time spent waiting for the channel to be replaced
+            # subtract the time spent waiting for the channel to be replaced
             next_refresh = random.uniform(refresh_interval_min, refresh_interval_max)
             next_sleep = max(next_refresh - (time.monotonic() - start_timestamp), 0)
 
@@ -415,7 +476,8 @@ class BigtableDataClientAsync(ClientWithProject):
     async def _register_instance(
         self,
         instance_id: str,
-        owner: _DataApiTargetAsync | ExecuteQueryIteratorAsync,
+        app_profile_id: Optional[str],
+        owner_id: int,
     ) -> None:
         """
         Registers an instance with the client, and warms the channel for the instance
@@ -425,13 +487,15 @@ class BigtableDataClientAsync(ClientWithProject):
 
         Args:
           instance_id: id of the instance to register.
-          owner: table that owns the instance. Owners will be tracked in
+          app_profile_id: id of the app profile calling the instance.
+          owner_id: integer id of the object owning the instance. Owners will be tracked in
               _instance_owners, and instances will only be unregistered when all
-              owners call _remove_instance_registration
+              owners call _remove_instance_registration. Can be obtained by calling
+              `id` identity funcion, using `id(owner)`
         """
         instance_name = self._gapic_client.instance_path(self.project, instance_id)
-        instance_key = _WarmedInstanceKey(instance_name, owner.app_profile_id)
-        self._instance_owners.setdefault(instance_key, set()).add(id(owner))
+        instance_key = _WarmedInstanceKey(instance_name, app_profile_id)
+        self._instance_owners.setdefault(instance_key, set()).add(owner_id)
         if instance_key not in self._active_instances:
             self._active_instances.add(instance_key)
             if self._channel_refresh_task:
@@ -449,10 +513,11 @@ class BigtableDataClientAsync(ClientWithProject):
             "_DataApiTargetAsync": "_DataApiTarget",
         }
     )
-    async def _remove_instance_registration(
+    def _remove_instance_registration(
         self,
         instance_id: str,
-        owner: _DataApiTargetAsync | ExecuteQueryIteratorAsync,
+        app_profile_id: Optional[str],
+        owner_id: int,
     ) -> bool:
         """
         Removes an instance from the client's registered instances, to prevent
@@ -462,17 +527,17 @@ class BigtableDataClientAsync(ClientWithProject):
 
         Args:
             instance_id: id of the instance to remove
-            owner: table that owns the instance. Owners will be tracked in
-              _instance_owners, and instances will only be unregistered when all
-              owners call _remove_instance_registration
+            app_profile_id: id of the app profile calling the instance.
+            owner_id: integer id of the object owning the instance. Can be
+                obtained by the `id` identity funcion, using `id(owner)`.
         Returns:
             bool: True if instance was removed, else False
         """
         instance_name = self._gapic_client.instance_path(self.project, instance_id)
-        instance_key = _WarmedInstanceKey(instance_name, owner.app_profile_id)
+        instance_key = _WarmedInstanceKey(instance_name, app_profile_id)
         owner_list = self._instance_owners.get(instance_key, set())
         try:
-            owner_list.remove(id(owner))
+            owner_list.remove(owner_id)
             if len(owner_list) == 0:
                 self._active_instances.remove(instance_key)
             return True
@@ -625,6 +690,7 @@ class BigtableDataClientAsync(ClientWithProject):
             DeadlineExceeded,
             ServiceUnavailable,
         ),
+        column_info: dict[str, Message | EnumTypeWrapper] | None = None,
     ) -> "ExecuteQueryIteratorAsync":
         """
         Executes an SQL query on an instance.
@@ -673,6 +739,62 @@ class BigtableDataClientAsync(ClientWithProject):
                 If None, defaults to prepare_operation_timeout.
             prepare_retryable_errors: a list of errors that will be retried if encountered during prepareQuery.
                 Defaults to 4 (DeadlineExceeded) and 14 (ServiceUnavailable)
+            column_info: (Optional) A dictionary mapping column names to Protobuf message classes or EnumTypeWrapper objects.
+                This dictionary provides the necessary type information for deserializing PROTO and
+                ENUM column values from the query results. When an entry is provided
+                for a PROTO or ENUM column, the client library will attempt to deserialize the raw data.
+
+                    - For PROTO columns: The value in the dictionary should be the
+                      Protobuf Message class (e.g., ``my_pb2.MyMessage``).
+                    - For ENUM columns: The value should be the Protobuf EnumTypeWrapper
+                      object (e.g., ``my_pb2.MyEnum``).
+
+                Example::
+
+                    import my_pb2
+
+                    column_info = {
+                        "my_proto_column": my_pb2.MyMessage,
+                        "my_enum_column": my_pb2.MyEnum
+                    }
+
+                If ``column_info`` is not provided, or if a specific column name is not found
+                in the dictionary:
+
+                    - PROTO columns will be returned as raw bytes.
+                    - ENUM columns will be returned as integers.
+
+                Note for Nested PROTO or ENUM Fields:
+
+                    To specify types for PROTO or ENUM fields within STRUCTs or MAPs, use a dot-separated
+                    path from the top-level column name.
+
+                        - For STRUCTs: ``struct_column_name.field_name``
+                        - For MAPs: ``map_column_name.key`` or ``map_column_name.value`` to specify types
+                          for the map keys or values, respectively.
+
+                    Example::
+
+                        import my_pb2
+
+                        column_info = {
+                            # Top-level column
+                            "my_proto_column": my_pb2.MyMessage,
+                            "my_enum_column": my_pb2.MyEnum,
+
+                            # Nested field in a STRUCT column named 'my_struct'
+                            "my_struct.nested_proto_field": my_pb2.OtherMessage,
+                            "my_struct.nested_enum_field": my_pb2.AnotherEnum,
+
+                            # Nested field in a MAP column named 'my_map'
+                            "my_map.key": my_pb2.MapKeyEnum, # If map keys were enums
+                            "my_map.value": my_pb2.MapValueMessage,
+
+                            # PROTO field inside a STRUCT, where the STRUCT is the value in a MAP column
+                            "struct_map.value.nested_proto_field": my_pb2.DeeplyNestedProto,
+                            "struct_map.value.nested_enum_field": my_pb2.DeeplyNestedEnum
+                        }
+
         Returns:
             ExecuteQueryIteratorAsync: an asynchronous iterator that yields rows returned by the query
         Raises:
@@ -682,6 +804,7 @@ class BigtableDataClientAsync(ClientWithProject):
             google.api_core.exceptions.GoogleAPIError: raised if the request encounters an unrecoverable error
             google.cloud.bigtable.data.exceptions.ParameterTypeInferenceFailed: Raised if
                 a parameter is passed without an explicit type, and the type cannot be infered
+            google.protobuf.message.DecodeError: raised if the deserialization of a PROTO/ENUM value fails.
         """
         instance_name = self._gapic_client.instance_path(self.project, instance_id)
         converted_param_types = _to_param_types(parameters, parameter_types)
@@ -739,6 +862,7 @@ class BigtableDataClientAsync(ClientWithProject):
             attempt_timeout,
             operation_timeout,
             retryable_excs=retryable_excs,
+            column_info=column_info,
         )
 
     @CrossSync.convert(sync_name="__enter__")
@@ -863,30 +987,39 @@ class _DataApiTargetAsync(abc.ABC):
         self.table_name = self.client._gapic_client.table_path(
             self.client.project, instance_id, table_id
         )
-        self.app_profile_id = app_profile_id
+        self.app_profile_id: str | None = app_profile_id
 
-        self.default_operation_timeout = default_operation_timeout
-        self.default_attempt_timeout = default_attempt_timeout
-        self.default_read_rows_operation_timeout = default_read_rows_operation_timeout
-        self.default_read_rows_attempt_timeout = default_read_rows_attempt_timeout
-        self.default_mutate_rows_operation_timeout = (
+        self.default_operation_timeout: float = default_operation_timeout
+        self.default_attempt_timeout: float | None = default_attempt_timeout
+        self.default_read_rows_operation_timeout: float = (
+            default_read_rows_operation_timeout
+        )
+        self.default_read_rows_attempt_timeout: float | None = (
+            default_read_rows_attempt_timeout
+        )
+        self.default_mutate_rows_operation_timeout: float = (
             default_mutate_rows_operation_timeout
         )
-        self.default_mutate_rows_attempt_timeout = default_mutate_rows_attempt_timeout
+        self.default_mutate_rows_attempt_timeout: float | None = (
+            default_mutate_rows_attempt_timeout
+        )
 
-        self.default_read_rows_retryable_errors = (
+        self.default_read_rows_retryable_errors: Sequence[type[Exception]] = (
             default_read_rows_retryable_errors or ()
         )
-        self.default_mutate_rows_retryable_errors = (
+        self.default_mutate_rows_retryable_errors: Sequence[type[Exception]] = (
             default_mutate_rows_retryable_errors or ()
         )
-        self.default_retryable_errors = default_retryable_errors or ()
+        self.default_retryable_errors: Sequence[type[Exception]] = (
+            default_retryable_errors or ()
+        )
 
         try:
             self._register_instance_future = CrossSync.create_task(
                 self.client._register_instance,
                 self.instance_id,
-                self,
+                self.app_profile_id,
+                id(self),
                 sync_executor=self.client._executor,
             )
         except RuntimeError as e:
@@ -1597,7 +1730,9 @@ class _DataApiTargetAsync(abc.ABC):
         """
         if self._register_instance_future:
             self._register_instance_future.cancel()
-        await self.client._remove_instance_registration(self.instance_id, self)
+        self.client._remove_instance_registration(
+            self.instance_id, self.app_profile_id, id(self)
+        )
 
     @CrossSync.convert(sync_name="__enter__")
     async def __aenter__(self):

google_cloud_bigtable-2.33.0/google/cloud/bigtable/data/_sync_autogen/_swappable_channel.py

@@ -0,0 +1,96 @@
+# 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 Channel
+
+
+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):
+        return self._channel.unary_unary(*args, **kwargs)
+
+    def unary_stream(self, *args, **kwargs):
+        return self._channel.unary_stream(*args, **kwargs)
+
+    def stream_unary(self, *args, **kwargs):
+        return self._channel.stream_unary(*args, **kwargs)
+
+    def stream_stream(self, *args, **kwargs):
+        return self._channel.stream_stream(*args, **kwargs)
+
+    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)
+
+    def close(self, grace=None):
+        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 SwappableChannel(_WrappedChannel):
+    """
+    Provides a grpc channel wrapper, that allows the internal channel to be swapped out
+
+    Args:
+      - channel_fn: a nullary function that returns a new channel instance.
+            It should be a partial with all channel configuration arguments built-in
+    """
+
+    def __init__(self, channel_fn: Callable[[], Channel]):
+        self._channel_fn = channel_fn
+        self._channel = channel_fn()
+
+    def create_channel(self) -> Channel:
+        """Create a fresh channel using the stored `channel_fn` partial"""
+        new_channel = self._channel_fn()
+        return new_channel
+
+    def swap_channel(self, new_channel: Channel) -> Channel:
+        """Replace the wrapped channel with a new instance. Typically created using `create_channel`"""
+        old_channel = self._channel
+        self._channel = new_channel
+        return old_channel