From 164848fbcc361d0e012b7dca6969c9e8a25d2ab7 Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Wed, 4 Oct 2023 13:42:44 -0400 Subject: [PATCH] Conditional imports to support pydantic>2 --- prefect_airbyte/connections.py | 8 +++++++- prefect_airbyte/server.py | 7 ++++++- 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/prefect_airbyte/connections.py b/prefect_airbyte/connections.py index beba9fb..786d199 100644 --- 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 index a54529d..d04d672 100644 --- 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