feat(copy): on_secrets hook to forward CREATE SECRET creds to COPY writers/readers

CopyToFunction/CopyFromFunction had a @Final on_bind with no seam to request
secrets, so a worker writing to / reading from secret-backed cloud storage
(S3/GCS/HTTP) could not receive the caller's credentials.

Add an overridable on_secrets(params) hook, called from the @Final on_bind on
both base classes. An author calls params.secrets.get(type, scope=...) — typically
scoping by the COPY path — and the framework's existing two-phase secret bind
resolves it from the caller's SecretManager and surfaces it on params.secrets at
write/close (TO) / read (FROM) time. Default is a no-op, so existing formats are
unaffected.

Add SecretLinesCopyToFunction / SecretLinesCopyFromFunction fixtures + unit tests.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: rustyconover

tests/test_copy_to_function.py

@@ -9,8 +9,14 @@
 
 import pyarrow as pa
 
-from vgi._test_fixtures.copy_to import ExampleLinesCopyToArgs, ExampleLinesCopyToFunction
+from vgi._test_fixtures.copy_to import (
+    ExampleLinesCopyToArgs,
+    ExampleLinesCopyToFunction,
+    SecretLinesCopyToArgs,
+    SecretLinesCopyToFunction,
+)
 from vgi._test_fixtures.worker import ExampleCatalog
+from vgi.table_function import ResolvedSecrets, SecretsAccessor
 
 SCHEMA = pa.schema([("a", pa.int64()), ("b", pa.string())])
 
@@ -100,6 +106,81 @@ def test_close_empty_input_with_header_writes_header_only() -> None:
     assert _read(out_name) == "a,b\n"
 
 
+# ---------------------------------------------------------------------------
+# Secret forwarding (CopyToFunction.on_secrets hook)
+# ---------------------------------------------------------------------------
+
+
+def _secret_params(store: _Store, secrets: ResolvedSecrets) -> types.SimpleNamespace:
+    """A TableBufferingParams-shaped stub carrying resolved secrets for write/close."""
+    bind_call = types.SimpleNamespace(input_schema=SCHEMA)
+    init_call = types.SimpleNamespace(bind_call=bind_call)
+    return types.SimpleNamespace(storage=store, init_call=init_call, execution_id=b"x", secrets=secrets)
+
+
+def test_on_secrets_requests_destination_scoped_secret() -> None:
+    """on_secrets() registers a pending lookup scoped to the COPY destination path."""
+    accessor = SecretsAccessor(None)  # nothing resolved yet → first-call behavior
+    params = types.SimpleNamespace(
+        args=SecretLinesCopyToArgs(),  # default secret_type='vgi_example'
+        bind_call=types.SimpleNamespace(copy_to=types.SimpleNamespace(file_path="s3://bucket/out.bin")),
+        secrets=accessor,
+    )
+    SecretLinesCopyToFunction.on_secrets(params)
+    # The hook must have asked the framework for a scoped two-phase resolution.
+    assert accessor.needs_resolution
+    pending = accessor.pending_lookups
+    assert len(pending) == 1
+    assert pending[0].secret_type == "vgi_example"
+    assert pending[0].scope == "s3://bucket/out.bin"
+
+
+def test_close_forwards_resolved_secret_api_key() -> None:
+    """close() writes the resolved (destination-scoped) secret's api_key + row count."""
+    store = _Store()
+    out_name = _tmp_path()
+    secrets = ResolvedSecrets(
+        {
+            "writer_creds": {
+                "type": pa.scalar("vgi_example"),
+                "scope": pa.scalar(out_name),
+                "api_key": pa.scalar("WRITER_KEY"),
+            }
+        }
+    )
+    params = _secret_params(store, secrets)
+    opts = SecretLinesCopyToArgs()
+
+    SecretLinesCopyToFunction.write(
+        batch=pa.record_batch({"a": [1, 2], "b": ["x", "y"]}, schema=SCHEMA),
+        options=opts,
+        file_path=out_name,
+        params=params,
+    )
+    n = SecretLinesCopyToFunction.close(options=opts, file_path=out_name, params=params)
+    assert n == 2
+    assert _read(out_name) == "api_key=WRITER_KEY\nrows=2\n"
+
+
+def test_close_writes_none_when_secret_absent() -> None:
+    """A genuinely missing secret resolves to 'NONE' (silent miss, not an error)."""
+    store = _Store()
+    out_name = _tmp_path()
+    params = _secret_params(store, ResolvedSecrets())
+    opts = SecretLinesCopyToArgs()
+    n = SecretLinesCopyToFunction.close(options=opts, file_path=out_name, params=params)
+    assert n == 0
+    assert _read(out_name) == "api_key=NONE\nrows=0\n"
+
+
+def test_catalog_advertises_secret_lines_out_format() -> None:
+    """The secret-forwarding writer is advertised like any other COPY TO format."""
+    formats = ExampleCatalog().copy_from_formats(attach_opaque_data=b"", transaction_opaque_data=None)
+    by = {(f.direction, f.format_name): f for f in formats}
+    assert ("to", "secret_lines_out") in by
+    assert by[("to", "secret_lines_out")].handler == "secret_lines_writer"
+
+
 def test_catalog_advertises_copy_to_format() -> None:
     """The example catalog advertises example_lines_out with direction='to'."""
     formats = ExampleCatalog().copy_from_formats(attach_opaque_data=b"", transaction_opaque_data=None)

