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

gooddata_flight_server/health/__init__.py

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

gooddata_flight_server/health/health_check_http_server.py

@@ -0,0 +1,103 @@
+# (C) 2024 GoodData Corporation
+from functools import partial
+from http import HTTPStatus
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from threading import Thread
+from typing import Any
+
+import structlog
+
+from gooddata_flight_server.health.server_health_monitor import (
+    ModuleHealthStatus,
+    ServerHealthMonitor,
+)
+
+LIVENESS_ENDPOINT_PATH = "/live"
+READINESS_ENDPOINT_PATH = "/ready"
+
+SERVER_MODULE_DEBUG_NAME = "server"
+
+LOGGER = structlog.get_logger("gooddata_flight_server.health_check_http_server")
+
+
+class _HealthCheckHandler(BaseHTTPRequestHandler):
+    def __init__(
+        self,
+        *args: Any,
+        server_health_monitor: ServerHealthMonitor,
+        **kwargs: Any,
+    ):
+        self._server_health_monitor = server_health_monitor
+        super().__init__(*args, **kwargs)
+
+    def log_message(self, format: str, *args: list) -> None:
+        pass
+
+    def do_GET(self) -> None:  # noqa: N802
+        if self.path == READINESS_ENDPOINT_PATH:
+            self.send_response(HTTPStatus.NO_CONTENT if self.is_ready() else HTTPStatus.INTERNAL_SERVER_ERROR)
+        elif self.path == LIVENESS_ENDPOINT_PATH:
+            self.send_response(HTTPStatus.NO_CONTENT if self.is_alive() else HTTPStatus.INTERNAL_SERVER_ERROR)
+        else:
+            self.send_response(HTTPStatus.NOT_FOUND)
+        self.end_headers()
+
+    def is_ready(self) -> bool:
+        """
+        Readiness check inspecting flight server startup state.
+
+        :return: True if instance is ready otherwise False
+        """
+        LOGGER.debug("checking_flight_server_ready")
+        if (
+            SERVER_MODULE_DEBUG_NAME in self._server_health_monitor.module_statuses
+            and self._server_health_monitor.module_statuses[SERVER_MODULE_DEBUG_NAME] == ModuleHealthStatus.OK
+        ):
+            LOGGER.debug("flight_server_ready")
+            return True
+        else:
+            LOGGER.warning("flight_server_not_ready")
+            return False
+
+    def is_alive(self) -> bool:
+        """
+        Liveness check inspecting all running modules statuses.
+
+        :return: True if server is healthy otherwise False
+        """
+        LOGGER.debug("checking_flight_server_healthy")
+
+        for (
+            module,
+            status,
+        ) in self._server_health_monitor.module_statuses.items():
+            if status == ModuleHealthStatus.NOT_OK:
+                LOGGER.warning("unhealthy_module", module=module)
+                return False
+
+        LOGGER.debug("flight_server_healthy")
+        return True
+
+
+class HealthCheckHttpServer:
+    """
+    HTTP server implementation for health checks mainly usable in k8s environment for readiness
+    and liveness probes. Exposes $LIVENESS_ENDPOINT_PATH and $READINESS_ENDPOINT_PATH returning HTTP 200 in
+    case of success otherwise HTTP 500.
+    """
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        server_health_monitor: ServerHealthMonitor,
+    ):
+        handler = partial(_HealthCheckHandler, server_health_monitor=server_health_monitor)
+        httpd = HTTPServer((host, port), handler)
+
+        def serve_forever(httpd_instance: HTTPServer) -> None:
+            with httpd_instance:
+                LOGGER.info("health_check_started", host=host, port=port)
+                httpd_instance.serve_forever()
+
+        Thread(target=serve_forever, args=(httpd,), daemon=True).start()

gooddata_flight_server/health/server_health_monitor.py

