adding many tests for validation and tool creation in general, fix some issues that came up

Author: bcallender

src/fenic/core/mcp/_tools.py

@@ -13,7 +13,7 @@
 from fenic.core._logical_plan.plans.base import LogicalPlan
 from fenic.core._utils.type_inference import infer_pytype_from_dtype
 from fenic.core.error import PlanError
-from fenic.core.mcp._validators import get_param_validator, maybe_get_param_validator
+from fenic.core.mcp._validators import get_param_validator
 from fenic.core.mcp.types import (
     BoundToolParam,
     TableFormat,
@@ -85,9 +85,10 @@ def bind_tool(
                 try:
                     validator = get_param_validator(validator_name)
                     if unresolved_expr.data_type not in validator.data_types():
+                        supported_data_types = ", ".join([str(dt) for dt in validator.data_types()])
                         raise PlanError(
-                            f"Param Validator {validator_name} supports data types {validator.data_types()}, "
-                            f"but the parameter {unresolved_expr_name} has data type {unresolved_expr.data_type}."
+                            f"Param Validator `{validator_name}` supports data types ({supported_data_types}), "
+                            f"but the parameter `{unresolved_expr_name}` has data type {unresolved_expr.data_type}."
                         )
                     validators.append(validator)
                 except KeyError:
@@ -132,19 +133,11 @@ def _infer_base_type(p: BoundToolParam):
             if isinstance(p.data_type, ArrayType):
                 return list[literal_type]  # type: ignore[valid-type]
             return literal_type
+        if isinstance(p.data_type, ArrayType):
+            inner_type = infer_pytype_from_dtype(p.data_type.element_type)
+            return list[inner_type]  # type: ignore[valid-type]
         return infer_pytype_from_dtype(p.data_type)
 
-    def _wrap_with_validator(base_t, validator_name: Optional[str]):
-        if not validator_name:
-            return base_t
-        pv = maybe_get_param_validator(validator_name)
-        if pv is None:
-            return base_t
-        def _wrap(v, _pv=pv):
-            _pv.validate(v)
-            return v
-        return TypingAnnotated[base_t, AfterValidator(_wrap)]  # type: ignore[valid-type]
-
     def _field_kwargs(p: BoundToolParam, include_default: bool) -> dict:
         kwargs: dict = {"description": p.description}
         constraints = p.constraints

src/fenic/core/mcp/_validators.py

@@ -1,6 +1,7 @@
 import re
-from typing import Dict, List, Optional, Protocol, Union, runtime_checkable
+from typing import Dict, List, Protocol, Union, runtime_checkable
 
+from fenic._polars_plugins import py_validate_regex  # noqa: F401
 from fenic.core.error import (
     ValidationError,
 )
@@ -21,7 +22,7 @@ def data_types(self) -> List[DataType]:
         """The data types that the validator operates on."""
         ...
 
-    def validate(self, value: Union[str, int, float, bool, list, dict]) -> bool:
+    def validate(self, value: Union[str, int, float, bool, list, dict]):
         """Validate an argument value.
 
         Args:
@@ -62,14 +63,6 @@ def validate(self, user_query: str):
         if len(query) > MAX_REGEX_LENGTH:
             raise ValidationError(f"Regex too long (>{MAX_REGEX_LENGTH} characters)")
 
-        # Support /pattern/flags and capture flags
-        query, flags = self._strip_slash_delimiters(query)
-        unsupported_flags = {f for f in flags if f not in {"i", "m", "s", "x"}}
-        if unsupported_flags:
-            raise ValidationError(
-                f"Unsupported regex flags: {''.join(sorted(unsupported_flags))}"
-            )
-
         # Strip inline flags at start like (?i), (?m), combined, to avoid duplication
         query = re.sub(r"^\(\?[aiLmsux]+\)", "", query)
 
@@ -89,15 +82,16 @@ def validate(self, user_query: str):
             except ValueError:
                 raise ValidationError("Invalid quantifier bounds") from None
             if m_val > MAX_QUANTIFIER_VALUE or n_val > MAX_QUANTIFIER_VALUE:
-                raise ValidationError("Quantifier bounds too large")
+                raise ValidationError(f"Quantifier bounds {m_val} or {n_val} > {MAX_QUANTIFIER_VALUE}")
             if n and n_val < m_val:
-                raise ValidationError("Quantifier upper bound less than lower bound")
+                raise ValidationError(f"Quantifier upper bound {n_val} < lower bound {m_val}")
 
         # Limit alternations
-        if query.count("|") > MAX_ALTERNATIONS:
-            raise ValidationError("Too many alternations in regex")
+        alternations = query.count("|")
+        if alternations > MAX_ALTERNATIONS:
+            raise ValidationError(f"Too many alternations ({alternations} > {MAX_ALTERNATIONS})")
 
-        # Disallow backreferences (simple and robust detection)
+        # Disallow backreferences
         if any(f"\\{d}" in query for d in "123456789"):
             raise ValidationError("Backreferences are not supported")
 
@@ -121,11 +115,11 @@ def validate(self, user_query: str):
         if re.search(r"\{\s*\d+\s*,\s*\d+\s*,", query):
             raise ValidationError("Invalid quantifier syntax")
 
-        # Ensure it compiles in Python as a basic sanity check
+        # Final check, ensure that the regex is valid for `rlike`
         try:
-            re.compile(query)
-        except re.error as err:
-            raise ValidationError(f"Invalid regex syntax: {err}") from None
+            py_validate_regex(query)
+        except Exception as err:
+            raise ValidationError(f"Invalid regex syntax: {query}") from err
 
         return
 
@@ -146,20 +140,6 @@ def _is_balanced(self, s: str, open_char: str, close_char: str) -> bool:
             i += 1
         return depth == 0
 
-
-    def _strip_slash_delimiters(self, pattern: str) -> tuple[str, set[str]]:
-        """Support /pattern/flags syntax; return (pattern, flags).
-
-        Only recognize i,m,s,x flags; others are rejected later.
-        """
-        if len(pattern) >= 2 and pattern.startswith("/") and pattern.rfind("/") > 0:
-            last = pattern.rfind("/")
-            core = pattern[1:last]
-            flags = set(pattern[last + 1 :].lower())
-            return core, flags
-        return pattern, set()
-
-
 # -- Registry for reusable ParamValidators --
 _PARAM_VALIDATOR_REGISTRY: Dict[str, ParamValidator] = {}
 
@@ -186,11 +166,5 @@ def get_param_validator(name: str) -> ParamValidator:
         raise KeyError(f"No ParamValidator registered under name '{name}'") from err
 
 
-def maybe_get_param_validator(name: Optional[str]) -> Optional[ParamValidator]:
-    if name is None:
-        return None
-    return get_param_validator(name)
-
-
 # Pre-register common validators
 register_param_validator("regex", RegexValidator())

tests/api/mcp/test_server.py

@@ -5,12 +5,58 @@
 
 from fenic import SystemTool, SystemToolConfig
 from fenic.api.mcp._tool_generation_utils import auto_generate_system_tools_from_tables
+from fenic.api.functions import col, tool_param
 from fenic.api.mcp.server import create_mcp_server
 from fenic.api.session.session import Session
 from fenic.core._utils.misc import to_snake_case
+from fenic.core.mcp._tools import bind_tool
+from fenic.core.mcp.types import ToolParam, ToolParamConstraints
+from fenic.core.types.datatypes import ArrayType, IntegerType, StringType
 from tests.api.mcp.utils import create_table_with_rows
 
 
+def test_server_generation_with_parameterized_tools(local_session: Session):
+    pytest.importorskip("fastmcp")
+    df = local_session.create_dataframe({"city": ["SF"], "age": [10], "user_name": ["Alice"]})
+    query = df.filter(
+        (col("city") == tool_param("city_name", StringType))
+        & (col("age") >= tool_param("age", IntegerType))
+        & (col("user_name").is_in(tool_param("user_names", ArrayType(StringType))))
+    )._logical_plan
+
+    parameterized_tool = bind_tool(
+        name="tool_x",
+        description="table one",
+        params=[
+            ToolParam(name="city_name", description="City name", constraints=ToolParamConstraints(pattern="^SF$")),
+            ToolParam(name="age", description="Age", constraints=ToolParamConstraints(gt=0, lt=120, multiple_of=2)),
+            ToolParam(name="user_names", description="User names", constraints=ToolParamConstraints(min_length=1, max_length=5)),
+        ],
+        result_limit=10,
+        query=query,
+    )
+
+    server = create_mcp_server(local_session, "Test Server", parameterized_tools=[parameterized_tool])
+    server_tools = asyncio.run(server.mcp.get_tools())
+    assert len(server_tools) == 1
+    parameter_schema = server_tools["tool_x"].parameters['properties']
+    city_name_param = parameter_schema['city_name']
+    assert city_name_param['type'] == 'string'
+    assert city_name_param['pattern'] == '^SF$'
+    assert city_name_param['description'] == "City name"
+    age_param = parameter_schema['age']
+    assert age_param['type'] == 'integer'
+    assert age_param['exclusiveMinimum'] == 0
+    assert age_param['exclusiveMaximum'] == 120
+    assert age_param['multipleOf'] == 2
+    assert age_param['description'] == "Age"
+    user_names_param = parameter_schema['user_names']
+    assert user_names_param['type'] == 'array'
+    assert user_names_param['items']['type'] == 'string'
+    assert user_names_param['maxItems'] == 5
+    assert user_names_param['minItems'] == 1
+    assert user_names_param['description'] == "User names"
+
 def test_server_generation(local_session: Session):
     pytest.importorskip("fastmcp")
     create_table_with_rows(local_session, "t1", [1, 2, 3], description="table one")

tests/core/mcp/test_tools.py

@@ -1,13 +1,14 @@
+import re
 
 import pytest
 from pydantic import BaseModel
 from pydantic import ValidationError as PydValidationError
 
 from fenic.api.functions import col, tool_param
-from fenic.core.error import PlanError
+from fenic.core.error import PlanError, ValidationError
 from fenic.core.mcp._tools import bind_tool, create_pydantic_model_for_tool
-from fenic.core.mcp.types import ToolParam
-from fenic.core.types.datatypes import IntegerType, StringType
+from fenic.core.mcp.types import ToolParam, ToolParamConstraints
+from fenic.core.types.datatypes import ArrayType, IntegerType, StringType
 
 
 def test_toolparam_required_and_default_validation():
@@ -51,6 +52,103 @@ def test_resolve_tool_validates_unresolved_params(local_session):
             query=query,
         )
 
+def test_resolve_tool_validates_mistyped_validators(local_session):
+    df = local_session.create_dataframe({"name": ["Alice", "Bob"], "age": [25, 30], "city": ["SF", "SEA"]})
+    query = df.filter((col("age") >= tool_param("min_age", IntegerType)) & (col("city") == tool_param("city_name", StringType)))._logical_plan
+
+    with pytest.raises(PlanError, match="Param Validator `regex` supports data types \(StringType\), but the parameter `min_age` has data type IntegerType."):
+        bind_tool(
+            name="users_by_city",
+            description="Filter users",
+            params=[
+                ToolParam(name="min_age", description="Minimum age", validator_names=["regex"]),
+                ToolParam(name="city_name", description="City name", validator_names=["regex"]),
+            ],
+            result_limit=50,
+            query=query,
+        )
+
+def test_resolve_tool_validates_missing_validators(local_session):
+    df = local_session.create_dataframe({"name": ["Alice", "Bob"], "age": [25, 30], "city": ["SF", "SEA"]})
+    query = df.filter((col("age") >= tool_param("min_age", IntegerType)) & (col("city") == tool_param("city_name", StringType)))._logical_plan
+
+    with pytest.raises(PlanError, match="Could not find a ParamValidator for the following validator names: \['non_existent'\]"):
+        bind_tool(
+            name="users_by_city",
+            description="Filter users",
+            params=[
+                ToolParam(name="min_age", description="Minimum age"),
+                ToolParam(name="city_name", description="City name", validator_names=["non_existent"]),
+            ],
+            result_limit=50,
+            query=query,
+        )
+
+def test_create_pydantic_model_for_tool_applies_validators(local_session):
+    df = local_session.create_dataframe({"name": ["Alice", "Bob"], "age": [25, 30], "city": ["SF", "SEA"]})
+    query = df.filter(
+        (col("age") >= tool_param("min_age", IntegerType)) &
+        (col("city") == tool_param("city_name", StringType))
+    )._logical_plan
+
+    tool = bind_tool(
+        name="users_by_city",
+        description="Filter users",
+        params=[
+            ToolParam(name="min_age", description="Minimum age"),
+            ToolParam(name="city_name", description="City name", validator_names=["regex"]),
+        ],
+        result_limit=50,
+        query=query,
+    )
+
+    Model: type[BaseModel] = create_pydantic_model_for_tool(tool)
+
+    with pytest.raises(ValidationError, match="Unbalanced curly braces"):
+        Model(city_name="{+---", min_age=25)
+
+    with pytest.raises(ValidationError, match="Too many alternations \(21 > 20\)"):
+        Model(city_name="SF|SEA|OAK|PHX|LAS|ORD|XRD|PRD|IAD|CRD|FRA|LON|UMEA|BOS|YYZ|DOG|BAT|BAN|LAP|LAX|TYO|HND", min_age=25)
+
+
+def test_create_pydantic_model_for_tool_applies_field_validators(local_session):
+    df = local_session.create_dataframe({"city": ["SF"], "age": [10], "user_name": ["Alice"]})
+    query = df.filter(
+        (col("city") == tool_param("city_name", StringType))
+        & (col("age") >= tool_param("age", IntegerType))
+        & (col("user_name").is_in(tool_param("user_names", ArrayType(StringType))))
+    )._logical_plan
+
+    tool = bind_tool(
+        name="tool_x",
+        description="",
+        params=[
+            ToolParam(name="city_name", description="City name", constraints=ToolParamConstraints(pattern="^SF$")),
+            ToolParam(name="age", description="Age", constraints=ToolParamConstraints(gt=0, lt=120, multiple_of=2)),
+            ToolParam(name="user_names", description="User names", constraints=ToolParamConstraints(min_length=1, max_length=5)),
+        ],
+        result_limit=10,
+        query=query,
+    )
+
+    Model: type[BaseModel] = create_pydantic_model_for_tool(tool)
+    #should pass validation
+    Model(city_name="SF", age=10, user_names=["Alice", "Bob"])
+    with pytest.raises(PydValidationError, match=re.escape("String should match pattern '^SF$'")):
+        Model(city_name="SEA", age=10, user_names=["Alice", "Bob"])
+
+    with pytest.raises(PydValidationError, match=re.escape("Input should be greater than 0")):
+        Model(city_name="SF", age=0, user_names=["Alice", "Bob"])
+
+    with pytest.raises(PydValidationError, match=re.escape("Input should be a multiple of 2")):
+        Model(city_name="SF", age=11, user_names=["Alice", "Bob"])
+
+    with pytest.raises(PydValidationError, match=re.escape("List should have at most 5 items after validation, not 6")):
+        Model(city_name="SF", age=10, user_names=["Alice", "Bob", "Charlie", "David", "Eve", "Frank"])
+
+    with pytest.raises(PydValidationError, match=re.escape("List should have at least 1 item after validation, not 0")):
+        Model(city_name="SF", age=10, user_names=[])
+
 
 def test_create_pydantic_model_for_tool_defaults_and_required(local_session):
     df = local_session.create_dataframe({"city": ["SF"], "age": [10]})

tests/core/mcp/test_validators.py

@@ -0,0 +1,65 @@
+import pytest
+
+from fenic.core.error import ValidationError
+from fenic.core.mcp._validators import (
+    RegexValidator,
+    get_param_validator,
+    register_param_validator,
+)
+
+
+def test_regex_validator_accepts_simple_pattern():
+    v = RegexValidator()
+    v.validate("foo|bar")
+
+
+def test_regex_validator_supports_slash_delimiters_and_flags():
+    v = RegexValidator()
+    v.validate("/foo.*/i")
+
+
+@pytest.mark.parametrize(
+    "pattern",
+    [
+        "   ",  # whitespace-only
+        "(",  # unbalanced paren
+        "[a-",  # unbalanced bracket
+        "{1,2,3}",  # invalid quantifier syntax
+    ],
+)
+def test_regex_validator_rejects_basic_invalid_patterns(pattern):
+    v = RegexValidator()
+    with pytest.raises(ValidationError):
+        v.validate(pattern)
+
+
+@pytest.mark.parametrize(
+    "pattern",
+    [
+        r"(.+)+",  # nested quantifier
+        r"(.*)+",  # nested quantifier
+        r"(?:.+){1001}",  # excessive bounded quantifier
+        r"(a|b|c|d|e|f|g|h|i|j|k|l|m|n|o|p|q|r|s|t|u|v|w|x|y|z|aa|ab|ac)",  # too many alternations
+        r"\1",  # backreference
+        r"(?<=a)b",  # lookbehind
+        r"(\.\*){2,1}",  # m > n for matching between m and n repeats of a character
+    ],
+)
+def test_regex_validator_rejects_redos_like_and_unsupported_constructs(pattern):
+    v = RegexValidator()
+    with pytest.raises(ValidationError):
+        v.validate(pattern)
+
+
+def test_registry_has_default_regex_validator():
+    v = get_param_validator("regex")
+    assert isinstance(v, RegexValidator)
+
+
+def test_registry_register_and_lookup_custom():
+    import uuid
+
+    unique_name = f"custom_regex_{uuid.uuid4().hex}"
+    register_param_validator(unique_name, RegexValidator())
+    v = get_param_validator(unique_name)
+    assert isinstance(v, RegexValidator)