gooddata-flight-server 1.34.1.dev1__py3-none-any.whl

gooddata_flight_server/server/flight_rpc/flight_middleware.py

@@ -0,0 +1,162 @@
+#  (C) 2024 GoodData Corporation
+from typing import Any, Callable, Optional
+
+import opentelemetry.context as otelctx
+import opentelemetry.propagate as otelpropagate
+import pyarrow.flight
+import structlog
+from opentelemetry import trace
+from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace import SpanKind, StatusCode, use_span
+from typing_extensions import TypeAlias
+
+from gooddata_flight_server.utils.otel_tracing import SERVER_TRACER
+
+_LOGGER = structlog.get_logger("gooddata_flight_server.rpc")
+
+
+class CallInfo(pyarrow.flight.ServerMiddleware):
+    """
+    Call Info middleware holds information about the current call:
+
+    - RPC method info
+    - headers used on RPC Call
+    """
+
+    MiddlewareName = "call_info"
+
+    def __init__(self, info: pyarrow.flight.CallInfo, headers: dict[str, list[str]]):
+        super().__init__()
+
+        self._info = info
+        self._headers = headers
+
+    @property
+    def info(self) -> pyarrow.flight.CallInfo:
+        """
+        :return: Flight's CallInfo with detail about the current call
+        """
+        return self._info
+
+    @property
+    def headers(self) -> dict[str, list[str]]:
+        """
+        :return: headers provided by the caller
+        """
+        return self._headers
+
+
+OnEndCallbackFn: TypeAlias = Callable[[Optional[pyarrow.ArrowException]], Any]
+
+
+class CallFinalizer(pyarrow.flight.ServerMiddleware):
+    """
+    Call Finalizer middleware can be used by method implementations to register
+    functions that should be called after the entire Flight RPC call finishes.
+
+    This may be important especially for the DoGet and DoExchange methods where the
+    method implementation prepares the stream and hands it over to PyArrow's Flight RPC.
+
+    It is often the case that the stream is backed by data which has its lifecycle
+    managed by the server. In these cases, the server needs to know the data is actually not
+    used anymore (so that it can release locks, free up the data or do whatever it needs
+    to do).
+    """
+
+    MiddlewareName = "call_finalizer"
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._on_end: list[OnEndCallbackFn] = []
+
+    def register_on_end(self, fun: OnEndCallbackFn) -> None:
+        """
+        Register a function that should be called once the call completes. It is possible to call this
+        multiple times and register multiple functions.
+
+        IMPORTANT: the function that you register here should be quick, do minimal blocking and have low chance
+        of hanging. The function will be called out from the gRPC server's thread; if the call hangs, server's thread
+        will be blocked.
+
+        :param fun: function to register, it will be called with one argument: exception,
+         which is either None on success or pyarrow.ArrowException on failure
+        :return: nothing
+        """
+        self._on_end.append(fun)
+
+    def call_completed(self, exception: Optional[pyarrow.lib.ArrowException]) -> None:
+        try:
+            for fun in self._on_end:
+                fun(exception)
+        except Exception:
+            _LOGGER.critical("call_finalization_failed", exc_info=True)
+
+
+class OtelMiddleware(pyarrow.flight.ServerMiddleware):
+    MiddlewareName = "otel_middleware"
+
+    def __init__(
+        self, info: pyarrow.flight.CallInfo, headers: dict[str, list[str]], extract_context: bool = False
+    ) -> None:
+        super().__init__()
+        method_name = info.method.name
+
+        if extract_context:
+            self._otel_ctx = otelpropagate.extract(headers)
+        else:
+            self._otel_ctx = otelctx.get_current()
+
+        self._otel_span = SERVER_TRACER.start_span(
+            f"{method_name}",
+            kind=SpanKind.SERVER,
+            context=self._otel_ctx,
+            attributes={
+                SpanAttributes.RPC_SYSTEM: "grpc",
+                SpanAttributes.RPC_SERVICE: "FlightRPC",
+                SpanAttributes.RPC_METHOD: method_name,
+            },
+        )
+
+        # note: the code does not set context / use span at this point because
+        # the middleware creation is done in a separate call, before the actual
+        # method invocation.
+        #
+        # the context has to be set & span used during the actual Flight RPC
+        # method handling
+
+    @property
+    def call_tracing(self) -> tuple[otelctx.Context, trace.Span]:
+        """
+        :return: tracing context & span for the current call
+        """
+        return self._otel_ctx, self._otel_span
+
+    def call_completed(self, exception: Optional[pyarrow.lib.ArrowException]) -> None:
+        # OpenTelemetry context/span restore;
+        #
+        # this has to happen because this method is done from thread managed
+        # by gRPC server / Flight & during a separate call after the request handling
+        # completes - any context juggling done previously is lost.
+        #
+        # - attach context extracted over the wire (see constructor)
+        # - use current method call's span (created during middleware init)
+        #
+        # code can then complete the span accordingly and finally detach the context
+        old_ctx = otelctx.attach(self._otel_ctx)
+        try:
+            with use_span(
+                self._otel_span,
+                end_on_exit=False,
+                record_exception=False,
+                set_status_on_exception=False,
+            ):
+                if exception is None:
+                    self._otel_span.set_status(StatusCode.OK)
+                else:
+                    self._otel_span.set_status(StatusCode.ERROR)
+                    self._otel_span.record_exception(exception)
+
+                self._otel_span.end()
+        finally:
+            otelctx.detach(old_ctx)