@@ -0,0 +1,83 @@
+# (C) 2024 GoodData Corporation
+import enum
+import threading
+import time
+
+import structlog
+
+from gooddata_flight_server.metrics import ServerMetrics
+from gooddata_flight_server.utils.libc_utils import LibcUtils
+
+
+class ModuleHealthStatus(enum.Enum):
+    """
+    This enum lists health status of a module
+    """
+
+    OK = "ok"
+    NOT_OK = "not_ok"
+
+
+class ServerHealthMonitor:
+    """
+    Server health monitor and maintenance.
+
+    The monitor includes a thread doing regular maintenance - namely periodically performing
+    malloc trim() to make system throw away the garbage. This is essential to survive in runtime environments
+    that impose memory (RSS) limits and kill the server if it exceeds it - the malloc does not
+    free the used memory back to the system; the RSS keeps growing and growing until the server
+    gets killed. The trim() call makes malloc drop all unneeded allocations.
+    """
+
+    def __init__(
+        self,
+        trim_interval: int = 30,
+    ) -> None:
+        self._logger = structlog.get_logger("gooddata_flight_server.maintenance")
+        self._libc = LibcUtils()
+        self._trim_interval = trim_interval
+        self._module_statuses: dict[str, ModuleHealthStatus] = {}
+
+        self._thread = threading.Thread(
+            name="gooddata_flight_server.maintenance",
+            target=self._maintenance,
+            daemon=True,
+        )
+        self._thread.start()
+
+        self._logger.info("server_health_monitor_started")
+
+    def _maintenance(self) -> None:
+        last_trim = time.time()
+
+        while True:
+            if time.time() - last_trim > self._trim_interval:
+                try:
+                    trim_start = time.perf_counter()
+                    self._libc.malloc_trim()
+                    duration = time.perf_counter() - trim_start
+
+                    ServerMetrics.TRIM_SUMMARY.observe(duration)
+                    self._logger.debug(
+                        "server_maintenance",
+                        op="malloc_trim",
+                        duration=duration,
+                    )
+                except Exception:
+                    ServerMetrics.TRIM_ERROR_COUNT.inc()
+                    self._logger.error("malloc_trim_failed", exc_info=True)
+
+                last_trim = time.time()
+
+            time.sleep(1.0)
+
+    @property
+    def module_statuses(self) -> dict[str, ModuleHealthStatus]:
+        return self._module_statuses
+
+    def set_module_status(self, name: str, status: ModuleHealthStatus) -> None:
+        """
+        :param name: name of the module to which the status belongs
+        :param status: health status of the module
+        """
+        self._module_statuses[name] = status

gooddata_flight_server/metrics.py

@@ -0,0 +1,16 @@
+#  (C) 2024 GoodData Corporation
+
+from prometheus_client import Counter, Summary
+
+
+# TODO: metric prefix should be configurable
+class ServerMetrics:
+    TRIM_SUMMARY = Summary(
+        "gdfs_malloc_trim",
+        "Summary of malloc trim call durations.",
+    )
+
+    TRIM_ERROR_COUNT = Counter(
+        "gdfs_malloc_trim_error",
+        "Number of times malloc_trim has failed. Repeated failures means big trouble incoming.",
+    )

gooddata_flight_server/py.typed

@@ -0,0 +1 @@
+Mark package as supporting typing. See https://www.python.org/dev/peps/pep-0561/ for details.

gooddata_flight_server/server/__init__.py

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

gooddata_flight_server/server/auth/__init__.py

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

gooddata_flight_server/server/auth/auth_middleware.py

@@ -0,0 +1,83 @@
+#  (C) 2024 GoodData Corporation
+from typing import Any, Optional
+
+import pyarrow.flight
+import structlog
+
+from gooddata_flight_server.server.auth.token_verifier import (
+    TokenVerificationStrategy,
+)
+
+
+class TokenAuthMiddleware(pyarrow.flight.ServerMiddleware):
+    MiddlewareName = "auth_token"
+
+    def __init__(self, token: str, token_data: Any) -> None:
+        super().__init__()
+
+        self._token = token
+        self._token_data = token_data
+
+    @property
+    def token_data(self) -> Any:
+        return self._token_data
+
+
+_DEFAULT_AUTH_TOKEN_HEADER = "Authorization"
+_LOGGER = structlog.get_logger("gooddata_flight_server.auth")
+_BEARER_END_IDX = 7
+
+
+class TokenAuthMiddlewareFactory(pyarrow.flight.ServerMiddlewareFactory):
+    def __init__(
+        self,
+        token_header_name: Optional[str],
+        strategy: TokenVerificationStrategy,
+    ):
+        super().__init__()
+
+        self._token_header_name = token_header_name
+        self._strategy = strategy
+
+    def _extract_token(self, headers: dict[str, list[str]]) -> str:
+        def _auth_header_value(lookup: str) -> str:
+            _lookup = lookup.lower()
+            values = [value for header, values in headers.items() if header.lower() == _lookup for value in values]
+
+            if len(values) > 1:
+                raise pyarrow.flight.FlightUnauthenticatedError(
+                    f"Authentication failed because the authentication header '{lookup}' is specified multiple times."
+                )
+
+            if not len(values):
+                raise pyarrow.flight.FlightUnauthenticatedError(
+                    "Authentication failed because the authentication header bearing the token was not included "
+                    "on the call."
+                )
+
+            return values[0]
+
+        if self._token_header_name is None:
+            token = _auth_header_value(_DEFAULT_AUTH_TOKEN_HEADER)
+            if not token.startswith("Bearer "):
+                raise pyarrow.flight.FlightUnauthenticatedError(
+                    "Authentication failed because the 'Authorization' header does not start with 'Bearer '"
+                )
+
+            return token[_BEARER_END_IDX:].strip()
+
+        token = _auth_header_value(self._token_header_name)
+        return token.strip()
+
+    def start_call(self, info: pyarrow.flight.CallInfo, headers: dict[str, list[str]]) -> Optional[TokenAuthMiddleware]:
+        try:
+            token = self._extract_token(headers)
+            result = self._strategy.verify(call_info=info, token=token)
+
+            return TokenAuthMiddleware(token=token, token_data=result)
+        except pyarrow.flight.FlightUnauthenticatedError as e:
+            _LOGGER.info("authentication_failed", reason=str(e))
+            raise
+        except pyarrow.flight.FlightUnauthorizedError as e:
+            _LOGGER.info("authorization_failed", reason=str(e))
+            raise