vgi/_test_fixtures/copy_from.py

@@ -27,9 +27,9 @@
 if TYPE_CHECKING:
     from vgi_rpc.rpc import OutputCollector
 
-    from vgi.table_function import ProcessParams
+    from vgi.table_function import BindParams, ProcessParams
 
-__all__ = ["ExampleLinesCopyFromFunction"]
+__all__ = ["ExampleLinesCopyFromFunction", "SecretLinesCopyFromFunction"]
 
 
 @dataclass(slots=True, frozen=True, kw_only=True)
@@ -97,3 +97,59 @@ def read(
             raw = [None if v == options.null_string else v for v in columns[idx]]
             arrays.append(pa.array(raw, type=pa.string()).cast(field.type))
         out.emit(pa.RecordBatch.from_arrays(arrays, schema=expected_schema))
+
+
+@dataclass(slots=True, frozen=True, kw_only=True)
+class SecretLinesCopyFromArgs:
+    """Options for the ``secret_lines_in`` COPY format."""
+
+    secret_type: Annotated[
+        str,
+        Arg("secret_type", default="vgi_example", doc="Secret type to fetch, scoped by the source path"),
+    ] = "vgi_example"
+
+
+class SecretLinesCopyFromFunction(CopyFromFunction[SecretLinesCopyFromArgs]):
+    """COPY ... FROM reader that forwards a ``CREATE SECRET`` credential.
+
+    Exercises the COPY-FROM secret-bind hook (:meth:`CopyFromFunction.on_secrets`):
+    it requests the ``secret_type`` secret scoped to the source path during bind,
+    and :meth:`read` emits a single VARCHAR row holding the resolved secret's
+    ``api_key`` (or ``NONE``) — so a test can assert that the caller's secret
+    reached the reader for a secret-backed cloud source.
+    """
+
+    COPY_FROM_FORMAT: ClassVar[str] = "secret_lines_in"
+    COPY_FROM_COMMENT: ClassVar[str | None] = "Reader that forwards a CREATE SECRET credential (test fixture)"
+
+    class Meta:
+        name = "secret_lines_reader"
+        description = "Emit the resolved secret's api_key as a single VARCHAR row"
+        categories = ["copy", "test", "secret"]
+        tags = {"category": "copy_from", "stability": "test"}
+
+    @classmethod
+    def on_secrets(cls, params: BindParams[SecretLinesCopyFromArgs]) -> None:
+        """Request the source-scoped secret; framework two-phase resolves it."""
+        cf = params.bind_call.copy_from
+        scope = cf.file_path if cf is not None else None
+        params.secrets.get(params.args.secret_type, scope=scope)
+
+    @classmethod
+    def read(
+        cls,
+        *,
+        path: str,
+        options: SecretLinesCopyFromArgs,
+        expected_schema: pa.Schema,
+        params: ProcessParams[SecretLinesCopyFromArgs],
+        out: OutputCollector,
+    ) -> None:
+        """Emit one row carrying the forwarded secret's api_key (or NONE)."""
+        secret = params.secrets.for_scope_of_type(path, options.secret_type) or {}
+        api_key = secret.get("api_key")
+        if api_key is not None and hasattr(api_key, "as_py"):
+            api_key = api_key.as_py()
+        value = "NONE" if api_key is None else str(api_key)
+        arrays = [pa.array([value], type=pa.string()).cast(expected_schema.field(0).type)]
+        out.emit(pa.RecordBatch.from_arrays(arrays, schema=expected_schema))

vgi/_test_fixtures/copy_to.py

@@ -29,8 +29,13 @@
 
 if TYPE_CHECKING:
     from vgi.table_buffering_function import TableBufferingParams
+    from vgi.table_function import BindParams
 
-__all__ = ["ExampleLinesCopyToFunction", "ExampleLinesOrderedCopyToFunction"]
+__all__ = [
+    "ExampleLinesCopyToFunction",
+    "ExampleLinesOrderedCopyToFunction",
+    "SecretLinesCopyToFunction",
+]
 
 _SHARD_NS = b"copy_to_shard"
 
@@ -158,3 +163,75 @@ class Meta:
         categories = ["copy", "test"]
         tags = {"category": "copy_to", "stability": "test"}
         sink_order_dependent = True  # ordered COPY TO → single-thread sink
+
+
+_SECRET_NS = b"copy_to_secret_shard"
+
+
+@dataclass(slots=True, frozen=True, kw_only=True)
+class SecretLinesCopyToArgs:
+    """Options for the ``secret_lines_out`` COPY format."""
+
+    secret_type: Annotated[
+        str,
+        Arg("secret_type", default="vgi_example", doc="Secret type to fetch, scoped by the destination path"),
+    ] = "vgi_example"
+
+
+class SecretLinesCopyToFunction(CopyToFunction[SecretLinesCopyToArgs]):
+    """COPY ... TO writer that forwards a ``CREATE SECRET`` credential.
+
+    Exercises the COPY-TO secret-bind hook (:meth:`CopyToFunction.on_secrets`):
+    it requests the ``secret_type`` secret scoped to the destination path during
+    bind, and ``close()`` writes the resolved secret's ``api_key`` (or ``NONE``)
+    plus the row count into the destination — so a test can assert that the
+    caller's secret reached the writer for a secret-backed cloud write.
+    """
+
+    COPY_TO_FORMAT: ClassVar[str] = "secret_lines_out"
+    COPY_TO_COMMENT: ClassVar[str | None] = "Writer that forwards a CREATE SECRET credential (test fixture)"
+
+    class Meta:
+        name = "secret_lines_writer"
+        description = "Write the resolved secret's api_key + row count to the destination"
+        categories = ["copy", "test", "secret"]
+        tags = {"category": "copy_to", "stability": "test"}
+
+    @classmethod
+    def on_secrets(cls, params: BindParams[SecretLinesCopyToArgs]) -> None:
+        """Request the destination-scoped secret; framework two-phase resolves it."""
+        cf = params.bind_call.copy_to
+        scope = cf.file_path if cf is not None else None
+        # Scoped lookup → longest-prefix match against the caller's CREATE SECRETs.
+        params.secrets.get(params.args.secret_type, scope=scope)
+
+    @classmethod
+    def write(
+        cls,
+        *,
+        batch: pa.RecordBatch,
+        options: SecretLinesCopyToArgs,
+        file_path: str,
+        params: TableBufferingParams[SecretLinesCopyToArgs],
+    ) -> None:
+        """Record this shard's row count (cross-process-safe append)."""
+        params.storage.state_append(_SECRET_NS, b"", str(batch.num_rows).encode())
+
+    @classmethod
+    def close(
+        cls,
+        *,
+        options: SecretLinesCopyToArgs,
+        file_path: str,
+        params: TableBufferingParams[SecretLinesCopyToArgs],
+    ) -> int:
+        """Write the forwarded secret's api_key + total row count, once."""
+        secret = params.secrets.for_scope_of_type(file_path, options.secret_type) or {}
+        api_key = secret.get("api_key")
+        if api_key is not None and hasattr(api_key, "as_py"):
+            api_key = api_key.as_py()
+        total = sum(int(blob) for _id, blob in params.storage.state_log_scan(_SECRET_NS, b"", after_id=-1))
+        with open(file_path, "w", encoding="utf-8") as fh:
+            fh.write(f"api_key={'NONE' if api_key is None else api_key}\n")
+            fh.write(f"rows={total}\n")
+        return total

vgi/_test_fixtures/worker.py

@@ -59,8 +59,12 @@
     SlowCancellableFunction,
     SlowCancellableInOutFunction,
 )
