vgi-python 0.8.0__py3-none-any.whl

vgi/schema_utils.py

@@ -0,0 +1,234 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Schema building utilities for VGI functions.
+
+This module provides helpers for creating and modifying Arrow schemas with
+minimal boilerplate, making output_schema definitions more concise.
+
+FUNCTIONS
+---------
+schema(**fields)
+    Build a schema from keyword arguments mapping names to types.
+
+schema_like(source, add, remove, rename, replace)
+    Derive a new schema from an existing one with modifications.
+
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+import pyarrow as pa
+
+# A field spec is either a bare DataType or a (DataType, metadata) tuple.
+FieldSpec = pa.DataType | tuple[pa.DataType, dict[bytes | str, bytes | str]]
+
+__all__ = [
+    "FieldSpec",
+    "VGI_PARTITION_COLUMN_KEY",
+    "partition_field",
+    "schema",
+    "schema_like",
+]
+
+#: KeyValueMetadata key that marks an Arrow field as a partition column
+#: for VGI's Hive-style partitioning (PartitionColumns mode). Workers
+#: opt in by setting ``Meta.partition_kind`` to a non-default
+#: :class:`vgi.metadata.PartitionKind` AND annotating at least one
+#: field of their bind schema with this key.
+VGI_PARTITION_COLUMN_KEY: bytes = b"vgi.partition_column"
+
+
+def partition_field(
+    name: str,
+    type: pa.DataType,
+    *,
+    nullable: bool = True,
+    metadata: dict[bytes | str, bytes | str] | None = None,
+) -> pa.Field[Any]:
+    """Build a ``pa.Field`` marked as a VGI partition column.
+
+    Equivalent to::
+
+        pa.field(name, type, nullable=nullable,
+                 metadata={VGI_PARTITION_COLUMN_KEY: b"true",
+                           **(metadata or {})})
+
+    Use in a bind schema when the function opts into PartitionColumns
+    mode by setting ``Meta.partition_kind`` to a non-default
+    :class:`vgi.metadata.PartitionKind`. Per-field metadata round-trips
+    through Arrow IPC, so the C++ extension can identify partition
+    columns from ``bind_result.output_schema`` without a parallel
+    list-of-names.
+
+    Args:
+        name: Column name.
+        type: Arrow data type.
+        nullable: Whether the column can contain nulls.
+        metadata: Extra field-level metadata to merge with the
+            partition-column marker. Useful for extension types
+            (e.g. geoarrow.wkb's ``ARROW:extension:name`` key).
+
+    Returns:
+        A ``pa.Field`` carrying ``{VGI_PARTITION_COLUMN_KEY: b"true"}``
+        in its metadata.
+
+    """
+    merged: dict[bytes, bytes] = {VGI_PARTITION_COLUMN_KEY: b"true"}
+    if metadata:
+        for k, v in metadata.items():
+            key = k if isinstance(k, bytes) else k.encode()
+            val = v if isinstance(v, bytes) else v.encode()
+            merged[key] = val
+    return pa.field(name, type, nullable=nullable, metadata=merged)
+
+
+def schema(
+    __fields: Mapping[str, FieldSpec] | None = None,
+    /,
+    **kwargs: FieldSpec,
+) -> pa.Schema:
+    """Build an Arrow schema from field definitions.
+
+    Creates a schema with fields in the order specified. Field names are
+    the keys and values are either Arrow data types or ``(type, metadata)``
+    tuples for attaching field-level metadata.
+
+    Args:
+        __fields: Optional mapping of field names to specs (for programmatic use).
+        **kwargs: Field names mapped to Arrow data types or ``(type, metadata)`` tuples.
+
+    Returns:
+        Arrow schema with the specified fields.
+
+    Raises:
+        TypeError: If a value is not a valid Arrow data type or field spec.
+
+    Examples::
+
+        schema(id=pa.int64(), name=pa.string())
+        schema(row_id=(pa.int64(), {b"is_row_id": b""}), id=pa.int64())
+
+    """
+    # Combine __fields dict with kwargs
+    all_fields: dict[str, FieldSpec] = {}
+    if __fields is not None:
+        all_fields.update(__fields)
+    all_fields.update(kwargs)
+
+    # Validate and build schema
+    pa_fields: list[pa.Field[Any]] = []
+    for name, spec in all_fields.items():
+        if isinstance(spec, tuple):
+            dtype, metadata = spec
+            if not isinstance(dtype, pa.DataType):
+                raise TypeError(
+                    f"Field '{name}': expected pa.DataType as first tuple element, "
+                    f"got {type(dtype).__name__}. Use pa.int64(), pa.string(), etc."
+                )
+            pa_fields.append(pa.field(name, dtype, metadata=metadata))
+        elif isinstance(spec, pa.DataType):
+            pa_fields.append(pa.field(name, spec))
+        else:
+            raise TypeError(
+                f"Field '{name}': expected pa.DataType or (pa.DataType, metadata) tuple, "
+                f"got {type(spec).__name__}. Use pa.int64(), pa.string(), etc."
+            )
+
+    return pa.schema(pa_fields)
+
+
+def schema_like(
+    source: pa.Schema,
+    *,
+    add: Mapping[str, pa.DataType] | None = None,
+    remove: list[str] | None = None,
+    rename: Mapping[str, str] | None = None,
+    replace: Mapping[str, pa.DataType] | None = None,
+) -> pa.Schema:
+    """Derive a new schema from an existing one with modifications.
+
+    Creates a modified copy of the source schema. Operations are applied
+    in this order: remove -> rename -> replace -> add.
+
+    Args:
+        source: The source schema to derive from.
+        add: Fields to add at the end. Dict mapping names to types.
+        remove: Field names to remove from the schema.
+        rename: Field name mappings (old_name -> new_name).
+        replace: Fields to replace with new types (keeps position).
+
+    Returns:
+        New schema with the specified modifications.
+
+    Raises:
+        KeyError: If a field to remove, rename, or replace doesn't exist.
+        ValueError: If trying to add a field that already exists.
+
+    """
+    # Start with source field names for tracking
+    field_names = set(source.names)
+
+    # Validate remove fields exist
+    if remove:
+        for name in remove:
+            if name not in field_names:
+                raise KeyError(f"Cannot remove field '{name}': not found in schema. Available fields: {source.names}")
+
+    # Validate rename fields exist
+    if rename:
+        for old_name in rename:
+            if old_name not in field_names:
+                raise KeyError(
+                    f"Cannot rename field '{old_name}': not found in schema. Available fields: {source.names}"
+                )
+
+    # Validate replace fields exist
+    if replace:
+        for name in replace:
+            if name not in field_names:
+                raise KeyError(f"Cannot replace field '{name}': not found in schema. Available fields: {source.names}")
+
+    # Build the new schema
+    # Step 1: Remove fields
+    remove_set = set(remove) if remove else set()
+
+    # Step 2 & 3: Process remaining fields (rename and replace)
+    rename_map = rename or {}
+    replace_map = replace or {}
+
+    new_fields: list[pa.Field[Any]] = []
+    final_names: set[str] = set()
+
+    for field in source:
+        # Skip removed fields
+        if field.name in remove_set:
+            continue
+
+        # Get the (possibly renamed) name
+        new_name = rename_map.get(field.name, field.name)
+
+        # Get the (possibly replaced) type
+        new_type = replace_map.get(field.name, field.type)
+
+        new_fields.append(pa.field(new_name, new_type, metadata=field.metadata))
+        final_names.add(new_name)
+
+    # Step 4: Add new fields
+    if add:
+        for name, dtype in add.items():
+            if name in final_names:
+                raise ValueError(
+                    f"Cannot add field '{name}': already exists in schema. "
+                    f"Use 'replace' to change an existing field's type."
+                )
+            if not isinstance(dtype, pa.DataType):
+                raise TypeError(
+                    f"Field '{name}': expected pa.DataType, "
+                    f"got {type(dtype).__name__}. Use pa.int64(), pa.string(), etc."
+                )
+            new_fields.append(pa.field(name, dtype))
+
+    return pa.schema(new_fields)

