From cbe3380aff66ea0b9a2f4a122d4a103577cd9a6b Mon Sep 17 00:00:00 2001
From: Chris Guidry <chris.g@prefect.io>
Date: Wed, 25 Oct 2023 15:17:34 -0400
Subject: [PATCH] Refining how we validate flow parameters under Pydantic 1
 versus 2

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:
* https://github.com/PrefectHQ/prefect-airbyte/issues/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.
---
 src/prefect/flows.py | 37 ++++++++++++++++++++++++++++++++-----
 1 file changed, 32 insertions(+), 5 deletions(-)

diff --git a/src/prefect/flows.py b/src/prefect/flows.py
index 10c0c316c1b9..b3bfa2028ee4 100644
--- a/src/prefect/flows.py
+++ b/src/prefect/flows.py
@@ -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:
@@ -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: