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# Copyright 2017 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.
15"""Helpers for :mod:`grpc`."""
17import collections
18import functools
19import warnings
20from typing import Generic, Iterator, Optional, TypeVar
22import google.auth
23import google.auth.credentials
24import google.auth.transport.grpc
25import google.auth.transport.requests
26import google.protobuf
27import grpc
29from google.api_core import exceptions, general_helpers
31# The list of gRPC Callable interfaces that return iterators.
32_STREAM_WRAP_CLASSES = (grpc.UnaryStreamMultiCallable, grpc.StreamStreamMultiCallable)
34# denotes the proto response type for grpc calls
35P = TypeVar("P")
38def _patch_callable_name(callable_):
39 """Fix-up gRPC callable attributes.
41 gRPC callable lack the ``__name__`` attribute which causes
42 :func:`functools.wraps` to error. This adds the attribute if needed.
43 """
44 if not hasattr(callable_, "__name__"):
45 callable_.__name__ = callable_.__class__.__name__
48def _wrap_unary_errors(callable_):
49 """Map errors for Unary-Unary and Stream-Unary gRPC callables."""
50 _patch_callable_name(callable_)
52 @functools.wraps(callable_)
53 def error_remapped_callable(*args, **kwargs):
54 try:
55 return callable_(*args, **kwargs)
56 except grpc.RpcError as exc:
57 raise exceptions.from_grpc_error(exc) from exc
59 return error_remapped_callable
62class _StreamingResponseIterator(Generic[P], grpc.Call):
63 def __init__(self, wrapped, prefetch_first_result=True):
64 self._wrapped = wrapped
66 # This iterator is used in a retry context, and returned outside after init.
67 # gRPC will not throw an exception until the stream is consumed, so we need
68 # to retrieve the first result, in order to fail, in order to trigger a retry.
69 try:
70 if prefetch_first_result:
71 self._stored_first_result = next(self._wrapped)
72 except TypeError:
73 # It is possible the wrapped method isn't an iterable (a grpc.Call
74 # for instance). If this happens don't store the first result.
75 pass
76 except StopIteration:
77 # ignore stop iteration at this time. This should be handled outside of retry.
78 pass
80 def __iter__(self) -> Iterator[P]:
81 """This iterator is also an iterable that returns itself."""
82 return self
84 def __next__(self) -> P:
85 """Get the next response from the stream.
87 Returns:
88 protobuf.Message: A single response from the stream.
89 """
90 try:
91 if hasattr(self, "_stored_first_result"):
92 result = self._stored_first_result
93 del self._stored_first_result
94 return result
95 return next(self._wrapped)
96 except grpc.RpcError as exc:
97 # If the stream has already returned data, we cannot recover here.
98 raise exceptions.from_grpc_error(exc) from exc
100 # grpc.Call & grpc.RpcContext interface
102 def add_callback(self, callback):
103 return self._wrapped.add_callback(callback)
105 def cancel(self):
106 return self._wrapped.cancel()
108 def code(self):
109 return self._wrapped.code()
111 def details(self):
112 return self._wrapped.details()
114 def initial_metadata(self):
115 return self._wrapped.initial_metadata()
117 def is_active(self):
118 return self._wrapped.is_active()
120 def time_remaining(self):
121 return self._wrapped.time_remaining()
123 def trailing_metadata(self):
124 return self._wrapped.trailing_metadata()
127# public type alias denoting the return type of streaming gapic calls
128GrpcStream = _StreamingResponseIterator[P]
131def _wrap_stream_errors(callable_):
132 """Wrap errors for Unary-Stream and Stream-Stream gRPC callables.
134 The callables that return iterators require a bit more logic to re-map
135 errors when iterating. This wraps both the initial invocation and the
136 iterator of the return value to re-map errors.
137 """
138 _patch_callable_name(callable_)
140 @functools.wraps(callable_)
141 def error_remapped_callable(*args, **kwargs):
142 try:
143 result = callable_(*args, **kwargs)
144 # Auto-fetching the first result causes PubSub client's streaming pull
145 # to hang when re-opening the stream, thus we need examine the hacky
146 # hidden flag to see if pre-fetching is disabled.
147 # https://github.com/googleapis/python-pubsub/issues/93#issuecomment-630762257
148 prefetch_first = getattr(callable_, "_prefetch_first_result_", True)
149 return _StreamingResponseIterator(
150 result, prefetch_first_result=prefetch_first
151 )
152 except grpc.RpcError as exc:
153 raise exceptions.from_grpc_error(exc) from exc
155 return error_remapped_callable
158def wrap_errors(callable_):
159 """Wrap a gRPC callable and map :class:`grpc.RpcErrors` to friendly error
160 classes.
162 Errors raised by the gRPC callable are mapped to the appropriate
163 :class:`google.api_core.exceptions.GoogleAPICallError` subclasses.
164 The original `grpc.RpcError` (which is usually also a `grpc.Call`) is
165 available from the ``response`` property on the mapped exception. This
166 is useful for extracting metadata from the original error.
168 Args:
169 callable_ (Callable): A gRPC callable.
171 Returns:
172 Callable: The wrapped gRPC callable.
173 """
174 if isinstance(callable_, _STREAM_WRAP_CLASSES):
175 return _wrap_stream_errors(callable_)
176 else:
177 return _wrap_unary_errors(callable_)
180def _create_composite_credentials(
181 credentials=None,
182 credentials_file=None,
183 default_scopes=None,
184 scopes=None,
185 ssl_credentials=None,
186 quota_project_id=None,
187 default_host=None,
188):
189 """Create the composite credentials for secure channels.
191 Args:
192 credentials (google.auth.credentials.Credentials): The credentials. If
193 not specified, then this function will attempt to ascertain the
194 credentials from the environment using :func:`google.auth.default`.
195 credentials_file (str): Deprecated. A file with credentials that can be loaded with
196 :func:`google.auth.load_credentials_from_file`. This argument is
197 mutually exclusive with credentials. This argument will be
198 removed in the next major version of `google-api-core`.
200 .. warning::
201 Important: If you accept a credential configuration (credential JSON/File/Stream)
202 from an external source for authentication to Google Cloud Platform, you must
203 validate it before providing it to any Google API or client library. Providing an
204 unvalidated credential configuration to Google APIs or libraries can compromise
205 the security of your systems and data. For more information, refer to
206 `Validate credential configurations from external sources`_.
208 .. _Validate credential configurations from external sources:
210 https://cloud.google.com/docs/authentication/external/externally-sourced-credentials
211 default_scopes (Sequence[str]): A optional list of scopes needed for this
212 service. These are only used when credentials are not specified and
213 are passed to :func:`google.auth.default`.
214 scopes (Sequence[str]): A optional list of scopes needed for this
215 service. These are only used when credentials are not specified and
216 are passed to :func:`google.auth.default`.
217 ssl_credentials (grpc.ChannelCredentials): Optional SSL channel
218 credentials. This can be used to specify different certificates.
219 quota_project_id (str): An optional project to use for billing and quota.
220 default_host (str): The default endpoint. e.g., "pubsub.googleapis.com".
222 Returns:
223 grpc.ChannelCredentials: The composed channel credentials object.
225 Raises:
226 google.api_core.DuplicateCredentialArgs: If both a credentials object and credentials_file are passed.
227 """
228 if credentials_file is not None:
229 warnings.warn(general_helpers._CREDENTIALS_FILE_WARNING, DeprecationWarning)
231 if credentials and credentials_file:
232 raise exceptions.DuplicateCredentialArgs(
233 "'credentials' and 'credentials_file' are mutually exclusive."
234 )
236 if credentials_file:
237 credentials, _ = google.auth.load_credentials_from_file(
238 credentials_file, scopes=scopes, default_scopes=default_scopes
239 )
240 elif credentials:
241 credentials = google.auth.credentials.with_scopes_if_required(
242 credentials, scopes=scopes, default_scopes=default_scopes
243 )
244 else:
245 credentials, _ = google.auth.default(
246 scopes=scopes, default_scopes=default_scopes
247 )
249 if quota_project_id and isinstance(
250 credentials, google.auth.credentials.CredentialsWithQuotaProject
251 ):
252 credentials = credentials.with_quota_project(quota_project_id)
254 request = google.auth.transport.requests.Request()
256 # Create the metadata plugin for inserting the authorization header.
257 try:
258 metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin(
259 credentials,
260 request,
261 default_host=default_host,
262 suppress_metrics_header=True,
263 )
264 except TypeError:
265 # Support older versions of google-auth that do not accept suppress_metrics_header
266 metadata_plugin = google.auth.transport.grpc.AuthMetadataPlugin(
267 credentials,
268 request,
269 default_host=default_host,
270 )
272 # Create a set of grpc.CallCredentials using the metadata plugin.
273 google_auth_credentials = grpc.metadata_call_credentials(metadata_plugin)
275 # if `ssl_credentials` is set, use `grpc.composite_channel_credentials` instead of
276 # `grpc.compute_engine_channel_credentials` as the former supports passing
277 # `ssl_credentials` via `channel_credentials` which is needed for mTLS.
278 if ssl_credentials:
279 # Combine the ssl credentials and the authorization credentials.
280 # See https://grpc.github.io/grpc/python/grpc.html#grpc.composite_channel_credentials
281 return grpc.composite_channel_credentials(
282 ssl_credentials, google_auth_credentials
283 )
284 else:
285 # Use grpc.compute_engine_channel_credentials in order to support Direct Path.
286 # See https://grpc.github.io/grpc/python/grpc.html#grpc.compute_engine_channel_credentials
287 # TODO(https://github.com/googleapis/python-api-core/issues/598):
288 # Although `grpc.compute_engine_channel_credentials` returns channel credentials
289 # outside of a Google Compute Engine environment (GCE), we should determine if
290 # there is a way to reliably detect a GCE environment so that
291 # `grpc.compute_engine_channel_credentials` is not called outside of GCE.
292 return grpc.compute_engine_channel_credentials(google_auth_credentials)
295def create_channel(
296 target,
297 credentials=None,
298 scopes=None,
299 ssl_credentials=None,
300 credentials_file=None,
301 quota_project_id=None,
302 default_scopes=None,
303 default_host=None,
304 compression=None,
305 attempt_direct_path: Optional[bool] = False,
306 **kwargs,
307):
308 """Create a secure channel with credentials.
310 Args:
311 target (str): The target service address in the format 'hostname:port'.
312 credentials (google.auth.credentials.Credentials): The credentials. If
313 not specified, then this function will attempt to ascertain the
314 credentials from the environment using :func:`google.auth.default`.
315 scopes (Sequence[str]): A optional list of scopes needed for this
316 service. These are only used when credentials are not specified and
317 are passed to :func:`google.auth.default`.
318 ssl_credentials (grpc.ChannelCredentials): Optional SSL channel
319 credentials. This can be used to specify different certificates.
320 credentials_file (str): A file with credentials that can be loaded with
321 :func:`google.auth.load_credentials_from_file`. This argument is
322 mutually exclusive with credentials.
324 .. warning::
325 Important: If you accept a credential configuration (credential JSON/File/Stream)
326 from an external source for authentication to Google Cloud Platform, you must
327 validate it before providing it to any Google API or client library. Providing an
328 unvalidated credential configuration to Google APIs or libraries can compromise
329 the security of your systems and data. For more information, refer to
330 `Validate credential configurations from external sources`_.
332 .. _Validate credential configurations from external sources:
334 https://cloud.google.com/docs/authentication/external/externally-sourced-credentials
335 quota_project_id (str): An optional project to use for billing and quota.
336 default_scopes (Sequence[str]): Default scopes passed by a Google client
337 library. Use 'scopes' for user-defined scopes.
338 default_host (str): The default endpoint. e.g., "pubsub.googleapis.com".
339 compression (grpc.Compression): An optional value indicating the
340 compression method to be used over the lifetime of the channel.
341 attempt_direct_path (Optional[bool]): If set, Direct Path will be attempted
342 when the request is made. Direct Path is only available within a Google
343 Compute Engine (GCE) environment and provides a proxyless connection
344 which increases the available throughput, reduces latency, and increases
345 reliability. Note:
347 - This argument should only be set in a GCE environment and for Services
348 that are known to support Direct Path.
349 - If this argument is set outside of GCE, then this request will fail
350 unless the back-end service happens to have configured fall-back to DNS.
351 - If the request causes a `ServiceUnavailable` response, it is recommended
352 that the client repeat the request with `attempt_direct_path` set to
353 `False` as the Service may not support Direct Path.
354 - Using `ssl_credentials` with `attempt_direct_path` set to `True` will
355 result in `ValueError` as this combination is not yet supported.
357 kwargs: Additional key-word args passed to
358 :func:`grpc.secure_channel`.
360 Returns:
361 grpc.Channel: The created channel.
363 Raises:
364 google.api_core.DuplicateCredentialArgs: If both a credentials object and credentials_file are passed.
365 ValueError: If `ssl_credentials` is set and `attempt_direct_path` is set to `True`.
366 """
368 # If `ssl_credentials` is set and `attempt_direct_path` is set to `True`,
369 # raise ValueError as this is not yet supported.
370 # See https://github.com/googleapis/python-api-core/issues/590
371 if ssl_credentials and attempt_direct_path:
372 raise ValueError("Using ssl_credentials with Direct Path is not supported")
374 composite_credentials = _create_composite_credentials(
375 credentials=credentials,
376 credentials_file=credentials_file,
377 default_scopes=default_scopes,
378 scopes=scopes,
379 ssl_credentials=ssl_credentials,
380 quota_project_id=quota_project_id,
381 default_host=default_host,
382 )
384 if attempt_direct_path:
385 target = _modify_target_for_direct_path(target)
387 return grpc.secure_channel(
388 target, composite_credentials, compression=compression, **kwargs
389 )
392def _modify_target_for_direct_path(target: str) -> str:
393 """
394 Given a target, return a modified version which is compatible with Direct Path.
396 Args:
397 target (str): The target service address in the format 'hostname[:port]' or
398 'dns://hostname[:port]'.
400 Returns:
401 target (str): The target service address which is converted into a format compatible with Direct Path.
402 If the target contains `dns:///` or does not contain `:///`, the target will be converted in
403 a format compatible with Direct Path; otherwise the original target will be returned as the
404 original target may already denote Direct Path.
405 """
407 # A DNS prefix may be included with the target to indicate the endpoint is living in the Internet,
408 # outside of Google Cloud Platform.
409 dns_prefix = "dns:///"
410 # Remove "dns:///" if `attempt_direct_path` is set to True as
411 # the Direct Path prefix `google-c2p:///` will be used instead.
412 target = target.replace(dns_prefix, "")
414 direct_path_separator = ":///"
415 if direct_path_separator not in target:
416 target_without_port = target.split(":")[0]
417 # Modify the target to use Direct Path by adding the `google-c2p:///` prefix
418 target = f"google-c2p{direct_path_separator}{target_without_port}"
419 return target
422_MethodCall = collections.namedtuple(
423 "_MethodCall", ("request", "timeout", "metadata", "credentials", "compression")
424)
426_ChannelRequest = collections.namedtuple("_ChannelRequest", ("method", "request"))
429class _CallableStub(object):
430 """Stub for the grpc.*MultiCallable interfaces."""
432 def __init__(self, method, channel):
433 self._method = method
434 self._channel = channel
435 self.response = None
436 """Union[protobuf.Message, Callable[protobuf.Message], exception]:
437 The response to give when invoking this callable. If this is a
438 callable, it will be invoked with the request protobuf. If it's an
439 exception, the exception will be raised when this is invoked.
440 """
441 self.responses = None
442 """Iterator[
443 Union[protobuf.Message, Callable[protobuf.Message], exception]]:
444 An iterator of responses. If specified, self.response will be populated
445 on each invocation by calling ``next(self.responses)``."""
446 self.requests = []
447 """List[protobuf.Message]: All requests sent to this callable."""
448 self.calls = []
449 """List[Tuple]: All invocations of this callable. Each tuple is the
450 request, timeout, metadata, compression, and credentials."""
452 def __call__(
453 self, request, timeout=None, metadata=None, credentials=None, compression=None
454 ):
455 self._channel.requests.append(_ChannelRequest(self._method, request))
456 self.calls.append(
457 _MethodCall(request, timeout, metadata, credentials, compression)
458 )
459 self.requests.append(request)
461 response = self.response
462 if self.responses is not None:
463 if response is None:
464 response = next(self.responses)
465 else:
466 raise ValueError(
467 "{method}.response and {method}.responses are mutually "
468 "exclusive.".format(method=self._method)
469 )
471 if callable(response):
472 return response(request)
474 if isinstance(response, Exception):
475 raise response
477 if response is not None:
478 return response
480 raise ValueError('Method stub for "{}" has no response.'.format(self._method))
483def _simplify_method_name(method):
484 """Simplifies a gRPC method name.
486 When gRPC invokes the channel to create a callable, it gives a full
487 method name like "/google.pubsub.v1.Publisher/CreateTopic". This
488 returns just the name of the method, in this case "CreateTopic".
490 Args:
491 method (str): The name of the method.
493 Returns:
494 str: The simplified name of the method.
495 """
496 return method.rsplit("/", 1).pop()
499class ChannelStub(grpc.Channel):
500 """A testing stub for the grpc.Channel interface.
502 This can be used to test any client that eventually uses a gRPC channel
503 to communicate. By passing in a channel stub, you can configure which
504 responses are returned and track which requests are made.
506 For example:
508 .. code-block:: python
510 channel_stub = grpc_helpers.ChannelStub()
511 client = FooClient(channel=channel_stub)
513 channel_stub.GetFoo.response = foo_pb2.Foo(name='bar')
515 foo = client.get_foo(labels=['baz'])
517 assert foo.name == 'bar'
518 assert channel_stub.GetFoo.requests[0].labels = ['baz']
520 Each method on the stub can be accessed and configured on the channel.
521 Here's some examples of various configurations:
523 .. code-block:: python
525 # Return a basic response:
527 channel_stub.GetFoo.response = foo_pb2.Foo(name='bar')
528 assert client.get_foo().name == 'bar'
530 # Raise an exception:
531 channel_stub.GetFoo.response = NotFound('...')
533 with pytest.raises(NotFound):
534 client.get_foo()
536 # Use a sequence of responses:
537 channel_stub.GetFoo.responses = iter([
538 foo_pb2.Foo(name='bar'),
539 foo_pb2.Foo(name='baz'),
540 ])
542 assert client.get_foo().name == 'bar'
543 assert client.get_foo().name == 'baz'
545 # Use a callable
547 def on_get_foo(request):
548 return foo_pb2.Foo(name='bar' + request.id)
550 channel_stub.GetFoo.response = on_get_foo
552 assert client.get_foo(id='123').name == 'bar123'
553 """
555 def __init__(self, responses=[]):
556 self.requests = []
557 """Sequence[Tuple[str, protobuf.Message]]: A list of all requests made
558 on this channel in order. The tuple is of method name, request
559 message."""
560 self._method_stubs = {}
562 def _stub_for_method(self, method):
563 method = _simplify_method_name(method)
564 self._method_stubs[method] = _CallableStub(method, self)
565 return self._method_stubs[method]
567 def __getattr__(self, key):
568 try:
569 return self._method_stubs[key]
570 except KeyError:
571 raise AttributeError
573 def unary_unary(
574 self,
575 method,
576 request_serializer=None,
577 response_deserializer=None,
578 _registered_method=False,
579 ):
580 """grpc.Channel.unary_unary implementation."""
581 return self._stub_for_method(method)
583 def unary_stream(
584 self,
585 method,
586 request_serializer=None,
587 response_deserializer=None,
588 _registered_method=False,
589 ):
590 """grpc.Channel.unary_stream implementation."""
591 return self._stub_for_method(method)
593 def stream_unary(
594 self,
595 method,
596 request_serializer=None,
597 response_deserializer=None,
598 _registered_method=False,
599 ):
600 """grpc.Channel.stream_unary implementation."""
601 return self._stub_for_method(method)
603 def stream_stream(
604 self,
605 method,
606 request_serializer=None,
607 response_deserializer=None,
608 _registered_method=False,
609 ):
610 """grpc.Channel.stream_stream implementation."""
611 return self._stub_for_method(method)
613 def subscribe(self, callback, try_to_connect=False):
614 """grpc.Channel.subscribe implementation."""
615 pass
617 def unsubscribe(self, callback):
618 """grpc.Channel.unsubscribe implementation."""
619 pass
621 def close(self):
622 """grpc.Channel.close implementation."""
623 pass