vgi-python 0.8.0__py3-none-any.whl

vgi/_test_fixtures/writable/__init__.py

@@ -0,0 +1,8 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Writable test-fixture worker and helpers.
+
+These fixtures depend on ``sqlglot`` (via ``vgi.transactor``) and live behind
+the ``vgi[test-fixtures-writable]`` extra. Tests that exercise the write
+subsystem (INSERT/UPDATE/DELETE/DDL) import from here.
+"""

vgi/_test_fixtures/writable/generic.py

@@ -0,0 +1,236 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Generic writable functions parameterized by table_name.
+
+These functions extend the base writable scan/insert/update/delete classes
+from ``writable_table`` but determine the table name dynamically from the
+first positional argument instead of a hardcoded class variable. This allows
+the same function classes to serve any table in the transactor's DuckDB
+database, which is essential for DDL-created tables.
+
+The table_name is passed as the first positional argument via the
+``ScanFunctionResult`` when the catalog dispatches scan/insert/update/delete
+for dynamically discovered tables.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pyarrow as pa
+from vgi_rpc import AnnotatedBatch
+from vgi_rpc.rpc import OutputCollector
+
+from vgi._test_fixtures.writable.table import (
+    _COUNT_SCHEMA,
+    WritableScanState,
+    _get_attach_opaque_data,
+    _get_pushdown_filters,
+    _get_tx_id,
+    _is_returning,
+    transactor_proxy,
+)
+from vgi.invocation import BindResponse, GlobalInitResponse
+from vgi.table_function import BindParams, InitParams, ProcessParams, TableFunctionGenerator
+from vgi.table_in_out_function import TableInOutGenerator
+
+__all__ = [
+    "GenericTableDelete",
+    "GenericTableInsert",
+    "GenericTableScan",
+    "GenericTableUpdate",
+]
+
+
+def _get_table_name_from_bind(params: BindParams[None]) -> str:
+    """Extract the table name from the first positional argument at bind time."""
+    args = params.bind_call.arguments
+    if not args.positional or args.positional[0] is None:
+        msg = "table_name positional argument is required"
+        raise ValueError(msg)
+    return str(args.positional[0].as_py())
+
+
+def _get_table_name_from_process(params: ProcessParams[None]) -> str:
+    """Extract the table name from the first positional argument at process time."""
+    assert params.init_call is not None
+    args = params.init_call.bind_call.arguments
+    if not args.positional or args.positional[0] is None:
+        msg = "table_name positional argument is required"
+        raise ValueError(msg)
+    return str(args.positional[0].as_py())
+
+
+def _get_table_schema_from_transactor(table_name: str, attach_opaque_data: bytes, tx_id: bytes) -> pa.Schema:
+    """Query the transactor for the table's Arrow schema (returned as IPC bytes)."""
+    proxy = transactor_proxy._get_proxy()
+    schema_bytes = proxy.table_schema(attach_opaque_data=attach_opaque_data, table_name=table_name, tx_id=tx_id)
+    return pa.ipc.read_schema(pa.BufferReader(schema_bytes))  # type: ignore[arg-type]
+
+
+# ============================================================================
+# Generic scan — dynamic table name + schema from transactor
+# ============================================================================
+
+
+class GenericTableScan(TableFunctionGenerator[None, WritableScanState]):
+    """Scan function for any table — determines table name from first positional arg."""
+
+    class Meta:
+        """Metadata for GenericTableScan."""
+
+        name = "generic_writable_scan"
+        projection_pushdown = True
+        filter_pushdown = True
+
+    @classmethod
+    def on_bind(cls, params: BindParams[None]) -> BindResponse:
+        """Bind: query transactor for table schema (already includes rowid)."""
+        table_name = _get_table_name_from_bind(params)
+        attach_opaque_data = params.attach_opaque_data  # unwrapped plaintext
+        tx_id = params.bind_call.transaction_opaque_data
+        assert attach_opaque_data is not None and tx_id is not None
+        table_schema = _get_table_schema_from_transactor(table_name, attach_opaque_data, tx_id)
+        return BindResponse(output_schema=table_schema)
+
+    @classmethod
+    def on_init(cls, params: InitParams[None]) -> GlobalInitResponse:
+        """Limit to a single worker."""
+        return GlobalInitResponse(max_workers=1)
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[None]) -> WritableScanState:
+        """Open the transactor scan stream once before processing begins."""
+        table_name = _get_table_name_from_process(params)
+        attach_opaque_data = _get_attach_opaque_data(params)
+        tx_id = _get_tx_id(params)
+        proxy = transactor_proxy._get_proxy()
+        columns = list(params.output_schema.names)
+        scan_iter = iter(
+            proxy.scan(
+                attach_opaque_data=attach_opaque_data,
+                tx_id=tx_id,
+                schema_name="",
+                table_name=table_name,
+                columns=columns,
+                pushdown_filters=_get_pushdown_filters(params),
+            )
+        )
+        return WritableScanState(scan_iter=scan_iter)
+
+    @classmethod
+    def process(cls, params: ProcessParams[None], state: WritableScanState, out: OutputCollector) -> None:
+        """Read the next batch from the scan stream."""
+        assert state.scan_iter is not None
+        try:
+            batch = next(state.scan_iter)
+            out.emit(batch.batch)
+        except StopIteration:
+            out.finish()
+
+
+# ============================================================================
+# Generic write base — shared INSERT/UPDATE/DELETE logic
+# ============================================================================
+
+
+class _GenericWriteBase(TableInOutGenerator[None, None]):
+    """Base for generic write handlers. Subclasses set _operation."""
+
+    _operation: str  # "insert" | "update" | "delete"
+
+    @classmethod
+    def on_bind(cls, params: BindParams[None]) -> BindResponse:
+        """Bind: query transactor for table schema to use for RETURNING."""
+        table_name = _get_table_name_from_bind(params)
+        if _is_returning(params):
+            attach_opaque_data = params.attach_opaque_data  # unwrapped plaintext
+            tx_id = params.bind_call.transaction_opaque_data
+            assert attach_opaque_data is not None and tx_id is not None
+            table_schema = _get_table_schema_from_transactor(table_name, attach_opaque_data, tx_id)
+            user_fields = [f for f in table_schema if f.name not in ("rowid", "row_id")]
+            return BindResponse(output_schema=pa.schema(user_fields))
+        return BindResponse(output_schema=_COUNT_SCHEMA)
+
+    @classmethod
+    def _open_stream(
+        cls,
+        proxy: Any,
+        attach_opaque_data: bytes,
+        tx_id: bytes,
+        table_name: str,
+        returning: bool,
+        batch: pa.RecordBatch,
+    ) -> Any:
+        """Open a write stream. Override for operations needing extra args."""
+        return getattr(proxy, cls._operation)(
+            attach_opaque_data=attach_opaque_data,
+            tx_id=tx_id,
+            table_name=table_name,
+            returning=returning,
+        )
+
+    @classmethod
+    def process(cls, params: ProcessParams[None], state: None, batch: pa.RecordBatch, out: OutputCollector) -> None:
+        """Forward batch to transactor write stream."""
+        table_name = _get_table_name_from_process(params)
+        attach_opaque_data = _get_attach_opaque_data(params)
+        tx_id = _get_tx_id(params)
+        returning = params.output_schema != _COUNT_SCHEMA
+        proxy = transactor_proxy._get_proxy()
+        with cls._open_stream(proxy, attach_opaque_data, tx_id, table_name, returning, batch) as stream:
+            response = stream.exchange(AnnotatedBatch(batch=batch))
+            out.emit(response.batch)
+
+
+class GenericTableInsert(_GenericWriteBase):
+    """INSERT handler for any table — determines table name from first positional arg."""
+
+    _operation = "insert"
+
+    class Meta:
+        """Metadata for GenericTableInsert."""
+
+        name = "generic_writable_insert"
+
+
+class GenericTableUpdate(_GenericWriteBase):
+    """UPDATE handler for any table — determines table name from first positional arg."""
+
+    _operation = "update"
+
+    class Meta:
+        """Metadata for GenericTableUpdate."""
+
+        name = "generic_writable_update"
+
+    @classmethod
+    def _open_stream(
+        cls,
+        proxy: Any,
+        attach_opaque_data: bytes,
+        tx_id: bytes,
+        table_name: str,
+        returning: bool,
+        batch: pa.RecordBatch,
+    ) -> Any:
+        """Open an update stream with column list derived from the batch."""
+        update_cols = [name for name in batch.schema.names if name != "rowid"]
+        return proxy.update(
+            attach_opaque_data=attach_opaque_data,
+            tx_id=tx_id,
+            table_name=table_name,
+            columns=update_cols,
+            returning=returning,
+        )
+
+
+class GenericTableDelete(_GenericWriteBase):
+    """DELETE handler for any table — determines table name from first positional arg."""
+
+    _operation = "delete"
+
+    class Meta:
+        """Metadata for GenericTableDelete."""
+
+        name = "generic_writable_delete"

