Add example of Flight/Python middleware

python/examples/flight/middleware.py

@@ -0,0 +1,165 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Example of invisibly propagating a request ID with middleware."""
+
+import argparse
+import sys
+import threading
+import time
+import uuid
+
+import pyarrow as pa
+import pyarrow.flight as flight
+
+
+class TraceContext:
+    _locals = threading.local()
+    _locals.trace_id = None
+
+    @classmethod
+    def current_trace_id(cls):
+        if not getattr(cls._locals, "trace_id", None):
+            cls.set_trace_id(uuid.uuid4().hex)
+        return cls._locals.trace_id
+
+    @classmethod
+    def set_trace_id(cls, trace_id):
+        cls._locals.trace_id = trace_id
+
+
+TRACE_HEADER = "x-tracing-id"
+
+
+class TracingServerMiddleware(flight.ServerMiddleware):
+    def __init__(self, trace_id):
+        self.trace_id = trace_id
+
+    def sending_headers(self):
+        return {
+            TRACE_HEADER: self.trace_id,
+        }
+
+
+class TracingServerMiddlewareFactory(flight.ServerMiddlewareFactory):
+    def start_call(self, info, headers):
+        print("Starting new call:", info)
+        if TRACE_HEADER in headers:
+            trace_id = headers[TRACE_HEADER][0]
+            print("Found trace header with value:", trace_id)
+            TraceContext.set_trace_id(trace_id)
+        return TracingServerMiddleware(TraceContext.current_trace_id())
+
+
+class TracingClientMiddleware(flight.ClientMiddleware):
+    def sending_headers(self):
+        print("Sending trace ID:", TraceContext.current_trace_id())
+        return {
+            "x-tracing-id": TraceContext.current_trace_id(),
+        }
+
+    def received_headers(self, headers):
+        if TRACE_HEADER in headers:
+            trace_id = headers[TRACE_HEADER][0]
+            print("Found trace header with value:", trace_id)
+            # Don't overwrite our trace ID
+
+
+class TracingClientMiddlewareFactory(flight.ClientMiddlewareFactory):
+    def start_call(self, info):
+        print("Starting new call:", info)
+        return TracingClientMiddleware()
+
+
+class FlightServer(flight.FlightServerBase):
+    def __init__(self, delegate, **kwargs):
+        super().__init__(**kwargs)
+        if delegate:
+            self.delegate = flight.connect(
+                delegate,
+                middleware=(TracingClientMiddlewareFactory(),))
+        else:
+            self.delegate = None
+
+    def list_actions(self, context):
+        return [
+            ("get-trace-id", "Get the trace context ID."),
+        ]
+
+    def do_action(self, context, action):
+        if action.type == "get-trace-id":
+            if self.delegate:
+                for result in self.delegate.do_action(action):
+                    yield result
+            else:
+                trace_id = TraceContext.current_trace_id().encode("utf-8")
+                print("Returning trace ID:", trace_id)
+                buf = pa.py_buffer(trace_id)
+                yield pa.flight.Result(buf)
+        else:
+            raise KeyError(f"Unknown action {action.type!r}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+
+    subparsers = parser.add_subparsers(dest="command")
+    client = subparsers.add_parser("client", help="Run the client.")
+    client.add_argument("server")
+    client.add_argument("--request-id", default=None)
+
+    server = subparsers.add_parser("server", help="Run the server.")
+    server.add_argument(
+        "--listen",
+        required=True,
+        help="The location to listen on (example: grpc://localhost:5050)",
+    )
+    server.add_argument(
+        "--delegate",
+        required=False,
+        default=None,
+        help=("A location to delegate to. That is, this server will "
+              "simply call the given server for the response. Demonstrates "
+              "propagation of the trace ID between servers."),
+    )
+
+    args = parser.parse_args()
+    if not getattr(args, "command"):
+        parser.print_help()
+        return 1
+
+    if args.command == "server":
+        server = FlightServer(
+            args.delegate,
+            location=args.listen,
+            middleware=[TracingServerMiddlewareFactory()])
+        server.serve()
+    elif args.command == "client":
+        client = flight.connect(
+            args.server,
+            middleware=(TracingClientMiddlewareFactory(),))
+        if args.request_id:
+            TraceContext.set_trace_id(args.request_id)
+        else:
+            TraceContext.set_trace_id("client-chosen-id")
+
+        for result in client.do_action(flight.Action("get-trace-id", b"")):
+            print(result.body.to_pybytes())
+
+
+if __name__ == "__main__":
+    sys.exit(main() or 0)

python/pyarrow/_flight.pyx

@@ -1825,19 +1825,74 @@ cdef wrap_call_info(const CCallInfo& c_info):
 
 
 cdef class ClientMiddlewareFactory:
+    """A factory for new middleware instances.
+
+    All middleware methods will be called from the same thread as the
+    RPC method implementation. That is, thread-locals set in the
+    client are accessible from the middleware itself.
+
+    """
+
     def start_call(self, info):
-        pass
+        """Called at the start of an RPC.
+
+        This must be thread-safe and must not raise exceptions.
+
+        Parameters
+        ----------
+        info : CallInfo
+            Information about the call.
+
+        Returns
+        -------
+        instance : ClientMiddleware optional
+            An instance of ClientMiddleware (the instance to use for
+            the call), or None if this call is not intercepted.
+
+        """
 
 
 cdef class ClientMiddleware:
+    """Client-side middleware for a call, instantiated per RPC.
+
+    Methods here should be fast and must be infalliable: they should
+    not raise exceptions or stall indefinitely.
+
+    """
+
     def sending_headers(self):
-        pass
+        """A callback before headers are sent.
+
+        Returns
+        -------
+        headers : dict optional
+            A dictionary of header values to add to the response, or
+            None if no headers are to be added. The dictionary should
+            have string keys and string or list-of-string values.
+
+        """
 
     def received_headers(self, headers):
-        pass
+        """A callback when headers are received.
+
+        Parameters
+        ----------
+        headers : dict
+            A dictionary of headers from the client. Keys are strings
+            and values are lists of strings.
+
+        """
 
     def call_completed(self, exception):
-        pass
+        """A callback when the call finishes.
+
+        Parameters
+        ----------
+        exception : ArrowException
+            If the call errored, this is the equivalent
+            exception. Will be None if the call succeeded.
+
+        """
 
     @staticmethod
     cdef void wrap(object py_middleware,
@@ -1850,16 +1905,73 @@ cdef class ClientMiddleware:
 
 
 cdef class ServerMiddlewareFactory:
+    """A factory for new middleware instances.
+
+    All middleware methods will be called from the same thread as the
+    RPC method implementation. That is, thread-locals set in the
+    middleware are accessible from the method itself.
+
+    """
+
     def start_call(self, info, headers):
-        pass
+        """Called at the start of an RPC.
+
+        This must be thread-safe.
+
+        Parameters
+        ----------
+        info : CallInfo
+            Information about the call.
+
+        headers : dict
+            A dictionary of headers from the client. Keys are strings
+            and values are lists of strings.
+
+        Returns
+        -------
+        instance : ServerMiddleware optional
+            An instance of ServerMiddleware (the instance to use for
+            the call), or None if this call is not intercepted.
+
+        Raises
+        ------
+        ArrowException
+            If an exception is raised, the call will be rejected with
+            the given error.
+
+        """
 
 
 cdef class ServerMiddleware:
+    """Server-side middleware for a call, instantiated per RPC.
+
+    Methods here should be fast and must be infalliable: they should
+    not raise exceptions or stall indefinitely.
+
+    """
+
     def sending_headers(self):
-        pass
+        """A callback before headers are sent.
+
+        Returns
+        -------
+        headers : dict optional
+            A dictionary of header values to add to the response, or
+            None if no headers are to be added. The dictionary should
+            have string keys and string or list-of-string values.
+
+        """
 
     def call_completed(self, exception):
-        pass
+        """A callback when the call finishes.
+
+        Parameters
+        ----------
+        exception : ArrowException
+            If the call errored, this is the equivalent
+            exception. Will be None if the call succeeded.
+
+        """
 
     @staticmethod
     cdef void wrap(object py_middleware,
@@ -1871,18 +1983,20 @@ cdef class ServerMiddleware:
 
 
 cdef class _ServerMiddlewareFactoryWrapper(ServerMiddlewareFactory):
-    """
+    """A wrapper that bundles server middleware into a single C++ one.
 
-    gRPC runs RPCs on C++ threads. We use PyGilState_Ensure/Release to
-    interact with the GIL. However, this has a major shortcoming: it
-    breaks thread locals, since PyGilState_Release will clean up
-    Python's thread state (and clean up any thread locals). This makes
+    gRPC runs RPCs on C++ threads, so we use PyGilState_Ensure/Release
+    to interact with the GIL before running Python code. However, this
+    has a major shortcoming: it breaks thread locals. When we enter
+    Python, it creates new thread state. When we exit, that thread
+    state is freed (and Python thread locals cleaned up), since the
+    thread isn't persistent from Python's point of view. This makes
     development of middleware harder than it needs to be.
 
-    Instead, we'll always register a middleware that always runs
-    first/last that artifically makes sure the GIL state stays valid.
-
-    TODO: after this, remove Status from middleware (it just mucks things up)
+    This middleware solves this by first incrementing the reference
+    count on the thread state, then running all Python middleware. It
+    then doesn't decrement the reference count until the call is
+    finished.
 
     """
 
@@ -1902,6 +2016,8 @@ cdef class _ServerMiddlewareFactoryWrapper(ServerMiddlewareFactory):
                 if instance:
                     instances.append(instance)
         except:
+            # Don't artifically keep thread state around if an error
+            # happens
             PyGILState_Release(state)
             raise
 
@@ -1923,8 +2039,16 @@ cdef class _ServerMiddlewareWrapper(ServerMiddleware):
         self.middleware = middleware
 
     def sending_headers(self):
-        # TODO: need to merge all results
-        pass
+        headers = collections.defaultdict(list)
+        for instance in self.middleware:
+            more_headers = instance.sending_headers()
+            # Manually merge with existing headers (since headers are
+            # multi-valued)
+            for key, values in more_headers.items():
+                if isinstance(values, (six.text_type, six.binary_type)):
+                    values = (values,)
+                headers[key].extend(values)
+        return headers
 
     def call_completed(self, exception):
         try:

python/pyarrow/flight.py

@@ -49,6 +49,7 @@ from pyarrow._flight import (  # noqa
     Result,
     ClientAuthHandler,
     ServerAuthHandler,
+    CallInfo,
     ClientMiddleware,
     ClientMiddlewareFactory,
     ServerMiddleware,