vgi-python 0.8.0__py3-none-any.whl

vgi/_test_fixtures/table/settings.py

@@ -0,0 +1,426 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Settings/secrets fixtures: settings_aware, struct_settings, secret_demo, scoped_secret_demo."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Annotated, Any, ClassVar
+
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass
+from vgi_rpc.rpc import OutputCollector
+
+from vgi._test_fixtures.table._common import (
+    _cardinality_from_count,
+)
+from vgi.arguments import Arg, Secret, Setting
+from vgi.invocation import BindResponse
+from vgi.metadata import FunctionExample
+from vgi.schema_utils import schema
+from vgi.table_function import (
+    BindParams,
+    ProcessParams,
+    TableFunctionGenerator,
+    init_single_worker,
+)
+
+
+@dataclass(slots=True, frozen=True)
+class SettingsAwareFunctionArguments:
+    """Arguments for SettingsAwareFunction."""
+
+    count: Annotated[int, Arg(0, doc="Number of rows to generate", ge=0)]
+
+
+@dataclass(kw_only=True)
+class SettingsAwareState(ArrowSerializableDataclass):
+    """Mutable state for SettingsAwareFunction with typed settings."""
+
+    remaining: int
+    current_index: int = 0
+    verbose: bool = False
+    greeting: str = "Hello"
+    multiplier: int = 1
+
+
+@init_single_worker
+@_cardinality_from_count
+class SettingsAwareFunction(TableFunctionGenerator[SettingsAwareFunctionArguments, SettingsAwareState]):
+    """Generates data demonstrating that settings are passed to functions.
+
+    USE CASE
+    --------
+    Demonstrates how functions can declare required settings via
+    Setting() annotations and access them via state (resolved once
+    in initial_state()). The output includes columns showing the actual
+    setting values that were passed.
+
+    This function uses three settings:
+    - vgi_verbose_mode: bool - when true, adds a details column
+    - greeting: str - a custom greeting message echoed in output
+    - multiplier: int - multiplies the value column
+
+    Settings are typed: the C++ extension sends Arrow scalars with proper
+    types (bool, int64, string). For backward compatibility, string values
+    like "true" are also accepted for vgi_verbose_mode.
+
+    SCHEMA
+    ------
+    Base output: {"id": int64, "greeting": string, "value": float64}
+    With vgi_verbose_mode=true: adds "details": string column
+
+    Example:
+    -------
+    With settings={vgi_verbose_mode: true, greeting: "Hi", multiplier: 2}:
+    Returns: [{"id": 0, "greeting": "Hi", "value": 0.0, "details": "row_0"}, ...]
+
+    """
+
+    class Meta:
+        """Metadata for SettingsAwareFunction."""
+
+        name = "settings_aware"
+        description = "Generates data demonstrating settings are passed"
+        categories = ["generator", "settings"]
+        examples = [
+            FunctionExample(
+                sql="SELECT * FROM settings_aware(5)",
+                description="Generate 5 rows showing setting values",
+            )
+        ]
+
+    BATCH_SIZE: ClassVar[int] = 1000
+
+    @staticmethod
+    def _is_verbose(val: object) -> bool:
+        """Check if verbose mode is enabled, handling both bool and string values."""
+        return val is True or val == "true"
+
+    @classmethod
+    def on_bind(
+        cls,
+        params: BindParams[SettingsAwareFunctionArguments],
+        *,
+        vgi_verbose_mode: Annotated[pa.Scalar[Any] | None, Setting()] = None,
+        greeting: Annotated[pa.Scalar[Any] | None, Setting()] = None,
+        multiplier: Annotated[pa.Scalar[Any] | None, Setting()] = None,
+    ) -> BindResponse:
+        """Return output schema based on vgi_verbose_mode setting.
+
+        Always includes id, greeting (from setting), and value (multiplied).
+        When vgi_verbose_mode is true, includes an extra "details" column.
+        """
+        fields: list[pa.Field[pa.DataType]] = [
+            pa.field("id", pa.int64()),
+            pa.field("greeting", pa.string()),
+            pa.field("value", pa.float64()),
+        ]
+
+        # Add details column if verbose mode is enabled (handles bool and string)
+        if vgi_verbose_mode is not None and cls._is_verbose(vgi_verbose_mode.as_py()):
+            fields.append(pa.field("details", pa.string()))
+
+        return BindResponse(output_schema=pa.schema(fields))
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[SettingsAwareFunctionArguments]) -> SettingsAwareState:
+        """Create initial state with typed settings resolved once."""
+        verbose_val = params.settings.get("vgi_verbose_mode", pa.scalar(False)).as_py()
+        greeting_val = params.settings.get("greeting", pa.scalar("Hello")).as_py()
+        multiplier_val = params.settings.get("multiplier", pa.scalar(1)).as_py()
+
+        return SettingsAwareState(
+            remaining=params.args.count,
+            verbose=cls._is_verbose(verbose_val),
+            greeting=str(greeting_val),
+            multiplier=int(multiplier_val),
+        )
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[SettingsAwareFunctionArguments],
+        state: SettingsAwareState,
+        out: OutputCollector,
+    ) -> None:
+        """Generate data based on settings stored in state."""
+        if state.remaining <= 0:
+            out.finish()
+            return
+
+        size = min(state.remaining, cls.BATCH_SIZE)
+        ids = list(range(state.current_index, state.current_index + size))
+
+        data: dict[str, list[int] | list[float] | list[str]] = {
+            "id": ids,
+            "greeting": [state.greeting] * size,
+            "value": [float(i) * 2.5 * state.multiplier for i in ids],
+        }
+
+        if state.verbose:
+            data["details"] = [f"row_{i}" for i in ids]
+
+        out.emit(pa.RecordBatch.from_pydict(data, schema=params.output_schema))
+
+        state.current_index += size
+        state.remaining -= size
+
+
+@dataclass(slots=True, frozen=True)
+class StructSettingsFunctionArguments:
+    """Arguments for StructSettingsFunction."""
+
+    count: Annotated[int, Arg(0, doc="Number of rows to generate", ge=0)]
+
+
+@dataclass(kw_only=True)
+class StructSettingsState(ArrowSerializableDataclass):
+    """Mutable state for StructSettingsFunction."""
+
+    remaining: int
+    current_index: int = 0
+    start: int = 0
+    step: int = 1
+    label: str = "item"
+
+
+@init_single_worker
+@_cardinality_from_count
+class StructSettingsFunction(TableFunctionGenerator[StructSettingsFunctionArguments, StructSettingsState]):
+    """Generates a sequence configured by a struct setting.
+
+    USE CASE
+    --------
+    Demonstrates how a single struct setting can configure multiple aspects
+    of a function's behavior. The config setting is a struct with fields:
+    - start: int64 - starting value for the sequence
+    - step: int64 - step between values
+    - label: string - prefix for label column
+
+    SCHEMA
+    ------
+    Output: {"n": int64, "label": string}
+
+    Example:
+    -------
+    With config={'start': 10, 'step': 5, 'label': 'item'} and count=3:
+    Returns: [{"n": 10, "label": "item_0"}, {"n": 15, "label": "item_1"}, {"n": 20, "label": "item_2"}]
+
+    """
+
+    class Meta:
+        """Metadata for StructSettingsFunction."""
+
+        name = "struct_settings"
+        description = "Generate a sequence configured by a struct setting"
+        categories = ["generator", "settings"]
+        examples = [
+            FunctionExample(
+                sql="SELECT * FROM struct_settings(5)",
+                description="Generate 5 rows configured by the config setting",
+            )
+        ]
+
+    FIXED_SCHEMA: ClassVar[pa.Schema] = schema({"n": pa.int64(), "label": pa.string()})
+
+    @classmethod
+    def on_bind(
+        cls,
+        params: BindParams[StructSettingsFunctionArguments],
+        *,
+        config: Annotated[pa.Scalar[Any] | None, Setting()] = None,
+    ) -> BindResponse:
+        """Return output schema. Config declared here for required_settings registration."""
+        return BindResponse(output_schema=cls.FIXED_SCHEMA)
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[StructSettingsFunctionArguments]) -> StructSettingsState:
+        """Create initial state with struct setting values resolved once."""
+        config = params.settings["config"]  # pa.StructScalar
+        cfg = config.as_py()  # dict
+        return StructSettingsState(
+            remaining=params.args.count,
+            start=cfg["start"],
+            step=cfg["step"],
+            label=cfg["label"],
+        )
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[StructSettingsFunctionArguments],
+        state: StructSettingsState,
+        out: OutputCollector,
+    ) -> None:
+        """Generate rows with values derived from the struct setting."""
+        if state.remaining <= 0:
+            out.finish()
+            return
+
+        size = min(state.remaining, 1000)
+        data: dict[str, list[int] | list[str]] = {
+            "n": [state.start + (state.current_index + i) * state.step for i in range(size)],
+            "label": [f"{state.label}_{state.current_index + i}" for i in range(size)],
+        }
+        out.emit(pa.RecordBatch.from_pydict(data, schema=params.output_schema))
+        state.current_index += size
+        state.remaining -= size
+
+
+# =============================================================================
+
+
+@dataclass(kw_only=True)
+class SecretDemoState(ArrowSerializableDataclass):
+    """State for SecretDemoFunction."""
+
+    keys: list[str] = field(default_factory=list)
+    values: list[str] = field(default_factory=list)
+    types: list[str] = field(default_factory=list)
+
+
+@init_single_worker
+class SecretDemoFunction(TableFunctionGenerator[None, SecretDemoState]):
+    """Table function that outputs secret key-value pairs as rows.
+
+    Demonstrates basic secret access via Secret() annotation.
+    """
+
+    class Meta:
+        """Metadata for SecretDemoFunction."""
+
+        name = "secret_demo"
+        description = "Outputs secret contents as key-value rows"
+
+    @classmethod
+    def on_bind(
+        cls,
+        params: BindParams[None],
+    ) -> BindResponse:
+        """Bind with secret request via SecretsAccessor."""
+        # Request the secret via the accessor — triggers two-phase bind
+        # so the resolved secret is available in initial_state().
+        params.secrets.get("vgi_example")
+        return BindResponse(
+            output_schema=schema(
+                {
+                    "key": pa.string(),
+                    "value": pa.string(),
+                    "arrow_type": pa.string(),
+                }
+            )
+        )
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[None]) -> SecretDemoState:
+        """Build initial state from secret key-value pairs."""
+        secret = params.secrets.get("vgi_example", {})
+        keys = list(secret.keys())
+        values = [str(v.as_py()) for v in secret.values()]
+        types = [str(v.type) for v in secret.values()]
+        return SecretDemoState(keys=keys, values=values, types=types)
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[None],
+        state: SecretDemoState,
+        out: OutputCollector,
+    ) -> None:
+        """Emit secret entries as rows."""
+        if not state.keys:
+            out.finish()
+            return
+        batch = pa.RecordBatch.from_pydict(
+            {"key": state.keys, "value": state.values, "arrow_type": state.types},
+            schema=params.output_schema,
+        )
+        out.emit(batch)
+        state.keys = []
+        state.values = []
+        state.types = []
+
+
+@dataclass(frozen=True)
+class ScopedSecretDemoArgs:
+    """Arguments for ScopedSecretDemoFunction."""
+
+    path: Annotated[str, Arg(0, doc="Scope path for secret lookup")]
+
+
+@dataclass(kw_only=True)
+class ScopedSecretDemoState(ArrowSerializableDataclass):
+    """State for ScopedSecretDemoFunction."""
+
+    found: bool = False
+    secret_keys: str = ""
+
+
+@init_single_worker
+class ScopedSecretDemoFunction(TableFunctionGenerator[ScopedSecretDemoArgs, ScopedSecretDemoState]):
+    """Demonstrates automatic two-phase bind with scoped secrets.
+
+    Requests a secret with a dynamic scope computed from the function argument.
+    The framework automatically handles the two-phase bind retry.
+    """
+
+    class Meta:
+        """Metadata for ScopedSecretDemoFunction."""
+
+        name = "scoped_secret_demo"
+        description = "Demo: resolves scoped secret based on argument"
+
+    @classmethod
+    def on_bind(
+        cls,
+        params: BindParams[ScopedSecretDemoArgs],
+        *,
+        vgi_example: Annotated[dict[str, pa.Scalar[Any]] | None, Secret("vgi_example")] = None,
+    ) -> BindResponse:
+        """Bind with dynamic scoped secret lookup."""
+        # Request secret with dynamic scope — framework handles retry automatically.
+        # The get() call registers a pending scoped lookup; the return value is
+        # unused because the framework will trigger a two-phase bind retry.
+        params.secrets.get("vgi_example", scope=params.args.path)
+
+        # On first call: secret is None (pending), framework triggers retry
+        # On retry: secret is dict (found) or None (genuinely not found)
+
+        return BindResponse(
+            output_schema=schema(
+                {
+                    "scope": pa.string(),
+                    "found": pa.bool_(),
+                    "secret_keys": pa.string(),
+                }
+            )
+        )
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[ScopedSecretDemoArgs]) -> ScopedSecretDemoState:
+        """Build state from resolved secrets."""
+        secret = params.secrets.get("vgi_example", {})
+        return ScopedSecretDemoState(
+            found=bool(secret),
+            secret_keys=",".join(secret.keys()) if secret else "",
+        )
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[ScopedSecretDemoArgs],
+        state: ScopedSecretDemoState,
+        out: OutputCollector,
+    ) -> None:
+        """Emit scope info and resolved secret keys."""
+        batch = pa.RecordBatch.from_pydict(
+            {
+                "scope": [params.args.path],
+                "found": [state.found],
+                "secret_keys": [state.secret_keys],
+            },
+            schema=params.output_schema,
+        )
+        out.emit(batch)
+        out.finish()

