Skip to content

Commit bf190fe

Browse files
rustyconoverclaude
andcommitted
Add validation for schema_contents type parameter in worker
- Validate type parameter is a valid SchemaObjectType value - Validate all returned objects match the expected type - Provide clear error messages when catalog returns wrong types This prevents confusing errors when a catalog's schema_contents() implementation doesn't properly filter by the type parameter. Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent fa2f61b commit bf190fe

1 file changed

Lines changed: 35 additions & 2 deletions

File tree

vgi/worker.py

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1004,13 +1004,37 @@ def _handle_catalog_invocation(
10041004
)
10051005

10061006
if method_name == "schema_contents":
1007-
type_param = SchemaObjectType(kwargs["type"])
1007+
# Validate the type parameter
1008+
type_value = kwargs.get("type")
1009+
valid_types = {e.value for e in SchemaObjectType}
1010+
if type_value not in valid_types:
1011+
raise ValueError(
1012+
f"Invalid schema_contents type: {type_value!r}. "
1013+
f"Must be one of: {sorted(valid_types)}"
1014+
)
1015+
type_param = SchemaObjectType(type_value)
1016+
1017+
# Determine expected class and schema based on type
1018+
expected_class: type[TableInfo | ViewInfo | FunctionInfo]
10081019
if type_param == SchemaObjectType.TABLE:
1020+
expected_class = TableInfo
10091021
schema = TableInfo.ARROW_SCHEMA
10101022
elif type_param == SchemaObjectType.VIEW:
1023+
expected_class = ViewInfo
10111024
schema = ViewInfo.ARROW_SCHEMA
10121025
else: # SCALAR_FUNCTION or TABLE_FUNCTION
1026+
expected_class = FunctionInfo
10131027
schema = FunctionInfo.ARROW_SCHEMA
1028+
1029+
# Validate all results match the expected type
1030+
for i, item in enumerate(result):
1031+
if not isinstance(item, expected_class):
1032+
raise TypeError(
1033+
f"schema_contents returned wrong type at index {i}: "
1034+
f"expected {expected_class.__name__}, "
1035+
f"got {type(item).__name__}. "
1036+
f"The catalog's schema_contents() must filter by type."
1037+
)
10141038
elif method_name == "schemas":
10151039
schema = SchemaInfo.ARROW_SCHEMA
10161040
else:
@@ -1034,7 +1058,16 @@ def _handle_catalog_invocation(
10341058
)
10351059

10361060
if method_name == "schema_contents":
1037-
type_param = SchemaObjectType(kwargs["type"])
1061+
# Validate the type parameter
1062+
type_value = kwargs.get("type")
1063+
valid_types = {e.value for e in SchemaObjectType}
1064+
if type_value not in valid_types:
1065+
raise ValueError(
1066+
f"Invalid schema_contents type: {type_value!r}. "
1067+
f"Must be one of: {sorted(valid_types)}"
1068+
)
1069+
type_param = SchemaObjectType(type_value)
1070+
10381071
if type_param == SchemaObjectType.TABLE:
10391072
schema = TableInfo.ARROW_SCHEMA
10401073
elif type_param == SchemaObjectType.VIEW:

0 commit comments

Comments
 (0)