gooddata_flight_server/server/auth/token_verifier.py

@@ -0,0 +1,62 @@
+#  (C) 2024 GoodData Corporation
+import abc
+from typing import Any
+
+import pyarrow.flight
+
+from gooddata_flight_server.server.base import ServerContext
+
+
+class TokenVerificationStrategy(abc.ABC):
+    """
+    Token verification strategy is used by server's token authentication
+    middleware to perform the actual verification:
+
+    - The middleware is responsible for extracting the token
+    - The strategy is responsible for verifying that the token is valid and
+      if applicable return any information carried in the token.
+    - The middleware is responsible for holding onto value returned by the
+      verifier and make it available to call handling code.
+    """
+
+    @abc.abstractmethod
+    def verify(self, call_info: pyarrow.flight.CallInfo, token: str) -> Any:
+        """
+        Perform token verification.
+
+        - If the token is not valid, this method should raise either pyarrow.flight.FlightUnauthenticated
+        - If the token is valid but fails authorization for current call, this method
+          should raise pyarrow.Flight.FlightUnauthorized
+
+        Otherwise, the method should return either nothing or a custom value describing
+        the contents of the token - as it sees fit.
+
+        For example a JWT Token Verification strategy may extract claims and return them in a
+        dictionary.
+
+        :param call_info: current call info
+        :param token: token to verify
+        :return: either nothing or a custom value with info extracted from the token
+        :raises pyarrow.flight.FlightUnauthenticated - when the token is invalid / fails validation
+        :raises pyarrow.flight.FlightUnauthorized - when the token is valid but the caller holding
+         the token is not authorized to make the call
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def create(cls, ctx: ServerContext) -> "TokenVerificationStrategy":
+        """
+        This method is called exactly once during the server startup to obtain
+        a singleton of the concrete verification strategy to be used by
+        the server.
+
+        A typical use case here is to inspect the settings included in the server
+        context and obtain any essential configuration from the `settings`.
+
+        The strategy _may_ use this opportunity to perform any long-running
+        and blocking one-time initialization.
+
+        :param ctx: server context where the verification strategy will be used
+        :return: an instance of the strategy.
+        """
+        return cls()

gooddata_flight_server/server/auth/token_verifier_factory.py

@@ -0,0 +1,55 @@
+#  (C) 2024 GoodData Corporation
+import importlib
+
+import structlog
+
+from gooddata_flight_server.exceptions import ServerStartupInterrupted
+from gooddata_flight_server.server.auth.token_verifier import TokenVerificationStrategy
+from gooddata_flight_server.server.auth.token_verifier_impl import EnumeratedTokenVerification
+from gooddata_flight_server.server.base import ServerContext
+
+_LOGGER = structlog.get_logger("gooddata_flight_server.auth")
+
+
+def _import_verification_strategy(module_name: str) -> type[TokenVerificationStrategy]:
+    _LOGGER.info("load_token_verification", module_name=module_name)
+    module = importlib.import_module(module_name)
+
+    for member in module.__dict__.values():
+        if not isinstance(member, type) or not issubclass(member, TokenVerificationStrategy):
+            # filter out module members which are not classes that implement the
+            # TokenVerificationStrategy interface
+            continue
+
+        if member == TokenVerificationStrategy:
+            # the TokenVerificationStrategy class is likely imported in the module -
+            # don't want that to interfere
+            continue
+
+        return member
+
+    raise ServerStartupInterrupted(
+        f"The module '{module_name}' specified in 'token_verification' setting does not "
+        f"include an implementation of {TokenVerificationStrategy.__name__}."
+    )
+
+
+def _find_token_verification_class(ctx: ServerContext) -> type[TokenVerificationStrategy]:
+    # config reader must ensure that there is always some value here
+    assert ctx.config.token_verification is not None
+
+    if ctx.config.token_verification == "EnumeratedTokenVerification":
+        return EnumeratedTokenVerification
+
+    return _import_verification_strategy(ctx.config.token_verification)
+
+
+def create_token_verification_strategy(ctx: ServerContext) -> TokenVerificationStrategy:
+    try:
+        cls = _find_token_verification_class(ctx)
+
+        _LOGGER.info("auth_token_strategy_init", cls=cls.__name__)
+        return cls.create(ctx)
+    except Exception:
+        _LOGGER.critical("auth_token_init_failed", exc_info=True)
+        raise

gooddata_flight_server/server/auth/token_verifier_impl.py

@@ -0,0 +1,41 @@
+#  (C) 2024 GoodData Corporation
+from typing import Any
+
+import pyarrow.flight
+from dynaconf import ValidationError
+
+from gooddata_flight_server.server.auth.token_verifier import TokenVerificationStrategy
+from gooddata_flight_server.server.base import ServerContext
+
+_TOKEN_ENUMERATION_SECTION = "enumerated_tokens"
+_TOKEN_SETTING = "tokens"
+
+
+class EnumeratedTokenVerification(TokenVerificationStrategy):
+    """
+    A simple token verification strategy that successfully verifies
+    a token if it matches one of the tokens specified in the settings.
+    """
+
+    def __init__(self, allowed_tokens: set[str]) -> None:
+        self._tokens = allowed_tokens
+
+    def verify(self, call_info: pyarrow.flight.CallInfo, token: str) -> Any:
+        if token not in self._tokens:
+            raise pyarrow.flight.FlightUnauthenticatedError("Authentication token is not valid.")
+
+        return None
+
+    @classmethod
+    def create(cls, ctx: ServerContext) -> "TokenVerificationStrategy":
+        tokens = list(ctx.settings.get(f"{_TOKEN_ENUMERATION_SECTION}.{_TOKEN_SETTING}") or [])
+        if not len(tokens):
+            raise ValidationError(
+                f"The 'EnumeratedTokenVerification' requires that you configure "
+                f"which tokens are allowed to use. You have to include section "
+                f"[{EnumeratedTokenVerification}] with a '{_TOKEN_SETTING}' setting that contains "
+                f"list of tokens. Alternatively, you can specify environment variable "
+                f"GOODDATA_FLIGHT_ENUMERATED_TOKENS__TOKENS."
+            )
+
+        return EnumeratedTokenVerification(set(tokens))

gooddata_flight_server/server/base.py

@@ -0,0 +1,63 @@
+#  (C) 2024 GoodData Corporation
+from dataclasses import dataclass
+from typing import Protocol
+
+import pyarrow.flight
+from dynaconf import Dynaconf
+
+from gooddata_flight_server.config.config import ServerConfig
+from gooddata_flight_server.health.server_health_monitor import (
+    ServerHealthMonitor,
+)
+from gooddata_flight_server.server.flight_rpc.server_methods import (
+    FlightServerMethods,
+)
+from gooddata_flight_server.tasks.task_executor import TaskExecutor
+
+
+@dataclass(frozen=True)
+class ServerContext:
+    """
+    Server's context.
+    """
+
+    settings: Dynaconf
+    """
+    All settings parsed from configuration files and/or environment variables provided at server startup.
+    """
+
+    config: ServerConfig
+    """
+    Server's configuration
+    """
+
+    location: pyarrow.flight.Location
+    """
+    Server's Flight RPC location - this location connectable by clients.
+    """
+
+    health: ServerHealthMonitor
+    """
+    Server's health monitor. Components may register their health into the monitor.
+
+    This monitor is integrated with health check server where the overall status
+    is reported using liveness/readiness endpoints.
+
+    If a component is unhealthy, the whole server will be reported as such.
+    """
+
+    task_executor: TaskExecutor
+    """
+    Task Executor can and should be used to implement long running Tasks which generate
+    Flight data.
+    """
+
+
+class FlightServerMethodsFactory(Protocol):
+    """
+    Factory function for server methods. This can be provided to the
+    GoodDataFlightServer - the server will invoke the function at the right
+    point and then integrate the FlightServerMethods into its Flight RPC service.
+    """
+
+    def __call__(self, ctx: ServerContext) -> FlightServerMethods: ...

gooddata_flight_server/server/default.logging.ini

@@ -0,0 +1,28 @@
+# (C) 2024 GoodData Corporation
+[loggers]
+keys = root, gooddata_flight_server
+
+[handlers]
+keys = gooddata_flight_server_stream_handler
+
+[formatters]
+keys = gooddata_flight_server_formatter
+
+[logger_root]
+level = INFO
+handlers = gooddata_flight_server_stream_handler
+
+[logger_gooddata_flight_server]
+level = INFO
+qualname = gooddata_flight_server
+handlers =
+
+
+[handler_gooddata_flight_server_stream_handler]
+class = StreamHandler
+level = DEBUG
+formatter = gooddata_flight_server_formatter
+args = (sys.stderr,)
+
+[formatter_gooddata_flight_server_formatter]
+format = %(message)s

gooddata_flight_server/server/flight_rpc/__init__.py

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