gooddata_flight_server/server/flight_rpc/flight_server.py

@@ -0,0 +1,228 @@
+# (C) 2024 GoodData Corporation
+"""
+This is a thin wrapper around the PyArrow FlightServerBase. It exists to provide a typed interface.
+
+There are two main pieces:
+
+-  FlightServer - this is an extension of pyarrow.flight.FlightServerBase; the sole purpose here is to decouple
+   implementation of the technical parts and the actual handling of Flight RPC Methods
+-  FlightServerMethods - base class containing typed definitions of all Flight RPC Methods
+"""
+
+from collections.abc import Generator
+from typing import Any, Callable, Optional, TypeVar, Union
+
+import opentelemetry.context as otelctx
+import opentelemetry.trace as oteltrace
+import pyarrow.flight
+from typing_extensions import Concatenate, ParamSpec, TypeAlias
+
+from gooddata_flight_server.server.flight_rpc.flight_middleware import OtelMiddleware
+from gooddata_flight_server.server.flight_rpc.server_methods import (
+    FlightServerMethods,
+)
+
+FlightServerLocation: TypeAlias = Union[str, bytes, Optional[tuple[str, int]], pyarrow.flight.Location]
+FlightTlsCertificates: TypeAlias = list[tuple[bytes, bytes]]
+FlightMiddlewares: TypeAlias = dict[str, pyarrow.flight.ServerMiddlewareFactory]
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+def rpc_decorator() -> Callable[
+    [Callable[Concatenate[Any, pyarrow.flight.ServerCallContext, P], T]],
+    Callable[Concatenate[Any, pyarrow.flight.ServerCallContext, P], T],
+]:
+    def _factory(
+        fun: Callable[Concatenate[Any, pyarrow.flight.ServerCallContext, P], T],
+    ) -> Callable[Concatenate[Any, pyarrow.flight.ServerCallContext, P], T]:
+        def _decorator(
+            self: Any,
+            context: pyarrow.flight.ServerCallContext,
+            *args: P.args,
+            **kwargs: P.kwargs,
+        ) -> T:
+            otel_middleware = context.get_middleware(OtelMiddleware.MiddlewareName)
+            if otel_middleware is not None:
+                otel_ctx, otel_span = otel_middleware.call_tracing
+            else:
+                otel_ctx = otelctx.get_current()
+                otel_span = oteltrace.INVALID_SPAN
+
+            old_otel_ctx = otelctx.attach(otel_ctx)
+            try:
+                with oteltrace.use_span(
+                    otel_span,
+                    end_on_exit=False,
+                    record_exception=False,
+                    set_status_on_exception=False,
+                ):
+                    return fun(self, context, *args, **kwargs)
+            finally:
+                otelctx.detach(old_otel_ctx)
+
+        return _decorator
+
+    return _factory
+
+
+class FlightServer(pyarrow.flight.FlightServerBase):
+    """
+    Creates AND starts the Flight RPC Server. To handle the RPC methods, the server will call out to the
+    `methods` object. The `methods` object can be switched from outside, thus allowing caller code to switch
+    implementation of the RPC methods at its own discretion. Typical use case for this is to keep returning
+    `UNAVAILABLE` until the entire server starts, then switch to the real implementation.
+
+    Please note, that the Flight RPC Server actually starts as soon as this class is created. The `serve` and
+    `serve_with_signals` are merely a synchronization mechanism to wait until the server gets stopped.
+
+    When managing the lifecycle of the Flight RPC from outside, it is better not to use the serve() methods.
+    """
+
+    def __init__(
+        self,
+        methods: Optional[FlightServerMethods] = None,
+        location: Optional[FlightServerLocation] = None,
+        auth_handler: Optional[pyarrow.flight.ServerAuthHandler] = None,
+        tls_certificates: Optional[FlightTlsCertificates] = None,
+        verify_client: Optional[bool] = None,
+        root_certificates: Optional[bytes] = None,
+        middleware: Optional[FlightMiddlewares] = None,
+    ):
+        """
+        Create Flight server
+
+        :param methods: implementation of Flight RPC methods to use; when None provided, will use implementation
+           that raises NotImplementedError for any method call
+        :param location : str, tuple or Location optional, default None
+           Location to serve on. Either a gRPC URI like `grpc://localhost:port`,
+           a tuple of (host, port) pair, or a Location instance.
+           If None is passed then the server will be started on localhost with a
+           system provided random port.
+        :param auth_handler : ServerAuthHandler optional, default None
+            An authentication mechanism to use. May be None.
+        :param tls_certificates : list optional, default None
+            A list of (certificate, key) pairs.
+        :param verify_client : boolean optional, default False
+            If True, then enable mutual TLS: require the client to present
+            a client certificate, and validate the certificate.
+        :param root_certificates : bytes optional, default None
+            If enabling mutual TLS, this specifies the PEM-encoded root
+            certificate used to validate client certificates.
+        :param middleware : list optional, default None
+            A dictionary of :class:`ServerMiddlewareFactory` items. The
+            keys are used to retrieve the middleware instance during calls
+            (see :meth:`ServerCallContext.get_middleware`).
+        """
+        super().__init__(
+            location=location,
+            auth_handler=auth_handler,
+            tls_certificates=tls_certificates,
+            verify_client=verify_client,
+            root_certificates=root_certificates,
+            middleware=middleware,
+        )
+
+        self._methods = methods or FlightServerMethods()
+
+    def switch_methods(self, methods: FlightServerMethods) -> None:
+        """
+        Switch methods used to handle Flight RPC requests. New requests will be handled using these methods.
+
+        :param methods: new set of methods to use
+        :return:
+        """
+        self._methods = methods
+
+    def serve(self) -> None:
+        """
+        Serve waits until the underlying server has stopped. This blocks current thread.
+
+        Note that the underlying Flight RPC Service starts up immediately when this class is constructed so this is
+        really just a synchronization mechanism.
+
+        :return: nothing
+        """
+        self.wait()
+
+    def serve_with_signals(self) -> None:
+        """
+        Serve waits until the underlying server has stopped OR until INT or TERM signal.
+        This blocks current thread.
+
+        When signal is received, the server will shut down as soon as current requests are
+        done and this method will return.
+
+        Note that the underlying Flight RPC Service starts up immediately when this class is
+        constructed so this is really just a synchronization mechanism.
+
+        :return: nothing
+        """
+        return self.serve()
+
+    #
+    # Delegates to methods impl
+    #
+
+    @rpc_decorator()
+    def list_flights(
+        self, context: pyarrow.flight.ServerCallContext, criteria: bytes
+    ) -> Generator[pyarrow.flight.FlightInfo, None, None]:
+        return self._methods.list_flights(context, criteria)
+
+    @rpc_decorator()
+    def get_flight_info(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        descriptor: pyarrow.flight.FlightDescriptor,
+    ) -> pyarrow.flight.FlightInfo:
+        return self._methods.get_flight_info(context, descriptor)
+
+    @rpc_decorator()
+    def get_schema(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        descriptor: pyarrow.flight.FlightDescriptor,
+    ) -> pyarrow.flight.SchemaResult:
+        return self._methods.get_schema(context, descriptor)
+
+    @rpc_decorator()
+    def do_put(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        descriptor: pyarrow.flight.FlightDescriptor,
+        reader: pyarrow.flight.MetadataRecordBatchReader,
+        writer: pyarrow.flight.FlightMetadataWriter,
+    ) -> None:
+        return self._methods.do_put(context, descriptor, reader, writer)
+
+    @rpc_decorator()
+    def do_get(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        ticket: pyarrow.flight.Ticket,
+    ) -> pyarrow.flight.FlightDataStream:
+        return self._methods.do_get(context, ticket)
+
+    @rpc_decorator()
+    def do_exchange(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        descriptor: pyarrow.flight.FlightDescriptor,
+        reader: pyarrow.flight.MetadataRecordBatchReader,
+        writer: pyarrow.flight.MetadataRecordBatchWriter,
+    ) -> None:
+        return self._methods.do_exchange(context, descriptor, reader, writer)
+
+    @rpc_decorator()
+    def list_actions(self, context: pyarrow.flight.ServerCallContext) -> list[tuple[str, str]]:
+        return self._methods.list_actions(context)
+
+    @rpc_decorator()
+    def do_action(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        action: pyarrow.flight.Action,
+    ) -> Generator[pyarrow.flight.Result, None, None]:
+        return self._methods.do_action(context, action)

