Skip to content

Commit

Permalink
Refining how we validate flow parameters under Pydantic 1 versus 2
Browse files Browse the repository at this point in the history
We have seen a few edge cases popping up where we're trying to use the
pydantic v2 validator to validate v1 models as flow function arguments.  This
leads to the opaque error:

```
TypeError: BaseModel.validate() takes 2 positional arguments but 3 were given
```

Examples:
* PrefectHQ/prefect-airbyte#61
* https://github.com/PrefectHQ/prefect-kubernetes/actions/runs/6628004996/job/18004160832?pr=92

The issue here is that we were trying to use our custom v2 port of the v1
`ValidatedFunction`, without inspecting the types of models involved.  Here, we
are looking at the types of the passed arguments, validating that they aren't
mixing v1 and v2 models, and then choosing the right validation implementation
from there.
  • Loading branch information
chrisguidry committed Oct 25, 2023
1 parent c6c2541 commit cbe3380
Showing 1 changed file with 32 additions and 5 deletions.
37 changes: 32 additions & 5 deletions src/prefect/flows.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,11 +42,14 @@

if HAS_PYDANTIC_V2:
import pydantic.v1 as pydantic
from pydantic import BaseModel as V2BaseModel
from pydantic import ValidationError as V2ValidationError
from pydantic.v1.decorator import ValidatedFunction
from pydantic.v1 import BaseModel as V1BaseModel
from pydantic.v1.decorator import ValidatedFunction as V1ValidatedFunction

from ._internal.pydantic.v2_validated_func import V2ValidatedFunction
from ._internal.pydantic.v2_validated_func import (
V2ValidatedFunction as ValidatedFunction, # noqa: F811
V2ValidatedFunction as ValidatedFunction,
)

else:
Expand Down Expand Up @@ -470,11 +473,35 @@ def validate_parameters(self, parameters: Dict[str, Any]) -> Dict[str, Any]:
Raises:
ParameterTypeError: if the provided parameters are not valid
"""
validated_fn = ValidatedFunction(
self.fn, config={"arbitrary_types_allowed": True}
)
args, kwargs = parameters_to_args_kwargs(self.fn, parameters)

if HAS_PYDANTIC_V2:
has_v1_models = any(isinstance(o, V1BaseModel) for o in args) or any(
isinstance(o, V1BaseModel) for o in kwargs.values()
)
has_v2_models = any(isinstance(o, V2BaseModel) for o in args) or any(
isinstance(o, V2BaseModel) for o in kwargs.values()
)

if has_v1_models and has_v2_models:
raise ParameterTypeError(
"Cannot mix Pydantic v1 and v2 models as arguments to a flow."
)

if has_v1_models:
validated_fn = V1ValidatedFunction(
self.fn, config={"arbitrary_types_allowed": True}
)
else:
validated_fn = V2ValidatedFunction(
self.fn, config={"arbitrary_types_allowed": True}
)

else:
validated_fn = ValidatedFunction(
self.fn, config={"arbitrary_types_allowed": True}
)

try:
model = validated_fn.init_model_instance(*args, **kwargs)
except pydantic.ValidationError as exc:
Expand Down

0 comments on commit cbe3380

Please sign in to comment.