Qualify successor protocol tests under strict typing

Author: rustyconover

tests/catalog/column_native_reference.py

@@ -27,13 +27,14 @@
 from vgi.worker import Worker
 
 from .column_paths_reference import ColumnPathsCatalog
+from .typing_helpers import present
 
 COUNT = pa.schema([pa.field("count", pa.int64(), nullable=False)])
 
 
 def _args(params: Any) -> dict[str, Any]:
     call = params.bind_call if isinstance(params, BindParams) else params.init_call.bind_call
-    target = json.loads(call.arguments.positional[0].as_py())
+    target = json.loads(present(call.arguments.positional[0]).as_py())
     return dict(
         attach_opaque_data=params.attach_opaque_data,
         transaction_opaque_data=params.transaction_opaque_data,
@@ -62,16 +63,20 @@ def on_bind(cls, params: BindParams[None]) -> BindResponse:
         if info is None:
             raise ValueError("table does not exist")
 
-        def data_field(field: pa.Field) -> pa.Field:
+        def data_field(field: pa.Field[Any]) -> pa.Field[Any]:
             # Column definitions live on TableInfo; the scan's data schema must
             # not ask generic function consumers to interpret catalog defaults.
             field = with_column_expressions(field, ColumnExpressions())
             if pa.types.is_struct(field.type):
                 field = field.with_type(pa.struct([data_field(child) for child in field.type]))
             return field
 
-        schema = pa.ipc.read_schema(pa.BufferReader(info.columns))
-        return BindResponse(output_schema=pa.schema([data_field(field) for field in schema], metadata=schema.metadata))
+        schema = pa.ipc.read_schema(pa.py_buffer(info.columns))
+        return BindResponse(
+            output_schema=pa.schema(
+                [data_field(field) for field in schema], metadata={k: v for k, v in (schema.metadata or {}).items()}
+            )
+        )
 
     @classmethod
     def on_init(cls, params: InitParams[None]) -> GlobalInitResponse:

tests/catalog/column_paths_reference.py

@@ -39,6 +39,7 @@
     CONSTRAINT_RECORDS_READ,
     CatalogCapabilities,
 )
+from vgi.catalog.catalog_interface import SerializedSchema
 from vgi.catalog.column_metadata import column_metadata
 from vgi.catalog.constraints import (
     ConstraintDefinition,
@@ -58,7 +59,9 @@
 from vgi.column_path import column_definition, column_path_key, replace_column_field, resolve_column_path
 from vgi.worker import Worker
 
-from .table_layout_reference import TABLE_PROFILES, TableLayoutCatalog, _table_key, _validate_expression
+from .expression_reference import validate_expression as _validate_expression
+from .table_layout_reference import TABLE_PROFILES, TableLayoutCatalog, _table_key
+from .typing_helpers import present
 
 COLUMN_PROFILES = CatalogCapabilities(
     read=sorted([*TABLE_PROFILES.read, CONSTRAINT_RECORDS_READ]),
@@ -86,7 +89,7 @@ def _eval(
     path: list[str],
     schema: pa.Schema | None = None,
     rows: list[dict[str, Any]] | None = None,
-) -> pa.Array:
+) -> pa.Array[Any]:
     expression.validate_owner(path)
     with duckdb.connect(config={"enable_external_access": False}) as connection:
         connection.execute("SET TimeZone='UTC'")
@@ -176,15 +179,15 @@ class ColumnPathsCatalog(TableLayoutCatalog):
     profiles = COLUMN_PROFILES
 
     def table_create(self, **kwargs: Any) -> None:
-        _fields(pa.ipc.read_schema(pa.BufferReader(kwargs["columns"])), kwargs["schema_path"])
+        _fields(pa.ipc.read_schema(pa.py_buffer(kwargs["columns"])), kwargs["schema_path"])
         super().table_create(**kwargs)
 
     def _create_constraints(self, args: dict[str, Any]) -> tuple[list[bytes], bytes | None]:
         table = CatalogObjectRef(kind="table", schema_path=args["schema_path"], name=args["name"])
         definitions = constraint_definitions(args["constraint_definitions"], table)
         if any(definition.kind != ConstraintKind.NOT_NULL for definition in definitions):
             raise NotImplementedError("column reference only enforces NOT_NULL constraints")
-        validate_constraint_schema(pa.ipc.read_schema(pa.BufferReader(args["columns"])), definitions)
+        validate_constraint_schema(pa.ipc.read_schema(pa.py_buffer(args["columns"])), definitions)
         return [
             self._not_null_info(table, definition).serialize_to_bytes() for definition in definitions
         ], uuid.uuid4().bytes
@@ -201,7 +204,7 @@ def _not_null_info(table: CatalogObjectRef, definition: ConstraintDefinition) ->
         )
 
     def _changed_constraints(
-        self, info: TableInfo, schema: pa.Schema, operation: str, path: list[str], replacement: pa.Field | None
+        self, info: TableInfo, schema: pa.Schema, operation: str, path: list[str], replacement: pa.Field[Any] | None
     ) -> list[bytes]:
         existing = {}
         for constraint in info.constraint_infos():
