1# Copyright 2019 gRPC authors.
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"""Interceptors implementation of gRPC Asyncio Python."""
15from __future__ import annotations
16
17from abc import ABCMeta
18from abc import abstractmethod
19import asyncio
20import collections
21import functools
22from typing import (
23 AsyncIterable,
24 AsyncIterator,
25 Awaitable,
26 Callable,
27 List,
28 Optional,
29 Sequence,
30 Union,
31)
32
33import grpc
34from grpc._cython import cygrpc
35
36from . import _base_call
37from ._call import AioRpcError
38from ._call import StreamStreamCall
39from ._call import StreamUnaryCall
40from ._call import UnaryStreamCall
41from ._call import UnaryUnaryCall
42from ._call import _API_STYLE_ERROR
43from ._call import _RPC_ALREADY_FINISHED_DETAILS
44from ._call import _RPC_HALF_CLOSED_DETAILS
45from ._metadata import Metadata
46from ._typing import DeserializingFunction
47from ._typing import DoneCallbackType
48from ._typing import EOFType
49from ._typing import RequestIterableType
50from ._typing import RequestType
51from ._typing import ResponseIterableType
52from ._typing import ResponseType
53from ._typing import SerializingFunction
54from ._utils import _timeout_to_deadline
55
56_LOCAL_CANCELLATION_DETAILS = "Locally cancelled by application!"
57
58
59class ServerInterceptor(metaclass=ABCMeta):
60 """Affords intercepting incoming RPCs on the service-side.
61
62 This is an EXPERIMENTAL API.
63 """
64
65 @abstractmethod
66 async def intercept_service(
67 self,
68 continuation: Callable[
69 [grpc.HandlerCallDetails], Awaitable[grpc.RpcMethodHandler]
70 ],
71 handler_call_details: grpc.HandlerCallDetails,
72 ) -> grpc.RpcMethodHandler:
73 """Intercepts incoming RPCs before handing them over to a handler.
74
75 State can be passed from an interceptor to downstream interceptors
76 via contextvars. The first interceptor is called from an empty
77 contextvars.Context, and the same Context is used for downstream
78 interceptors and for the final handler call. Note that there are no
79 guarantees that interceptors and handlers will be called from the
80 same thread.
81
82 Args:
83 continuation: A function that takes a HandlerCallDetails and
84 proceeds to invoke the next interceptor in the chain, if any,
85 or the RPC handler lookup logic, with the call details passed
86 as an argument, and returns an RpcMethodHandler instance if
87 the RPC is considered serviced, or None otherwise.
88 handler_call_details: A HandlerCallDetails describing the RPC.
89
90 Returns:
91 An RpcMethodHandler with which the RPC may be serviced if the
92 interceptor chooses to service this RPC, or None otherwise.
93 """
94
95
96class ClientCallDetails(
97 collections.namedtuple(
98 "ClientCallDetails",
99 ("method", "timeout", "metadata", "credentials", "wait_for_ready"),
100 ),
101 grpc.ClientCallDetails,
102):
103 """Describes an RPC to be invoked.
104
105 This is an EXPERIMENTAL API.
106
107 Args:
108 method: The method name of the RPC.
109 timeout: An optional duration of time in seconds to allow for the RPC.
110 metadata: Optional metadata to be transmitted to the service-side of
111 the RPC.
112 credentials: An optional CallCredentials for the RPC.
113 wait_for_ready: An optional flag to enable :term:`wait_for_ready` mechanism.
114 """
115
116 method: bytes
117 timeout: Optional[float]
118 metadata: Optional[Metadata]
119 credentials: Optional[grpc.CallCredentials]
120 wait_for_ready: Optional[bool]
121
122
123class ClientInterceptor(metaclass=ABCMeta):
124 """Base class used for all Aio Client Interceptor classes"""
125
126
127class UnaryUnaryClientInterceptor(ClientInterceptor, metaclass=ABCMeta):
128 """Affords intercepting unary-unary invocations."""
129
130 @abstractmethod
131 async def intercept_unary_unary(
132 self,
133 continuation: Callable[
134 [ClientCallDetails, RequestType], UnaryUnaryCall
135 ],
136 client_call_details: ClientCallDetails,
137 request: RequestType,
138 ) -> Union[UnaryUnaryCall, ResponseType]:
139 """Intercepts a unary-unary invocation asynchronously.
140
141 Args:
142 continuation: A coroutine that proceeds with the invocation by
143 executing the next interceptor in the chain or invoking the
144 actual RPC on the underlying Channel. It is the interceptor's
145 responsibility to call it if it decides to move the RPC forward.
146 The interceptor can use
147 `call = await continuation(client_call_details, request)`
148 to continue with the RPC. `continuation` returns the call to the
149 RPC.
150 client_call_details: A ClientCallDetails object describing the
151 outgoing RPC.
152 request: The request value for the RPC.
153
154 Returns:
155 An object with the RPC response.
156
157 Raises:
158 AioRpcError: Indicating that the RPC terminated with non-OK status.
159 asyncio.CancelledError: Indicating that the RPC was canceled.
160 """
161
162
163class UnaryStreamClientInterceptor(ClientInterceptor, metaclass=ABCMeta):
164 """Affords intercepting unary-stream invocations."""
165
166 @abstractmethod
167 async def intercept_unary_stream(
168 self,
169 continuation: Callable[
170 [ClientCallDetails, RequestType], UnaryStreamCall
171 ],
172 client_call_details: ClientCallDetails,
173 request: RequestType,
174 ) -> Union[ResponseIterableType, UnaryStreamCall]:
175 """Intercepts a unary-stream invocation asynchronously.
176
177 The function could return the call object or an asynchronous
178 iterator, in case of being an asyncrhonous iterator this will
179 become the source of the reads done by the caller.
180
181 Args:
182 continuation: A coroutine that proceeds with the invocation by
183 executing the next interceptor in the chain or invoking the
184 actual RPC on the underlying Channel. It is the interceptor's
185 responsibility to call it if it decides to move the RPC forward.
186 The interceptor can use
187 `call = await continuation(client_call_details, request)`
188 to continue with the RPC. `continuation` returns the call to the
189 RPC.
190 client_call_details: A ClientCallDetails object describing the
191 outgoing RPC.
192 request: The request value for the RPC.
193
194 Returns:
195 The RPC Call or an asynchronous iterator.
196
197 Raises:
198 AioRpcError: Indicating that the RPC terminated with non-OK status.
199 asyncio.CancelledError: Indicating that the RPC was canceled.
200 """
201
202
203class StreamUnaryClientInterceptor(ClientInterceptor, metaclass=ABCMeta):
204 """Affords intercepting stream-unary invocations."""
205
206 @abstractmethod
207 async def intercept_stream_unary(
208 self,
209 continuation: Callable[
210 [ClientCallDetails, RequestType], StreamUnaryCall
211 ],
212 client_call_details: ClientCallDetails,
213 request_iterator: RequestIterableType,
214 ) -> StreamUnaryCall:
215 """Intercepts a stream-unary invocation asynchronously.
216
217 Within the interceptor the usage of the call methods like `write` or
218 even awaiting the call should be done carefully, since the caller
219 could be expecting an untouched call, for example for start writing
220 messages to it.
221
222 Args:
223 continuation: A coroutine that proceeds with the invocation by
224 executing the next interceptor in the chain or invoking the
225 actual RPC on the underlying Channel. It is the interceptor's
226 responsibility to call it if it decides to move the RPC forward.
227 The interceptor can use
228 `call = await continuation(client_call_details, request_iterator)`
229 to continue with the RPC. `continuation` returns the call to the
230 RPC.
231 client_call_details: A ClientCallDetails object describing the
232 outgoing RPC.
233 request_iterator: The request iterator that will produce requests
234 for the RPC.
235
236 Returns:
237 The RPC Call.
238
239 Raises:
240 AioRpcError: Indicating that the RPC terminated with non-OK status.
241 asyncio.CancelledError: Indicating that the RPC was canceled.
242 """
243
244
245class StreamStreamClientInterceptor(ClientInterceptor, metaclass=ABCMeta):
246 """Affords intercepting stream-stream invocations."""
247
248 @abstractmethod
249 async def intercept_stream_stream(
250 self,
251 continuation: Callable[
252 [ClientCallDetails, RequestType], StreamStreamCall
253 ],
254 client_call_details: ClientCallDetails,
255 request_iterator: RequestIterableType,
256 ) -> Union[ResponseIterableType, StreamStreamCall]:
257 """Intercepts a stream-stream invocation asynchronously.
258
259 Within the interceptor the usage of the call methods like `write` or
260 even awaiting the call should be done carefully, since the caller
261 could be expecting an untouched call, for example for start writing
262 messages to it.
263
264 The function could return the call object or an asynchronous
265 iterator, in case of being an asyncrhonous iterator this will
266 become the source of the reads done by the caller.
267
268 Args:
269 continuation: A coroutine that proceeds with the invocation by
270 executing the next interceptor in the chain or invoking the
271 actual RPC on the underlying Channel. It is the interceptor's
272 responsibility to call it if it decides to move the RPC forward.
273 The interceptor can use
274 `call = await continuation(client_call_details, request_iterator)`
275 to continue with the RPC. `continuation` returns the call to the
276 RPC.
277 client_call_details: A ClientCallDetails object describing the
278 outgoing RPC.
279 request_iterator: The request iterator that will produce requests
280 for the RPC.
281
282 Returns:
283 The RPC Call or an asynchronous iterator.
284
285 Raises:
286 AioRpcError: Indicating that the RPC terminated with non-OK status.
287 asyncio.CancelledError: Indicating that the RPC was canceled.
288 """
289
290
291class InterceptedCall:
292 """Base implementation for all intercepted call arities.
293
294 Interceptors might have some work to do before the RPC invocation with
295 the capacity of changing the invocation parameters, and some work to do
296 after the RPC invocation with the capacity for accessing to the wrapped
297 `UnaryUnaryCall`.
298
299 It handles also early and later cancellations, when the RPC has not even
300 started and the execution is still held by the interceptors or when the
301 RPC has finished but again the execution is still held by the interceptors.
302
303 Once the RPC is finally executed, all methods are finally done against the
304 intercepted call, being at the same time the same call returned to the
305 interceptors.
306
307 As a base class for all of the interceptors implements the logic around
308 final status, metadata and cancellation.
309 """
310
311 _interceptors_task: asyncio.Task
312 _pending_add_done_callbacks: Sequence[DoneCallbackType]
313
314 def __init__(self, interceptors_task: asyncio.Task) -> None:
315 self._interceptors_task = interceptors_task
316 self._pending_add_done_callbacks = []
317 self._interceptors_task.add_done_callback(
318 self._fire_or_add_pending_done_callbacks
319 )
320
321 def __del__(self):
322 self.cancel()
323
324 def _fire_or_add_pending_done_callbacks(
325 self, interceptors_task: asyncio.Task
326 ) -> None:
327 if not self._pending_add_done_callbacks:
328 return
329
330 call_completed = False
331
332 if interceptors_task.cancelled() or (
333 interceptors_task.done()
334 and interceptors_task.exception() is not None
335 ):
336 call_completed = True
337 else:
338 call = interceptors_task.result()
339 if call.done():
340 call_completed = True
341
342 if call_completed:
343 for callback in self._pending_add_done_callbacks:
344 callback(self)
345 else:
346 for callback in self._pending_add_done_callbacks:
347 callback = functools.partial(
348 self._wrap_add_done_callback, callback
349 )
350 call.add_done_callback(callback)
351
352 self._pending_add_done_callbacks = []
353
354 def _wrap_add_done_callback(
355 self, callback: DoneCallbackType, unused_call: _base_call.Call
356 ) -> None:
357 callback(self)
358
359 def cancel(self) -> bool:
360 if self._interceptors_task.cancelled():
361 return False
362
363 if not self._interceptors_task.done():
364 # There is no yet the intercepted call available,
365 # Trying to cancel it by using the generic Asyncio
366 # cancellation method.
367 return self._interceptors_task.cancel()
368
369 if self._interceptors_task.exception() is not None:
370 return False
371
372 call = self._interceptors_task.result()
373
374 return call.cancel()
375
376 def cancelled(self) -> bool:
377 if self._interceptors_task.cancelled():
378 return True
379 if not self._interceptors_task.done():
380 return False
381
382 exc = self._interceptors_task.exception()
383 if exc is not None:
384 if isinstance(exc, AioRpcError):
385 return exc.code() == grpc.StatusCode.CANCELLED
386 return False
387
388 call = self._interceptors_task.result()
389
390 return call.cancelled()
391
392 def done(self) -> bool:
393 if self._interceptors_task.cancelled():
394 return True
395 if not self._interceptors_task.done():
396 return False
397 if self._interceptors_task.exception() is not None:
398 return True
399
400 call = self._interceptors_task.result()
401
402 return call.done()
403
404 def add_done_callback(self, callback: DoneCallbackType) -> None:
405 if not self._interceptors_task.done():
406 self._pending_add_done_callbacks.append(callback)
407 return
408
409 if (
410 self._interceptors_task.cancelled()
411 or self._interceptors_task.exception() is not None
412 ):
413 callback(self)
414 return
415
416 call = self._interceptors_task.result()
417
418 if call.done():
419 callback(self)
420 else:
421 callback = functools.partial(self._wrap_add_done_callback, callback)
422 call.add_done_callback(callback)
423
424 def time_remaining(self) -> Optional[float]:
425 raise NotImplementedError()
426
427 async def initial_metadata(self) -> Optional[Metadata]:
428 try:
429 call = await self._interceptors_task
430 except AioRpcError as err:
431 return err.initial_metadata()
432 except asyncio.CancelledError:
433 return None
434
435 return await call.initial_metadata()
436
437 async def trailing_metadata(self) -> Optional[Metadata]:
438 try:
439 call = await self._interceptors_task
440 except AioRpcError as err:
441 return err.trailing_metadata()
442 except asyncio.CancelledError:
443 return None
444
445 return await call.trailing_metadata()
446
447 async def code(self) -> grpc.StatusCode:
448 try:
449 call = await self._interceptors_task
450 except AioRpcError as err:
451 return err.code()
452 except asyncio.CancelledError:
453 return grpc.StatusCode.CANCELLED
454
455 return await call.code()
456
457 async def details(self) -> str:
458 try:
459 call = await self._interceptors_task
460 except AioRpcError as err:
461 return err.details()
462 except asyncio.CancelledError:
463 return _LOCAL_CANCELLATION_DETAILS
464
465 return await call.details()
466
467 async def debug_error_string(self) -> Optional[str]:
468 try:
469 call = await self._interceptors_task
470 except AioRpcError as err:
471 return err.debug_error_string()
472 except asyncio.CancelledError:
473 return ""
474
475 return await call.debug_error_string()
476
477 async def wait_for_connection(self) -> None:
478 call = await self._interceptors_task
479 return await call.wait_for_connection()
480
481
482class _InterceptedUnaryResponseMixin:
483 def __await__(self):
484 call = yield from self._interceptors_task.__await__()
485 response = yield from call.__await__()
486 return response
487
488
489class _InterceptedStreamResponseMixin:
490 _response_aiter: Optional[AsyncIterable[ResponseType]]
491
492 def _init_stream_response_mixin(self) -> None:
493 # Is initialized later, otherwise if the iterator is not finally
494 # consumed a logging warning is emitted by Asyncio.
495 self._response_aiter = None
496
497 async def _wait_for_interceptor_task_response_iterator(
498 self,
499 ) -> ResponseType:
500 call = await self._interceptors_task
501 async for response in call:
502 yield response
503
504 def __aiter__(self) -> AsyncIterator[ResponseType]:
505 if self._response_aiter is None:
506 self._response_aiter = (
507 self._wait_for_interceptor_task_response_iterator()
508 )
509 return self._response_aiter
510
511 async def read(self) -> Union[EOFType, ResponseType]:
512 if self._response_aiter is None:
513 self._response_aiter = (
514 self._wait_for_interceptor_task_response_iterator()
515 )
516 try:
517 return await self._response_aiter.asend(None)
518 except StopAsyncIteration:
519 return cygrpc.EOF
520
521
522class _InterceptedStreamRequestMixin:
523 _write_to_iterator_async_gen: Optional[AsyncIterable[RequestType]]
524 _write_to_iterator_queue: Optional[asyncio.Queue]
525 _status_code_task: Optional[asyncio.Task]
526
527 _FINISH_ITERATOR_SENTINEL = object()
528
529 def _init_stream_request_mixin(
530 self, request_iterator: Optional[RequestIterableType]
531 ) -> RequestIterableType:
532 if request_iterator is None:
533 # We provide our own request iterator which is a proxy
534 # of the futures writes that will be done by the caller.
535 self._write_to_iterator_queue = asyncio.Queue(maxsize=1)
536 self._write_to_iterator_async_gen = (
537 self._proxy_writes_as_request_iterator()
538 )
539 self._status_code_task = None
540 request_iterator = self._write_to_iterator_async_gen
541 else:
542 self._write_to_iterator_queue = None
543
544 return request_iterator
545
546 async def _proxy_writes_as_request_iterator(self):
547 await self._interceptors_task
548
549 while True:
550 value = await self._write_to_iterator_queue.get()
551 if (
552 value
553 is _InterceptedStreamRequestMixin._FINISH_ITERATOR_SENTINEL
554 ):
555 break
556 yield value
557
558 async def _write_to_iterator_queue_interruptible(
559 self,
560 request: RequestType,
561 call: _base_call.Call,
562 ):
563 # Write the specified 'request' to the request iterator queue using the
564 # specified 'call' to allow for interruption of the write in the case
565 # of abrupt termination of the call.
566 if self._status_code_task is None:
567 self._status_code_task = self._loop.create_task(call.code())
568
569 await asyncio.wait(
570 (
571 self._loop.create_task(
572 self._write_to_iterator_queue.put(request)
573 ),
574 self._status_code_task,
575 ),
576 return_when=asyncio.FIRST_COMPLETED,
577 )
578
579 async def write(self, request: RequestType) -> None:
580 # If no queue was created it means that requests
581 # should be expected through an iterators provided
582 # by the caller.
583 if self._write_to_iterator_queue is None:
584 raise cygrpc.UsageError(_API_STYLE_ERROR)
585
586 try:
587 call = await self._interceptors_task
588 except (asyncio.CancelledError, AioRpcError):
589 raise asyncio.InvalidStateError(_RPC_ALREADY_FINISHED_DETAILS)
590
591 if call.done():
592 raise asyncio.InvalidStateError(_RPC_ALREADY_FINISHED_DETAILS)
593 if call._done_writing_flag:
594 raise asyncio.InvalidStateError(_RPC_HALF_CLOSED_DETAILS)
595
596 await self._write_to_iterator_queue_interruptible(request, call)
597
598 if call.done():
599 raise asyncio.InvalidStateError(_RPC_ALREADY_FINISHED_DETAILS)
600
601 async def done_writing(self) -> None:
602 """Signal peer that client is done writing.
603
604 This method is idempotent.
605 """
606 # If no queue was created it means that requests
607 # should be expected through an iterators provided
608 # by the caller.
609 if self._write_to_iterator_queue is None:
610 raise cygrpc.UsageError(_API_STYLE_ERROR)
611
612 try:
613 call = await self._interceptors_task
614 except asyncio.CancelledError:
615 raise asyncio.InvalidStateError(_RPC_ALREADY_FINISHED_DETAILS)
616
617 await self._write_to_iterator_queue_interruptible(
618 _InterceptedStreamRequestMixin._FINISH_ITERATOR_SENTINEL, call
619 )
620
621
622class InterceptedUnaryUnaryCall(
623 _InterceptedUnaryResponseMixin, InterceptedCall, _base_call.UnaryUnaryCall
624):
625 """Used for running a `UnaryUnaryCall` wrapped by interceptors.
626
627 For the `__await__` method is it is proxied to the intercepted call only when
628 the interceptor task is finished.
629 """
630
631 _loop: asyncio.AbstractEventLoop
632 _channel: cygrpc.AioChannel
633
634 # pylint: disable=too-many-arguments
635 def __init__(
636 self,
637 interceptors: Sequence[UnaryUnaryClientInterceptor],
638 request: RequestType,
639 timeout: Optional[float],
640 metadata: Metadata,
641 credentials: Optional[grpc.CallCredentials],
642 wait_for_ready: Optional[bool],
643 channel: cygrpc.AioChannel,
644 method: bytes,
645 request_serializer: Optional[SerializingFunction],
646 response_deserializer: Optional[DeserializingFunction],
647 loop: asyncio.AbstractEventLoop,
648 ) -> None:
649 self._loop = loop
650 self._channel = channel
651 interceptors_task = loop.create_task(
652 self._invoke(
653 interceptors,
654 method,
655 timeout,
656 metadata,
657 credentials,
658 wait_for_ready,
659 request,
660 request_serializer,
661 response_deserializer,
662 )
663 )
664 super().__init__(interceptors_task)
665
666 # pylint: disable=too-many-arguments
667 async def _invoke(
668 self,
669 interceptors: Sequence[UnaryUnaryClientInterceptor],
670 method: bytes,
671 timeout: Optional[float],
672 metadata: Optional[Metadata],
673 credentials: Optional[grpc.CallCredentials],
674 wait_for_ready: Optional[bool],
675 request: RequestType,
676 request_serializer: Optional[SerializingFunction],
677 response_deserializer: Optional[DeserializingFunction],
678 ) -> Union[UnaryUnaryCall, UnaryUnaryCallResponse]:
679 """Run the RPC call wrapped in interceptors"""
680
681 async def _run_interceptor(
682 interceptors: List[UnaryUnaryClientInterceptor],
683 client_call_details: ClientCallDetails,
684 request: RequestType,
685 ) -> Union[UnaryUnaryCall, UnaryUnaryCallResponse]:
686 if interceptors:
687 continuation = functools.partial(
688 _run_interceptor, interceptors[1:]
689 )
690 call_or_response = await interceptors[0].intercept_unary_unary(
691 continuation, client_call_details, request
692 )
693
694 if isinstance(call_or_response, _base_call.UnaryUnaryCall):
695 return call_or_response
696 return UnaryUnaryCallResponse(call_or_response)
697
698 return UnaryUnaryCall(
699 request,
700 _timeout_to_deadline(client_call_details.timeout),
701 client_call_details.metadata,
702 client_call_details.credentials,
703 client_call_details.wait_for_ready,
704 self._channel,
705 client_call_details.method,
706 request_serializer,
707 response_deserializer,
708 self._loop,
709 )
710
711 client_call_details = ClientCallDetails(
712 method, timeout, metadata, credentials, wait_for_ready
713 )
714 return await _run_interceptor(
715 list(interceptors), client_call_details, request
716 )
717
718 def time_remaining(self) -> Optional[float]:
719 raise NotImplementedError()
720
721
722class InterceptedUnaryStreamCall(
723 _InterceptedStreamResponseMixin, InterceptedCall, _base_call.UnaryStreamCall
724):
725 """Used for running a `UnaryStreamCall` wrapped by interceptors."""
726
727 _loop: asyncio.AbstractEventLoop
728 _channel: cygrpc.AioChannel
729 _last_returned_call_from_interceptors = Optional[_base_call.UnaryStreamCall]
730
731 # pylint: disable=too-many-arguments
732 def __init__(
733 self,
734 interceptors: Sequence[UnaryStreamClientInterceptor],
735 request: RequestType,
736 timeout: Optional[float],
737 metadata: Metadata,
738 credentials: Optional[grpc.CallCredentials],
739 wait_for_ready: Optional[bool],
740 channel: cygrpc.AioChannel,
741 method: bytes,
742 request_serializer: Optional[SerializingFunction],
743 response_deserializer: Optional[DeserializingFunction],
744 loop: asyncio.AbstractEventLoop,
745 ) -> None:
746 self._loop = loop
747 self._channel = channel
748 self._init_stream_response_mixin()
749 self._last_returned_call_from_interceptors = None
750 interceptors_task = loop.create_task(
751 self._invoke(
752 interceptors,
753 method,
754 timeout,
755 metadata,
756 credentials,
757 wait_for_ready,
758 request,
759 request_serializer,
760 response_deserializer,
761 )
762 )
763 super().__init__(interceptors_task)
764
765 # pylint: disable=too-many-arguments
766 async def _invoke(
767 self,
768 interceptors: Sequence[UnaryStreamClientInterceptor],
769 method: bytes,
770 timeout: Optional[float],
771 metadata: Optional[Metadata],
772 credentials: Optional[grpc.CallCredentials],
773 wait_for_ready: Optional[bool],
774 request: RequestType,
775 request_serializer: Optional[SerializingFunction],
776 response_deserializer: Optional[DeserializingFunction],
777 ) -> Union[UnaryStreamCall, UnaryStreamCallResponseIterator]:
778 """Run the RPC call wrapped in interceptors"""
779
780 async def _run_interceptor(
781 interceptors: List[UnaryStreamClientInterceptor],
782 client_call_details: ClientCallDetails,
783 request: RequestType,
784 ) -> Union[UnaryStreamCall, UnaryStreamCallResponseIterator]:
785 if interceptors:
786 continuation = functools.partial(
787 _run_interceptor, interceptors[1:]
788 )
789
790 call_or_response_iterator = await interceptors[
791 0
792 ].intercept_unary_stream(
793 continuation, client_call_details, request
794 )
795
796 if isinstance(
797 call_or_response_iterator, _base_call.UnaryStreamCall
798 ):
799 self._last_returned_call_from_interceptors = (
800 call_or_response_iterator
801 )
802 else:
803 self._last_returned_call_from_interceptors = (
804 UnaryStreamCallResponseIterator(
805 self._last_returned_call_from_interceptors,
806 call_or_response_iterator,
807 )
808 )
809 return self._last_returned_call_from_interceptors
810 self._last_returned_call_from_interceptors = UnaryStreamCall(
811 request,
812 _timeout_to_deadline(client_call_details.timeout),
813 client_call_details.metadata,
814 client_call_details.credentials,
815 client_call_details.wait_for_ready,
816 self._channel,
817 client_call_details.method,
818 request_serializer,
819 response_deserializer,
820 self._loop,
821 )
822
823 return self._last_returned_call_from_interceptors
824
825 client_call_details = ClientCallDetails(
826 method, timeout, metadata, credentials, wait_for_ready
827 )
828 return await _run_interceptor(
829 list(interceptors), client_call_details, request
830 )
831
832 def time_remaining(self) -> Optional[float]:
833 raise NotImplementedError()
834
835
836class InterceptedStreamUnaryCall(
837 _InterceptedUnaryResponseMixin,
838 _InterceptedStreamRequestMixin,
839 InterceptedCall,
840 _base_call.StreamUnaryCall,
841):
842 """Used for running a `StreamUnaryCall` wrapped by interceptors.
843
844 For the `__await__` method is it is proxied to the intercepted call only when
845 the interceptor task is finished.
846 """
847
848 _loop: asyncio.AbstractEventLoop
849 _channel: cygrpc.AioChannel
850
851 # pylint: disable=too-many-arguments
852 def __init__(
853 self,
854 interceptors: Sequence[StreamUnaryClientInterceptor],
855 request_iterator: Optional[RequestIterableType],
856 timeout: Optional[float],
857 metadata: Metadata,
858 credentials: Optional[grpc.CallCredentials],
859 wait_for_ready: Optional[bool],
860 channel: cygrpc.AioChannel,
861 method: bytes,
862 request_serializer: Optional[SerializingFunction],
863 response_deserializer: Optional[DeserializingFunction],
864 loop: asyncio.AbstractEventLoop,
865 ) -> None:
866 self._loop = loop
867 self._channel = channel
868 request_iterator = self._init_stream_request_mixin(request_iterator)
869 interceptors_task = loop.create_task(
870 self._invoke(
871 interceptors,
872 method,
873 timeout,
874 metadata,
875 credentials,
876 wait_for_ready,
877 request_iterator,
878 request_serializer,
879 response_deserializer,
880 )
881 )
882 super().__init__(interceptors_task)
883
884 # pylint: disable=too-many-arguments
885 async def _invoke(
886 self,
887 interceptors: Sequence[StreamUnaryClientInterceptor],
888 method: bytes,
889 timeout: Optional[float],
890 metadata: Optional[Metadata],
891 credentials: Optional[grpc.CallCredentials],
892 wait_for_ready: Optional[bool],
893 request_iterator: RequestIterableType,
894 request_serializer: Optional[SerializingFunction],
895 response_deserializer: Optional[DeserializingFunction],
896 ) -> StreamUnaryCall:
897 """Run the RPC call wrapped in interceptors"""
898
899 async def _run_interceptor(
900 interceptors: Sequence[StreamUnaryClientInterceptor],
901 client_call_details: ClientCallDetails,
902 request_iterator: RequestIterableType,
903 ) -> _base_call.StreamUnaryCall:
904 if interceptors:
905 continuation = functools.partial(
906 _run_interceptor, interceptors[1:]
907 )
908
909 return await interceptors[0].intercept_stream_unary(
910 continuation, client_call_details, request_iterator
911 )
912 return StreamUnaryCall(
913 request_iterator,
914 _timeout_to_deadline(client_call_details.timeout),
915 client_call_details.metadata,
916 client_call_details.credentials,
917 client_call_details.wait_for_ready,
918 self._channel,
919 client_call_details.method,
920 request_serializer,
921 response_deserializer,
922 self._loop,
923 )
924
925 client_call_details = ClientCallDetails(
926 method, timeout, metadata, credentials, wait_for_ready
927 )
928 return await _run_interceptor(
929 list(interceptors), client_call_details, request_iterator
930 )
931
932 def time_remaining(self) -> Optional[float]:
933 raise NotImplementedError()
934
935
936class InterceptedStreamStreamCall(
937 _InterceptedStreamResponseMixin,
938 _InterceptedStreamRequestMixin,
939 InterceptedCall,
940 _base_call.StreamStreamCall,
941):
942 """Used for running a `StreamStreamCall` wrapped by interceptors."""
943
944 _loop: asyncio.AbstractEventLoop
945 _channel: cygrpc.AioChannel
946 _last_returned_call_from_interceptors = Optional[
947 _base_call.StreamStreamCall
948 ]
949
950 # pylint: disable=too-many-arguments
951 def __init__(
952 self,
953 interceptors: Sequence[StreamStreamClientInterceptor],
954 request_iterator: Optional[RequestIterableType],
955 timeout: Optional[float],
956 metadata: Metadata,
957 credentials: Optional[grpc.CallCredentials],
958 wait_for_ready: Optional[bool],
959 channel: cygrpc.AioChannel,
960 method: bytes,
961 request_serializer: Optional[SerializingFunction],
962 response_deserializer: Optional[DeserializingFunction],
963 loop: asyncio.AbstractEventLoop,
964 ) -> None:
965 self._loop = loop
966 self._channel = channel
967 self._init_stream_response_mixin()
968 request_iterator = self._init_stream_request_mixin(request_iterator)
969 self._last_returned_call_from_interceptors = None
970 interceptors_task = loop.create_task(
971 self._invoke(
972 interceptors,
973 method,
974 timeout,
975 metadata,
976 credentials,
977 wait_for_ready,
978 request_iterator,
979 request_serializer,
980 response_deserializer,
981 )
982 )
983 super().__init__(interceptors_task)
984
985 # pylint: disable=too-many-arguments
986 async def _invoke(
987 self,
988 interceptors: Sequence[StreamStreamClientInterceptor],
989 method: bytes,
990 timeout: Optional[float],
991 metadata: Optional[Metadata],
992 credentials: Optional[grpc.CallCredentials],
993 wait_for_ready: Optional[bool],
994 request_iterator: RequestIterableType,
995 request_serializer: Optional[SerializingFunction],
996 response_deserializer: Optional[DeserializingFunction],
997 ) -> Union[StreamStreamCall, StreamStreamCallResponseIterator]:
998 """Run the RPC call wrapped in interceptors"""
999
1000 async def _run_interceptor(
1001 interceptors: List[StreamStreamClientInterceptor],
1002 client_call_details: ClientCallDetails,
1003 request_iterator: RequestIterableType,
1004 ) -> Union[StreamStreamCall, StreamStreamCallResponseIterator]:
1005 if interceptors:
1006 continuation = functools.partial(
1007 _run_interceptor, interceptors[1:]
1008 )
1009
1010 call_or_response_iterator = await interceptors[
1011 0
1012 ].intercept_stream_stream(
1013 continuation, client_call_details, request_iterator
1014 )
1015
1016 if isinstance(
1017 call_or_response_iterator, _base_call.StreamStreamCall
1018 ):
1019 self._last_returned_call_from_interceptors = (
1020 call_or_response_iterator
1021 )
1022 else:
1023 self._last_returned_call_from_interceptors = (
1024 StreamStreamCallResponseIterator(
1025 self._last_returned_call_from_interceptors,
1026 call_or_response_iterator,
1027 )
1028 )
1029 return self._last_returned_call_from_interceptors
1030 self._last_returned_call_from_interceptors = StreamStreamCall(
1031 request_iterator,
1032 _timeout_to_deadline(client_call_details.timeout),
1033 client_call_details.metadata,
1034 client_call_details.credentials,
1035 client_call_details.wait_for_ready,
1036 self._channel,
1037 client_call_details.method,
1038 request_serializer,
1039 response_deserializer,
1040 self._loop,
1041 )
1042 return self._last_returned_call_from_interceptors
1043
1044 client_call_details = ClientCallDetails(
1045 method, timeout, metadata, credentials, wait_for_ready
1046 )
1047 return await _run_interceptor(
1048 list(interceptors), client_call_details, request_iterator
1049 )
1050
1051 def time_remaining(self) -> Optional[float]:
1052 raise NotImplementedError()
1053
1054
1055class UnaryUnaryCallResponse(_base_call.UnaryUnaryCall):
1056 """Final UnaryUnaryCall class finished with a response."""
1057
1058 _response: ResponseType
1059
1060 def __init__(self, response: ResponseType) -> None:
1061 self._response = response
1062
1063 def cancel(self) -> bool:
1064 return False
1065
1066 def cancelled(self) -> bool:
1067 return False
1068
1069 def done(self) -> bool:
1070 return True
1071
1072 def add_done_callback(self, unused_callback) -> None:
1073 raise NotImplementedError()
1074
1075 def time_remaining(self) -> Optional[float]:
1076 raise NotImplementedError()
1077
1078 async def initial_metadata(self) -> Optional[Metadata]:
1079 return None
1080
1081 async def trailing_metadata(self) -> Optional[Metadata]:
1082 return None
1083
1084 async def code(self) -> grpc.StatusCode:
1085 return grpc.StatusCode.OK
1086
1087 async def details(self) -> str:
1088 return ""
1089
1090 async def debug_error_string(self) -> Optional[str]:
1091 return None
1092
1093 def __await__(self):
1094 if False: # pylint: disable=using-constant-test
1095 # This code path is never used, but a yield statement is needed
1096 # for telling the interpreter that __await__ is a generator.
1097 yield None
1098 return self._response
1099
1100 async def wait_for_connection(self) -> None:
1101 pass
1102
1103
1104class _StreamCallResponseIterator:
1105 _call: Union[_base_call.UnaryStreamCall, _base_call.StreamStreamCall]
1106 _response_iterator: AsyncIterable[ResponseType]
1107
1108 def __init__(
1109 self,
1110 call: Union[_base_call.UnaryStreamCall, _base_call.StreamStreamCall],
1111 response_iterator: AsyncIterable[ResponseType],
1112 ) -> None:
1113 self._response_iterator = response_iterator
1114 self._call = call
1115
1116 def cancel(self) -> bool:
1117 return self._call.cancel()
1118
1119 def cancelled(self) -> bool:
1120 return self._call.cancelled()
1121
1122 def done(self) -> bool:
1123 return self._call.done()
1124
1125 def add_done_callback(self, callback) -> None:
1126 self._call.add_done_callback(callback)
1127
1128 def time_remaining(self) -> Optional[float]:
1129 return self._call.time_remaining()
1130
1131 async def initial_metadata(self) -> Optional[Metadata]:
1132 return await self._call.initial_metadata()
1133
1134 async def trailing_metadata(self) -> Optional[Metadata]:
1135 return await self._call.trailing_metadata()
1136
1137 async def code(self) -> grpc.StatusCode:
1138 return await self._call.code()
1139
1140 async def details(self) -> str:
1141 return await self._call.details()
1142
1143 async def debug_error_string(self) -> Optional[str]:
1144 return await self._call.debug_error_string()
1145
1146 def __aiter__(self):
1147 return self._response_iterator.__aiter__()
1148
1149 async def wait_for_connection(self) -> None:
1150 return await self._call.wait_for_connection()
1151
1152
1153class UnaryStreamCallResponseIterator(
1154 _StreamCallResponseIterator, _base_call.UnaryStreamCall
1155):
1156 """UnaryStreamCall class which uses an alternative response iterator."""
1157
1158 async def read(self) -> Union[EOFType, ResponseType]:
1159 # Behind the scenes everything goes through the
1160 # async iterator. So this path should not be reached.
1161 raise NotImplementedError()
1162
1163
1164class StreamStreamCallResponseIterator(
1165 _StreamCallResponseIterator, _base_call.StreamStreamCall
1166):
1167 """StreamStreamCall class which uses an alternative response iterator."""
1168
1169 async def read(self) -> Union[EOFType, ResponseType]:
1170 # Behind the scenes everything goes through the
1171 # async iterator. So this path should not be reached.
1172 raise NotImplementedError()
1173
1174 async def write(self, request: RequestType) -> None:
1175 # Behind the scenes everything goes through the
1176 # async iterator provided by the InterceptedStreamStreamCall.
1177 # So this path should not be reached.
1178 raise NotImplementedError()
1179
1180 async def done_writing(self) -> None:
1181 # Behind the scenes everything goes through the
1182 # async iterator provided by the InterceptedStreamStreamCall.
1183 # So this path should not be reached.
1184 raise NotImplementedError()
1185
1186 @property
1187 def _done_writing_flag(self) -> bool:
1188 return self._call._done_writing_flag