Skip to content

Commit

Permalink
Keep client interceptors in sync with grpc client interceptors
Browse files Browse the repository at this point in the history
  • Loading branch information
Mihir Gore committed Apr 15, 2021
1 parent e7d26a4 commit 759b52b
Show file tree
Hide file tree
Showing 2 changed files with 156 additions and 16 deletions.
130 changes: 114 additions & 16 deletions ...metry-instrumentation-grpc/src/opentelemetry/instrumentation/grpc/grpcext/_interceptor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,32 +48,71 @@ def __init__(self, method, base_callable, interceptor):
self._base_callable = base_callable
self._interceptor = interceptor

def __call__(self, request, timeout=None, metadata=None, credentials=None):
def __call__(
self,
request,
timeout=None,
metadata=None,
credentials=None,
wait_for_ready=None,
compression=None,
):
def invoker(request, metadata):
return self._base_callable(request, timeout, metadata, credentials)
return self._base_callable(
request,
timeout,
metadata,
credentials,
wait_for_ready,
compression,
)

client_info = _UnaryClientInfo(self._method, timeout)
return self._interceptor.intercept_unary(
request, metadata, client_info, invoker
)

def with_call(
self, request, timeout=None, metadata=None, credentials=None
self,
request,
timeout=None,
metadata=None,
credentials=None,
wait_for_ready=None,
compression=None,
):
def invoker(request, metadata):
return self._base_callable.with_call(
request, timeout, metadata, credentials
request,
timeout,
metadata,
credentials,
wait_for_ready,
compression,
)

client_info = _UnaryClientInfo(self._method, timeout)
return self._interceptor.intercept_unary(
request, metadata, client_info, invoker
)

def future(self, request, timeout=None, metadata=None, credentials=None):
def future(
self,
request,
timeout=None,
metadata=None,
credentials=None,
wait_for_ready=None,
compression=None,
):
def invoker(request, metadata):
return self._base_callable.future(
request, timeout, metadata, credentials
request,
timeout,
metadata,
credentials,
wait_for_ready,
compression,
)

client_info = _UnaryClientInfo(self._method, timeout)
Expand All @@ -88,9 +127,24 @@ def __init__(self, method, base_callable, interceptor):
self._base_callable = base_callable
self._interceptor = interceptor

def __call__(self, request, timeout=None, metadata=None, credentials=None):
def __call__(
self,
request,
timeout=None,
metadata=None,
credentials=None,
wait_for_ready=None,
compression=None,
):
def invoker(request, metadata):
return self._base_callable(request, timeout, metadata, credentials)
return self._base_callable(
request,
timeout,
metadata,
credentials,
wait_for_ready,
compression,
)

client_info = _StreamClientInfo(self._method, False, True, timeout)
return self._interceptor.intercept_stream(
Expand All @@ -105,11 +159,22 @@ def __init__(self, method, base_callable, interceptor):
self._interceptor = interceptor

def __call__(
self, request_iterator, timeout=None, metadata=None, credentials=None
self,
request_iterator,
timeout=None,
metadata=None,
credentials=None,
wait_for_ready=None,
compression=None,
):
def invoker(request_iterator, metadata):
return self._base_callable(
request_iterator, timeout, metadata, credentials
request_iterator,
timeout,
metadata,
credentials,
wait_for_ready,
compression,
)

client_info = _StreamClientInfo(self._method, True, False, timeout)
Expand All @@ -118,11 +183,22 @@ def invoker(request_iterator, metadata):
)

def with_call(
self, request_iterator, timeout=None, metadata=None, credentials=None
self,
request_iterator,
timeout=None,
metadata=None,
credentials=None,
wait_for_ready=None,
compression=None,
):
def invoker(request_iterator, metadata):
return self._base_callable.with_call(
request_iterator, timeout, metadata, credentials
request_iterator,
timeout,
metadata,
credentials,
wait_for_ready,
compression,
)

client_info = _StreamClientInfo(self._method, True, False, timeout)
Expand All @@ -131,11 +207,22 @@ def invoker(request_iterator, metadata):
)

def future(
self, request_iterator, timeout=None, metadata=None, credentials=None
self,
request_iterator,
timeout=None,
metadata=None,
credentials=None,
wait_for_ready=None,
compression=None,
):
def invoker(request_iterator, metadata):
return self._base_callable.future(
request_iterator, timeout, metadata, credentials
request_iterator,
timeout,
metadata,
credentials,
wait_for_ready,
compression,
)

client_info = _StreamClientInfo(self._method, True, False, timeout)
Expand All @@ -151,11 +238,22 @@ def __init__(self, method, base_callable, interceptor):
self._interceptor = interceptor

def __call__(
self, request_iterator, timeout=None, metadata=None, credentials=None
self,
request_iterator,
timeout=None,
metadata=None,
credentials=None,
wait_for_ready=None,
compression=None,
):
def invoker(request_iterator, metadata):
return self._base_callable(
request_iterator, timeout, metadata, credentials
request_iterator,
timeout,
metadata,
credentials,
wait_for_ready,
compression,
)

client_info = _StreamClientInfo(self._method, True, True, timeout)
Expand Down
42 changes: 42 additions & 0 deletions instrumentation/opentelemetry-instrumentation-grpc/tests/test_client_interceptor.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,13 +40,55 @@
from .protobuf.test_server_pb2 import Request


class Interceptor(
grpc.UnaryUnaryClientInterceptor,
grpc.UnaryStreamClientInterceptor,
grpc.StreamUnaryClientInterceptor,
grpc.StreamStreamClientInterceptor,
):
def __init__(self):
pass

def intercept_unary_unary(
self, continuation, client_call_details, request
):
return self._intercept_call(continuation, client_call_details, request)

def intercept_unary_stream(
self, continuation, client_call_details, request
):
return self._intercept_call(continuation, client_call_details, request)

def intercept_stream_unary(
self, continuation, client_call_details, request_iterator
):
return self._intercept_call(
continuation, client_call_details, request_iterator
)

def intercept_stream_stream(
self, continuation, client_call_details, request_iterator
):
return self._intercept_call(
continuation, client_call_details, request_iterator
)

@staticmethod
def _intercept_call(
continuation, client_call_details, request_or_iterator
):
return continuation(client_call_details, request_or_iterator)


class TestClientProto(TestBase):
def setUp(self):
super().setUp()
GrpcInstrumentorClient().instrument()
self.server = create_test_server(25565)
self.server.start()
interceptors = [Interceptor()]
self.channel = grpc.insecure_channel("localhost:25565")
self.channel = grpc.intercept_channel(self.channel, *interceptors)
self._stub = test_server_pb2_grpc.GRPCTestServerStub(self.channel)

def tearDown(self):
Expand Down

0 comments on commit 759b52b

Please sign in to comment.