@@ -234,7 +237,7 @@ def collect(fields: Any, prefix: list[str]) -> None:
         return result
 
     def _arrange(self, rows: list[dict[str, Any]], info: TableInfo) -> list[dict[str, Any]]:
-        schema = pa.ipc.read_schema(pa.BufferReader(info.columns))
+        schema = pa.ipc.read_schema(pa.py_buffer(info.columns))
         # Resolve defaults from the same locked snapshot that stores the rows.
         # Reading table metadata before _change would race concurrent DDL.
         rows = [_defaulted(schema, row, info.schema_path) for row in rows]
@@ -250,7 +253,7 @@ def _column_change(self, operation: str, args: dict[str, Any]) -> ConstraintSnap
             version, data = self._snapshot(db, args["attach_opaque_data"], args["transaction_opaque_data"])
             if key not in data:
                 if args.get("ignore_not_found"):
-                    return
+                    return None
                 raise ValueError("table does not exist")
             info = self._info(data[key])
             request = args.get("constraint_request")
@@ -264,12 +267,14 @@ def _column_change(self, operation: str, args: dict[str, Any]) -> ConstraintSnap
                         raise ValueError("NOT_NULL constraint already exists")
                     operation = "not_null_set"
                 else:
-                    matches = [item for item in info.constraint_infos() if item.ref.key() == change.ref.key()]
-                    if len(matches) != 1 or matches[0].definition.kind != ConstraintKind.NOT_NULL:
+                    constraint_matches = [
+                        item for item in info.constraint_infos() if item.ref.key() == change.ref.key()
+                    ]
+                    if len(constraint_matches) != 1 or constraint_matches[0].definition.kind != ConstraintKind.NOT_NULL:
                         raise ValueError("unknown NOT_NULL constraint identity")
-                    path = matches[0].definition.column_paths[0]
+                    path = constraint_matches[0].definition.column_paths[0]
                     operation = "not_null_drop"
-            schema = pa.ipc.read_schema(pa.BufferReader(info.columns))
+            schema = pa.ipc.read_schema(pa.py_buffer(info.columns))
             parent_fields = schema if len(path) == 1 else resolve_column_path(schema, path[:-1]).type
             if not isinstance(parent_fields, pa.Schema) and not pa.types.is_struct(parent_fields):
                 raise ValueError("column path has a non-struct parent")
@@ -279,11 +284,11 @@ def _column_change(self, operation: str, args: dict[str, Any]) -> ConstraintSnap
             existing = matches[0] if matches else None
             if operation == "add" and existing is not None:
                 if args.get("if_column_not_exists"):
-                    return
+                    return None
                 raise ValueError("column already exists")
             if operation != "add" and existing is None:
                 if operation == "drop" and args.get("if_column_exists"):
