Skip to content

Commit

Permalink
feat: Add OperationsRestAsyncTransport to support long running operat…
Browse files Browse the repository at this point in the history
…ions
  • Loading branch information
parthea committed Sep 20, 2024
1 parent 58516ef commit b08c465
Show file tree
Hide file tree
Showing 6 changed files with 580 additions and 15 deletions.
14 changes: 14 additions & 0 deletions google/api_core/operations_v1/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,23 @@
from google.api_core.operations_v1.operations_client import OperationsClient
from google.api_core.operations_v1.transports.rest import OperationsRestTransport

try:
from google.api_core.operations_v1.transports.rest_asyncio import (
OperationsRestAsyncTransport,
)

HAS_ASYNC_TRANSPORT = True
except ImportError as e:
# Don't raise an exception if `OperationsRestAsyncTransport` cannot be imported
# so that OperationsRestTransport can still be imported
HAS_ASYNC_TRANSPORT = False

__all__ = [
"AbstractOperationsClient",
"OperationsAsyncClient",
"OperationsClient",
"OperationsRestTransport",
]

if HAS_ASYNC_TRANSPORT:
__all__.append("OperationsRestAsyncTransport")
14 changes: 14 additions & 0 deletions google/api_core/operations_v1/transports/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,26 @@
from .base import OperationsTransport
from .rest import OperationsRestTransport

try:
from .rest_asyncio import OperationsRestAsyncTransport

HAS_ASYNC_TRANSPORT = True
except ImportError as e:
# Don't raise an exception if `OperationsRestAsyncTransport` cannot be imported
# so that OperationsRestTransport can still be imported
HAS_ASYNC_TRANSPORT = False

# Compile a registry of transports.
_transport_registry = OrderedDict()
_transport_registry["rest"] = OperationsRestTransport

if HAS_ASYNC_TRANSPORT:
_transport_registry["rest_asyncio"] = OperationsRestAsyncTransport

__all__ = (
"OperationsTransport",
"OperationsRestTransport",
)

if HAS_ASYNC_TRANSPORT:
__all__.append("OperationsRestAsyncTransport")
12 changes: 7 additions & 5 deletions google/api_core/operations_v1/transports/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@
from google.api_core import exceptions as core_exceptions # type: ignore
from google.api_core import gapic_v1 # type: ignore
from google.api_core import retry as retries # type: ignore
from google.api_core import retry_async as retries_async # type: ignore
from google.api_core import version
import google.auth # type: ignore
from google.auth import credentials as ga_credentials # type: ignore
Expand Down Expand Up @@ -115,12 +116,13 @@ def __init__(
# Save the credentials.
self._credentials = credentials

def _prep_wrapped_messages(self, client_info):
def _prep_wrapped_messages(self, client_info, is_async=False):
# Precompute the wrapped methods.
retry_class = retries_async.AsyncRetry if is_async else retries.Retry
self._wrapped_methods = {
self.list_operations: gapic_v1.method.wrap_method(
self.list_operations,
default_retry=retries.Retry(
default_retry=retry_class(
initial=0.5,
maximum=10.0,
multiplier=2.0,
Expand All @@ -135,7 +137,7 @@ def _prep_wrapped_messages(self, client_info):
),
self.get_operation: gapic_v1.method.wrap_method(
self.get_operation,
default_retry=retries.Retry(
default_retry=retry_class(
initial=0.5,
maximum=10.0,
multiplier=2.0,
Expand All @@ -150,7 +152,7 @@ def _prep_wrapped_messages(self, client_info):
),
self.delete_operation: gapic_v1.method.wrap_method(
self.delete_operation,
default_retry=retries.Retry(
default_retry=retry_class(
initial=0.5,
maximum=10.0,
multiplier=2.0,
Expand All @@ -165,7 +167,7 @@ def _prep_wrapped_messages(self, client_info):
),
self.cancel_operation: gapic_v1.method.wrap_method(
self.cancel_operation,
default_retry=retries.Retry(
default_retry=retry_class(
initial=0.5,
maximum=10.0,
multiplier=2.0,
Expand Down
Loading

0 comments on commit b08c465

Please sign in to comment.