gooddata-flight-server 1.29.2.dev2__py3-none-any.whl

gooddata_flight_server/server/flight_rpc/server_methods.py

@@ -0,0 +1,200 @@
+# (C) 2024 GoodData Corporation
+from collections.abc import Generator
+from typing import Optional
+
+import pyarrow.flight
+import structlog
+
+from gooddata_flight_server.errors.error_info import ErrorCode, ErrorInfo
+from gooddata_flight_server.server.flight_rpc.flight_middleware import (
+    CallFinalizer,
+    CallInfo,
+)
+from gooddata_flight_server.tasks.task_executor import TaskExecutor
+from gooddata_flight_server.tasks.task_result import FlightDataTaskResult
+
+_LOGGER = structlog.get_logger("gooddata_flight_server.rpc")
+
+
+class FlightServerMethods:
+    """
+    Base class for implementations of Flight RPC server methods. This class contains a couple of utility
+    methods that may be useful in subclasses.
+
+    Typings reverse-engineered from PyArrow's Cython code.
+    """
+
+    @staticmethod
+    def call_info_middleware(
+        from_context: pyarrow.flight.ServerCallContext,
+    ) -> CallInfo:
+        """
+        Utility method to obtain CallInfo middleware from the call context. The
+        CallInfo middleware can be used to access Flight's CallInfo AND all headers
+        that were passed during the call.
+
+        :param from_context: server call context
+        :return: middleware
+        """
+        mw = from_context.get_middleware(CallInfo.MiddlewareName)
+        assert isinstance(mw, CallInfo)
+
+        return mw
+
+    @staticmethod
+    def call_finalizer_middleware(
+        from_context: pyarrow.flight.ServerCallContext,
+    ) -> CallFinalizer:
+        """
+        Utility method to obtain CallFinalizer middleware from the call context. The
+        CallFinalizer middleware can be used to register functions that should be
+        called after the entire Flight RPC completes.
+
+        See CallFinalizer documentation for more details..
+
+        :param from_context: server call context
+        :return: middleware
+        """
+        mw = from_context.get_middleware(CallFinalizer.MiddlewareName)
+        assert isinstance(mw, CallFinalizer)
+
+        return mw
+
+    @staticmethod
+    def do_get_task_result(
+        context: pyarrow.flight.ServerCallContext, task_executor: TaskExecutor, task_id: str
+    ) -> pyarrow.flight.FlightDataStream:
+        """
+        Utility method that creates a FlightDataStream from a result of a task that was
+        previously executed. You typically want to use this in implementation of `do_get`.
+
+        This method ensures that once the data is sent out, all necessary locks it previously
+        acquired to protect the data are freed. Single-use results will be closed once they
+        are sent out. The method uses current's call finalizer middleware to accomplish this.
+
+        :param context: server call context
+        :param task_executor: task executor where the task run
+        :param task_id: task identifier
+        :return: FlightDataStream, can be returned as-is as result of do_get
+        """
+        try:
+            task_result = task_executor.wait_for_result(task_id)
+            if task_result is None:
+                raise ErrorInfo.for_reason(
+                    ErrorCode.INVALID_TICKET,
+                    f"Unable to serve data for task '{task_id}'. The task result is not present.",
+                ).to_user_error()
+
+            if task_result.error is not None:
+                raise task_result.error.as_flight_error()
+
+            if task_result.cancelled:
+                raise ErrorInfo.for_reason(
+                    ErrorCode.COMMAND_CANCELLED,
+                    f"FlexFun invocation was cancelled. Invocation task was: '{task_result.task_id}'.",
+                ).to_server_error()
+
+            result = task_result.result
+            if not isinstance(result, FlightDataTaskResult):
+                raise ErrorInfo.for_reason(
+                    ErrorCode.INTERNAL_ERROR,
+                    f"An internal error has occurred while attempting read result for '{task_id}'."
+                    f"While the result exists, it is of an unexpected type: {type(result).__name__} ",
+                ).to_internal_error()
+
+            rlock, data = result.acquire_data()
+
+            def _on_end(_: Optional[pyarrow.ArrowException]) -> None:
+                """
+                Once the request that streams the data out is done, make sure
+                to release the read-lock. Single-use results are closed at
+                this point because the data cannot be read again anyway.
+                """
+                rlock.release()
+
+                if result.single_use_data:
+                    # note: results with single-use data can only ever have one active
+                    #  reader (e.g. this one). since the rlock is now released the
+                    #  close will proceed without chance of being blocked
+                    try:
+                        result.close()
+                    except Exception:
+                        # log and sink these Exceptions - not much to do
+                        _LOGGER.error("do_get_close_failed", exc_info=True)
+
+            finalizer = FlightServerMethods.call_finalizer_middleware(context)
+            finalizer.register_on_end(_on_end)
+
+            if isinstance(data, pyarrow.Table):
+                _LOGGER.info("do_get_table", task_id=task_id, num_rows=data.num_rows)
+
+                return pyarrow.flight.RecordBatchStream(data)
+            elif isinstance(data, pyarrow.RecordBatchReader):
+                _LOGGER.info("do_get_reader", task_id=task_id)
+
+                return pyarrow.flight.RecordBatchStream(data)
+
+            _LOGGER.info("do_get_generator", task_id=task_id)
+            return pyarrow.flight.GeneratorStream(data)
+        except Exception:
+            _LOGGER.error("do_get_failed", exc_info=True)
+            raise
+
+    ###################################################################
+    # Flight RPC methods - to be implemented as needed by
+    # subclasses.
+    ###################################################################
+
+    def list_flights(
+        self, context: pyarrow.flight.ServerCallContext, criteria: bytes
+    ) -> Generator[pyarrow.flight.FlightInfo, None, None]:
+        raise NotImplementedError
+
+    def get_flight_info(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        descriptor: pyarrow.flight.FlightDescriptor,
+    ) -> pyarrow.flight.FlightInfo:
+        raise NotImplementedError
+
+    def get_schema(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        descriptor: pyarrow.flight.FlightDescriptor,
+    ) -> pyarrow.flight.SchemaResult:
+        raise NotImplementedError
+
+    def do_put(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        descriptor: pyarrow.flight.FlightDescriptor,
+        reader: pyarrow.flight.MetadataRecordBatchReader,
+        writer: pyarrow.flight.FlightMetadataWriter,
+    ) -> None:
+        raise NotImplementedError
+
+    def do_get(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        ticket: pyarrow.flight.Ticket,
+    ) -> pyarrow.flight.FlightDataStream:
+        raise NotImplementedError
+
+    def do_exchange(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        descriptor: pyarrow.flight.FlightDescriptor,
+        reader: pyarrow.flight.MetadataRecordBatchReader,
+        writer: pyarrow.flight.MetadataRecordBatchWriter,
+    ) -> None:
+        raise NotImplementedError
+
+    def list_actions(self, context: pyarrow.flight.ServerCallContext) -> list[tuple[str, str]]:
+        raise NotImplementedError
+
+    def do_action(
+        self,
+        context: pyarrow.flight.ServerCallContext,
+        action: pyarrow.flight.Action,
+    ) -> Generator[pyarrow.flight.Result, None, None]:
+        raise NotImplementedError

