|
9 | 9 | from typing import Annotated, Any |
10 | 10 |
|
11 | 11 | import pyarrow as pa |
| 12 | +import pytest |
12 | 13 |
|
13 | 14 | from tests.conftest import FIXTURE_WORKER |
14 | 15 | from vgi.arguments import ( |
@@ -418,6 +419,16 @@ def test_bind_params_secrets_type(self) -> None: |
418 | 419 | assert len(outputs) > 0 |
419 | 420 |
|
420 | 421 |
|
| 422 | +def _wire_secrets(*secrets: dict[str, str]) -> pa.RecordBatch: |
| 423 | + """Build a secrets batch the way the DuckDB extension does. |
| 424 | +
|
| 425 | + One row, one struct column per secret, the column named by the secret's |
| 426 | + ``name``; ``type``, ``provider``, ``name`` and ``scope`` are fields. |
| 427 | + """ |
| 428 | + columns = {s["name"]: [{"provider": "config", **s}] for s in secrets} |
| 429 | + return pa.RecordBatch.from_pydict(columns) |
| 430 | + |
| 431 | + |
421 | 432 | # --------------------------------------------------------------------------- |
422 | 433 | # Tests for SecretsAccessor |
423 | 434 | # --------------------------------------------------------------------------- |
@@ -461,57 +472,157 @@ def test_scoped_get_registers_pending(self) -> None: |
461 | 472 | assert entry.secret_type == "s3" |
462 | 473 | assert entry.scope == "s3://bucket/" |
463 | 474 |
|
464 | | - def test_to_dict_combines_unscoped_and_scoped(self) -> None: |
465 | | - """to_dict() includes both unscoped and scoped entries.""" |
466 | | - # Build a batch with an unscoped column and a scoped column |
467 | | - struct_type = pa.struct([("k", pa.string())]) |
468 | | - unscoped_arr = pa.array([{"k": "unscoped_val"}], type=struct_type) |
469 | | - |
470 | | - scoped_field = pa.field( |
471 | | - "secret_0", |
472 | | - struct_type, |
473 | | - metadata={"secret_type": "scoped_type", "scope": "s3://x/"}, |
| 475 | + def test_to_dict_is_keyed_by_secret_name(self) -> None: |
| 476 | + """to_dict() keeps every resolved secret under its own name, nulls dropped.""" |
| 477 | + batch = _wire_secrets( |
| 478 | + {"type": "s3", "name": "prod", "scope": "s3://prod/", "key_id": "P"}, |
| 479 | + {"type": "s3", "name": "dev", "scope": "s3://dev/", "key_id": "D"}, |
474 | 480 | ) |
475 | | - scoped_arr = pa.array([{"k": "scoped_val"}], type=struct_type) |
| 481 | + struct_type = pa.struct([("type", pa.string())]) |
| 482 | + batch = batch.append_column(pa.field("gone", struct_type), pa.array([None], type=struct_type)) |
| 483 | + result = SecretsAccessor(batch).to_dict() |
| 484 | + assert sorted(result) == ["dev", "prod"] |
| 485 | + assert result["prod"]["key_id"].as_py() == "P" |
476 | 486 |
|
477 | | - batch = pa.RecordBatch.from_arrays( |
478 | | - [unscoped_arr, scoped_arr], |
479 | | - schema=pa.schema([pa.field("my_type", struct_type), scoped_field]), |
| 487 | + def test_a_secret_named_like_the_old_wire_columns_is_an_ordinary_secret(self) -> None: |
| 488 | + """A user may name a secret ``secret_0``; it is not a special column.""" |
| 489 | + batch = _wire_secrets({"type": "s3", "name": "secret_0", "scope": "s3://", "key_id": "K"}) |
| 490 | + accessor = SecretsAccessor(batch, is_retry=True) |
| 491 | + assert accessor.to_dict()["secret_0"]["key_id"].as_py() == "K" |
| 492 | + found = accessor.get("s3", name="secret_0") |
| 493 | + assert found is not None |
| 494 | + assert found["key_id"].as_py() == "K" |
| 495 | + |
| 496 | + def test_get_by_name_on_retry_returns_the_secret(self) -> None: |
| 497 | + """The two-phase bind: a by-name lookup registers, then the retry returns it. |
| 498 | +
|
| 499 | + Regression: the retry looked for ``secret_N`` columns that no client |
| 500 | + sends, so the lookup came back ``None`` although the secret was in the |
| 501 | + batch, and ``required=True`` raised. |
| 502 | + """ |
| 503 | + first = SecretsAccessor(None) |
| 504 | + assert first.get("s3", name="prod", required=True) is None |
| 505 | + assert first.pending_lookups == [SecretLookupEntry(secret_type="s3", secret_name="prod")] |
| 506 | + |
| 507 | + batch = _wire_secrets({"type": "s3", "name": "prod", "scope": "s3://", "key_id": "P"}) |
| 508 | + retry = SecretsAccessor(batch, is_retry=True) |
| 509 | + found = retry.get("s3", name="prod", required=True) |
| 510 | + assert found is not None |
| 511 | + assert found["key_id"].as_py() == "P" |
| 512 | + assert retry.all_resolved |
| 513 | + |
| 514 | + def test_get_by_name_checks_the_type(self) -> None: |
| 515 | + """A secret with the right name but another type is not a match.""" |
| 516 | + batch = _wire_secrets({"type": "gcs", "name": "prod", "scope": "gcs://"}) |
| 517 | + accessor = SecretsAccessor(batch, is_retry=True) |
| 518 | + assert accessor.get("s3", name="prod") is None |
| 519 | + with pytest.raises(ValueError, match="Required secret 's3' not found"): |
| 520 | + accessor.get("s3", name="prod", required=True) |
| 521 | + |
| 522 | + def test_get_by_name_already_resolved_needs_no_retry(self) -> None: |
| 523 | + """A by-name lookup the client pre-resolved is answered on the first bind.""" |
| 524 | + batch = _wire_secrets({"type": "s3", "name": "prod", "scope": "s3://", "key_id": "P"}) |
| 525 | + accessor = SecretsAccessor(batch) |
| 526 | + assert accessor.get("s3", name="prod") is not None |
| 527 | + assert not accessor.needs_resolution |
| 528 | + |
| 529 | + def test_get_by_scope_on_retry_picks_the_longest_prefix(self) -> None: |
| 530 | + """A scope lookup returns the secret whose scope is the longest prefix.""" |
| 531 | + batch = _wire_secrets( |
| 532 | + {"type": "s3", "name": "any", "scope": "s3://", "key_id": "ANY"}, |
| 533 | + {"type": "s3", "name": "prod", "scope": "s3://prod/\ns3n://prod/", "key_id": "PROD"}, |
| 534 | + {"type": "gcs", "name": "other", "scope": "s3://prod/data/", "key_id": "WRONG-TYPE"}, |
480 | 535 | ) |
481 | | - result = SecretsAccessor(batch).to_dict() |
482 | | - assert "my_type" in result |
483 | | - assert result["my_type"]["k"].as_py() == "unscoped_val" |
484 | | - assert "scoped_type" in result |
485 | | - assert result["scoped_type"]["k"].as_py() == "scoped_val" |
486 | | - |
487 | | - def test_to_dict_skips_null_scoped(self) -> None: |
488 | | - """to_dict() skips scoped entries with null values.""" |
489 | | - struct_type = pa.struct([("k", pa.string())]) |
490 | | - scoped_field = pa.field( |
491 | | - "secret_0", |
492 | | - struct_type, |
493 | | - metadata={"secret_type": "null_type"}, |
| 536 | + accessor = SecretsAccessor(batch, is_retry=True) |
| 537 | + prod = accessor.get("s3", scope="s3://prod/data/x.parquet", required=True) |
| 538 | + assert prod is not None and prod["key_id"].as_py() == "PROD" |
| 539 | + other = accessor.get("s3", scope="s3://elsewhere/x.parquet", required=True) |
| 540 | + assert other is not None and other["key_id"].as_py() == "ANY" |
| 541 | + assert accessor.get("s3", scope="az://nope/") is None |
| 542 | + |
| 543 | + def test_get_by_scope_always_asks_the_client_first(self) -> None: |
| 544 | + """Only the client knows whether a longer-scoped secret exists.""" |
| 545 | + batch = _wire_secrets({"type": "s3", "name": "any", "scope": "s3://", "key_id": "ANY"}) |
| 546 | + accessor = SecretsAccessor(batch) |
| 547 | + assert accessor.get("s3", scope="s3://prod/x") is None |
| 548 | + assert accessor.pending_lookups == [SecretLookupEntry(secret_type="s3", scope="s3://prod/x")] |
| 549 | + |
| 550 | + def test_get_by_type_finds_a_secret_whatever_its_name(self) -> None: |
| 551 | + """A by-type lookup matches the ``type`` field, not the column name.""" |
| 552 | + batch = _wire_secrets( |
| 553 | + {"type": "vgi_example", "name": "scoped", "scope": "x://a/", "v": "scoped"}, |
| 554 | + {"type": "vgi_example", "name": "__default_vgi_example", "scope": "", "v": "default"}, |
| 555 | + {"type": "s3", "name": "aaa", "scope": "", "v": "s3"}, |
494 | 556 | ) |
495 | | - scoped_arr = pa.array([None], type=struct_type) |
496 | | - batch = pa.RecordBatch.from_arrays( |
497 | | - [scoped_arr], |
498 | | - schema=pa.schema([scoped_field]), |
| 557 | + accessor = SecretsAccessor(batch) |
| 558 | + found = accessor.get("vgi_example") |
| 559 | + assert found is not None |
| 560 | + assert found["v"].as_py() == "default" |
| 561 | + assert not accessor.needs_resolution |
| 562 | + |
| 563 | + def test_get_by_type_falls_back_to_a_scoped_secret(self) -> None: |
| 564 | + """With no unscoped secret of the type, the first one of that type is used.""" |
| 565 | + batch = _wire_secrets({"type": "s3", "name": "prod", "scope": "s3://", "v": "p"}) |
| 566 | + found = SecretsAccessor(batch).get("s3") |
| 567 | + assert found is not None |
| 568 | + assert found["v"].as_py() == "p" |
| 569 | + |
| 570 | + def test_annotated_secret_is_selected_by_name_and_scope(self) -> None: |
| 571 | + """A ``Secret()`` annotation on on_bind gets the secret it names, through a real bind. |
| 572 | +
|
| 573 | + Regression: the table-function path looked for a column named by the |
| 574 | + secret *type*, so an annotated secret arrived as ``None`` unless the |
| 575 | + user had named the secret after its type. |
| 576 | + """ |
| 577 | + from vgi.arguments import Arguments |
| 578 | + from vgi.invocation import FunctionType |
| 579 | + from vgi.protocol import BindRequest |
| 580 | + |
| 581 | + seen: dict[str, Any] = {} |
| 582 | + |
| 583 | + @dataclass(slots=True, frozen=True) |
| 584 | + class _Args: |
| 585 | + count: Annotated[int, Arg(0, doc="Count")] = 10 |
| 586 | + |
| 587 | + class _TestFunc(TableFunctionGenerator[_Args]): |
| 588 | + @classmethod |
| 589 | + def on_bind( |
| 590 | + cls, |
| 591 | + params: BindParams[_Args], |
| 592 | + *, |
| 593 | + by_name: Annotated[dict[str, pa.Scalar[Any]] | None, Secret("s3", name="dev")] = None, |
| 594 | + by_scope: Annotated[dict[str, pa.Scalar[Any]] | None, Secret("s3", scope="s3://prod/x")] = None, |
| 595 | + missing: Annotated[dict[str, pa.Scalar[Any]] | None, Secret("gcs")] = None, |
| 596 | + ) -> BindResponse: |
| 597 | + seen.update(by_name=by_name, by_scope=by_scope, missing=missing) |
| 598 | + return BindResponse(output_schema=pa.schema([pa.field("x", pa.int64())])) |
| 599 | + |
| 600 | + @classmethod |
| 601 | + def process(cls, params: Any, state: Any, out: Any) -> None: ... |
| 602 | + |
| 603 | + request = BindRequest( |
| 604 | + function_name="_test_func", |
| 605 | + arguments=Arguments(positional=(pa.scalar(1),)), |
| 606 | + function_type=FunctionType.TABLE, |
| 607 | + secrets=_wire_secrets( |
| 608 | + {"type": "s3", "name": "dev", "scope": "s3://dev/", "key_id": "DEV"}, |
| 609 | + {"type": "s3", "name": "prod", "scope": "s3://prod/", "key_id": "PROD"}, |
| 610 | + ), |
499 | 611 | ) |
500 | | - result = SecretsAccessor(batch).to_dict() |
501 | | - assert "null_type" not in result |
| 612 | + _TestFunc.bind(request) |
| 613 | + |
| 614 | + assert seen["by_name"]["key_id"].as_py() == "DEV" |
| 615 | + assert seen["by_scope"]["key_id"].as_py() == "PROD" |
| 616 | + assert seen["missing"] is None |
502 | 617 |
|
503 | 618 | def test_required_raises_on_retry(self) -> None: |
504 | 619 | """required=True raises ValueError when secret is genuinely missing on retry.""" |
505 | | - import pytest |
506 | | - |
507 | 620 | accessor = SecretsAccessor(None, is_retry=True) |
508 | 621 | with pytest.raises(ValueError, match="Required secret"): |
509 | 622 | accessor.get("missing_type", required=True) |
510 | 623 |
|
511 | 624 | def test_scoped_required_raises_on_retry(self) -> None: |
512 | 625 | """required=True with scope raises ValueError on retry when not found.""" |
513 | | - import pytest |
514 | | - |
515 | 626 | accessor = SecretsAccessor(None, is_retry=True) |
516 | 627 | with pytest.raises(ValueError, match="Required secret"): |
517 | 628 | accessor.get("s3", scope="s3://bucket/", required=True) |
@@ -557,7 +668,6 @@ def test_non_serializable_state_raises(self) -> None: |
557 | 668 | still rejected, and rejected at class-definition time rather than |
558 | 669 | silently falling back to initial_state() on every HTTP exchange. |
559 | 670 | """ |
560 | | - import pytest |
561 | 671 |
|
562 | 672 | @dataclass |
563 | 673 | class BadState: |
|
0 commit comments