-                    return
+                    return None
                 raise ValueError("column does not exist")
             if (
                 operation in ("drop", "rename", "type")
@@ -303,7 +308,7 @@ def _column_change(self, operation: str, args: dict[str, Any]) -> ConstraintSnap
                     parent = _parent(row, path)
                     if parent is not None:
                         parent[path[-1]] = (
-                            pc.cast(_eval(default, info.schema_path), replacement.type, safe=True)[0].as_py()
+                            pc.cast(_eval(default, info.schema_path), present(replacement).type, safe=True)[0].as_py()
                             if default is not None
                             else None
                         )
@@ -327,7 +332,12 @@ def _column_change(self, operation: str, args: dict[str, Any]) -> ConstraintSnap
                 # annotations when replacing the primitive physical type.
                 metadata = dict(existing.metadata or {})
                 metadata.update(target.metadata or {})
-                replacement = existing.with_type(target.type).with_metadata(metadata or None)
+                replacement = existing.with_type(target.type)
+                replacement = (
+                    replacement.with_metadata({k: v for k, v in metadata.items()})
+                    if metadata
+                    else replacement.remove_metadata()
+                )
                 _fields([replacement], info.schema_path)
                 if args["expression"] is not None:
                     # Bind before changes even if the table has no rows.
@@ -363,13 +373,17 @@ def _column_change(self, operation: str, args: dict[str, Any]) -> ConstraintSnap
                 else:
                     raise AssertionError(operation)
                 assert replacement is not None
-                replacement = replacement.with_metadata(metadata or None)
+                replacement = (
+                    replacement.with_metadata({k: v for k, v in metadata.items()})
+                    if metadata
+                    else replacement.remove_metadata()
+                )
             updated = replace_column_field(schema, path, replacement, add=operation == "add")
             constraints = self._changed_constraints(info, updated, operation, path, replacement)
             changed = not updated.equals(schema, check_metadata=True) or constraints != info.constraints
             info = replace(
                 info,
-                columns=updated.serialize().to_pybytes(),
+                columns=SerializedSchema(updated.serialize().to_pybytes()),
                 constraints=constraints,
                 definition_revision=uuid.uuid4().bytes
                 if changed and operation != "comment"
@@ -392,7 +406,9 @@ def _column_change(self, operation: str, args: dict[str, Any]) -> ConstraintSnap
                 db.execute("UPDATE transactions SET version=?, data=? WHERE id=?", (version + 1, json.dumps(data), tx))
             if request is not None:
                 return ConstraintSnapshot(
-                    table=request.table, definition_revision=info.definition_revision, constraints=info.constraints
+                    table=request.table,
+                    definition_revision=present(info.definition_revision),
+                    constraints=info.constraints,
                 )
         return None
 
@@ -406,8 +422,8 @@ def table_constraints_apply(self, request: ConstraintsApplyRequest) -> Constrain
         if change.operation == ConstraintOperation.ADD:
             definition = change.definition
             if (
-                definition.kind != ConstraintKind.NOT_NULL
-                or definition.declared_name is not None
+                present(definition).kind != ConstraintKind.NOT_NULL
+                or present(definition).declared_name is not None
                 or change.enforcement != ConstraintEnforcement.ENFORCED
                 or change.validation != ConstraintValidation.ALL_ROWS
             ):
@@ -458,4 +474,9 @@ def table_column_comment_set(self, **kwargs: Any) -> None:
 class ColumnPathsWorker(Worker):
     """Expose the durable column reference through the normal protocol worker."""
 
+    def _get_catalog(self) -> ColumnPathsCatalog:
+        catalog = super()._get_catalog()
+        assert isinstance(catalog, ColumnPathsCatalog)
+        return catalog
+
     catalog_interface = ColumnPathsCatalog

tests/catalog/constraints_protocol_reference.py

@@ -13,21 +13,26 @@
 import os
 import sqlite3
 import uuid
+from collections.abc import Sequence
 from contextlib import closing
 from dataclasses import replace
-from typing import Any, ClassVar
+from typing import Any, ClassVar, Literal, overload
 
 import pyarrow as pa
 
 from vgi.catalog import (
     AttachOpaqueData,
     CatalogAttachResult,
+    FunctionInfo,
+    IndexInfo,
+    MacroInfo,
     OnConflict,
     ReadOnlyCatalogInterface,
     SchemaInfo,
     SchemaObjectType,
     TableInfo,
     TransactionOpaqueData,
+    ViewInfo,
 )
 from vgi.catalog.capabilities import (
     CONSTRAINT_CREATE,
@@ -38,7 +43,7 @@
     TRANSACTION_DDL,
     CatalogCapabilities,
 )
-from vgi.catalog.catalog_interface import CatalogContentsResult, SchemaContentsInfo
+from vgi.catalog.catalog_interface import CatalogContentsResult, SchemaContentsInfo, SerializedSchema
 from vgi.catalog.constraints import (
     ConstraintChange,
     ConstraintDefinition,
@@ -50,6 +55,7 @@
     ConstraintValidation,
 )
 from vgi.catalog.dependencies import CatalogObjectRef
+from vgi.schema_path import SchemaPath
 from vgi.worker import Worker
 
 from .constraints_reference import ConstraintCatalog, _schema
@@ -111,7 +117,7 @@ def _info(record: dict[str, Any]) -> TableInfo:
         return TableInfo(
             name=snapshot.table.name,
             schema_path=snapshot.table.schema_path,
-            columns=_schema(record).serialize().to_pybytes(),
+            columns=SerializedSchema(_schema(record).serialize().to_pybytes()),
             constraints=snapshot.constraints,
             definition_revision=snapshot.definition_revision,
             comment=None,
@@ -144,14 +150,68 @@ def table_get(
                 return None
             return self._info(record)
 
+    @overload
+    def schema_contents(
+        self,
+        *,
+        attach_opaque_data: AttachOpaqueData,
+        transaction_opaque_data: TransactionOpaqueData | None,
+        path: SchemaPath,
+        type: Literal[SchemaObjectType.TABLE],
+    ) -> Sequence[TableInfo]: ...
+
+    @overload
+    def schema_contents(
+        self,
+        *,
+        attach_opaque_data: AttachOpaqueData,
+        transaction_opaque_data: TransactionOpaqueData | None,
+        path: SchemaPath,
+        type: Literal[SchemaObjectType.VIEW],
+    ) -> Sequence[ViewInfo]: ...
+
+    @overload
+    def schema_contents(
+        self,
+        *,
+        attach_opaque_data: AttachOpaqueData,
+        transaction_opaque_data: TransactionOpaqueData | None,
+        path: SchemaPath,
+        type: Literal[
+            SchemaObjectType.SCALAR_FUNCTION,
+            SchemaObjectType.TABLE_FUNCTION,
+            SchemaObjectType.AGGREGATE_FUNCTION,
+        ],
+    ) -> Sequence[FunctionInfo]: ...
+
+    @overload
+    def schema_contents(
+        self,
+        *,
+        attach_opaque_data: AttachOpaqueData,
+        transaction_opaque_data: TransactionOpaqueData | None,
+        path: SchemaPath,
+        type: Literal[SchemaObjectType.SCALAR_MACRO, SchemaObjectType.TABLE_MACRO],
+    ) -> Sequence[MacroInfo]: ...
+
+    @overload
+    def schema_contents(
+        self,
+        *,
+        attach_opaque_data: AttachOpaqueData,
+        transaction_opaque_data: TransactionOpaqueData | None,
+        path: SchemaPath,
+        type: Literal[SchemaObjectType.INDEX],
+    ) -> Sequence[IndexInfo]: ...
+
     def schema_contents(
         self,
         *,
         attach_opaque_data: bytes,
         transaction_opaque_data: bytes | None,
         path: list[str],
         type: SchemaObjectType,
-    ) -> list[TableInfo]:
+    ) -> Sequence[TableInfo | ViewInfo | FunctionInfo | MacroInfo | IndexInfo]:
         _, infos = self._snapshot_infos(attach_opaque_data, transaction_opaque_data)
         return [info for info in infos if info.schema_path == path] if type == SchemaObjectType.TABLE else []
 
@@ -188,7 +248,7 @@ def table_create(
         self.authority.create(
             attach_opaque_data,
             CatalogObjectRef(kind="table", schema_path=schema_path, name=name),
-            pa.ipc.read_schema(pa.BufferReader(columns)),
+            pa.ipc.read_schema(pa.py_buffer(columns)),
             [ConstraintDefinition.deserialize_from_bytes(value) for value in constraint_definitions],
             tx=transaction_opaque_data,
         )

tests/catalog/constraints_reference.py

@@ -88,7 +88,7 @@ def _table_ref(record: dict[str, Any]) -> CatalogObjectRef:
 
 
 def _base_schema(record: dict[str, Any]) -> pa.Schema:
-    return pa.ipc.read_schema(pa.BufferReader(_unb64(record["schema"])))
+    return pa.ipc.read_schema(pa.py_buffer(_unb64(record["schema"])))
 
 
 def _rows(record: dict[str, Any]) -> pa.Table:
@@ -104,7 +104,10 @@ def _schema(record: dict[str, Any]) -> pa.Schema:
         for path in info.definition.column_paths
     }
     base = _base_schema(record)
-    return pa.schema([field.with_nullable(field.name not in required) for field in base], metadata=base.metadata)
+    return pa.schema(
+        [field.with_nullable(field.name not in required) for field in base],
+        metadata=None if base.metadata is None else {key: value for key, value in (base.metadata).items()},
+    )
 
 
 def _touch(record: dict[str, Any]) -> None:
@@ -203,6 +206,7 @@ def _lookup(data: dict[str, Any], table: CatalogObjectRef) -> dict[str, Any]:
         record = data.get(_table_key(table))
         if record is None:
             raise ValueError("table does not exist")
+        assert isinstance(record, dict)
         return record
 
     @staticmethod
@@ -322,7 +326,12 @@ def create(
             key = _table_key(table)
             if key in data:
                 raise ValueError("table already exists")
-            base = pa.schema([field.with_nullable(True) for field in columns], metadata=columns.metadata)
+            base = pa.schema(
+                [field.with_nullable(True) for field in columns],
+                metadata=None
+                if columns.metadata is None
+                else {key: value for key, value in (columns.metadata).items()},
+            )
             record = {
                 "table": _b64(table.serialize_to_bytes()),
                 "schema": _b64(base.serialize().to_pybytes()),