Merge pull request #60 from PrefectHQ/pydantic-v2-compatibility

Conditional imports to support operating with `pydantic>2` installed
PrefectHQ · Oct 4, 2023 · d1760bb · d1760bb
2 parents 9079bbe + 164848f
commit d1760bb
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/prefect_airbyte/connections.py b/prefect_airbyte/connections.py
@@ -8,7 +8,13 @@
 from prefect import get_run_logger, task
 from prefect.blocks.abstract import JobBlock, JobRun
 from prefect.utilities.asyncutils import sync_compatible
-from pydantic import BaseModel, Field
+from pydantic import VERSION as PYDANTIC_VERSION
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import BaseModel, Field
+else:
+    from pydantic import BaseModel, Field
+
 from typing_extensions import Literal
 
 from prefect_airbyte import exceptions as err

diff --git a/prefect_airbyte/server.py b/prefect_airbyte/server.py
@@ -3,7 +3,12 @@
 from logging import Logger
 
 from prefect.blocks.core import Block
-from pydantic import Field, SecretStr
+from pydantic import VERSION as PYDANTIC_VERSION
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import Field, SecretStr
+else:
+    from pydantic import Field, SecretStr
 
 from prefect_airbyte.client import AirbyteClient