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