vgi/_test_fixtures/table/transaction_storage.py

@@ -0,0 +1,162 @@
+# Copyright 2025, 2026 Query Farm LLC - https://query.farm
+
+"""Demo of transaction-scoped storage (``BindParams.transaction_storage``).
+
+Backs the ``example.main.tx_cached_value(key, seed)`` function exposed by
+``vgi-fixture-worker``. The function uses ``BindParams.transaction_storage``
+to cache its ``seed`` argument per ``(transaction_opaque_data, key)``:
+
+* First call within a transaction for a given ``key``: stores ``seed`` and
+  emits it.
+* Subsequent calls within the **same** transaction for the **same** ``key``:
+  emit the originally-cached value and **ignore** the new ``seed``.
+* New transaction or different ``key``: produces a fresh cached value.
+* Without a transaction (``params.transaction_storage is None``): no caching;
+  every call emits its own ``seed``.
+
+The resolved value is shipped from ``on_bind`` to ``process`` via
+``BindResponse.opaque_data`` so any worker in the pool can produce the same
+answer — the value lives in shared storage (sqlite/CF DO/Azure SQL), not in
+the bind worker's local memory.
+"""
+
+from __future__ import annotations
+
+import struct
+from dataclasses import dataclass
+from typing import Annotated, ClassVar
+
+import pyarrow as pa
+from vgi_rpc import ArrowSerializableDataclass
+from vgi_rpc.rpc import OutputCollector
+
+from vgi.arguments import Arg
+from vgi.invocation import BindResponse, GlobalInitResponse
+from vgi.schema_utils import schema
+from vgi.table_function import (
+    BindParams,
+    InitParams,
+    ProcessParams,
+    TableCardinality,
+    TableFunctionGenerator,
+)
+
+__all__ = ["TxCachedValueFunction"]
+
+
+@dataclass(frozen=True)
+class TxCachedValueArgs:
+    """Arguments for ``tx_cached_value``."""
+
+    key: Annotated[str, Arg(0, doc="Cache key, scoped to the current transaction")]
+    seed: Annotated[int, Arg(1, doc="Value to cache on first call; ignored on cache hit")]
+
+
+@dataclass(kw_only=True)
+class _TxCachedValueState(ArrowSerializableDataclass):
+    """Mutable per-process state carried into ``process``."""
+
+    # Resolved value (cached or freshly-seeded). Carried from bind via
+    # opaque_data so process() doesn't need access to transaction_storage
+    # (which is only populated on BindParams).
+    value: int
+    emitted: bool = False
+
+
+class TxCachedValueFunction(TableFunctionGenerator[TxCachedValueArgs, _TxCachedValueState]):
+    """Returns a single-row table whose value is cached per (transaction, key).
+
+    The cache lives in ``BindParams.transaction_storage`` — a view over
+    ``FunctionStorage.transaction_state_*``. On a cache hit the stored value
+    is returned; on a miss, the supplied ``seed`` is written to storage and
+    returned.
+
+    Without a transaction (``params.transaction_storage is None``) every
+    bind acts as a cache miss and emits the caller's ``seed`` verbatim —
+    so the same SQL run inside vs. outside a ``BEGIN``/``COMMIT`` block
+    visibly differs.
+    """
+
+    FunctionArguments = TxCachedValueArgs
+    State = _TxCachedValueState
+
+    class Meta:
+        """Metadata for tx_cached_value."""
+
+        name = "tx_cached_value"
+        description = "Return a value cached per (transaction_opaque_data, key) via transaction_storage."
+        categories = ["test", "transaction-storage"]
+        tags = {"category": "test"}
+
+    OUTPUT_SCHEMA: ClassVar[pa.Schema] = schema({"v": pa.int64()})
+
+    @staticmethod
+    def _storage_key(user_key: str) -> bytes:
+        """Storage key — namespaced so unrelated demos can share one transaction."""
+        return f"vgi-fixture:tx_cached_value:{user_key}".encode()
+
+    @classmethod
+    def on_bind(cls, params: BindParams[TxCachedValueArgs]) -> BindResponse:
+        """Resolve the value via transaction_storage, ship it via opaque_data."""
+        storage = params.transaction_storage
+        if storage is not None:
+            key = cls._storage_key(params.args.key)
+            cached = storage.get_one(key)
+            if cached is not None:
+                value = struct.unpack(">q", cached)[0]
+            else:
+                value = params.args.seed
+                storage.put_one(key, struct.pack(">q", value))
+        else:
+            # No transaction → no caching possible. Every call is a fresh
+            # bind that uses the caller's seed verbatim.
+            value = params.args.seed
+
+        return BindResponse(
+            output_schema=cls.OUTPUT_SCHEMA,
+            opaque_data=struct.pack(">q", value),
+        )
+
+    @classmethod
+    def cardinality(cls, params: BindParams[TxCachedValueArgs]) -> TableCardinality:
+        """One row, always."""
+        del params
+        return TableCardinality(estimate=1, max=1)
+
+    @classmethod
+    def on_init(cls, params: InitParams[TxCachedValueArgs]) -> GlobalInitResponse:
+        """Pass the resolved value through to process().
+
+        ``max_workers=1`` because this function emits exactly one row and
+        does no work-queue coordination — running it in parallel would
+        cause every secondary worker to re-emit the same row.
+        """
+        return GlobalInitResponse(
+            max_workers=1,
+            opaque_data=params.init_call.bind_opaque_data,
+        )
+
+    @classmethod
+    def initial_state(cls, params: ProcessParams[TxCachedValueArgs]) -> _TxCachedValueState:
+        """Decode the opaque_data shipped from on_bind()."""
+        assert params.init_response is not None
+        opaque = params.init_response.opaque_data
+        assert opaque is not None and len(opaque) == 8, (
+            "tx_cached_value: bind must populate opaque_data with an 8-byte int"
+        )
+        return _TxCachedValueState(value=struct.unpack(">q", opaque)[0])
+
+    @classmethod
+    def process(
+        cls,
+        params: ProcessParams[TxCachedValueArgs],
+        state: _TxCachedValueState,
+        out: OutputCollector,
+    ) -> None:
+        """Emit the resolved value as a single-row batch, then finish."""
+        del params
+        if state.emitted:
+            out.finish()
+            return
+        out.emit(pa.RecordBatch.from_pydict({"v": [state.value]}, schema=cls.OUTPUT_SCHEMA))
+        state.emitted = True