vgi/_test_fixtures/writable/table.py

@@ -0,0 +1,149 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Writable table infrastructure — transactor proxy and shared helpers.
+
+Provides the ``TransactorProxy`` for connecting to the db-transactor subprocess,
+and helper functions used by the generic writable functions in ``writable_generic.py``.
+All tables are created dynamically via CREATE TABLE DDL at the client side.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Annotated
+
+import pyarrow as pa
+from vgi_rpc import AnnotatedBatch, ArrowSerializableDataclass, Transient
+
+from vgi.schema_utils import schema
+from vgi.table_function import BindParams, ProcessParams
+
+if TYPE_CHECKING:
+    from vgi.protocol import BindRequest
+
+from vgi.transactor.client import TransactorClient
+from vgi.transactor.protocol import TransactorProtocol
+
+__all__ = [
+    "TransactorProxy",
+    "WritableScanState",
+    "transactor_proxy",
+]
+
+# Output schema for write functions returning affected row counts.
+_COUNT_SCHEMA = schema(count=pa.int64())
+
+# DuckDB's native rowid pseudocolumn, marked with is_row_id metadata so the
+# C++ extension knows which column carries the physical row identifier.
+_ROWID_FIELD = pa.field("rowid", pa.int64(), metadata={b"is_row_id": b""})
+
+
+def _parse_write_options(bind_call: BindRequest) -> dict[str, bool | str | list[str]]:
+    """Parse the write_options RecordBatch from the bind call's named arguments."""
+    defaults: dict[str, bool | str | list[str]] = {
+        "return_chunks": False,
+        "on_conflict": "throw",
+        "on_conflict_columns": [],
+    }
+    if not (bind_call.arguments and bind_call.arguments.named):
+        return defaults
+    val = bind_call.arguments.named.get("write_options")
+    if val is None:
+        return defaults
+    from vgi_rpc.utils import deserialize_record_batch
+
+    options_bytes = val.as_py()
+    batch, _ = deserialize_record_batch(options_bytes)
+    result = dict(defaults)
+    if "return_chunks" in batch.schema.names:
+        result["return_chunks"] = batch.column("return_chunks")[0].as_py()
+    if "on_conflict" in batch.schema.names:
+        result["on_conflict"] = batch.column("on_conflict")[0].as_py()
+    if "on_conflict_columns" in batch.schema.names:
+        result["on_conflict_columns"] = batch.column("on_conflict_columns")[0].as_py()
+    return result
+
+
+def _is_returning(params: BindParams[None]) -> bool:
+    """Check if the C++ operator requested RETURNING rows."""
+    opts = _parse_write_options(params.bind_call)
+    return bool(opts.get("return_chunks", False))
+
+
+def _get_tx_id(params: ProcessParams[None]) -> bytes:
+    """Get transaction_opaque_data from the bind request."""
+    assert params.init_call is not None
+    tx_id = params.init_call.bind_call.transaction_opaque_data
+    if tx_id:
+        return tx_id
+    msg = "transaction_opaque_data is required but was not provided in the bind request"
+    raise ValueError(msg)
+
+
+def _get_attach_opaque_data(params: ProcessParams[None]) -> bytes:
+    """Get the unwrapped plaintext attach (storage shards on the sealed form)."""
+    attach_opaque_data = params.attach_opaque_data
+    if attach_opaque_data:
+        return attach_opaque_data
+    msg = "attach_opaque_data is required but was not provided in the bind request"
+    raise ValueError(msg)
+
+
+def _get_pushdown_filters(params: ProcessParams[None]) -> bytes | None:
+    """Get pushdown_filters as serialized IPC bytes from params (or None)."""
+    assert params.init_call is not None
+    pf_batch = params.init_call.pushdown_filters
+    if pf_batch is None:
+        return None
+    sink = pa.BufferOutputStream()
+    writer = pa.ipc.new_stream(sink, pf_batch.schema)
+    writer.write_batch(pf_batch)
+    writer.close()
+    return sink.getvalue().to_pybytes()
+
+
+@dataclass(kw_only=True)
+class WritableScanState(ArrowSerializableDataclass):
+    """State for writable table scans — holds the live transactor scan iterator."""
+
+    scan_iter: Annotated[Iterator[AnnotatedBatch] | None, Transient()] = None
+
+
+# ============================================================================
+# TransactorProxy — manages the db-transactor connection
+# ============================================================================
+
+
+class TransactorProxy:
+    """Manages connections to the shared db-transactor subprocess.
+
+    The transactor manages multiple databases internally (one per attach_opaque_data).
+    DDL statements are run during register() for each new catalog attachment.
+    """
+
+    def __init__(self, ddl_statements: list[str] | None = None) -> None:
+        """Initialize the proxy."""
+        self._ddl = ddl_statements or []
+        self._client: TransactorClient | None = None
+
+    def _get_proxy(self) -> TransactorProtocol:
+        """Get the transactor RPC proxy (auto-spawn if needed)."""
+        if self._client is None:
+            self._client = TransactorClient()
+        return self._client.get_proxy()  # type: ignore[no-any-return]
+
+    def register(self, attach_opaque_data: bytes, catalog_name: str = "") -> None:
+        """Register a new database for this attach_opaque_data and run initial DDL."""
+        proxy = self._get_proxy()
+        proxy.register(attach_opaque_data=attach_opaque_data, catalog_name=catalog_name, ddl_statements=self._ddl)
+
+    def close(self) -> None:
+        """Close the transactor connection."""
+        if self._client is not None:
+            self._client.close()
+            self._client = None
+
+
+# Module-level proxy — all tables created dynamically via DDL.
+transactor_proxy = TransactorProxy()