gooddata_flight_server/server/flight_rpc/flight_service.py

@@ -0,0 +1,279 @@
+#  (C) 2024 GoodData Corporation
+#
+# mypy: no-strict-optional
+
+from threading import Thread
+from typing import Optional
+
+import pyarrow.flight
+import structlog
+
+from gooddata_flight_server.config.config import AuthenticationMethod, ServerConfig
+from gooddata_flight_server.errors.error_code import ErrorCode
+from gooddata_flight_server.errors.error_info import ErrorInfo
+from gooddata_flight_server.server.auth.auth_middleware import TokenAuthMiddleware, TokenAuthMiddlewareFactory
+from gooddata_flight_server.server.auth.token_verifier_factory import create_token_verification_strategy
+from gooddata_flight_server.server.base import ServerContext
+from gooddata_flight_server.server.flight_rpc.flight_middleware import (
+    CallFinalizer,
+    CallInfo,
+    OtelMiddleware,
+)
+from gooddata_flight_server.server.flight_rpc.flight_server import FlightServer
+from gooddata_flight_server.server.flight_rpc.server_methods import (
+    FlightServerMethods,
+)
+
+
+def _get_flight_server_locations(config: ServerConfig) -> tuple[str, str]:
+    transport = "grpc+tls" if config.use_tls else "grpc"
+
+    return (
+        f"{transport}://{config.listen_host}:{config.listen_port}",
+        f"{transport}://{config.advertise_host}:{config.advertise_port}",
+    )
+
+
+class _AvailabilityMiddlewareFactory(pyarrow.flight.ServerMiddlewareFactory):
+    """
+    Optionally rejects calls with FlightUnavailableError & some reason.
+
+    If unavailable_reason is set -> reject. Otherwise, let the request through.
+    """
+
+    def __init__(self, unavailable_reason: Optional[ErrorInfo] = None):
+        super().__init__()
+
+        self.unavailable_reason: Optional[ErrorInfo] = unavailable_reason
+
+    def start_call(
+        self, info: pyarrow.flight.CallInfo, headers: dict[str, list[str]]
+    ) -> Optional[pyarrow.flight.ServerMiddleware]:
+        if self.unavailable_reason is None:
+            return None
+
+        raise self.unavailable_reason.to_unavailable_error()
+
+
+class _CallInfoMiddlewareFactory(pyarrow.flight.ServerMiddlewareFactory):
+    def start_call(
+        self, info: pyarrow.flight.CallInfo, headers: dict[str, list[str]]
+    ) -> Optional[pyarrow.flight.ServerMiddleware]:
+        return CallInfo(info, headers)
+
+
+class _CallFinalizerMiddlewareFactory(pyarrow.flight.ServerMiddlewareFactory):
+    def start_call(
+        self, info: pyarrow.flight.CallInfo, headers: dict[str, list[str]]
+    ) -> Optional[pyarrow.flight.ServerMiddleware]:
+        return CallFinalizer()
+
+
+class _OtelMiddlewareFactory(pyarrow.flight.ServerMiddlewareFactory):
+    def __init__(self, extract_context: bool) -> None:
+        super().__init__()
+        self._extract_context = extract_context
+
+    def start_call(
+        self, info: pyarrow.flight.CallInfo, headers: dict[str, list[str]]
+    ) -> Optional[pyarrow.flight.ServerMiddleware]:
+        return OtelMiddleware(info, headers, self._extract_context)
+
+
+class FlightRpcService:
+    """
+    Service that exposes Flight RPC.
+
+    Handles Flight server start/stop & importantly also allows switching running server from
+    available to unavailable state; typical use case is to reject any RPC during startup/shutdown
+    of the server or loss of connection to the cluster.
+    """
+
+    def __init__(
+        self,
+        config: ServerConfig,
+        methods: FlightServerMethods = FlightServerMethods(),
+    ):
+        self._config = config
+        self._methods = methods
+
+        (
+            self._listen_url,
+            self._client_url,
+        ) = _get_flight_server_locations(config=config)
+
+        self._logger = structlog.get_logger("gooddata_flight_server.rpc")
+
+        self._availability = _AvailabilityMiddlewareFactory(
+            unavailable_reason=ErrorInfo.for_reason(ErrorCode.NOT_READY, "Try again later.")
+        )
+        # internal mutable state
+        # server starts immediately when constructed (PyArrow stuff); thus defer
+        # construction until start() is called
+        self._server: Optional[FlightServer] = None
+        self._flight_shutdown_thread: Optional[Thread] = None
+        self._stopped = False
+
+    def _initialize_authentication(
+        self, ctx: ServerContext
+    ) -> Optional[tuple[str, pyarrow.flight.ServerMiddlewareFactory]]:
+        if self._config.authentication_method == AuthenticationMethod.NoAuth:
+            if self._config.use_mutual_tls:
+                return None
+
+            if "127.0.0.1" not in self._config.listen_host and "localhost" not in self._config.listen_host:
+                print("!" * 72)
+                print("!!! Your server is configured without authentication and ")
+                print("!!! it seems it is listening on a non-loopback interface. ")
+                print("!!! The server may be reachable from public network. ")
+                print(f"!!! Listening on: {self._config.listen_host}. ")
+                print("!" * 72)
+
+                self._logger.warning("insecure_warning", listen_url=self._config.listen_host)
+
+            return None
+
+        verification = create_token_verification_strategy(ctx)
+
+        return TokenAuthMiddleware.MiddlewareName, TokenAuthMiddlewareFactory(
+            ctx.config.token_header_name, verification
+        )
+
+    def _initialize_otel_tracing(
+        self, ctx: ServerContext
+    ) -> Optional[tuple[str, pyarrow.flight.ServerMiddlewareFactory]]:
+        if self._config.otel_config.exporter_type is None:
+            return None
+
+        return OtelMiddleware.MiddlewareName, _OtelMiddlewareFactory(
+            self._config.otel_config.extract_context_from_headers
+        )
+
+    def start(self, ctx: ServerContext) -> None:
+        """
+        Starts the server. This will start the Flight RPC Server bound to configured host and port.
+        The server will be returning UNAVAILABLE for all methods until it is switched to serving.
+
+        See `switch_to_serving`.
+
+        :return:  nothing
+        """
+        # massaging before sending these out to PyArrow
+        tls_certificates = (
+            [self._config.tls_cert_and_key]
+            if self._config.use_tls and self._config.tls_cert_and_key is not None
+            else None
+        )
+
+        self._logger.info(
+            "flight_service_start",
+            listen_url=self._listen_url,
+            client_url=self._client_url,
+            tls=tls_certificates is not None,
+        )
+
+        middleware = {
+            "_availability": self._availability,
+            CallInfo.MiddlewareName: _CallInfoMiddlewareFactory(),
+            CallFinalizer.MiddlewareName: _CallFinalizerMiddlewareFactory(),
+        }
+
+        auth_middleware = self._initialize_authentication(ctx)
+        if auth_middleware is not None:
+            middleware[auth_middleware[0]] = auth_middleware[1]
+
+        otel_middleware = self._initialize_otel_tracing(ctx)
+        if otel_middleware is not None:
+            middleware[otel_middleware[0]] = otel_middleware[1]
+
+        # server starts right as it is constructed
+        # the serve() method does not have to be called; moreover, it should not be called by the server
+        # as it makes PyArrow to install signal handlers that interfere with quiver server's handlers
+        #
+        # see: https://github.com/apache/arrow/issues/11932
+        self._server = FlightServer(
+            methods=self._methods,
+            location=self._listen_url,
+            tls_certificates=tls_certificates,
+            verify_client=self._config.use_mutual_tls,
+            root_certificates=self._config.tls_root_cert,
+            middleware=middleware,
+        )
+
+    def switch_to_serving(self, methods: FlightServerMethods) -> None:
+        """
+        Switches the Flight RPC server to serving mode.
+
+        :param methods: implementation of the Flight RPC methods
+        :return: nothing
+        """
+        if self._server is None:
+            raise AssertionError("Flight server was never started")
+
+        self._methods = methods
+        self._server.switch_methods(methods)
+        self._availability.unavailable_reason = None
+
+    def switch_to_unavailable(self, err: ErrorInfo) -> None:
+        """
+        Switches the Flight RPC server to unavailable mode. All new Flight RPC method
+        calls will be returning UNAVAILABLE from now on.
+
+        :param err: error info to include in the FlightUnavailableError
+        :return: nothing
+        """
+        if self._server is None:
+            raise AssertionError("Flight server was not started")
+
+        self._availability.unavailable_reason = err
+
+    def stop(self) -> None:
+        """
+        Stops service. This method will switch the Flight RPC server to return
+        'UNAVAILABLE' for all following calls and then initiates shutdown of the server.
+
+        The shutdown is done asynchronously. Use `wait_for_stop` to sync for completion.
+        Since the server shutdown will wait until existing clients finish it may take longer.
+
+        :return:
+        """
+        if self._server is None:
+            raise AssertionError("Flight server was never started")
+
+        if self._flight_shutdown_thread is not None:
+            return
+
+        # do not handle any more new requests
+        self.switch_to_unavailable(ErrorInfo.for_reason(ErrorCode.SHUTTING_DOWN, "Server is shutting down."))
+
+        # trigger server shutdown in separate thread because shutdown blocks and violates the contract
+        # for quiver's long-running services
+        self._flight_shutdown_thread = Thread(
+            name="flight_service_shutdown",
+            daemon=True,
+            target=self._shutdown_server,
+        )
+        self._flight_shutdown_thread.start()
+
+    def _shutdown_server(self) -> None:
+        self._logger.info("flight_service_shutdown")
+        # this will block until server stops
+        self._server.shutdown()
+        self._logger.info("flight_service_finished")
+
+    def wait_for_stop(self, timeout: Optional[float] = None) -> bool:
+        if self._flight_shutdown_thread is None:
+            # this is really some mess in the caller code.. did not call stop() but tries to wait for it..
+            raise AssertionError("Flight server stop() was not issued yet attempting to wait for the server to stop.")
+
+        if self._flight_shutdown_thread.is_alive():
+            self._flight_shutdown_thread.join(timeout=timeout)
+
+        return not self._flight_shutdown_thread.is_alive()
+
+    @property
+    def client_url(self) -> str:
+        """
+        :return: location URL to advertise to clients
+        """
+        return self._client_url