Skip to content
This repository was archived by the owner on Apr 1, 2026. It is now read-only.

Commit 57ecce0

Browse files
gkevinzhenggkevinzheng
committed
feat: Async consistency polling harness
1 parent 3d714e2 commit 57ecce0

6 files changed

Lines changed: 561 additions & 1 deletion

File tree

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

Lines changed: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,228 @@
1+
# Copyright 2025 Google LLC.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# -*- coding: utf-8 -*-
16+
# Copyright 2025 Google LLC
17+
#
18+
# Licensed under the Apache License, Version 2.0 (the "License");
19+
# you may not use this file except in compliance with the License.
20+
# You may obtain a copy of the License at
21+
#
22+
# http://www.apache.org/licenses/LICENSE-2.0
23+
#
24+
# Unless required by applicable law or agreed to in writing, software
25+
# distributed under the License is distributed on an "AS IS" BASIS,
26+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
27+
# See the License for the specific language governing permissions and
28+
# limitations under the License.
29+
30+
import copy
31+
import functools
32+
33+
from typing import Callable, Optional, Sequence, Tuple, Union
34+
from google.api_core import gapic_v1
35+
from google.api_core import retry as retries
36+
37+
try:
38+
OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None]
39+
except AttributeError: # pragma: NO COVER
40+
OptionalRetry = Union[retries.Retry, object, None] # type: ignore
41+
42+
from google.api_core import client_options as client_options_lib
43+
from google.auth import credentials as ga_credentials # type: ignore
44+
45+
from google.cloud.bigtable.admin_v2.types import bigtable_table_admin
46+
47+
from google.cloud.bigtable.admin_v2.services.bigtable_table_admin import (
48+
async_client as base_client,
49+
)
50+
from google.cloud.bigtable.admin_v2.services.bigtable_table_admin.transports.base import (
51+
BigtableTableAdminTransport,
52+
)
53+
from google.cloud.bigtable.admin_v2.overlay.types import async_consistency
54+
55+
from google.cloud.bigtable.gapic_version import __version__ as bigtable_version
56+
57+
58+
DEFAULT_CLIENT_INFO = copy.copy(base_client.DEFAULT_CLIENT_INFO)
59+
DEFAULT_CLIENT_INFO.client_library_version = f"{bigtable_version}-admin-overlay-async"
60+
61+
class BigtableTableAdminAsyncClient(base_client.BaseBigtableTableAdminAsyncClient):
62+
def __init__(
63+
self,
64+
*,
65+
credentials: Optional[ga_credentials.Credentials] = None,
66+
transport: Optional[
67+
Union[
68+
str,
69+
BigtableTableAdminTransport,
70+
Callable[..., BigtableTableAdminTransport],
71+
]
72+
] = "grpc_asyncio",
73+
client_options: Optional[client_options_lib.ClientOptions] = None,
74+
client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO,
75+
) -> None:
76+
"""Instantiates the Bigtable table admin async client.
77+
78+
Args:
79+
credentials (Optional[google.auth.credentials.Credentials]): The
80+
authorization credentials to attach to requests. These
81+
credentials identify the application to the service; if none
82+
are specified, the client will attempt to ascertain the
83+
credentials from the environment.
84+
transport (Optional[Union[str,BigtableTableAdminTransport,Callable[..., BigtableTableAdminTransport]]]):
85+
The transport to use, or a Callable that constructs and returns a new transport to use.
86+
If a Callable is given, it will be called with the same set of initialization
87+
arguments as used in the BigtableTableAdminTransport constructor.
88+
If set to None, a transport is chosen automatically.
89+
client_options (Optional[Union[google.api_core.client_options.ClientOptions, dict]]):
90+
Custom options for the client.
91+
92+
1. The ``api_endpoint`` property can be used to override the
93+
default endpoint provided by the client when ``transport`` is
94+
not explicitly provided. Only if this property is not set and
95+
``transport`` was not explicitly provided, the endpoint is
96+
determined by the GOOGLE_API_USE_MTLS_ENDPOINT environment
97+
variable, which have one of the following values:
98+
"always" (always use the default mTLS endpoint), "never" (always
99+
use the default regular endpoint) and "auto" (auto-switch to the
100+
default mTLS endpoint if client certificate is present; this is
101+
the default value).
102+
103+
2. If the GOOGLE_API_USE_CLIENT_CERTIFICATE environment variable
104+
is "true", then the ``client_cert_source`` property can be used
105+
to provide a client certificate for mTLS transport. If
106+
not provided, the default SSL client certificate will be used if
107+
present. If GOOGLE_API_USE_CLIENT_CERTIFICATE is "false" or not
108+
set, no client certificate will be used.
109+
110+
3. The ``universe_domain`` property can be used to override the
111+
default "googleapis.com" universe. Note that ``api_endpoint``
112+
property still takes precedence; and ``universe_domain`` is
113+
currently not supported for mTLS.
114+
115+
client_info (google.api_core.gapic_v1.client_info.ClientInfo):
116+
The client info used to send a user-agent string along with
117+
API requests. If ``None``, then default info will be used.
118+
Generally, you only need to set this if you're developing
119+
your own client library.
120+
121+
Raises:
122+
google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport
123+
creation failed for any reason.
124+
"""
125+
super(BigtableTableAdminAsyncClient, self).__init__(
126+
credentials=credentials,
127+
transport=transport,
128+
client_options=client_options,
129+
client_info=client_info,
130+
)
131+
132+
133+
async def wait_for_consistency(
134+
self,
135+
request: Optional[
136+
Union[bigtable_table_admin.CheckConsistencyRequest, dict]
137+
] = None,
138+
*,
139+
name: Optional[str] = None,
140+
consistency_token: Optional[str] = None,
141+
retry: OptionalRetry = gapic_v1.method.DEFAULT,
142+
timeout: Union[float, object] = gapic_v1.method.DEFAULT,
143+
metadata: Sequence[Tuple[str, Union[str, bytes]]] = (),
144+
) -> async_consistency.AsyncCheckConsistencyPollingFuture:
145+
r"""Creates a polling future that periodically checks replication
146+
consistency based on a consistency token, that is, if replication
147+
has caught up based on the conditions specified in the token and
148+
the check request. The future will stop checking once the underlying
149+
:meth:`check_consistency` request involving that token returns True.
150+
151+
.. code-block:: python
152+
153+
# This snippet should be regarded as a code template only.
154+
# It will require modifications to work:
155+
# - It may require correct/in-range values for request initialization.
156+
# - It may require specifying regional endpoints when creating the service
157+
# client as shown in:
158+
# https://googleapis.dev/python/google-api-core/latest/client_options.html
159+
from google.cloud.bigtable import admin_v2
160+
161+
def sample_wait_for_consistency():
162+
# Create a client
163+
client = admin_v2.BigtableTableAdminClient()
164+
165+
# Initialize request argument(s)
166+
request = admin_v2.CheckConsistencyRequest(
167+
name="name_value",
168+
consistency_token="consistency_token_value",
169+
)
170+
171+
# Make the request
172+
future = client.wait_for_consistency(request=request)
173+
174+
# Wait for the table to become consistent
175+
print("Waiting for operation to complete...")
176+
177+
response = future.result()
178+
179+
# Handle the response
180+
print(response)
181+
182+
Args:
183+
request (Union[google.cloud.bigtable.admin_v2.types.CheckConsistencyRequest, dict]):
184+
The request object. Request message for
185+
[google.bigtable.admin.v2.BigtableTableAdmin.CheckConsistency][google.bigtable.admin.v2.BigtableTableAdmin.CheckConsistency]
186+
name (str):
187+
Required. The unique name of the Table for which to
188+
check replication consistency. Values are of the form
189+
``projects/{project}/instances/{instance}/tables/{table}``.
190+
191+
This corresponds to the ``name`` field
192+
on the ``request`` instance; if ``request`` is provided, this
193+
should not be set.
194+
consistency_token (str):
195+
Required. The token created using
196+
GenerateConsistencyToken for the Table.
197+
198+
This corresponds to the ``consistency_token`` field
199+
on the ``request`` instance; if ``request`` is provided, this
200+
should not be set.
201+
retry (google.api_core.retry.Retry): Designation of what errors, if any,
202+
should be retried.
203+
timeout (float): The timeout for this request.
204+
metadata (Sequence[Tuple[str, Union[str, bytes]]]): Key/value pairs which should be
205+
sent along with the request as metadata. Normally, each value must be of type `str`,
206+
but for metadata keys ending with the suffix `-bin`, the corresponding values must
207+
be of type `bytes`.
208+
209+
Returns:
210+
google.cloud.bigtable.admin_v2.overlay.types.async_consistency.AsyncCheckConsistencyPollingFuture:
211+
An object representing a polling future.
212+
213+
The result type for the operation will be `bool`, and will return True when the
214+
consistency check involving the given consistency token returns True.
215+
"""
216+
api_call = functools.partial(
217+
self.check_consistency,
218+
request,
219+
name=name,
220+
consistency_token=consistency_token,
221+
timeout=timeout,
222+
metadata=metadata,
223+
)
224+
225+
return async_consistency.AsyncCheckConsistencyPollingFuture(
226+
api_call,
227+
check_consistency_call_retry=retry
228+
)

