Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# -*- coding: utf-8 -*-
2# Copyright 2025 Google LLC
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8# http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16import inspect
17import json
18import logging as std_logging
19import pickle
20from typing import Awaitable, Callable, Dict, Optional, Sequence, Tuple, Union
21import warnings
23from google.api_core import exceptions as core_exceptions
24from google.api_core import gapic_v1, grpc_helpers_async
25from google.api_core import retry_async as retries
26from google.auth import credentials as ga_credentials # type: ignore
27from google.auth.transport.grpc import SslCredentials # type: ignore
28from google.protobuf.json_format import MessageToJson
29import google.protobuf.message
30import grpc # type: ignore
31from grpc.experimental import aio # type: ignore
32import proto # type: ignore
34from google.cloud.bigquery_storage_v1.types import storage, stream
36from .base import DEFAULT_CLIENT_INFO, BigQueryReadTransport
37from .grpc import BigQueryReadGrpcTransport
39try:
40 from google.api_core import client_logging # type: ignore
42 CLIENT_LOGGING_SUPPORTED = True # pragma: NO COVER
43except ImportError: # pragma: NO COVER
44 CLIENT_LOGGING_SUPPORTED = False
46_LOGGER = std_logging.getLogger(__name__)
49class _LoggingClientAIOInterceptor(
50 grpc.aio.UnaryUnaryClientInterceptor
51): # pragma: NO COVER
52 async def intercept_unary_unary(self, continuation, client_call_details, request):
53 logging_enabled = CLIENT_LOGGING_SUPPORTED and _LOGGER.isEnabledFor(
54 std_logging.DEBUG
55 )
56 if logging_enabled: # pragma: NO COVER
57 request_metadata = client_call_details.metadata
58 if isinstance(request, proto.Message):
59 request_payload = type(request).to_json(request)
60 elif isinstance(request, google.protobuf.message.Message):
61 request_payload = MessageToJson(request)
62 else:
63 request_payload = f"{type(request).__name__}: {pickle.dumps(request)}"
65 request_metadata = {
66 key: value.decode("utf-8") if isinstance(value, bytes) else value
67 for key, value in request_metadata
68 }
69 grpc_request = {
70 "payload": request_payload,
71 "requestMethod": "grpc",
72 "metadata": dict(request_metadata),
73 }
74 _LOGGER.debug(
75 f"Sending request for {client_call_details.method}",
76 extra={
77 "serviceName": "google.cloud.bigquery.storage.v1.BigQueryRead",
78 "rpcName": str(client_call_details.method),
79 "request": grpc_request,
80 "metadata": grpc_request["metadata"],
81 },
82 )
83 response = await continuation(client_call_details, request)
84 if logging_enabled: # pragma: NO COVER
85 response_metadata = await response.trailing_metadata()
86 # Convert gRPC metadata `<class 'grpc.aio._metadata.Metadata'>` to list of tuples
87 metadata = (
88 dict([(k, str(v)) for k, v in response_metadata])
89 if response_metadata
90 else None
91 )
92 result = await response
93 if isinstance(result, proto.Message):
94 response_payload = type(result).to_json(result)
95 elif isinstance(result, google.protobuf.message.Message):
96 response_payload = MessageToJson(result)
97 else:
98 response_payload = f"{type(result).__name__}: {pickle.dumps(result)}"
99 grpc_response = {
100 "payload": response_payload,
101 "metadata": metadata,
102 "status": "OK",
103 }
104 _LOGGER.debug(
105 f"Received response to rpc {client_call_details.method}.",
106 extra={
107 "serviceName": "google.cloud.bigquery.storage.v1.BigQueryRead",
108 "rpcName": str(client_call_details.method),
109 "response": grpc_response,
110 "metadata": grpc_response["metadata"],
111 },
112 )
113 return response
116class BigQueryReadGrpcAsyncIOTransport(BigQueryReadTransport):
117 """gRPC AsyncIO backend transport for BigQueryRead.
119 BigQuery Read API.
121 The Read API can be used to read data from BigQuery.
123 This class defines the same methods as the primary client, so the
124 primary client can load the underlying transport implementation
125 and call it.
127 It sends protocol buffers over the wire using gRPC (which is built on
128 top of HTTP/2); the ``grpcio`` package must be installed.
129 """
131 _grpc_channel: aio.Channel
132 _stubs: Dict[str, Callable] = {}
134 @classmethod
135 def create_channel(
136 cls,
137 host: str = "bigquerystorage.googleapis.com",
138 credentials: Optional[ga_credentials.Credentials] = None,
139 credentials_file: Optional[str] = None,
140 scopes: Optional[Sequence[str]] = None,
141 quota_project_id: Optional[str] = None,
142 **kwargs,
143 ) -> aio.Channel:
144 """Create and return a gRPC AsyncIO channel object.
145 Args:
146 host (Optional[str]): The host for the channel to use.
147 credentials (Optional[~.Credentials]): The
148 authorization credentials to attach to requests. These
149 credentials identify this application to the service. If
150 none are specified, the client will attempt to ascertain
151 the credentials from the environment.
152 credentials_file (Optional[str]): A file with credentials that can
153 be loaded with :func:`google.auth.load_credentials_from_file`.
154 scopes (Optional[Sequence[str]]): A optional list of scopes needed for this
155 service. These are only used when credentials are not specified and
156 are passed to :func:`google.auth.default`.
157 quota_project_id (Optional[str]): An optional project to use for billing
158 and quota.
159 kwargs (Optional[dict]): Keyword arguments, which are passed to the
160 channel creation.
161 Returns:
162 aio.Channel: A gRPC AsyncIO channel object.
163 """
165 return grpc_helpers_async.create_channel(
166 host,
167 credentials=credentials,
168 credentials_file=credentials_file,
169 quota_project_id=quota_project_id,
170 default_scopes=cls.AUTH_SCOPES,
171 scopes=scopes,
172 default_host=cls.DEFAULT_HOST,
173 **kwargs,
174 )
176 def __init__(
177 self,
178 *,
179 host: str = "bigquerystorage.googleapis.com",
180 credentials: Optional[ga_credentials.Credentials] = None,
181 credentials_file: Optional[str] = None,
182 scopes: Optional[Sequence[str]] = None,
183 channel: Optional[Union[aio.Channel, Callable[..., aio.Channel]]] = None,
184 api_mtls_endpoint: Optional[str] = None,
185 client_cert_source: Optional[Callable[[], Tuple[bytes, bytes]]] = None,
186 ssl_channel_credentials: Optional[grpc.ChannelCredentials] = None,
187 client_cert_source_for_mtls: Optional[Callable[[], Tuple[bytes, bytes]]] = None,
188 quota_project_id: Optional[str] = None,
189 client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO,
190 always_use_jwt_access: Optional[bool] = False,
191 api_audience: Optional[str] = None,
192 ) -> None:
193 """Instantiate the transport.
195 Args:
196 host (Optional[str]):
197 The hostname to connect to (default: 'bigquerystorage.googleapis.com').
198 credentials (Optional[google.auth.credentials.Credentials]): The
199 authorization credentials to attach to requests. These
200 credentials identify the application to the service; if none
201 are specified, the client will attempt to ascertain the
202 credentials from the environment.
203 This argument is ignored if a ``channel`` instance is provided.
204 credentials_file (Optional[str]): A file with credentials that can
205 be loaded with :func:`google.auth.load_credentials_from_file`.
206 This argument is ignored if a ``channel`` instance is provided.
207 scopes (Optional[Sequence[str]]): A optional list of scopes needed for this
208 service. These are only used when credentials are not specified and
209 are passed to :func:`google.auth.default`.
210 channel (Optional[Union[aio.Channel, Callable[..., aio.Channel]]]):
211 A ``Channel`` instance through which to make calls, or a Callable
212 that constructs and returns one. If set to None, ``self.create_channel``
213 is used to create the channel. If a Callable is given, it will be called
214 with the same arguments as used in ``self.create_channel``.
215 api_mtls_endpoint (Optional[str]): Deprecated. The mutual TLS endpoint.
216 If provided, it overrides the ``host`` argument and tries to create
217 a mutual TLS channel with client SSL credentials from
218 ``client_cert_source`` or application default SSL credentials.
219 client_cert_source (Optional[Callable[[], Tuple[bytes, bytes]]]):
220 Deprecated. A callback to provide client SSL certificate bytes and
221 private key bytes, both in PEM format. It is ignored if
222 ``api_mtls_endpoint`` is None.
223 ssl_channel_credentials (grpc.ChannelCredentials): SSL credentials
224 for the grpc channel. It is ignored if a ``channel`` instance is provided.
225 client_cert_source_for_mtls (Optional[Callable[[], Tuple[bytes, bytes]]]):
226 A callback to provide client certificate bytes and private key bytes,
227 both in PEM format. It is used to configure a mutual TLS channel. It is
228 ignored if a ``channel`` instance or ``ssl_channel_credentials`` is provided.
229 quota_project_id (Optional[str]): An optional project to use for billing
230 and quota.
231 client_info (google.api_core.gapic_v1.client_info.ClientInfo):
232 The client info used to send a user-agent string along with
233 API requests. If ``None``, then default info will be used.
234 Generally, you only need to set this if you're developing
235 your own client library.
236 always_use_jwt_access (Optional[bool]): Whether self signed JWT should
237 be used for service account credentials.
239 Raises:
240 google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport
241 creation failed for any reason.
242 google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials``
243 and ``credentials_file`` are passed.
244 """
245 self._grpc_channel = None
246 self._ssl_channel_credentials = ssl_channel_credentials
247 self._stubs: Dict[str, Callable] = {}
249 if api_mtls_endpoint:
250 warnings.warn("api_mtls_endpoint is deprecated", DeprecationWarning)
251 if client_cert_source:
252 warnings.warn("client_cert_source is deprecated", DeprecationWarning)
254 if isinstance(channel, aio.Channel):
255 # Ignore credentials if a channel was passed.
256 credentials = None
257 self._ignore_credentials = True
258 # If a channel was explicitly provided, set it.
259 self._grpc_channel = channel
260 self._ssl_channel_credentials = None
261 else:
262 if api_mtls_endpoint:
263 host = api_mtls_endpoint
265 # Create SSL credentials with client_cert_source or application
266 # default SSL credentials.
267 if client_cert_source:
268 cert, key = client_cert_source()
269 self._ssl_channel_credentials = grpc.ssl_channel_credentials(
270 certificate_chain=cert, private_key=key
271 )
272 else:
273 self._ssl_channel_credentials = SslCredentials().ssl_credentials
275 else:
276 if client_cert_source_for_mtls and not ssl_channel_credentials:
277 cert, key = client_cert_source_for_mtls()
278 self._ssl_channel_credentials = grpc.ssl_channel_credentials(
279 certificate_chain=cert, private_key=key
280 )
282 # The base transport sets the host, credentials and scopes
283 super().__init__(
284 host=host,
285 credentials=credentials,
286 credentials_file=credentials_file,
287 scopes=scopes,
288 quota_project_id=quota_project_id,
289 client_info=client_info,
290 always_use_jwt_access=always_use_jwt_access,
291 api_audience=api_audience,
292 )
294 if not self._grpc_channel:
295 # initialize with the provided callable or the default channel
296 channel_init = channel or type(self).create_channel
297 self._grpc_channel = channel_init(
298 self._host,
299 # use the credentials which are saved
300 credentials=self._credentials,
301 # Set ``credentials_file`` to ``None`` here as
302 # the credentials that we saved earlier should be used.
303 credentials_file=None,
304 scopes=self._scopes,
305 ssl_credentials=self._ssl_channel_credentials,
306 quota_project_id=quota_project_id,
307 options=[
308 ("grpc.max_send_message_length", -1),
309 ("grpc.max_receive_message_length", -1),
310 ],
311 )
313 self._interceptor = _LoggingClientAIOInterceptor()
314 self._grpc_channel._unary_unary_interceptors.append(self._interceptor)
315 self._logged_channel = self._grpc_channel
316 self._wrap_with_kind = (
317 "kind" in inspect.signature(gapic_v1.method_async.wrap_method).parameters
318 )
319 # Wrap messages. This must be done after self._logged_channel exists
320 self._prep_wrapped_messages(client_info)
322 @property
323 def grpc_channel(self) -> aio.Channel:
324 """Create the channel designed to connect to this service.
326 This property caches on the instance; repeated calls return
327 the same channel.
328 """
329 # Return the channel from cache.
330 return self._grpc_channel
332 @property
333 def create_read_session(
334 self,
335 ) -> Callable[[storage.CreateReadSessionRequest], Awaitable[stream.ReadSession]]:
336 r"""Return a callable for the create read session method over gRPC.
338 Creates a new read session. A read session divides
339 the contents of a BigQuery table into one or more
340 streams, which can then be used to read data from the
341 table. The read session also specifies properties of the
342 data to be read, such as a list of columns or a
343 push-down filter describing the rows to be returned.
345 A particular row can be read by at most one stream. When
346 the caller has reached the end of each stream in the
347 session, then all the data in the table has been read.
349 Data is assigned to each stream such that roughly the
350 same number of rows can be read from each stream.
351 Because the server-side unit for assigning data is
352 collections of rows, the API does not guarantee that
353 each stream will return the same number or rows.
354 Additionally, the limits are enforced based on the
355 number of pre-filtered rows, so some filters can lead to
356 lopsided assignments.
358 Read sessions automatically expire 6 hours after they
359 are created and do not require manual clean-up by the
360 caller.
362 Returns:
363 Callable[[~.CreateReadSessionRequest],
364 Awaitable[~.ReadSession]]:
365 A function that, when called, will call the underlying RPC
366 on the server.
367 """
368 # Generate a "stub function" on-the-fly which will actually make
369 # the request.
370 # gRPC handles serialization and deserialization, so we just need
371 # to pass in the functions for each.
372 if "create_read_session" not in self._stubs:
373 self._stubs["create_read_session"] = self._logged_channel.unary_unary(
374 "/google.cloud.bigquery.storage.v1.BigQueryRead/CreateReadSession",
375 request_serializer=storage.CreateReadSessionRequest.serialize,
376 response_deserializer=stream.ReadSession.deserialize,
377 )
378 return self._stubs["create_read_session"]
380 @property
381 def read_rows(
382 self,
383 ) -> Callable[[storage.ReadRowsRequest], Awaitable[storage.ReadRowsResponse]]:
384 r"""Return a callable for the read rows method over gRPC.
386 Reads rows from the stream in the format prescribed
387 by the ReadSession. Each response contains one or more
388 table rows, up to a maximum of 100 MiB per response;
389 read requests which attempt to read individual rows
390 larger than 100 MiB will fail.
392 Each request also returns a set of stream statistics
393 reflecting the current state of the stream.
395 Returns:
396 Callable[[~.ReadRowsRequest],
397 Awaitable[~.ReadRowsResponse]]:
398 A function that, when called, will call the underlying RPC
399 on the server.
400 """
401 # Generate a "stub function" on-the-fly which will actually make
402 # the request.
403 # gRPC handles serialization and deserialization, so we just need
404 # to pass in the functions for each.
405 if "read_rows" not in self._stubs:
406 self._stubs["read_rows"] = self._logged_channel.unary_stream(
407 "/google.cloud.bigquery.storage.v1.BigQueryRead/ReadRows",
408 request_serializer=storage.ReadRowsRequest.serialize,
409 response_deserializer=storage.ReadRowsResponse.deserialize,
410 )
411 return self._stubs["read_rows"]
413 @property
414 def split_read_stream(
415 self,
416 ) -> Callable[
417 [storage.SplitReadStreamRequest], Awaitable[storage.SplitReadStreamResponse]
418 ]:
419 r"""Return a callable for the split read stream method over gRPC.
421 Splits a given ``ReadStream`` into two ``ReadStream`` objects.
422 These ``ReadStream`` objects are referred to as the primary and
423 the residual streams of the split. The original ``ReadStream``
424 can still be read from in the same manner as before. Both of the
425 returned ``ReadStream`` objects can also be read from, and the
426 rows returned by both child streams will be the same as the rows
427 read from the original stream.
429 Moreover, the two child streams will be allocated back-to-back
430 in the original ``ReadStream``. Concretely, it is guaranteed
431 that for streams original, primary, and residual, that
432 original[0-j] = primary[0-j] and original[j-n] = residual[0-m]
433 once the streams have been read to completion.
435 Returns:
436 Callable[[~.SplitReadStreamRequest],
437 Awaitable[~.SplitReadStreamResponse]]:
438 A function that, when called, will call the underlying RPC
439 on the server.
440 """
441 # Generate a "stub function" on-the-fly which will actually make
442 # the request.
443 # gRPC handles serialization and deserialization, so we just need
444 # to pass in the functions for each.
445 if "split_read_stream" not in self._stubs:
446 self._stubs["split_read_stream"] = self._logged_channel.unary_unary(
447 "/google.cloud.bigquery.storage.v1.BigQueryRead/SplitReadStream",
448 request_serializer=storage.SplitReadStreamRequest.serialize,
449 response_deserializer=storage.SplitReadStreamResponse.deserialize,
450 )
451 return self._stubs["split_read_stream"]
453 def _prep_wrapped_messages(self, client_info):
454 """Precompute the wrapped methods, overriding the base class method to use async wrappers."""
455 self._wrapped_methods = {
456 self.create_read_session: self._wrap_method(
457 self.create_read_session,
458 default_retry=retries.AsyncRetry(
459 initial=0.1,
460 maximum=60.0,
461 multiplier=1.3,
462 predicate=retries.if_exception_type(
463 core_exceptions.DeadlineExceeded,
464 core_exceptions.ServiceUnavailable,
465 ),
466 deadline=600.0,
467 ),
468 default_timeout=600.0,
469 client_info=client_info,
470 ),
471 self.read_rows: self._wrap_method(
472 self.read_rows,
473 default_retry=retries.AsyncRetry(
474 initial=0.1,
475 maximum=60.0,
476 multiplier=1.3,
477 predicate=retries.if_exception_type(
478 core_exceptions.ServiceUnavailable,
479 ),
480 deadline=86400.0,
481 ),
482 default_timeout=86400.0,
483 client_info=client_info,
484 ),
485 self.split_read_stream: self._wrap_method(
486 self.split_read_stream,
487 default_retry=retries.AsyncRetry(
488 initial=0.1,
489 maximum=60.0,
490 multiplier=1.3,
491 predicate=retries.if_exception_type(
492 core_exceptions.DeadlineExceeded,
493 core_exceptions.ServiceUnavailable,
494 ),
495 deadline=600.0,
496 ),
497 default_timeout=600.0,
498 client_info=client_info,
499 ),
500 }
502 def _wrap_method(self, func, *args, **kwargs):
503 if self._wrap_with_kind: # pragma: NO COVER
504 kwargs["kind"] = self.kind
505 return gapic_v1.method_async.wrap_method(func, *args, **kwargs)
507 def close(self):
508 return self._logged_channel.close()
510 @property
511 def kind(self) -> str:
512 return "grpc_asyncio"
515__all__ = ("BigQueryReadGrpcAsyncIOTransport",)