vgi/secret_protocol.py

@@ -0,0 +1,124 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""VGI secret protocol — the wire contract for Orchard's standalone secret service.
+
+Orchard is an independently-deployed microservice that brokers downstream
+credentials (S3/HTTP/GCS/…) for a single authenticated account. The DuckDB
+extension's ``VgiRemoteSecretStorage`` calls :meth:`VgiSecretProtocol.secret_lookup`
+lazily whenever a secret consumer (e.g. httpfs resolving an ``s3://`` path) asks
+the secret manager for a credential.
+
+This protocol is **versioned independently** of :class:`vgi.protocol.VgiProtocol`
+(the worker/catalog protocol). It has exactly one method and a tiny surface so it
+can evolve on its own cadence — see ``protocol_version`` below.
+
+Wire shape
+----------
+``secret_lookup`` takes the requested ``path`` and ``type`` as direct scalar
+parameters (not a wrapped ``request`` dataclass), so the generated C++ builder
+``BuildSecretLookupParams(path, type)`` is directly callable without a hand-coded
+inner serializer. The response is :class:`SecretLookupResponse`, IPC-serialized
+into the unary ``result`` envelope and validated C++-side against
+``SecretLookupResultSchema()``.
+
+Identity is carried entirely by the OAuth bearer token on the HTTP request (the
+same ``CatalogAuth`` the catalog established at ATTACH) — there is no
+account/storage identifier in the request body.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Annotated, Any, ClassVar, Protocol
+
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass, ArrowType
+
+
+def encode_secret_values(mapping: dict[str, Any]) -> pa.RecordBatch | None:
+    """Build the one-row ``values`` RecordBatch from a Python mapping.
+
+    Each key becomes a column; the cell at row 0 is the secret value. Types are
+    inferred by pyarrow (str→utf8, int→int64, bool→bool, dict→struct, list→list,
+    …). Pass a ``pa.array([...])`` as a value for explicit control over the type.
+    Returns ``None`` for an empty mapping (no values to ship).
+    """
+    if not mapping:
+        return None
+    columns: dict[str, pa.Array[Any]] = {}
+    for key, value in mapping.items():
+        columns[key] = value if isinstance(value, pa.Array) else pa.array([value])
+    return pa.RecordBatch.from_pydict(columns)
+
+
+@dataclass(frozen=True, slots=True, kw_only=True)
+class SecretLookupResponse(ArrowSerializableDataclass):
+    """Response for :meth:`VgiSecretProtocol.secret_lookup`.
+
+    ``values`` is the secret's key→value map carried as a **one-row RecordBatch**
+    (serialized to binary on the wire): each column is a secret key and its row-0
+    cell is the value. This lets values be any Arrow/DuckDB type — string, int64,
+    bool, struct, list, nested — not just strings. Build it with
+    :func:`encode_secret_values`. The C++ side converts each cell to a typed
+    DuckDB ``Value`` via the Arrow→DuckDB bridge. ``redact_keys`` lists the subset
+    of keys whose values must be redacted by ``duckdb_secrets()`` — honor it or
+    values leak.
+
+    ``ttl_seconds`` is the server's suggested cache lifetime. ``expires_at_unix``
+    is the *credential's own* hard expiry as a Unix timestamp (0 = no intrinsic
+    expiry); the client caches for ``min(ttl_seconds, expires_at_unix - now)`` so
+    a short-lived STS token is never served past its own expiry.
+
+    When ``found`` is False every other field is empty/zero and the client caches
+    a short-TTL negative entry.
+    """
+
+    found: bool
+    secret_type: str = ""
+    provider: str = ""
+    name: str = ""
+    scope: list[str] = field(default_factory=list)
+    values: Annotated[pa.RecordBatch | None, ArrowType(pa.binary())] = None
+    redact_keys: list[str] = field(default_factory=list)
+    ttl_seconds: int = 0
+    expires_at_unix: int = 0
+
+
+# ---------------------------------------------------------------------------
+# VGI Secret Protocol
+# ---------------------------------------------------------------------------
+
+
+class VgiSecretProtocol(Protocol):
+    """Wire protocol for Orchard's standalone secret service.
+
+    A single unary method, ``secret_lookup``. ``vgi_rpc.RpcServer(VgiSecretProtocol,
+    impl)`` handles serialization, dispatching, and version enforcement exactly as
+    it does for :class:`vgi.protocol.VgiProtocol`.
+
+    Application protocol surface version
+    ------------------------------------
+    ``protocol_version`` is the canonical semver (MAJOR.MINOR.PATCH) of this
+    contract, **independent** of ``VgiProtocol.protocol_version``. The framework
+    enforces an exact major+minor match (patch ignored) at the dispatch boundary.
+    The C++ extension reads ``VGI_SECRET_PROTOCOL_VERSION`` from
+    ``vgi/src/generated/vgi_secret_protocol_version.hpp`` (generated; sibling of
+    ``vgi_protocol_version.hpp``) and passes it as a per-call
+    ``protocol_version_override`` so it never collides with the worker protocol's
+    global version constant.
+
+    Bump rules mirror :class:`vgi.protocol.VgiProtocol`: major for any
+    backwards-incompatible change, minor for additive, patch for worker-side fixes.
+    """
+
+    protocol_version: ClassVar[str] = "1.0.0"
+
+    def secret_lookup(self, path: str, type: str) -> SecretLookupResponse:  # noqa: A002
+        """Resolve the credential for ``path`` of secret ``type``.
+
+        ``type`` is the lowercased DuckDB secret type the consumer probed for
+        (``s3`` / ``r2`` / ``gcs`` / ``aws`` / ``http`` / …). Identity comes from
+        the OAuth bearer on the transport. Return ``SecretLookupResponse(found=False)``
+        when the account has no matching credential.
+        """
+        ...

vgi/secret_service.py

@@ -0,0 +1,238 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Standalone serving harness for the VGI secret protocol (Orchard).
+
+Orchard's secret service is an *independently-deployed microservice* — separate
+from the worker/catalog ``vgi-serve`` deployable and speaking
+:class:`vgi.secret_protocol.VgiSecretProtocol`. This module provides:
+
+- :func:`create_secret_app` — build a WSGI app for any ``VgiSecretProtocol``
+  implementation (usable with gunicorn/waitress/uwsgi).
+- :func:`serve_secret_http` — run it under waitress (prints ``PORT:<n>`` for test
+  harnesses, mirroring :mod:`vgi.serve`).
+- :class:`ExampleOrchardSecretService` — a reference implementation that returns a
+  canned ``s3`` credential for ``s3://test-bucket*``; the C++ integration tests
+  point ``vgi-secret-serve`` at this class.
+- :func:`main` — the ``vgi-secret-serve`` CLI entry point.
+
+Auth: identity is carried by the HTTP bearer. Production deployments wire an
+``authenticate`` callback (reuse the ``VGI_BEARER_TOKENS`` / ``VGI_JWT_*`` env
+vars supported by :mod:`vgi.serve`). The example service ignores identity.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import os
+import sys
+import time
+from typing import TYPE_CHECKING, Any
+
+import pyarrow as pa
+
+from vgi.secret_protocol import SecretLookupResponse, VgiSecretProtocol, encode_secret_values
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    import falcon
+
+
+# --------------------------------------------------------------------------- #
+# Reference implementation (also the integration-test fixture)
+# --------------------------------------------------------------------------- #
+
+
+class ExampleOrchardSecretService:
+    """Reference :class:`VgiSecretProtocol` implementation for tests/demos.
+
+    Returns a canned ``s3`` credential for any path under ``s3://test-bucket``
+    with a short ``expires_at_unix`` (so the ``min(ttl, expiry)`` cache path is
+    exercised) and ``secret`` marked for redaction. Everything else is a miss.
+    """
+
+    #: Seconds until the canned credential's intrinsic expiry.
+    credential_lifetime_seconds: int = 30
+
+    def secret_lookup(self, path: str, type: str) -> SecretLookupResponse:  # noqa: A002
+        """Return the canned credential for known test paths; empty otherwise."""
+        if type == "s3" and path.startswith("s3://test-bucket"):
+            # Heterogeneous typed values: string, int64, bool, and a nested struct
+            # — exercising the full Arrow→DuckDB Value bridge, not just string→string.
+            values: dict[str, object] = {
+                "key_id": "AKIAEXAMPLEORCHARD",
+                "secret": "examplesecretvalue",
+                "region": "us-east-1",
+                "port": pa.array([9000], pa.int64()),
+                "use_ssl": True,
+                "endpoint_config": {"connect_timeout_ms": 5000, "max_retries": 3},
+            }
+            # When VGI_MOCK_S3_ENDPOINT is set, point httpfs at a local mock S3 so
+            # a real `SELECT … FROM 's3://…'` read exercises the null-ClientContext
+            # system-transaction lookup path end to end.
+            mock_endpoint = os.environ.get("VGI_MOCK_S3_ENDPOINT")
+            if mock_endpoint:
+                values["endpoint"] = mock_endpoint  # host:port, no scheme
+                values["use_ssl"] = False
+                values["url_style"] = "path"
+            return SecretLookupResponse(
+                found=True,
+                secret_type="s3",
+                provider="orchard",
+                name="orchard_test_bucket",
+                scope=["s3://test-bucket"],
+                values=encode_secret_values(values),
+                redact_keys=["secret"],
+                ttl_seconds=60,
+                expires_at_unix=int(time.time()) + self.credential_lifetime_seconds,
+            )
+        return SecretLookupResponse(found=False)
+
+
+# --------------------------------------------------------------------------- #
+# WSGI app + HTTP server
+# --------------------------------------------------------------------------- #
+
+
+def create_secret_app(
+    impl: object,
+    *,
+    prefix: str = "",
+    cors_origins: str = "*",
+    signing_key: bytes | None = None,
+    authenticate: Callable[[falcon.Request], Any] | None = None,
+    oauth_resource_metadata: Any = None,
+) -> falcon.App[Any, Any]:
+    """Build a WSGI app serving *impl* over :class:`VgiSecretProtocol`.
+
+    *impl* is any object implementing ``secret_lookup(path, type)``. The default
+    landing/describe pages are disabled — this is a credential endpoint, not a
+    browsable worker.
+    """
+    try:
+        from vgi_rpc.http import make_wsgi_app
+    except ImportError:
+        sys.stderr.write(
+            "Error: HTTP dependencies not installed.\nInstall with: pip install vgi[http]  (or: uv sync --extra http)\n"
+        )
+        sys.exit(1)
+
+    from vgi_rpc.rpc import RpcServer
+
+    if signing_key is None:
+        signing_key = os.urandom(32)
+
+    server = RpcServer(VgiSecretProtocol, impl, enable_describe=False)
+    return make_wsgi_app(
+        server,
+        prefix=prefix,
+        cors_origins=cors_origins,
+        token_key=signing_key,
+        authenticate=authenticate,
+        oauth_resource_metadata=oauth_resource_metadata,
+        enable_landing_page=False,
+        enable_describe_page=False,
+    )
+
+
+def serve_secret_http(
+    impl: object,
+    *,
+    host: str = "0.0.0.0",
+    port: int | None = None,
+    prefix: str = "",
+    cors_origins: str = "*",
+    signing_key: bytes | None = None,
+    authenticate: Callable[..., Any] | None = None,
+    oauth_resource_metadata: Any = None,
+) -> None:
+    """Serve *impl* over HTTP under waitress. Prints ``PORT:<n>`` once bound."""
+    import socket
+
+    try:
+        import waitress  # type: ignore[import-untyped]
+    except ImportError:
+        sys.stderr.write(
+            "Error: waitress not installed.\nInstall with: pip install vgi[http]  (or: uv sync --extra http)\n"
+        )
+        sys.exit(1)
+
+    if port is None:
+        env_port = os.environ.get("PORT")
+        port = int(env_port) if env_port else 8080
+    if port == 0:
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind((host, 0))
+            port = int(s.getsockname()[1])
+
+    wsgi_app = create_secret_app(
+        impl,
+        prefix=prefix,
+        cors_origins=cors_origins,
+        signing_key=signing_key,
+        authenticate=authenticate,
+        oauth_resource_metadata=oauth_resource_metadata,
+    )
+
+    print(f"PORT:{port}", flush=True)
+    sys.stderr.write(f"Serving {type(impl).__name__} (VgiSecretProtocol) on http://{host}:{port}{prefix}\n")
+    sys.stderr.flush()
+    waitress.serve(wsgi_app, host=host, port=port, _quiet=True)
+
+
+def _load_impl(reference: str) -> object:
+    """Instantiate a ``VgiSecretProtocol`` implementation from ``module:Class``."""
+    if ":" not in reference:
+        sys.stderr.write(f"Error: expected 'module:ClassName', got {reference!r}\n")
+        sys.exit(1)
+    module_ref, class_name = reference.rsplit(":", 1)
+    try:
+        module = importlib.import_module(module_ref)
+    except ImportError as exc:
+        sys.stderr.write(f"Error: could not import {module_ref!r}: {exc}\n")
+        sys.exit(1)
+    cls = getattr(module, class_name, None)
+    if cls is None or not isinstance(cls, type):
+        sys.stderr.write(f"Error: {class_name!r} not found in {module_ref!r}\n")
+        sys.exit(1)
+    return cls()
+
+
+def main() -> None:
+    """CLI entry point for ``vgi-secret-serve``."""
+    from vgi.serve import _resolve_authenticate, _resolve_oauth_resource_metadata, _resolve_signing_key
+
+    parser = argparse.ArgumentParser(
+        prog="vgi-secret-serve",
+        description="Serve a VgiSecretProtocol implementation over HTTP (Orchard secret service).",
+    )
+    parser.add_argument(
+        "impl",
+        nargs="?",
+        default="vgi.secret_service:ExampleOrchardSecretService",
+        help="Implementation reference: module:ClassName (default: the built-in ExampleOrchardSecretService fixture).",
+    )
+    parser.add_argument("--host", default="0.0.0.0", help="HTTP bind address")
+    parser.add_argument(
+        "--port", "-p", type=int, default=None, help="HTTP port (default: $PORT or 8080; 0 = ephemeral)"
+    )
+    parser.add_argument("--prefix", default="", help="URL prefix for RPC endpoints")
+    parser.add_argument("--cors-origins", default="*", help="Allowed CORS origins")
+    args = parser.parse_args()
+
+    impl = _load_impl(args.impl)
+    serve_secret_http(
+        impl,
+        host=args.host,
+        port=args.port,
+        prefix=args.prefix,
+        cors_origins=args.cors_origins,
+        signing_key=_resolve_signing_key(),
+        authenticate=_resolve_authenticate(),
+        oauth_resource_metadata=_resolve_oauth_resource_metadata(),
+    )
+
+
+if __name__ == "__main__":
+    main()