From aa58459b08b689cfff7acae63192cf902fda3948 Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Wed, 4 Oct 2023 15:50:07 -0400 Subject: [PATCH] Conditional imports to support pydantic>2 (#215) --- prefect_gcp/aiplatform.py | 8 +++++++- prefect_gcp/bigquery.py | 7 ++++++- prefect_gcp/cloud_run.py | 8 +++++++- prefect_gcp/cloud_storage.py | 7 ++++++- prefect_gcp/credentials.py | 7 ++++++- prefect_gcp/secret_manager.py | 7 ++++++- prefect_gcp/workers/cloud_run.py | 7 ++++++- prefect_gcp/workers/vertex.py | 8 +++++++- tests/test_cloud_run.py | 8 +++++++- tests/test_cloud_run_worker.py | 8 +++++++- tests/test_vertex_worker.py | 8 +++++++- 11 files changed, 72 insertions(+), 11 deletions(-) diff --git a/prefect_gcp/aiplatform.py b/prefect_gcp/aiplatform.py index 130754b2..edd9e181 100644 --- a/prefect_gcp/aiplatform.py +++ b/prefect_gcp/aiplatform.py @@ -62,7 +62,13 @@ from prefect.exceptions import InfrastructureNotFound from prefect.infrastructure import Infrastructure, InfrastructureResult from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible -from pydantic import Field +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import Field +else: + from pydantic import Field + from slugify import slugify from typing_extensions import Literal diff --git a/prefect_gcp/bigquery.py b/prefect_gcp/bigquery.py index 52c48895..45241545 100644 --- a/prefect_gcp/bigquery.py +++ b/prefect_gcp/bigquery.py @@ -10,7 +10,12 @@ from prefect.blocks.abstract import DatabaseBlock from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible from prefect.utilities.hashing import hash_objects -from pydantic import Field +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import Field +else: + from pydantic import Field from prefect_gcp.credentials import GcpCredentials diff --git a/prefect_gcp/cloud_run.py b/prefect_gcp/cloud_run.py index ff500fb0..faf55845 100644 --- a/prefect_gcp/cloud_run.py +++ b/prefect_gcp/cloud_run.py @@ -43,7 +43,13 @@ from prefect.exceptions import InfrastructureNotFound from prefect.infrastructure.base import Infrastructure, InfrastructureResult from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible -from pydantic import BaseModel, Field, root_validator, validator +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import BaseModel, Field, root_validator, validator +else: + from pydantic import BaseModel, Field, root_validator, validator + from typing_extensions import Literal from prefect_gcp.credentials import GcpCredentials diff --git a/prefect_gcp/cloud_storage.py b/prefect_gcp/cloud_storage.py index cc144022..228f58ae 100644 --- a/prefect_gcp/cloud_storage.py +++ b/prefect_gcp/cloud_storage.py @@ -13,7 +13,12 @@ from prefect.logging import disable_run_logger from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible from prefect.utilities.filesystem import filter_files -from pydantic import Field, validator +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import Field, validator +else: + from pydantic import Field, validator # cannot be type_checking only or else `fields = cls.schema()` raises # TypeError: issubclass() arg 1 must be a class diff --git a/prefect_gcp/credentials.py b/prefect_gcp/credentials.py index 8900fae5..8fa805ff 100644 --- a/prefect_gcp/credentials.py +++ b/prefect_gcp/credentials.py @@ -12,7 +12,12 @@ from prefect.blocks.abstract import CredentialsBlock from prefect.blocks.fields import SecretDict from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible -from pydantic import Field, root_validator, validator +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import Field, root_validator, validator +else: + from pydantic import Field, root_validator, validator try: from google.cloud.bigquery import Client as BigQueryClient diff --git a/prefect_gcp/secret_manager.py b/prefect_gcp/secret_manager.py index 1132c15a..0d77ff7b 100644 --- a/prefect_gcp/secret_manager.py +++ b/prefect_gcp/secret_manager.py @@ -6,7 +6,12 @@ from prefect import get_run_logger, task from prefect.blocks.abstract import SecretBlock from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible -from pydantic import Field +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import Field +else: + from pydantic import Field from prefect_gcp.credentials import GcpCredentials diff --git a/prefect_gcp/workers/cloud_run.py b/prefect_gcp/workers/cloud_run.py index e3824d9b..649423ac 100644 --- a/prefect_gcp/workers/cloud_run.py +++ b/prefect_gcp/workers/cloud_run.py @@ -163,7 +163,12 @@ BaseWorker, BaseWorkerResult, ) -from pydantic import Field, validator +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import Field, validator +else: + from pydantic import Field, validator from prefect_gcp.cloud_run import Execution, Job from prefect_gcp.credentials import GcpCredentials diff --git a/prefect_gcp/workers/vertex.py b/prefect_gcp/workers/vertex.py index cbb12e0f..d95027a4 100644 --- a/prefect_gcp/workers/vertex.py +++ b/prefect_gcp/workers/vertex.py @@ -36,7 +36,13 @@ BaseWorker, BaseWorkerResult, ) -from pydantic import Field, validator +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import Field, validator +else: + from pydantic import Field, validator + from slugify import slugify from prefect_gcp.credentials import GcpCredentials diff --git a/tests/test_cloud_run.py b/tests/test_cloud_run.py index 0dcc9004..9e8a8f46 100644 --- a/tests/test_cloud_run.py +++ b/tests/test_cloud_run.py @@ -1,7 +1,13 @@ from unittest.mock import Mock import anyio -import pydantic +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + import pydantic.v1 as pydantic +else: + import pydantic + import pytest from googleapiclient.errors import HttpError from prefect.exceptions import InfrastructureNotFound diff --git a/tests/test_cloud_run_worker.py b/tests/test_cloud_run_worker.py index 119e8f2f..171032e8 100644 --- a/tests/test_cloud_run_worker.py +++ b/tests/test_cloud_run_worker.py @@ -2,7 +2,13 @@ from unittest.mock import Mock import anyio -import pydantic +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + import pydantic.v1 as pydantic +else: + import pydantic + import pytest from googleapiclient.errors import HttpError from jsonschema.exceptions import ValidationError diff --git a/tests/test_vertex_worker.py b/tests/test_vertex_worker.py index 11d73a50..8faf2a0f 100644 --- a/tests/test_vertex_worker.py +++ b/tests/test_vertex_worker.py @@ -3,7 +3,13 @@ from unittest.mock import MagicMock import anyio -import pydantic +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + import pydantic.v1 as pydantic +else: + import pydantic + import pytest from google.cloud.aiplatform_v1.types.job_service import CancelCustomJobRequest from google.cloud.aiplatform_v1.types.job_state import JobState