‎google/cloud/bigtable/admin_v2/overlay/types/__init__.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,10 @@
1212
# See the License for the specific language governing permissions and
1313
# limitations under the License.
1414

15+
from .async_consistency import (
16+
AsyncCheckConsistencyPollingFuture,
17+
)
18+
1519
from .consistency import (
1620
CheckConsistencyPollingFuture,
1721
)
@@ -21,6 +25,7 @@
2125
)
2226

2327
__all__ = (
28+
"AsyncCheckConsistencyPollingFuture",
2429
"CheckConsistencyPollingFuture",
2530
"RestoreTableOperation",
2631
)

‎google/cloud/bigtable/admin_v2/overlay/types/async_consistency.py‎

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
# Copyright 2025 Google LLC.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from typing import Awaitable, Union, Callable
16+
17+
from google.api_core.future import async_future
18+
from google.api_core import gapic_v1
19+
from google.api_core import retry as retries
20+
from google.cloud.bigtable.admin_v2.types import bigtable_table_admin
21+
22+
try:
23+
OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None]
24+
except AttributeError: # pragma: NO COVER
25+
OptionalRetry = Union[retries.Retry, object, None] # type: ignore
26+
27+
28+
class AsyncCheckConsistencyPollingFuture(async_future.AsyncFuture):
29+
"""A Future that polls an underlying `check_consistency` operation until it returns True.
30+
31+
**This class should not be instantiated by users** and should only be instantiated by the admin
32+
client's
33+
:meth:`google.cloud.bigtable.admin_v2.overlay.services.bigtable_table_admin.AsyncBigtableTableAdminClient.wait_for_consistency`
34+
or
35+
:meth:`google.cloud.bigtable.admin_v2.overlay.services.bigtable_table_admin.AsyncBigtableTableAdminClient.wait_for_replication`
36+
methods.
37+
38+
Args:
39+
check_consistency_call(Callable[
40+
[Optional[google.api_core.retry.Retry],
41+
google.cloud.bigtable.admin_v2.types.CheckConsistencyResponse]):
42+
A :meth:`check_consistency
43+
<google.cloud.bigtable.admin_v2.overlay.services.bigtable_table_admin.AsyncBigtableTableAdminClient.check_consistency>`
44+
call from the admin client. The call should fix every user parameter except for retry,
45+
which will be done via :meth:`functools.partial`.
46+
default_retry(Optional[google.api_core.retry.Retry]): The `retry` parameter passed in to either
47+
:meth:`wait_for_consistency
48+
<google.cloud.bigtable.admin_v2.overlay.services.bigtable_table_admin.BigtableTableAdminClient.wait_for_consistency>`
49+
or :meth:`wait_for_replication
50+
<google.cloud.bigtable.admin_v2.overlay.services.bigtable_table_admin.BigtableTableAdminClient.wait_for_replication>`
51+
polling (google.api_core.retry.Retry): The configuration used for polling.
52+
This parameter controls how often :meth:`done` is polled. If the
53+
``timeout`` argument is specified in the :meth:`result
54+
<google.api_core.future.polling.PollingFuture.result>` method it will
55+
override the ``polling.timeout`` property.
56+
"""
57+
58+
def __init__(
59+
self,
60+
check_consistency_call: Callable[
61+
[OptionalRetry], Awaitable[bigtable_table_admin.CheckConsistencyResponse]
62+
],
63+
check_consistency_call_retry: OptionalRetry = gapic_v1.method.DEFAULT,
64+
future_retry: retries.Retry = async_future.DEFAULT_RETRY,
65+
**kwargs
66+
):
67+
super(AsyncCheckConsistencyPollingFuture, self).__init__(retry=future_retry, **kwargs)
68+
69+
# Done is called with two different scenarios, retry is specified or not specified.
70+
# API_call will be a functools partial with everything except retry specified because of
71+
# that.
72+
self._check_consistency_call = check_consistency_call
73+
self._check_consistency_call_retry = check_consistency_call_retry
74+
75+
async def done(self, retry: OptionalRetry = None):
76+
"""Polls the underlying `check_consistency` call to see if the future is complete.
77+
78+
This should not be used by the user to wait until the `check_consistency` call finishes;
79+
use the :meth:`result <google.api_core.future.async_future.AsyncFuture.result>` method of
80+
this class instead.
81+
82+
Args:
83+
retry (google.api_core.retry.Retry): (Optional) How to retry the
84+
polling RPC (to not be confused with polling configuration. See
85+
the documentation for :meth:`result <google.api_core.future.async_future.AsyncFuture.result>`
86+
for details).
87+
88+
Returns:
89+
bool: True if the future is complete, False otherwise.
90+
"""
91+
if self._future.done():
92+
return True
93+
94+
retry = retry or self._check_consistency_call_retry
95+
96+
try:
97+
check_consistency_response = await self._check_consistency_call(retry=retry)
98+
if check_consistency_response.consistent:
99+
self.set_result(True)
100+
101+
return check_consistency_response.consistent
102+
except Exception as e:
103+
self.set_exception(e)
104+
105+
def cancel(self):
106+
raise NotImplementedError("Cannot cancel consistency token operation")
107+
108+
def cancelled(self):
109+
raise NotImplementedError("Cannot cancel consistency token operation")

0 commit comments

Comments
 (0)