gooddata_flight_server/server/server_base.py

@@ -0,0 +1,321 @@
+# (C) 2024 GoodData Corporation
+import abc
+import os
+import platform
+import signal
+from abc import abstractmethod
+from threading import Condition, Thread
+from typing import Any, Optional
+
+import pyarrow
+import structlog
+from prometheus_client import start_http_server
+
+from gooddata_flight_server._version import __version__
+from gooddata_flight_server.config.config import ServerConfig
+from gooddata_flight_server.health.health_check_http_server import (
+    SERVER_MODULE_DEBUG_NAME,
+    HealthCheckHttpServer,
+)
+from gooddata_flight_server.health.server_health_monitor import (
+    ModuleHealthStatus,
+    ServerHealthMonitor,
+)
+from gooddata_flight_server.utils.otel_tracing import SERVER_TRACER
+
+# DEV ONLY - heap usage debugger
+#
+# uncomment to turn on, then uncomment the dump done at the time of server stop()
+# from guppy import hpy
+# h = hpy()
+
+DEFAULT_LOGGING_INI = os.path.join(os.path.dirname(__file__), "default.logging.ini")
+
+
+class ServerBase(abc.ABC):
+    """
+    Server base class. Template class which takes care of the infrastructure and other boring
+    stuff and lets subclasses focus on just the value added services.
+    """
+
+    def __init__(
+        self,
+        config: ServerConfig,
+    ):
+        self._logger = structlog.get_logger("gooddata_flight_server.server")
+        self._config = config
+        self._health = ServerHealthMonitor(
+            trim_interval=config.malloc_trim_interval_sec,
+        )
+
+        # main server thread; this is responsible for starting all sub-services,
+        # don't want the main thread to be blocked
+        self._main_thread = Thread(
+            name="gooddata_flight_server.server",
+            target=self._server_main,
+            daemon=True,
+        )
+
+        # main server waits for this condition, once notified, it will stop or abort all sub-services
+        self._stop_cond = Condition()
+        # main server notifies on this condition once all sub-services are started
+        self._start_cond = Condition()
+        self._started = False
+        self._startup_interrupted: Optional[Exception] = None
+        self._stop = False
+        self._abort = False
+
+    #
+    # Server template
+    #
+
+    def _sig_handler(self, sig: Any, frame: Any) -> None:
+        self._logger.info("server_shutdown_initiated")
+        self.stop()
+
+    def _health_check_http_server_start(self) -> None:
+        """
+        If configured, start health checks HTTP server.
+
+        :return: nothing
+        """
+        if self._config.health_check_host is not None:
+            self._logger.debug(
+                "health_check_starting",
+                host=self._config.health_check_host,
+                port=self._config.health_check_port,
+            )
+
+            HealthCheckHttpServer(
+                host=self._config.health_check_host,
+                port=self._config.health_check_port,
+                server_health_monitor=self._health,
+            )
+
+    def _metrics_server_start(self) -> None:
+        """
+        If configured, start Prometheus metric server.
+
+        :return: nothing
+        """
+        if self._config.metrics_host is not None:
+            self._logger.debug(
+                "metrics_starting",
+                host=self._config.metrics_host,
+                port=self._config.metrics_port,
+            )
+
+            start_http_server(
+                addr=self._config.metrics_host,
+                port=self._config.metrics_port,
+            )
+
+            self._logger.info(
+                "metrics_started",
+                host=self._config.metrics_host,
+                port=self._config.metrics_port,
+            )
+
+    def _server_main(self) -> None:
+        self._logger.info(
+            "server_startup",
+            platform=platform.platform(terse=True),
+            python_version=platform.python_version(),
+            arrow_version=pyarrow.__version__,
+            server_version=__version__,
+            config=self._config.without_tls(),
+        )
+
+        try:
+            self._pre_startup()
+        except Exception as e:
+            self._startup_interrupted = e
+            return
+
+        with self._stop_cond:
+            if self._stop:
+                return
+
+        try:
+            self._metrics_server_start()
+            self._health_check_http_server_start()
+
+            self._startup_services()
+            self._health.set_module_status(SERVER_MODULE_DEBUG_NAME, ModuleHealthStatus.OK)
+
+            with self._start_cond:
+                self._started = True
+                self._start_cond.notify_all()
+
+            with self._stop_cond:
+                while not self._stop:
+                    self._stop_cond.wait()
+
+            if not self._abort:
+                self._shutdown_services()
+            else:
+                self._abort_services()
+        except Exception as e:
+            self._startup_interrupted = e
+
+            # wake up anyone who may be waiting for the server to start
+            with self._start_cond:
+                self._start_cond.notify_all()
+
+            return
+        finally:
+            self._logger.info("server_main_finished")
+
+    #
+    # Context
+    #
+
+    @property
+    def logger(self) -> structlog.stdlib.BoundLogger:
+        """
+        :return: server's main logger
+        """
+        return self._logger
+
+    @property
+    def health(self) -> ServerHealthMonitor:
+        """
+        :return: server's health monitor, this is initialized from the very beginning of the server's life
+        """
+        return self._health
+
+    #
+    # Lifecycle
+    #
+    @SERVER_TRACER.start_as_current_span("server_start")
+    def start(self) -> None:
+        """
+        Starts the server. This spins of the main server thread where it all happens. The start() method
+        returns immediately.
+
+        :return: nothing
+        """
+        signal.signal(signal.SIGINT, self._sig_handler)
+        signal.signal(signal.SIGTERM, self._sig_handler)
+
+        self._main_thread.start()
+
+    def stop(self) -> None:
+        """
+        Gracefully stops the server. This method will not block; see wait_for_stop().
+
+        :return: nothing
+        """
+        # DEV ONLY - heap usage debugger; also uncomment the imports
+        # print(h.heap())
+
+        with self._stop_cond:
+            self._stop = True
+            self._stop_cond.notify_all()
+
+    def abort(self) -> None:
+        """
+        Triggers hard-stop of the server. This is typically done as a response to unrecoverable failure that
+        warrants for quick and dirty shutdown.
+
+        On hard stop, the server will only try to do bare minimum in an attempt to clean up bare essentials.
+
+        :return:
+        """
+        with self._stop_cond:
+            self._abort = True
+            self._stop = True
+            self._stop_cond.notify_all()
+
+    def aborted(self) -> bool:
+        """
+        :return: True if the server aborted due to some serious failure
+        """
+        return self._abort
+
+    def wait_for_start(self, timeout: Optional[float] = None) -> bool:
+        """
+        Waits until server and all its services are up and running.
+
+        :param timeout: time in fractions of seconds
+        :return: true if started, false if not
+        """
+        with self._start_cond:
+            completed = self._start_cond.wait_for(
+                lambda: self._started is True or self._startup_interrupted is not None,
+                timeout=timeout,
+            )
+            if not completed:
+                return False
+
+            if self._startup_interrupted is not None:
+                raise self._startup_interrupted
+
+            return True
+
+    def wait_for_stop(self, timeout: Optional[float] = None) -> bool:
+        """
+        Waits until the main server thread stops. If the server startup encountered error (and never started), then
+        that error will be raised.
+
+        :param timeout: time to wait
+        :return: True if the main server thread finished; false if it is still running
+        """
+        self._main_thread.join(timeout=timeout)
+
+        if self._startup_interrupted is not None:
+            self._logger.fatal("server_startup_interrupted", exc_info=self._startup_interrupted)
+
+            raise self._startup_interrupted
+
+        return not self._main_thread.is_alive()
+
+    #
+    # template methods to be implemented by subclasses
+    #
+
+    def _pre_startup(self) -> None:
+        """
+        Perform any work before the actual startup logic is initiated. This may gracefully interrupt the server
+        startup if you raise exception created using the `_make_interrupt_exc()` method.
+
+        :return: nothing
+        """
+
+    @abstractmethod
+    def _startup_services(self) -> None:
+        """
+        Create and start up any services that the concrete server type requires. After the method completes,
+        the server should be up and running, ready to serve clients.
+
+        If this method raises an exception, then the server start will be interrupted and the server will exit.
+        You may gracefully interrupt the server startup if you raise exception created using the `make_interrupt_exc()`
+        method.
+
+        :return: nothing
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _shutdown_services(self) -> None:
+        """
+        Gracefully shutdown any services that the concrete server type uses. This method should block until all
+        the necessary services are stopped.
+
+        :return: nothing
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def _abort_services(self) -> None:
+        """
+        Abort services.
+
+        This method is called on unrecoverable errors when server just needs to shut down as soon as possible.
+        During abort, the server should try and do most important sanity (if any) and exit asap.
+
+        Attempting more complex and possibly blocking operations is not a good idea.
+
+        :return: nothing
+        """
+        raise NotImplementedError

