Initial UDF server implementation

Author: BohuTANG

.gitignore

@@ -0,0 +1,4 @@
+.venv/
+__pycache__/
+.pytest_cache/
+*.egg-info/

README.md

@@ -1 +1,41 @@
-# ai-server
+# ai-server
+
+Stage aware Databend UDF server providing storage centric helpers used by the AI
+extensions. The initial release ships with:
+
+- `list_stage_files`: enumerate objects inside an external S3 stage via
+  [Apache OpenDAL](https://opendal.apache.org/docs/python).
+- `read_pdf`: pull down a PDF object from a stage and return its extracted text
+  as a `STRING`.
+- `read_docx`: fetch Microsoft Word files (`.docx`) and expose their textual
+  content as a `STRING`.
+
+## Getting started
+
+```bash
+python3 -m venv .venv
+source .venv/bin/activate
+pip install -e .[dev]
+ai-udf-server --port 8815 --metrics-port 9091
+```
+
+Databend can now connect to the running Flight server and call the registered
+functions. `list_stage_files` expects a `STAGE_LOCATION` argument supplied by
+Databend plus a numeric limit. Example:
+
+```sql
+SELECT
+    list_stage_files(
+        @stage_location,
+        50
+    );
+```
+
+To retrieve file contents call the document readers with an explicit path. The
+PDF reader optionally accepts `NULL` for the `max_pages` argument to stream the
+entire document:
+
+```sql
+SELECT read_pdf(@stage_location, 'inbox/manual.pdf', NULL);
+SELECT read_docx(@stage_location, 'reports/summary.docx');
+```

ai_server/__init__.py

@@ -0,0 +1,5 @@
+"""AI-enhanced Databend UDF server utilities."""
+
+from .server import create_server
+
+__all__ = ["create_server"]

ai_server/main.py

@@ -0,0 +1,76 @@
+"""Command line entrypoint for the Databend AI UDF server."""
+
+from __future__ import annotations
+
+import argparse
+import logging
+import signal
+import sys
+from contextlib import suppress
+from typing import Optional
+
+from prometheus_client import start_http_server as start_prometheus_server
+
+from ai_server.server import create_server
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_args(argv: Optional[list[str]] = None) -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Databend AI UDF server")
+    parser.add_argument("--host", default="0.0.0.0", help="Bind address for the gRPC server")
+    parser.add_argument("--port", type=int, default=8815, help="Port for the gRPC server")
+    parser.add_argument(
+        "--metrics-port",
+        type=int,
+        default=None,
+        help="Port for Prometheus metrics exporter (disabled when omitted).",
+    )
+    parser.add_argument(
+        "--log-level",
+        default="INFO",
+        help="Python logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL).",
+    )
+    return parser.parse_args(argv)
+
+
+def _configure_logging(level: str) -> None:
+    logging.basicConfig(
+        level=getattr(logging, level.upper(), logging.INFO),
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+
+
+def main(argv: Optional[list[str]] = None) -> int:
+    args = _parse_args(argv)
+    _configure_logging(args.log_level)
+
+    if args.metrics_port is not None:
+        start_prometheus_server(args.metrics_port)
+        logger.info("Prometheus metrics server started on port %s", args.metrics_port)
+
+    server = create_server(host=args.host, port=args.port, metric_port=args.metrics_port)
+    logger.info("Starting Databend AI UDF server on %s:%s", args.host, args.port)
+
+    # Handle shutdown gracefully to ensure we stop serving when receiving termination signals.
+    stop_event = getattr(server, "stopped", None)
+
+    def _handle_signal(signum, frame):  # noqa: ANN001 - signature dictated by signal library
+        logger.info("Received signal %s, shutting down.", signum)
+        with suppress(Exception):
+            server.shutdown()
+        if stop_event is not None:
+            stop_event.set()
+
+    signal.signal(signal.SIGINT, _handle_signal)
+    signal.signal(signal.SIGTERM, _handle_signal)
+
+    try:
+        server.serve()
+    except KeyboardInterrupt:
+        logger.info("Interrupted, stopping server.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

ai_server/server.py

@@ -0,0 +1,36 @@
+"""Entrypoint for the Databend AI UDF server."""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from databend_udf import UDFServer
+
+from ai_server.udfs import list_stage_files, read_docx, read_pdf
+
+
+def create_server(
+    host: str = "0.0.0.0", port: int = 8815, metric_port: Optional[int] = None
+) -> UDFServer:
+    """
+    Create a configured UDF server instance.
+
+    Parameters
+    ----------
+    host:
+        Bind address for the Flight server.
+    port:
+        Bind port for the Flight server.
+    metric_port:
+        Optional metrics port for Prometheus exporter. When provided the
+        databend-udf server will expose metrics via Prometheus.
+    """
+    location = f"{host}:{port}"
+    metric_location = (
+        f"{host}:{metric_port}" if metric_port is not None else None
+    )
+    server = UDFServer(location, metric_location=metric_location)
+    server.add_function(list_stage_files)
+    server.add_function(read_pdf)
+    server.add_function(read_docx)
+    return server

ai_server/stages/operator.py

@@ -0,0 +1,192 @@
+"""Helpers to translate Databend stage metadata into OpenDAL operators."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import logging
+import threading
+from typing import Any, Dict, Mapping, Tuple
+
+from databend_udf import StageLocation
+from opendal import Operator, exceptions as opendal_exceptions
+
+logger = logging.getLogger(__name__)
+
+_OPERATOR_CACHE: Dict[str, Operator] = {}
+_CACHE_LOCK = threading.Lock()
+
+
+class StageConfigurationError(RuntimeError):
+    """Raised when an unsupported or invalid stage configuration is encountered."""
+
+
+def _normalize_bool(value: Any) -> bool:
+    if isinstance(value, bool):
+        return value
+    if value is None:
+        return False
+    if isinstance(value, (int, float)):
+        return bool(value)
+    if isinstance(value, str):
+        normalized = value.strip().lower()
+        if normalized in {"1", "true", "t", "yes", "y", "on"}:
+            return True
+        if normalized in {"0", "false", "f", "no", "n", "off"}:
+            return False
+    return bool(value)
+
+
+def _first_present(storage: Mapping[str, Any], *keys: str) -> Any:
+    for key in keys:
+        if key in storage:
+            value = storage[key]
+            if value not in (None, "", {}):
+                return value
+    return None
+
+
+def _build_s3_options(storage: Mapping[str, Any]) -> Dict[str, Any]:
+    bucket = _first_present(storage, "bucket", "name")
+    if not bucket:
+        raise StageConfigurationError("S3 stage is missing bucket configuration")
+
+    region = _first_present(storage, "region")
+    endpoint = _first_present(storage, "endpoint", "endpoint_url")
+    access_key = _first_present(storage, "access_key_id", "aws_key_id")
+    secret_key = _first_present(storage, "secret_access_key", "aws_secret_key")
+    security_token = _first_present(storage, "security_token", "session_token", "aws_token")
+    master_key = _first_present(storage, "master_key")
+    root = _first_present(storage, "root")
+    role_arn = _first_present(storage, "role_arn", "aws_role_arn")
+    external_id = _first_present(storage, "external_id", "aws_external_id")
+    virtual_host_style = storage.get("enable_virtual_host_style")
+    disable_loader = storage.get("disable_credential_loader")
+
+    options: Dict[str, Any] = {"bucket": bucket}
+
+    if region:
+        options["region"] = region
+    else:
+        # Databend stages may skip region when working with S3 compatible endpoints.
+        # OpenDAL requires a region, default to us-east-1 if not provided.
+        options["region"] = "us-east-1"
+
+    if endpoint:
+        options["endpoint"] = endpoint
+    if access_key:
+        options["access_key_id"] = access_key
+    if secret_key:
+        options["secret_access_key"] = secret_key
+    if security_token:
+        options["security_token"] = security_token
+    if master_key:
+        options["master_key"] = master_key
+    if root:
+        options["root"] = root
+    if role_arn:
+        options["role_arn"] = role_arn
+    if external_id:
+        options["external_id"] = external_id
+    if virtual_host_style is not None:
+        options["enable_virtual_host_style"] = _normalize_bool(virtual_host_style)
+    if disable_loader is not None:
+        options["disable_credential_loader"] = _normalize_bool(disable_loader)
+
+    return options
+
+
+def _build_memory_options(_: Mapping[str, Any]) -> Dict[str, Any]:
+    # Useful for local testing; not a Databend production configuration.
+    return {}
+
+
+_STORAGE_BUILDERS: Dict[str, Any] = {"s3": _build_s3_options, "memory": _build_memory_options}
+
+
+def _cache_key(stage: StageLocation) -> str:
+    payload = {
+        "stage_name": stage.stage_name,
+        "stage_type": stage.stage_type,
+        "storage": stage.storage,
+    }
+    encoded = json.dumps(payload, sort_keys=True, default=str)
+    return hashlib.sha256(encoded.encode("utf-8")).hexdigest()
+
+
+def _build_operator(stage: StageLocation) -> Operator:
+    storage = stage.storage or {}
+    storage_type = str(storage.get("type", "")).lower()
+
+    if storage_type not in _STORAGE_BUILDERS:
+        raise StageConfigurationError(
+            f"Unsupported stage storage type '{storage_type or 'unknown'}'"
+        )
+
+    builder = _STORAGE_BUILDERS[storage_type]
+    options = builder(storage)
+    logger.debug(
+        "Creating OpenDAL operator for stage '%s' with backend '%s'",
+        stage.stage_name,
+        storage_type,
+    )
+    try:
+        return Operator(storage_type, **options)
+    except opendal_exceptions.Error as exc:
+        raise StageConfigurationError(
+            f"Failed to construct operator for stage '{stage.stage_name}': {exc}"
+        ) from exc
+
+
+def get_operator(stage: StageLocation) -> Operator:
+    """Return a cached OpenDAL operator for the given stage."""
+
+    cache_key = _cache_key(stage)
+    with _CACHE_LOCK:
+        operator = _OPERATOR_CACHE.get(cache_key)
+        if operator is None:
+            operator = _build_operator(stage)
+            _OPERATOR_CACHE[cache_key] = operator
+    return operator
+
+
+def clear_operator_cache() -> None:
+    """Utility that clears cached operators, primarily for testing."""
+
+    with _CACHE_LOCK:
+        _OPERATOR_CACHE.clear()
+
+
+def resolve_stage_subpath(stage: StageLocation, path: str | None = None) -> str:
+    """
+    Combine the stage's relative path with a user-provided path.
+
+    The resulting string is relative to the OpenDAL operator's configured root.
+    """
+
+    def _normalize(component: str | None) -> Tuple[str, ...]:
+        if not component:
+            return ()
+        parts = []
+        for part in component.split("/"):
+            chunk = part.strip()
+            if not chunk or chunk == ".":
+                continue
+            if chunk == "..":
+                raise ValueError("Stage paths must not contain '..'")
+            parts.append(chunk)
+        return tuple(parts)
+
+    base_parts = _normalize(stage.relative_path)
+    extra_parts = _normalize(path)
+    full_parts = base_parts + extra_parts
+    if not full_parts:
+        return ""
+    return "/".join(full_parts)
+
+
+def as_directory_path(path: str) -> str:
+    """Ensure the provided path represents a directory for list operations."""
+    if not path:
+        return ""
+    return path if path.endswith("/") else f"{path}/"

ai_server/udfs/__init__.py

@@ -0,0 +1,6 @@
+"""Collection of UDF implementations exposed by the AI server."""
+
+from .stage import list_stage_files
+from .files import read_docx, read_pdf
+
+__all__ = ["list_stage_files", "read_pdf", "read_docx"]