d5h-foss / grpc-interceptor

Simplified Python gRPC interceptors
MIT License
136 stars 21 forks source link

Support arbitrary messages in testing framework #39

Open andrewmwhite opened 1 year ago

andrewmwhite commented 1 year ago

I'm writing an interceptor for applying field masks to arbitrary messages, something like this:

import typing as t

import grpc
from grpc_interceptor.exceptions import GrpcException
from grpc_interceptor.server import AsyncServerInterceptor
from google.protobuf.field_mask_pb2 import FieldMask

class FieldMaskInterceptor(AsyncServerInterceptor):

    def __init__(self, *args, field_name: str = "field_mask", **kwargs):
        self.field_name = field_name
        super().__init__(*args, **kwargs)

    def apply_mask(self, request, response):
        if self.field_name not in request.DESCRIPTOR.fields_by_name:
            # Field mask field is not defined for message type, so ignore.
            return response

        if not request.HasField(self.field_name):
            # Field mask field is defined but not set for this message, so ignore.
            return response

        mask = FieldMask()
        mask.CanonicalFormFromMask(mask=request.field_mask)

        if len(mask.paths) == 0:
            # Field mask is default/empty, so ignore by gRPC convention.
            return response

        # Mask is set so apply and return the masked response.
        masked_response = response.__class__()
        mask.MergeMessage(response, masked_response)
        return masked_response

    async def intercept(
        self,
        method: t.Callable,
        request_or_iterator: t.Any,
        context: grpc.ServicerContext,
        method_name: str,
    ) -> t.Any:
        if hasattr(request_or_iterator, "__aiter__"):
            # see: https://grpc-interceptor.readthedocs.io/en/latest/#async-server-interceptors
            raise NotImplementedError("Client-streaming not yet supported")

        try:
            response_or_iterator = method(request_or_iterator, context)
            if hasattr(response_or_iterator, "__aiter__"):
                # see: https://grpc-interceptor.readthedocs.io/en/latest/#async-server-interceptors
                raise NotImplementedError("Server-streaming not yet supported")

            # Unary => await, apply the mask, return the response
            response = await response_or_iterator
            response = self.apply_mask(request_or_iterator, response)
            return response
        except GrpcException as exc:
            await context.set_code(exc.status_code)
            await context.set_details(exc.details)
            raise

I was hoping to leverage your lovely testing setup to test this, but it looks like DummyRequest and DummyResponse are baked in to the testing setup and can't be overridden. To test the field mask interceptor I'd need a new request type with a FieldMask field. Is the request message type something that can be made configurable? I'm not familiar enough with gRPC to tell if that's even possible.

d5h commented 1 year ago

Hi Andrew. Yea the testing service and hence request / response types are currently baked in. It shouldn't be hard to allow custom types though. I'll think about what the API would look like for that. In the meanwhile I'd suggest unit testing your apply_mask method directly.