feat: Restore Table LRO rework + async restore table

Author: gkevinzheng

google/cloud/bigtable/admin_v2/overlay/services/bigtable_table_admin/async_client.py

@@ -52,6 +52,7 @@
 )
 from google.cloud.bigtable.admin_v2.overlay.types import (
     async_consistency,
+    async_restore_table,
     wait_for_consistency_request,
 )
 
@@ -133,6 +134,100 @@ def __init__(
             client_info=client_info,
         )
 
+    async def restore_table(
+        self,
+        request: Optional[Union[bigtable_table_admin.RestoreTableRequest, dict]] = None,
+        *,
+        retry: OptionalRetry = gapic_v1.method.DEFAULT,
+        timeout: Union[float, object] = gapic_v1.method.DEFAULT,
+        metadata: Sequence[Tuple[str, Union[str, bytes]]] = (),
+    ) -> async_restore_table.AsyncRestoreTableOperation:
+        r"""Create a new table by restoring from a completed backup. The
+        returned table :class:`long-running operation
+        <google.cloud.bigtable.admin_v2.overlay.types.restore_table.RestoreTableOperation>`
+        can be used to track the progress of the operation, and to cancel it. The
+        :attr:`metadata <google.api_core.operation.Operation.metadata>` field type is
+        :class:`RestoreTableMetadata <google.cloud.bigtable.admin_v2.types.RestoreTableMetadata>`.
+        The :meth:`response <google.api_core.operation.Operation.result>` type is
+        :class:`google.cloud.bigtable.admin_v2.types.Table`, if successful.
+
+        Additionally, the returned :class:`long-running-operation <google.cloud.bigtable.admin_v2.overlay.types.async_restore_table.AsyncRestoreTableOperation>`
+        provides a method, :meth:`google.cloud.bigtable.admin_v2.overlay.types.async_restore_table.AsyncRestoreTableOperation.optimize_restore_table_operation` that
+        provides access to a :class:`google.api_core.operation_async.AsyncOperation` object representing the OptimizeRestoreTable long-running-operation
+        after the current one has completed.
+
+        .. code-block:: python
+
+            # This snippet should be regarded as a code template only.
+            #
+            # It will require modifications to work:
+            # - It may require correct/in-range values for request initialization.
+            # - It may require specifying regional endpoints when creating the service
+            #   client as shown in:
+            #   https://googleapis.dev/python/google-api-core/latest/client_options.html
+            from google.cloud.bigtable import admin_v2
+
+            async def sample_restore_table():
+                # Create a client
+                client = admin_v2.BigtableTableAdminAsyncClient()
+
+                # Initialize request argument(s)
+                request = admin_v2.RestoreTableRequest(
+                    backup="backup_value",
+                    parent="parent_value",
+                    table_id="table_id_value",
+                )
+
+                # Make the request
+                operation = await client.restore_table(request=request)
+
+                print("Waiting for operation to complete...")
+
+                response = await operation.result()
+
+                # Handle the response
+                print(response)
+
+                # Handle LRO2
+                optimize_operation = await operation.optimize_restore_table_operation()
+
+                if optimize_operation:
+                    print("Waiting for table optimization to complete...")
+
+                    response = await optimize_operation.result()
+
+        Args:
+            request (Union[google.cloud.bigtable.admin_v2.types.RestoreTableRequest, dict]):
+                The request object. The request for
+                [RestoreTable][google.bigtable.admin.v2.BigtableTableAdmin.RestoreTable].
+            retry (google.api_core.retry.Retry): Designation of what errors, if any,
+                should be retried.
+            timeout (float): The timeout for this request.
+            metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be
+                sent along with the request as metadata. Normally, each value must be of type `str`,
+                but for metadata keys ending with the suffix `-bin`, the corresponding values must
+                be of type `bytes`.
+
+        Returns:
+            google.cloud.bigtable.admin_v2.overlay.types.async_restore_table.AsyncRestoreTableOperation:
+                An object representing a long-running operation.
+
+                The result type for the operation will be :class:`google.cloud.bigtable.admin_v2.types.Table` A collection of user data indexed by row, column, and timestamp.
+                   Each table is served using the resources of its
+                   parent cluster.
+        """
+        operation = await self._restore_table(
+            request=request,
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
+        restore_table_operation = async_restore_table.AsyncRestoreTableOperation(
+            self._client._transport.operations_client, operation
+        )
+        return restore_table_operation
+
     async def wait_for_consistency(
         self,
         request: Optional[

google/cloud/bigtable/admin_v2/overlay/types/async_restore_table.py

@@ -0,0 +1,99 @@
+# Copyright 2025 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 typing import Optional
+
+from google.api_core import exceptions
+from google.api_core import operation_async
+from google.protobuf import empty_pb2
+
+from google.cloud.bigtable.admin_v2.types import OptimizeRestoredTableMetadata
+
+
+class AsyncRestoreTableOperation(operation_async.AsyncOperation):
+    """A Future for interacting with Bigtable Admin's RestoreTable Long-Running Operation.
+
+    This is needed to expose a potential long-running operation that might run after this operation
+    finishes, OptimizeRestoreTable. This is exposed via the the :meth:`optimize_restore_table_operation`
+    method.
+
+    **This class should not be instantiated by users** and should only be instantiated by the admin
+    client's :meth:`restore_table
+    <google.cloud.bigtable.admin_v2.overlay.services.bigtable_table_admin.BigtableTableAdminAsyncClient.restore_table>`
+    method.
+
+    Args:
+        operations_client (google.api_core.operations_v1.AbstractOperationsClient): The operations
+            client from the admin client class's transport.
+        restore_table_operation (google.api_core.operation_async.AsyncOperation): A
+            :class:`google.api_core.operation_async.AsyncOperation`
+            instance resembling a RestoreTable long-running operation
+    """
+
+    def __init__(
+        self, operations_client, restore_table_operation: operation_async.AsyncOperation
+    ):
+        self._operations_client = operations_client
+        self._optimize_restored_table_operation = None
+        super().__init__(
+            restore_table_operation._operation,
+            restore_table_operation._refresh,
+            restore_table_operation._cancel,
+            restore_table_operation._result_type,
+            restore_table_operation._metadata_type,
+            retry=restore_table_operation._retry,
+        )
+
+    async def optimize_restored_table_operation(
+        self,
+    ) -> Optional[operation_async.AsyncOperation]:
+        """Gets the OptimizeRestoredTable long-running operation that runs after this operation finishes.
+        The current operation might not trigger a follow-up OptimizeRestoredTable operation, in which case, this
+        method will return `None`.
+        This method must not be called before the parent restore_table operation is complete.
+        Returns:
+            An object representing a long-running operation, or None if there is no OptimizeRestoredTable operation
+                after this one.
+        Raises:
+            RuntimeError: raised when accessed before the restore_table operation is complete
+
+        Raises:
+            google.api_core.GoogleAPIError: raised when accessed before the restore_table operation is complete
+        """
+        if not await self.done():
+            raise exceptions.GoogleAPIError(
+                "optimize_restored_table operation can't be accessed until the restore_table operation is complete"
+            )
+
+        if self._optimize_restored_table_operation is not None:
+            return self._optimize_restored_table_operation
+
+        operation_name = self.metadata.optimize_table_operation_name
+
+        # When the RestoreTable operation finishes, it might not necessarily trigger
+        # an optimize operation.
+        if operation_name:
+            gapic_operation = await self._operations_client.get_operation(
+                name=operation_name
+            )
+            self._optimize_restored_table_operation = operation_async.from_gapic(
+                gapic_operation,
+                self._operations_client,
+                empty_pb2.Empty,
+                metadata_type=OptimizeRestoredTableMetadata,
+            )
+            return self._optimize_restored_table_operation
+        else:
+            # no optimize operation found
+            return None

google/cloud/bigtable/admin_v2/overlay/types/restore_table.py

@@ -12,22 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import enum
-from typing import Optional, Union
+from typing import Optional
 
-from google.api_core import retry
+from google.api_core import exceptions
 from google.api_core import operation
 from google.protobuf import empty_pb2
 
 from google.cloud.bigtable.admin_v2.types import OptimizeRestoredTableMetadata
 
 
-# This is needed due to a documentation issue with using object() as the default
-# value for a parameter.
-class Timeout(enum.Enum):
-    DEFAULT_TIMEOUT = "DEFAULT_TIMEOUT"
-
-
 class RestoreTableOperation(operation.Operation):
     """A Future for interacting with Bigtable Admin's RestoreTable Long-Running Operation.
 
@@ -59,16 +52,14 @@ def __init__(self, operations_client, restore_table_operation: operation.Operati
             polling=restore_table_operation._polling,
         )
 
-    def optimize_restored_table_operation(
-        self,
-        timeout: Optional[Union[int, Timeout]] = Timeout.DEFAULT_TIMEOUT,
-        retry: Optional[retry.Retry] = None,
-        polling: Optional[retry.Retry] = None,
-    ) -> Optional[operation.Operation]:
+    def optimize_restored_table_operation(self) -> Optional[operation.Operation]:
         """Gets the OptimizeRestoredTable long-running operation that runs after this operation finishes.
 
-        This is a blocking call that will return the operation after this current long-running operation
-        finishes, just like :meth:`google.api_core.operation.Operation.result`. The follow-up operation has
+        This must not be called before the parent restore_table operation is complete. You can guarantee
+        this happening by calling this function after this class's :meth:`google.api_core.operation.Operation.result`
+        method.
+
+        The follow-up operation has
         :attr:`metadata <google.api_core.operation.Operation.metadata>` type
         :class:`OptimizeRestoredTableMetadata
         <google.cloud.bigtable.admin_v2.types.bigtable_table_admin.OptimizeRestoredTableMetadata>`
@@ -77,51 +68,35 @@ def optimize_restored_table_operation(
         The current operation might not trigger a follow-up OptimizeRestoredTable operation, in which case, this
         method will return `None`.
 
-        Args:
-            timeout (Optional[int | google.cloud.bigtable.admin_v2.overlay.types.restore_table.Timeout]):
-                How long (in seconds) to wait for the operation to complete. If None, wait indefinitely. If
-                `Timeout.DEFAULT_TIMEOUT`, wait the default amount.
-            retry (Optional[google.api_core.retry.Retry]): How to retry the polling RPC. This defines ONLY
-                how the polling RPC call is retried (i.e. what to do if the RPC we used for polling returned
-                an error). It does NOT define how the polling is done (i.e. how frequently and for how long
-                to call the polling RPC).
-            polling (Optional[google.api_core.retry.Retry]): How often and for how long to call polling RPC
-                periodically. This parameter does NOT define how to retry each individual polling RPC call
-                (use the `retry` parameter for that).
-
         Returns:
             Optional[google.api_core.operation.Operation]:
                 An object representing a long-running operation, or None if there is no OptimizeRestoredTable operation
                 after this one.
-        """
-        if timeout == Timeout.DEFAULT_TIMEOUT:
-            timeout = operation.Operation._DEFAULT_VALUE
-
-        self._blocking_poll(timeout=timeout, retry=retry, polling=polling)
 
-        if self._exception is not None:
-            # pylint: disable=raising-bad-type
-            # Pylint doesn't recognize that this is valid in this case.
-            raise self._exception
+        Raises:
+            google.api_core.GoogleAPIError: raised when accessed before the restore_table operation is complete
+        """
+        if not self.done():
+            raise exceptions.GoogleAPIError(
+                "optimize_restored_table operation can't be accessed until the restore_table operation is complete"
+            )
 
-        return self._optimize_restored_table_operation
+        if self._optimize_restored_table_operation is not None:
+            return self._optimize_restored_table_operation
 
-    def set_result(self, response):
-        optimize_restored_table_operation_name = (
-            self.metadata.optimize_table_operation_name
-        )
+        operation_name = self.metadata.optimize_table_operation_name
 
         # When the RestoreTable operation finishes, it might not necessarily trigger
         # an optimize operation.
-        if optimize_restored_table_operation_name:
-            optimize_restore_table_operation = self._operations_client.get_operation(
-                name=optimize_restored_table_operation_name
-            )
+        if operation_name:
+            gapic_operation = self._operations_client.get_operation(name=operation_name)
             self._optimize_restored_table_operation = operation.from_gapic(
-                optimize_restore_table_operation,
+                gapic_operation,
                 self._operations_client,
                 empty_pb2.Empty,
                 metadata_type=OptimizeRestoredTableMetadata,
             )
-
-        super().set_result(response)
+            return self._optimize_restored_table_operation
+        else:
+            # no optimize operation found
+            return None