vgi-python 0.8.0__py3-none-any.whl

vgi/catalog/setting.py

@@ -0,0 +1,253 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Setting descriptor for declarative worker settings.
+
+This module provides the Setting descriptor class for defining worker settings
+using Python's Annotated type hints, similar to how Arg works for function arguments.
+
+"""
+
+from dataclasses import dataclass, field
+from typing import (
+    TYPE_CHECKING,
+    Annotated,
+    Any,
+    ClassVar,
+    cast,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+import pyarrow as pa
+from vgi_rpc.utils import deserialize_record_batch, serialize_record_batch_bytes
+
+from vgi.schema_utils import schema
+
+if TYPE_CHECKING:
+    from typing import Self
+
+__all__ = [
+    "Setting",
+    "SettingSpec",
+    "extract_setting_specs",
+]
+
+
+@dataclass(frozen=True)
+class SettingSpec:
+    """Extracted setting metadata for catalog serialization.
+
+    This is the resolved form of a Setting, with all types inferred and
+    ready for serialization.
+
+    Attributes:
+        name: The setting name (from the class attribute name).
+        desc: Human-readable description.
+        type: The Arrow data type for this setting.
+        default: The default value (Python object).
+
+    """
+
+    name: str
+    desc: str
+    type: pa.DataType
+    default: Any
+
+    ARROW_SCHEMA: ClassVar[pa.Schema] = pa.schema(
+        [
+            pa.field("name", pa.string(), nullable=False),
+            pa.field("description", pa.string(), nullable=False),
+            pa.field("type", pa.binary(), nullable=False),
+            pa.field("default_value", pa.binary(), nullable=True),
+        ]  # type: ignore[arg-type]
+    )
+
+    def serialize(self) -> bytes:
+        """Serialize to Arrow IPC bytes."""
+        # Serialize type as a single-field schema
+        type_schema = schema(value=self.type)
+        type_bytes = type_schema.serialize().to_pybytes()
+
+        # Serialize default value if present
+        default_bytes: bytes | None = None
+        if self.default is not None:
+            default_batch = pa.RecordBatch.from_pydict({"value": [self.default]}, schema=type_schema)
+            default_bytes = serialize_record_batch_bytes(default_batch)
+
+        batch = pa.RecordBatch.from_pylist(
+            [
+                {
+                    "name": self.name,
+                    "description": self.desc,
+                    "type": type_bytes,
+                    "default_value": default_bytes,
+                }
+            ],
+            schema=self.ARROW_SCHEMA,
+        )
+        return serialize_record_batch_bytes(batch)
+
+    @classmethod
+    def deserialize(cls, batch: pa.RecordBatch) -> "Self":
+        """Deserialize from Arrow RecordBatch."""
+        from vgi_rpc.utils import _validate_single_row_batch
+
+        row = _validate_single_row_batch(
+            batch,
+            cls.__name__,
+            required_fields=["name", "description", "type"],
+        )
+        # Deserialize type from schema bytes
+        type_schema = pa.ipc.read_schema(pa.py_buffer(cast(bytes, row["type"])))
+        data_type = type_schema.field("value").type
+
+        # Deserialize default value if present
+        default: Any = None
+        if row["default_value"] is not None:
+            default_batch, _ = deserialize_record_batch(cast(bytes, row["default_value"]))
+            default = default_batch.column("value")[0].as_py()
+
+        return cls(
+            name=cast(str, row["name"]),
+            desc=cast(str, row["description"]),
+            type=data_type,
+            default=default,
+        )
+
+
+# Python type to Arrow type mapping
+_PYTHON_TO_ARROW: dict[type, pa.DataType] = {
+    bool: pa.bool_(),
+    int: pa.int64(),
+    float: pa.float64(),
+    str: pa.string(),
+    bytes: pa.binary(),
+}
+
+
+def _resolve_arrow_type(type_hint: type | pa.DataType) -> pa.DataType:
+    """Resolve Arrow type from either a Python type or Arrow DataType.
+
+    Args:
+        type_hint: A Python type (bool, int, float, str, bytes) or Arrow DataType.
+
+    Returns:
+        The resolved Arrow DataType.
+
+    Raises:
+        TypeError: If the type cannot be resolved.
+
+    """
+    # If already an Arrow DataType, use it directly
+    if isinstance(type_hint, pa.DataType):
+        return type_hint
+
+    # Map Python types to Arrow types
+    if type_hint in _PYTHON_TO_ARROW:
+        return _PYTHON_TO_ARROW[type_hint]
+
+    raise TypeError(
+        f"Cannot resolve Arrow type from: {type_hint}. "
+        "Use a Python type (bool, int, float, str, bytes) or Arrow DataType."
+    )
+
+
+@dataclass
+class Setting:
+    """Descriptor for declarative setting definitions using Annotated.
+
+    Use with Annotated type hints to declare settings in a Worker's Settings class.
+    The Arrow type is resolved from the base type in the Annotated hint.
+
+    Attributes:
+        desc: Human-readable description of the setting.
+        arrow_type: Optional explicit Arrow type (overrides inference from annotation).
+
+    """
+
+    desc: str = ""
+    arrow_type: pa.DataType | None = None
+
+    # Internal fields set during class creation
+    _name: str = field(default="", init=False, repr=False)
+
+    def __set_name__(self, owner: type, name: str) -> None:
+        """Store the attribute name when assigned to a class."""
+        self._name = name
+
+    def __get__(self, obj: object | None, objtype: type | None = None) -> Any:
+        """Get the setting value.
+
+        When accessed on the class, returns the descriptor itself.
+        When accessed on an instance, returns the default value.
+        """
+        if obj is None:
+            return self
+        # Return the class-level default
+        return getattr(type(obj), self._name, None)
+
+
+def extract_setting_specs(settings_cls: type) -> list[SettingSpec]:
+    """Extract SettingSpec objects from a Settings class.
+
+    Parses a Settings class with Annotated type hints and extracts
+    SettingSpec objects for each setting definition.
+
+    Args:
+        settings_cls: A class with Annotated[type, Setting(...)] attributes.
+
+    Returns:
+        List of SettingSpec objects extracted from the class.
+
+    Raises:
+        TypeError: If a setting's Arrow type cannot be resolved.
+
+    """
+    specs: list[SettingSpec] = []
+
+    # Get type hints with extras (preserves Annotated)
+    try:
+        hints = get_type_hints(settings_cls, include_extras=True)
+    except Exception:
+        # If type hints can't be resolved, return empty list
+        return specs
+
+    for name, hint in hints.items():
+        # Skip non-Annotated hints
+        if get_origin(hint) is not Annotated:
+            continue
+
+        args = get_args(hint)
+        if len(args) < 2:
+            continue
+
+        base_type = args[0]
+
+        # Find Setting in the annotation args
+        setting = None
+        for arg in args[1:]:
+            if isinstance(arg, Setting):
+                setting = arg
+                break
+
+        if setting is None:
+            continue
+
+        # Get default value from class attribute
+        default = getattr(settings_cls, name, None)
+
+        # Resolve Arrow type: explicit Setting.type takes precedence,
+        # otherwise resolve from base_type (Python type or Arrow DataType)
+        arrow_type = setting.arrow_type if setting.arrow_type is not None else _resolve_arrow_type(base_type)
+
+        specs.append(
+            SettingSpec(
+                name=name,
+                desc=setting.desc,
+                type=arrow_type,
+                default=default,
+            )
+        )
+
+    return specs

vgi/catalog/storage.py

@@ -0,0 +1,372 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Storage for VGI catalog state.
+
+This module provides a storage protocol and implementation for persisting
+catalog attach_opaque_data and transaction_opaque_data state across worker processes.
+
+Protocol:
+    CatalogStorage: Protocol for catalog state persistence.
+
+Implementation:
+    CatalogStorageSqlite: SQLite-backed storage implementation.
+
+"""
+
+import random
+import sqlite3
+import uuid
+from typing import Any, Protocol
+
+from vgi.catalog.catalog_interface import AttachOpaqueData, TransactionOpaqueData
+
+__all__ = [
+    "CatalogStorage",
+    "CatalogStorageSqlite",
+]
+
+
+def _get_default_db_path() -> str:
+    """Return the default SQLite database path for catalog storage."""
+    from pathlib import Path
+
+    from platformdirs import user_state_dir
+
+    state_dir = Path(user_state_dir("vgi"))
+    state_dir.mkdir(parents=True, exist_ok=True)
+    return str((state_dir / "vgi_catalog.db").resolve())
+
+
+class CatalogStorage(Protocol):
+    """Storage protocol for VGI catalog state persistence.
+
+    Provides two access patterns for catalog state:
+
+    **Attachments** - Track catalog attachments with their options.
+    Stores the mapping from attach_opaque_data to catalog name and options.
+
+    **Transactions** - Track active transactions.
+    Stores transaction state for catalogs that support transactions.
+
+    """
+
+    # --- Attachment State ---
+
+    def attach_put(self, attach_opaque_data: AttachOpaqueData, catalog_name: str, options: dict[str, Any]) -> None:
+        """Store attachment state.
+
+        Args:
+            attach_opaque_data: Unique identifier for the attachment.
+            catalog_name: Name of the attached catalog.
+            options: Options passed during attachment.
+
+        """
+        ...
+
+    def attach_get(self, attach_opaque_data: AttachOpaqueData) -> tuple[str, dict[str, Any]] | None:
+        """Retrieve attachment state by attach_opaque_data.
+
+        Args:
+            attach_opaque_data: Unique identifier for the attachment.
+
+        Returns:
+            Tuple of (catalog_name, options), or None if not found.
+
+        """
+        ...
+
+    def attach_delete(self, attach_opaque_data: AttachOpaqueData) -> None:
+        """Delete attachment state.
+
+        Args:
+            attach_opaque_data: Unique identifier for the attachment.
+
+        """
+        ...
+
+    def attach_list(self) -> list[AttachOpaqueData]:
+        """List all active attachments.
+
+        Returns:
+            List of all attach opaque data values in storage.
+
+        """
+        ...
+
+    # --- Transaction State ---
+
+    def transaction_put(
+        self, transaction_opaque_data: TransactionOpaqueData, attach_opaque_data: AttachOpaqueData, state: bytes
+    ) -> None:
+        """Store transaction state.
+
+        Args:
+            transaction_opaque_data: Unique identifier for the transaction.
+            attach_opaque_data: Attachment the transaction belongs to.
+            state: Serialized transaction state.
+
+        """
+        ...
+
+    def transaction_get(self, transaction_opaque_data: TransactionOpaqueData) -> tuple[AttachOpaqueData, bytes] | None:
+        """Retrieve transaction state.
+
+        Args:
+            transaction_opaque_data: Unique identifier for the transaction.
+
+        Returns:
+            Tuple of (attach_opaque_data, state bytes), or None if not found.
+
+        """
+        ...
+
+    def transaction_delete(self, transaction_opaque_data: TransactionOpaqueData) -> None:
+        """Delete transaction state.
+
+        Args:
+            transaction_opaque_data: Unique identifier for the transaction.
+
+        """
+        ...
+
+
+class CatalogStorageSqlite:
+    """SQLite-backed storage for VGI catalog state.
+
+    This implementation uses SQLite with WAL mode to allow multiple worker
+    processes to share catalog state. It manages two tables:
+
+    - catalog_attachments: Maps attach_opaque_data to catalog name and options
+    - catalog_transactions: Tracks active transactions
+
+    """
+
+    def __init__(self, db_path: str | None = None) -> None:
+        """Initialize SQLite catalog storage.
+
+        Args:
+            db_path: Path to the SQLite database file. If None, uses a default
+                location in the user's state directory.
+
+        """
+        self.db_path = db_path if db_path is not None else _get_default_db_path()
+        self._ensure_tables()
+
+    def _connect(self) -> sqlite3.Connection:
+        """Create a new database connection."""
+        conn = sqlite3.connect(self.db_path, timeout=30.0)
+        conn.execute("PRAGMA journal_mode=WAL")
+        return conn
+
+    def _ensure_tables(self) -> None:
+        """Create all storage tables if they don't exist."""
+        conn = self._connect()
+        try:
+            # Attachment table
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS catalog_attachments (
+                    attach_opaque_data BLOB PRIMARY KEY,
+                    catalog_name TEXT NOT NULL,
+                    options_json TEXT NOT NULL,
+                    created_at REAL DEFAULT (julianday('now'))
+                )
+            """)
+            # Transaction table
+            conn.execute("""
+                CREATE TABLE IF NOT EXISTS catalog_transactions (
+                    transaction_opaque_data BLOB PRIMARY KEY,
+                    attach_opaque_data BLOB NOT NULL,
+                    state_data BLOB NOT NULL,
+                    created_at REAL DEFAULT (julianday('now')),
+                    FOREIGN KEY (attach_opaque_data) REFERENCES catalog_attachments(attach_opaque_data)
+                )
+            """)
+            conn.execute("""
+                CREATE INDEX IF NOT EXISTS idx_transactions_attach
+                ON catalog_transactions(attach_opaque_data)
+            """)
+            conn.commit()
+        finally:
+            conn.close()
+
+    # --- Attachment State ---
+
+    def attach_put(self, attach_opaque_data: AttachOpaqueData, catalog_name: str, options: dict[str, Any]) -> None:
+        """Store attachment state."""
+        import json
+
+        # Opportunistically clean old entries (1% of calls)
+        if random.random() < 0.01:
+            self.cleanup_old_entries(max_age_days=7.0)
+
+        options_json = json.dumps(options)
+
+        conn = self._connect()
+        try:
+            conn.execute(
+                """
+                INSERT OR REPLACE INTO catalog_attachments
+                (attach_opaque_data, catalog_name, options_json, created_at)
+                VALUES (?, ?, ?, julianday('now'))
+                """,
+                (attach_opaque_data, catalog_name, options_json),
+            )
+            conn.commit()
+        finally:
+            conn.close()
+
+    def attach_get(self, attach_opaque_data: AttachOpaqueData) -> tuple[str, dict[str, Any]] | None:
+        """Retrieve attachment state by attach_opaque_data."""
+        import json
+
+        conn = self._connect()
+        try:
+            cursor = conn.execute(
+                """SELECT catalog_name, options_json
+                FROM catalog_attachments WHERE attach_opaque_data = ?""",
+                (attach_opaque_data,),
+            )
+            row = cursor.fetchone()
+        finally:
+            conn.close()
+
+        if row is None:
+            return None
+
+        catalog_name: str = row[0]
+        options: dict[str, Any] = json.loads(row[1])
+        return (catalog_name, options)
+
+    def attach_delete(self, attach_opaque_data: AttachOpaqueData) -> None:
+        """Delete attachment state."""
+        conn = self._connect()
+        try:
+            # Delete associated transactions first
+            conn.execute(
+                "DELETE FROM catalog_transactions WHERE attach_opaque_data = ?",
+                (attach_opaque_data,),
+            )
+            conn.execute(
+                "DELETE FROM catalog_attachments WHERE attach_opaque_data = ?",
+                (attach_opaque_data,),
+            )
+            conn.commit()
+        finally:
+            conn.close()
+
+    def attach_list(self) -> list[AttachOpaqueData]:
+        """List all active attachment IDs."""
+        conn = self._connect()
+        try:
+            cursor = conn.execute("SELECT attach_opaque_data FROM catalog_attachments")
+            return [AttachOpaqueData(row[0]) for row in cursor.fetchall()]
+        finally:
+            conn.close()
+
+    # --- Transaction State ---
+
+    def transaction_put(
+        self, transaction_opaque_data: TransactionOpaqueData, attach_opaque_data: AttachOpaqueData, state: bytes
+    ) -> None:
+        """Store transaction state."""
+        # Opportunistically clean old entries (1% of calls)
+        if random.random() < 0.01:
+            self.cleanup_old_entries(max_age_days=7.0)
+
+        conn = self._connect()
+        try:
+            conn.execute(
+                """
+                INSERT OR REPLACE INTO catalog_transactions
+                (transaction_opaque_data, attach_opaque_data, state_data, created_at)
+                VALUES (?, ?, ?, julianday('now'))
+                """,
+                (transaction_opaque_data, attach_opaque_data, state),
+            )
+            conn.commit()
+        finally:
+            conn.close()
+
+    def transaction_get(self, transaction_opaque_data: TransactionOpaqueData) -> tuple[AttachOpaqueData, bytes] | None:
+        """Retrieve transaction state."""
+        conn = self._connect()
+        try:
+            cursor = conn.execute(
+                """SELECT attach_opaque_data, state_data
+                FROM catalog_transactions WHERE transaction_opaque_data = ?""",
+                (transaction_opaque_data,),
+            )
+            row = cursor.fetchone()
+        finally:
+            conn.close()
+
+        if row is None:
+            return None
+
+        return (AttachOpaqueData(row[0]), row[1])
+
+    def transaction_delete(self, transaction_opaque_data: TransactionOpaqueData) -> None:
+        """Delete transaction state."""
+        conn = self._connect()
+        try:
+            conn.execute(
+                "DELETE FROM catalog_transactions WHERE transaction_opaque_data = ?",
+                (transaction_opaque_data,),
+            )
+            conn.commit()
+        finally:
+            conn.close()
+
+    # --- Utility Methods ---
+
+    def generate_attach_opaque_data(self) -> AttachOpaqueData:
+        """Generate a new unique attach_opaque_data.
+
+        Returns:
+            A new AttachOpaqueData based on UUID4.
+
+        """
+        return AttachOpaqueData(uuid.uuid4().bytes)
+
+    def generate_transaction_opaque_data(self) -> TransactionOpaqueData:
+        """Generate a new unique transaction_opaque_data.
+
+        Returns:
+            A new TransactionOpaqueData based on UUID4.
+
+        """
+        return TransactionOpaqueData(uuid.uuid4().bytes)
+
+    # --- Maintenance ---
+
+    def cleanup_old_entries(self, max_age_days: float = 7.0) -> int:
+        """Remove entries older than the specified age from all tables.
+
+        Args:
+            max_age_days: Maximum age in days for entries to keep.
+
+        Returns:
+            Total number of entries deleted.
+
+        """
+        conn = self._connect()
+        try:
+            # Delete old transactions first (foreign key constraint)
+            cursor1 = conn.execute(
+                """
+                DELETE FROM catalog_transactions
+                WHERE julianday('now') - created_at > ?
+                """,
+                (max_age_days,),
+            )
+            cursor2 = conn.execute(
+                """
+                DELETE FROM catalog_attachments
+                WHERE julianday('now') - created_at > ?
+                """,
+                (max_age_days,),
+            )
+            conn.commit()
+            return int(cursor1.rowcount) + int(cursor2.rowcount)
+        finally:
+            conn.close()

vgi/client/__init__.py

@@ -0,0 +1,67 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""VGI client package for communicating with VGI workers.
+
+This package provides:
+- Client: A class for programmatic interaction with VGI workers, including
+  both function invocation and catalog operations
+- ClientError: Exception raised by Client function operations
+- CatalogClientMixin: Mixin class providing catalog operations
+- OutputWriter: Helper for writing output in various formats
+- main: CLI entry point
+
+Usage (API):
+    from vgi.client import Client, ClientError
+    from vgi.arguments import Arguments
+
+    with Client("./my_worker.py") as client:
+        for batch in client.table_in_out_function(
+            function_name="echo",
+            arguments=Arguments(positional=[], named={}),
+            input=input_batches,
+        ):
+            process(batch)
+
+Usage (Catalog API):
+    from vgi.client import Client
+
+    client = Client("./my_worker")
+    result = client.catalog_attach(
+        name="my_catalog", options={}, data_version_spec=None, implementation_version=None
+    )
+
+Usage (CLI):
+    vgi-client --input data.parquet --function echo
+    vgi-client --input data.parquet --function sum_all_columns
+
+"""
+
+from typing import TYPE_CHECKING, Any
+
+from vgi.client.catalog_mixin import CatalogClientMixin
+from vgi.client.client import Client, ClientError, ResumableTableScan, ResumeUnsupported
+
+if TYPE_CHECKING:
+    from vgi.client.cli import OutputWriter, main
+
+__all__ = [
+    "CatalogClientMixin",
+    "Client",
+    "ClientError",
+    "OutputWriter",
+    "ResumableTableScan",
+    "ResumeUnsupported",
+    "main",
+]
+
+
+# Lazy-load the CLI surface. ``vgi.client.cli`` transitively imports
+# ``pyarrow.parquet`` / ``pyarrow._s3fs`` / ``pyarrow._gcsfs`` etc., which add
+# ~2 seconds to the cold import path. Programmatic users of ``Client`` don't
+# need any of that; only the ``vgi-client`` CLI entry point does.
+def __getattr__(name: str) -> Any:
+    if name in {"OutputWriter", "main"}:
+        from vgi.client import cli
+
+        return getattr(cli, name)
+    raise AttributeError(f"module 'vgi.client' has no attribute {name!r}")