vgi-python 0.8.0__py3-none-any.whl

vgi/_test_fixtures/attach_options.py

@@ -0,0 +1,308 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Example worker that declares attach-time options of many types and echoes them back.
+
+This worker exists to exercise the attach-time options pipeline end-to-end:
+- Declared options (``AttachOptions`` inner class) advertised via the ``catalogs()``
+  RPC for pre-attach discovery.
+- Values received at ``catalog_attach`` are serialized into the returned
+  ``attach_opaque_data`` so they survive pooled-worker reuse (subprocess) and stateless
+  transports (HTTP). Nothing is stored on ``self``.
+- The ``echo_attach_options`` table function decodes ``attach_opaque_data`` on every
+  invocation and returns a one-row batch containing every declared option.
+
+Run directly as a worker::
+
+    vgi-fixture-attach-options-worker
+
+Or serve over HTTP via ``vgi-serve``.
+"""
+
+from __future__ import annotations
+
+import datetime
+import uuid
+from dataclasses import dataclass
+from decimal import Decimal
+from typing import Annotated, Any, ClassVar
+
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass
+from vgi_rpc.rpc import CallContext, OutputCollector
+from vgi_rpc.utils import deserialize_record_batch, serialize_record_batch_bytes
+
+from vgi.catalog.attach_option import AttachOption, AttachOptionSpec, extract_attach_option_specs
+from vgi.catalog.catalog_interface import (
+    AttachOpaqueData,
+    CatalogAttachResult,
+    CatalogInfo,
+    ReadOnlyCatalogInterface,
+    TransactionOpaqueData,
+)
+from vgi.catalog.descriptors import Catalog, Schema
+from vgi.invocation import BindResponse
+from vgi.schema_utils import schema
+from vgi.table_function import BindParams, ProcessParams, TableFunctionGenerator, init_single_worker
+from vgi.worker import Worker
+
+__all__ = [
+    "AttachOptionsWorker",
+    "EchoAttachOptionsFunction",
+    "main",
+]
+
+
+CATALOG_NAME = "attach_options"
+_ATTACH_ID_SEP = b"\x00"
+_UUID_BYTES = 16
+
+
+# ---------------------------------------------------------------------------
+# Declared attach-time options: one per supported type
+# ---------------------------------------------------------------------------
+
+
+class AttachOptions:
+    """Attach-time options covering the supported Arrow/DuckDB type space."""
+
+    # Scalar primitives
+    opt_bool: Annotated[bool, AttachOption(desc="Boolean option")] = True
+    opt_int8: Annotated[int, AttachOption(desc="int8", arrow_type=pa.int8())] = -8
+    opt_int16: Annotated[int, AttachOption(desc="int16", arrow_type=pa.int16())] = -16
+    opt_int32: Annotated[int, AttachOption(desc="int32", arrow_type=pa.int32())] = -32
+    opt_int64: Annotated[int, AttachOption(desc="int64")] = -64
+    opt_uint8: Annotated[int, AttachOption(desc="uint8", arrow_type=pa.uint8())] = 8
+    opt_uint16: Annotated[int, AttachOption(desc="uint16", arrow_type=pa.uint16())] = 16
+    opt_uint32: Annotated[int, AttachOption(desc="uint32", arrow_type=pa.uint32())] = 32
+    opt_uint64: Annotated[int, AttachOption(desc="uint64", arrow_type=pa.uint64())] = 64
+    opt_float32: Annotated[float, AttachOption(desc="float32", arrow_type=pa.float32())] = 1.5
+    opt_float64: Annotated[float, AttachOption(desc="float64")] = 2.5
+    opt_string: Annotated[str, AttachOption(desc="UTF-8 string")] = "hello"
+    opt_blob: Annotated[bytes, AttachOption(desc="Binary blob")] = b"\x00\x01\x02"
+
+    # Temporal
+    opt_date: Annotated[datetime.date, AttachOption(desc="Date", arrow_type=pa.date32())] = datetime.date(2026, 4, 24)
+    opt_time: Annotated[datetime.time, AttachOption(desc="Time of day", arrow_type=pa.time64("us"))] = datetime.time(
+        12, 34, 56
+    )
+    opt_timestamp: Annotated[datetime.datetime, AttachOption(desc="Naive timestamp", arrow_type=pa.timestamp("us"))] = (
+        datetime.datetime(2026, 4, 24, 12, 34, 56)
+    )
+    opt_timestamp_tz: Annotated[
+        datetime.datetime,
+        AttachOption(desc="Timestamp with UTC tz", arrow_type=pa.timestamp("us", tz="UTC")),
+    ] = datetime.datetime(2026, 4, 24, 12, 34, 56, tzinfo=datetime.UTC)
+
+    # Precision
+    opt_decimal: Annotated[Decimal, AttachOption(desc="Decimal(18,4)", arrow_type=pa.decimal128(18, 4))] = Decimal(
+        "123.4500"
+    )
+
+    # Nested
+    opt_list: Annotated[list[int], AttachOption(desc="List of int64", arrow_type=pa.list_(pa.int64()))] = [1, 2, 3]
+    opt_struct: Annotated[
+        dict[str, object],
+        AttachOption(
+            desc="Struct",
+            arrow_type=pa.struct([pa.field("a", pa.int64()), pa.field("b", pa.string())]),
+        ),
+    ] = {"a": 1, "b": "x"}
+
+
+# Resolve once at import time; used both to build the echo function's output schema
+# and to backfill defaults in catalog_attach.
+_ATTACH_OPTION_SPECS: list[AttachOptionSpec] = extract_attach_option_specs(AttachOptions)
+
+_ECHO_SCHEMA: pa.Schema = schema({spec.name: spec.type for spec in _ATTACH_OPTION_SPECS})
+
+
+# ---------------------------------------------------------------------------
+# attach_opaque_data encoding / decoding
+# ---------------------------------------------------------------------------
+
+
+def _build_echo_batch(received: dict[str, Any]) -> pa.RecordBatch:
+    """Merge received option values with declared defaults, return a one-row batch."""
+    row: dict[str, Any] = {spec.name: spec.default for spec in _ATTACH_OPTION_SPECS}
+    for name, value in received.items():
+        if name in row:
+            row[name] = value
+    return pa.RecordBatch.from_pylist([row], schema=_ECHO_SCHEMA)
+
+
+def _encode_attach_opaque_data(received: dict[str, Any]) -> AttachOpaqueData:
+    batch = _build_echo_batch(received)
+    ipc_bytes = serialize_record_batch_bytes(batch)
+    return AttachOpaqueData(uuid.uuid4().bytes + _ATTACH_ID_SEP + ipc_bytes)
+
+
+def _decode_attach_opaque_data(attach_opaque_data: bytes) -> pa.RecordBatch:
+    raw = bytes(attach_opaque_data)
+    if len(raw) <= _UUID_BYTES + 1 or raw[_UUID_BYTES : _UUID_BYTES + 1] != _ATTACH_ID_SEP:
+        raise ValueError("attach_opaque_data does not carry an options payload")
+    ipc_bytes = raw[_UUID_BYTES + 1 :]
+    batch, _ = deserialize_record_batch(ipc_bytes)
+    return batch
+
+
+# ---------------------------------------------------------------------------
+# Echo table function
+# ---------------------------------------------------------------------------
+
+
+@dataclass(slots=True, frozen=True)
+class _EchoArgs:
+    """No arguments — the echo function reads state from ``attach_opaque_data``."""
+
+
+@dataclass(kw_only=True)
+class _EchoState(ArrowSerializableDataclass):
+    emitted: bool = False
+
+
+@init_single_worker
+class EchoAttachOptionsFunction(TableFunctionGenerator[_EchoArgs, _EchoState]):
+    """Return the attach-time option values that were passed at ATTACH.
+
+    One row, one column per declared option. The values come from ``attach_opaque_data``
+    so the function is safe under pool reuse (subprocess) and stateless
+    dispatch (HTTP): no per-attach state lives on ``self``.
+    """
+
+    FunctionArguments = _EchoArgs
+
+    class Meta:
+        name = "echo_attach_options"
+        description = "Echo the attach-time option values carried in attach_opaque_data"
+        categories = ["generator", "testing"]
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = _ECHO_SCHEMA
+
+    @classmethod
+    def on_bind(cls, params: BindParams[_EchoArgs]) -> BindResponse:
+        return BindResponse(output_schema=cls.FIXED_SCHEMA)
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[_EchoArgs]) -> _EchoState:
+        return _EchoState()
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[_EchoArgs],
+        state: _EchoState,
+        out: OutputCollector,
+    ) -> None:
+        if state.emitted:
+            out.finish()
+            return
+
+        # params.attach_opaque_data is the unwrapped plaintext attach (storage
+        # shards on the sealed form via the request); see InitParams.
+        attach_opaque_data = params.attach_opaque_data
+        if attach_opaque_data is None:
+            raise ValueError("echo_attach_options requires an attach_opaque_data")
+
+        batch = _decode_attach_opaque_data(attach_opaque_data)
+        # Re-cast to the declared schema so column order matches what bind promised.
+        batch = batch.select(_ECHO_SCHEMA.names)
+        out.emit(batch)
+        state.emitted = True
+
+
+# ---------------------------------------------------------------------------
+# Catalog interface
+# ---------------------------------------------------------------------------
+
+
+_CATALOG_DESCRIPTOR = Catalog(
+    name=CATALOG_NAME,
+    schemas=[
+        Schema(
+            name="main",
+            tables=(),
+            views=(),
+            functions=(EchoAttachOptionsFunction,),
+        ),
+    ],
+)
+
+
+class AttachOptionsCatalog(ReadOnlyCatalogInterface):
+    """Catalog that advertises AttachOptions and echoes values via attach_opaque_data."""
+
+    catalog = _CATALOG_DESCRIPTOR
+    catalog_name = CATALOG_NAME
+
+    def catalog_attach(
+        self,
+        *,
+        name: str,
+        options: dict[str, Any],
+        data_version_spec: str | None,
+        implementation_version: str | None,
+        ctx: CallContext | None = None,
+    ) -> CatalogAttachResult:
+        del data_version_spec, implementation_version, ctx
+        if name != CATALOG_NAME:
+            raise ValueError(f"Unknown catalog: {name!r}. Available: {CATALOG_NAME}")
+
+        attach_opaque_data = _encode_attach_opaque_data(options)
+
+        return CatalogAttachResult(
+            attach_opaque_data=attach_opaque_data,
+            supports_transactions=False,
+            supports_time_travel=False,
+            catalog_version_frozen=True,
+            catalog_version=1,
+            attach_opaque_data_required=True,
+            default_schema="main",
+            settings=[],
+            resolved_data_version=None,
+            resolved_implementation_version=None,
+        )
+
+    def catalogs(self) -> list[CatalogInfo]:
+        return [
+            CatalogInfo(
+                name=CATALOG_NAME,
+                implementation_version=None,
+                data_version_spec=None,
+                attach_option_specs=[spec.serialize() for spec in _ATTACH_OPTION_SPECS],
+            ),
+        ]
+
+    def catalog_version(
+        self,
+        *,
+        attach_opaque_data: AttachOpaqueData,
+        transaction_opaque_data: TransactionOpaqueData | None,
+        ctx: CallContext | None = None,
+    ) -> int:
+        del attach_opaque_data, transaction_opaque_data, ctx
+        return 1
+
+
+# ---------------------------------------------------------------------------
+# Worker + entry point
+# ---------------------------------------------------------------------------
+
+
+class AttachOptionsWorker(Worker):
+    """Worker exposing :class:`AttachOptionsCatalog`."""
+
+    # The AttachOptions inner class is picked up by Worker.__init_subclass__,
+    # which extracts specs and injects them into the catalog interface.
+    AttachOptions = AttachOptions
+
+    catalog_interface = AttachOptionsCatalog
+    catalog_name = CATALOG_NAME
+    functions = [EchoAttachOptionsFunction]
+
+
+def main() -> None:
+    AttachOptionsWorker.main()
+
+
+if __name__ == "__main__":
+    main()

vgi/_test_fixtures/bad_protocol.py

@@ -0,0 +1,62 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Fixture worker that deliberately advertises a mismatched protocol_version.
+
+This fixture exists to exercise the vgi-rpc framework's *application protocol
+version* enforcement (added in vgi-rpc 0.18.0) end-to-end, across a real
+transport, from both the Python ``Client`` and the C++ DuckDB extension.
+
+The trick is entirely Python-side and needs no extension rebuild: the worker
+hands :class:`BadProtocol` (a :class:`~vgi.protocol.VgiProtocol` subclass that
+redeclares ``protocol_version`` to an impossible major version) to
+``RpcServer`` via :attr:`~vgi.worker.Worker.protocol_class`. vgi-rpc reads the
+version with ``vars(protocol).get("protocol_version")`` — reading the class's
+own ``__dict__``, not an inherited attribute — so the redeclaration on this
+subclass's body is what takes effect.
+
+A normal client (Python ``Client`` or the C++ extension) declares
+``protocol_version = "1.0.0"`` and sends it on every request. Because this
+worker enforces ``"99.0.0"``, the major versions differ and the dispatch
+boundary raises ``ProtocolVersionError`` with a directional "upgrade the
+client" message that round-trips back to the caller.
+
+Otherwise this is a drop-in replacement for ``vgi-fixture-worker``: it
+inherits every function and the catalog from :class:`ExampleWorker`, so any
+request reaches the dispatch boundary (and trips the version check) using the
+same SQL the example worker accepts.
+
+Registered as the ``vgi-fixture-bad-protocol-worker`` entry point.
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+from vgi._test_fixtures.worker import ExampleWorker
+from vgi.protocol import VgiProtocol
+
+# A major-version bump guarantees a mismatch against any real client's
+# "1.0.0" (vgi-rpc compares major+minor exactly, ignoring patch). Declared on
+# this class body so ``vars(BadProtocol)["protocol_version"]`` resolves to it.
+BAD_PROTOCOL_VERSION = "99.0.0"
+
+
+class BadProtocol(VgiProtocol):
+    """VgiProtocol surface with a deliberately incompatible version."""
+
+    protocol_version: ClassVar[str] = BAD_PROTOCOL_VERSION
+
+
+class BadProtocolWorker(ExampleWorker):
+    """ExampleWorker that serves the example catalog under a bad protocol version."""
+
+    protocol_class: ClassVar[type[VgiProtocol]] = BadProtocol  # type: ignore[type-abstract]
+
+
+def main() -> None:
+    """Run the mismatched-protocol fixture worker process."""
+    BadProtocolWorker.main()
+
+
+if __name__ == "__main__":
+    main()