From 525c9173a7db0e7cbe337a2252204b2d076124cf Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Wed, 4 Oct 2023 15:48:55 -0400 Subject: [PATCH] Conditional imports to support operating with `pydantic>2` installed (#317) --- prefect_aws/client_parameters.py | 7 ++++++- prefect_aws/credentials.py | 7 ++++++- prefect_aws/ecs.py | 8 +++++++- prefect_aws/s3.py | 7 ++++++- prefect_aws/secrets_manager.py | 7 ++++++- prefect_aws/workers/ecs_worker.py | 8 +++++++- requirements-dev.txt | 4 +++- tests/test_ecs.py | 7 ++++++- tests/workers/test_ecs_worker.py | 8 +++++++- 9 files changed, 54 insertions(+), 9 deletions(-) diff --git a/prefect_aws/client_parameters.py b/prefect_aws/client_parameters.py index 3f02b4b2..bf030590 100644 --- a/prefect_aws/client_parameters.py +++ b/prefect_aws/client_parameters.py @@ -5,7 +5,12 @@ from botocore import UNSIGNED from botocore.client import Config -from pydantic import BaseModel, Field, FilePath, root_validator, validator +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import BaseModel, Field, FilePath, root_validator, validator +else: + from pydantic import BaseModel, Field, FilePath, root_validator, validator class AwsClientParameters(BaseModel): diff --git a/prefect_aws/credentials.py b/prefect_aws/credentials.py index f390faef..7810f7f9 100644 --- a/prefect_aws/credentials.py +++ b/prefect_aws/credentials.py @@ -7,7 +7,12 @@ from mypy_boto3_s3 import S3Client from mypy_boto3_secretsmanager import SecretsManagerClient from prefect.blocks.abstract import CredentialsBlock -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_aws.client_parameters import AwsClientParameters diff --git a/prefect_aws/ecs.py b/prefect_aws/ecs.py index a9534332..7dbfa68a 100644 --- a/prefect_aws/ecs.py +++ b/prefect_aws/ecs.py @@ -121,7 +121,13 @@ from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible from prefect.utilities.dockerutils import get_prefect_image_name from prefect.utilities.pydantic import JsonPatch -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 + from slugify import slugify from typing_extensions import Literal, Self diff --git a/prefect_aws/s3.py b/prefect_aws/s3.py index 17fe8870..fe646ddc 100644 --- a/prefect_aws/s3.py +++ b/prefect_aws/s3.py @@ -14,7 +14,12 @@ from prefect.filesystems import WritableDeploymentStorage, WritableFileSystem from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible from prefect.utilities.filesystem import filter_files -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_aws import AwsCredentials, MinIOCredentials from prefect_aws.client_parameters import AwsClientParameters diff --git a/prefect_aws/secrets_manager.py b/prefect_aws/secrets_manager.py index 01bc9ba1..5eea5ba8 100644 --- a/prefect_aws/secrets_manager.py +++ b/prefect_aws/secrets_manager.py @@ -5,7 +5,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_aws import AwsCredentials diff --git a/prefect_aws/workers/ecs_worker.py b/prefect_aws/workers/ecs_worker.py index f4d97bef..1b1c8073 100644 --- a/prefect_aws/workers/ecs_worker.py +++ b/prefect_aws/workers/ecs_worker.py @@ -68,7 +68,13 @@ BaseWorker, BaseWorkerResult, ) -from pydantic import Field, root_validator +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import Field, root_validator +else: + from pydantic import Field, root_validator + from slugify import slugify from tenacity import retry, stop_after_attempt, wait_fixed, wait_random from typing_extensions import Literal diff --git a/requirements-dev.txt b/requirements-dev.txt index 9003b759..e22e9290 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -9,7 +9,9 @@ mkdocs-gen-files mkdocs-material mkdocstrings-python-legacy mock; python_version < '3.8' -moto >= 3.1.16 +# moto 4.2.5 broke something fairly deep in our test suite +# https://github.com/PrefectHQ/prefect-aws/issues/318 +moto >= 3.1.16, < 4.2.5 mypy pillow pre-commit diff --git a/tests/test_ecs.py b/tests/test_ecs.py index 9282a859..a27ff925 100644 --- a/tests/test_ecs.py +++ b/tests/test_ecs.py @@ -16,7 +16,12 @@ from prefect.server.schemas.core import Deployment, Flow, FlowRun from prefect.utilities.asyncutils import run_sync_in_worker_thread from prefect.utilities.dockerutils import get_prefect_image_name -from pydantic import ValidationError +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import ValidationError +else: + from pydantic import ValidationError from prefect_aws.ecs import ( ECS_DEFAULT_CPU, diff --git a/tests/workers/test_ecs_worker.py b/tests/workers/test_ecs_worker.py index ef57372b..1be6eb70 100644 --- a/tests/workers/test_ecs_worker.py +++ b/tests/workers/test_ecs_worker.py @@ -12,7 +12,13 @@ from moto.ec2.utils import generate_instance_identity_document from prefect.server.schemas.core import FlowRun from prefect.utilities.asyncutils import run_sync_in_worker_thread -from pydantic import ValidationError +from pydantic import VERSION as PYDANTIC_VERSION + +if PYDANTIC_VERSION.startswith("2."): + from pydantic.v1 import ValidationError +else: + from pydantic import ValidationError + from tenacity import RetryError from prefect_aws.workers.ecs_worker import (