-from vgi._test_fixtures.copy_from import ExampleLinesCopyFromFunction
-from vgi._test_fixtures.copy_to import ExampleLinesCopyToFunction, ExampleLinesOrderedCopyToFunction
+from vgi._test_fixtures.copy_from import ExampleLinesCopyFromFunction, SecretLinesCopyFromFunction
+from vgi._test_fixtures.copy_to import (
+    ExampleLinesCopyToFunction,
+    ExampleLinesOrderedCopyToFunction,
+    SecretLinesCopyToFunction,
+)
 from vgi._test_fixtures.nest_tensor import NestTensorFunction, UnnestTensorFunction, UnnestTensorRowsFunction
 from vgi._test_fixtures.scalar import (
     AddValuesFunction,
@@ -355,9 +359,11 @@ def _build_enum_stats() -> dict[str, ColumnStatisticsInput]:
                 SlowCancellableBufferingFunction,
                 # CopyFromFunction - custom COPY ... FROM format reader
                 ExampleLinesCopyFromFunction,
+                SecretLinesCopyFromFunction,
                 # CopyToFunction - custom COPY ... TO format writer
                 ExampleLinesCopyToFunction,
                 ExampleLinesOrderedCopyToFunction,
+                SecretLinesCopyToFunction,
                 # TableFunctionGenerator - generate output without input
                 ConstantColumnsFunction,
                 SlowCancellableFunction,

vgi/copy_from_function.py

@@ -98,10 +98,32 @@ def on_bind(cls, params: BindParams[TArgs]) -> BindResponse:
                 f"{cls.__name__} is a COPY FROM format reader; invoke it via "
                 f"COPY <table> FROM '<path>' (FORMAT {fmt}), not as a table function."
             )