gooddata_flight_server/server/server_main.py

@@ -0,0 +1,116 @@
+#  (C) 2024 GoodData Corporation
+from typing import Union
+
+import pyarrow.flight
+from dynaconf import Dynaconf
+
+from gooddata_flight_server.config.config import ServerConfig, read_config
+from gooddata_flight_server.exceptions import FlightMethodsModuleError
+from gooddata_flight_server.server.base import FlightServerMethodsFactory, ServerContext
+from gooddata_flight_server.server.flight_rpc.flight_service import FlightRpcService
+from gooddata_flight_server.server.flight_rpc.server_methods import FlightServerMethods
+from gooddata_flight_server.server.server_base import DEFAULT_LOGGING_INI, ServerBase
+from gooddata_flight_server.tasks.task_executor import TaskExecutor
+from gooddata_flight_server.tasks.thread_task_executor import ThreadTaskExecutor
+from gooddata_flight_server.utils.logging import init_logging
+from gooddata_flight_server.utils.otel_tracing import initialize_otel_tracing
+
+
+class GoodDataFlightServer(ServerBase):
+    def __init__(
+        self,
+        settings: Dynaconf,
+        config: ServerConfig,
+        methods: Union[FlightServerMethods, FlightServerMethodsFactory],
+    ):
+        super().__init__(config)
+
+        self._settings = settings
+        self._methods = methods if isinstance(methods, FlightServerMethods) else None
+        self._methods_factory = methods if not isinstance(methods, FlightServerMethods) else None
+
+        self._flight_service = FlightRpcService(config=config)
+        self._location = pyarrow.flight.Location(self._flight_service.client_url)
+
+        # TODO: make metric prefix configurable
+        self._task_executor = ThreadTaskExecutor(
+            metric_prefix="gdfs",
+            task_threads=config.task_threads,
+            result_close_threads=config.task_close_threads,
+            keep_results_for=config.task_result_ttl_sec,
+        )
+
+    @property
+    def location(self) -> pyarrow.flight.Location:
+        """
+        Server's location - this should be sent in all infos returned via Flight RPC.
+
+        :return: location
+        """
+        return self._location
+
+    @property
+    def task_executor(self) -> TaskExecutor:
+        """
+
+        :return:
+        """
+        return self._task_executor
+
+    def _startup_services(self) -> None:
+        server_ctx = ServerContext(
+            settings=self._settings,
+            config=self._config,
+            location=self._location,
+            task_executor=self._task_executor,
+            health=self.health,
+        )
+
+        self._flight_service.start(server_ctx)
+
+        if self._methods_factory is not None:
+            self.logger.info("flight_service_init_methods")
+
+            try:
+                self._methods = self._methods_factory(server_ctx)
+                if not isinstance(self._methods, FlightServerMethods):
+                    raise FlightMethodsModuleError(
+                        f"The provided FlightMethodsFactory has a valid signature but returned an invalid result of type "
+                        f"{type(self._methods)}. Make sure the factory function returns an instance of FlightServerMethods."
+                    )
+            except Exception as e:
+                self.logger.critical("flight_service_init_failed", exc_info=e)
+                raise
+
+        assert self._methods is not None
+        self._flight_service.switch_to_serving(self._methods)
+
+        self.logger.info("rpc_enabled", methods=type(self._methods).__name__)
+
+    def _shutdown_services(self) -> None:
+        self._flight_service.stop()
+        self._flight_service.wait_for_stop()
+
+    def _abort_services(self) -> None:
+        self._flight_service.stop()
+
+
+def create_server(
+    methods: Union[FlightServerMethods, FlightServerMethodsFactory],
+    config_files: tuple[str, ...] = (),
+    logging_config: str = DEFAULT_LOGGING_INI,
+    dev_log: bool = True,
+) -> "GoodDataFlightServer":
+    settings, config = read_config(files=config_files)
+
+    init_logging(
+        logging_config,
+        dev_log=dev_log,
+        event_key=config.log_event_key_name,
+        trace_ctx_keys=config.log_trace_keys,
+        add_trace_ctx=config.otel_config.exporter_type is not None,
+    )
+
+    initialize_otel_tracing(config=config.otel_config)
+
+    return GoodDataFlightServer(settings=settings, config=config, methods=methods)

gooddata_flight_server/tasks/__init__.py

@@ -0,0 +1 @@
+#  (C) 2024 GoodData Corporation

gooddata_flight_server/tasks/base.py

@@ -0,0 +1,21 @@
+#  (C) 2024 GoodData Corporation
+from typing import Any, Union
+
+import pyarrow
+from typing_extensions import TypeAlias
+
+# TODO: may be move to some more 'common' place
+ArrowData: TypeAlias = Union[pyarrow.lib.Table, pyarrow.lib.RecordBatchReader]
+
+
+class TaskWaitTimeoutError(TimeoutError):
+    """
+    This exception is thrown when TaskExecutor's wait_for_result() times out.
+
+    The exception includes task identifier and the payload that was used to
+    start the task.
+    """
+
+    def __init__(self, task_id: str, cmd: Any) -> None:
+        self.task_id = task_id
+        self.cmd = cmd