+        # Forward credentials for secret-backed cloud sources (S3/GCS/HTTP/...)
+        # via the framework's two-phase secret bind. on_bind is @final (the output
+        # schema is fixed to the COPY target), so on_secrets is the seam.
+        cls.on_secrets(params)
         # ``expected_schema`` is transparently a ``pa.Schema`` here — the
         # ArrowType(binary) annotation only governs the wire encoding.
         return BindResponse(output_schema=cf.expected_schema)
 
+    @classmethod
+    def on_secrets(cls, params: BindParams[TArgs]) -> None:
+        """Request the credentials this reader needs to reach its source.
+
+        Override to forward ``CREATE SECRET`` values to :meth:`read` for
+        secret-backed cloud sources. Call ``params.secrets.get(secret_type,
+        scope=..., name=...)`` — typically scoping by the source path
+        (``params.bind_call.copy_from.file_path``) so DuckDB resolves the
+        longest-prefix-matching secret. The framework issues a two-phase bind retry
+        to resolve every requested secret from the caller's secret store, then
+        surfaces the resolved values on ``params.secrets`` (a
+        :class:`ResolvedSecrets`) at :meth:`read` time. Pass ``required=True`` to
+        ``get()`` to make a missing secret fail the bind.
+
+        Default: request nothing (no secrets forwarded).
+        """
+        # Default no-op. Subclasses override to call params.secrets.get(...).
+
     @final
     @classmethod
     def initial_state(cls, params: ProcessParams[TArgs]) -> _CopyFromState:

vgi/copy_to_function.py

@@ -96,9 +96,40 @@ class CopyToFunction[TArgs](TableBufferingFunction[TArgs, None]):
     @final
     @classmethod
     def on_bind(cls, params: BindParams[TArgs]) -> BindResponse:
-        """A sink produces no rows — bind to an empty output schema."""
+        """A sink produces no rows — bind to an empty output schema.
+
+        ``on_bind`` is ``@final`` (a writer has no output schema to compute), but
+        a cloud writer still needs credentials. The seam is :meth:`on_secrets`,
+        called here so a subclass can request ``CREATE SECRET`` values via the
+        framework's two-phase secret bind without overriding ``on_bind`` itself.
+        """
+        cls.on_secrets(params)
         return BindResponse(output_schema=pa.schema([]))
 
+    @classmethod
+    def on_secrets(cls, params: BindParams[TArgs]) -> None:
+        """Request the credentials this writer needs to reach its destination.
+
+        Override to forward ``CREATE SECRET`` values to :meth:`write` / :meth:`close`
+        for secret-backed cloud writes (S3/GCS/HTTP/...). Call
+        ``params.secrets.get(secret_type, scope=..., name=...)`` — typically scoping
+        by the destination path (:meth:`copy_to_path` / ``params.bind_call.copy_to``)
+        so DuckDB resolves the longest-prefix-matching secret. The framework issues a
+        two-phase bind retry to resolve every requested secret from the caller's
+        secret store, then surfaces the resolved values on ``params.secrets`` (a
+        :class:`ResolvedSecrets`) at :meth:`write` / :meth:`close` time. Requested
+        secrets that don't exist resolve to "not found" rather than an error — pass
+        ``required=True`` to ``get()`` to make a missing secret fail the bind.
+
+        The destination path is available without an ``init_call`` here via
+        ``params.bind_call.copy_to.file_path`` (use :meth:`copy_to_path` only at
+        write/close time, where an ``init_call`` exists).
+
+        Default: request nothing (no secrets forwarded), so existing writers that
+        never touched credentials are unaffected.
+        """
+        # Default no-op. Subclasses override to call params.secrets.get(...).
+
     @final
     @classmethod
     def copy_to_path(cls, params: TableBufferingParams[TArgs]) -> str: