migrate prefect-aws to pydantic 2 (PrefectHQ#13726)

requirements-client.txt

@@ -4,6 +4,7 @@ cachetools >= 5.3, < 6.0
 cloudpickle >= 2.0, < 4.0
 coolname >= 1.0.4, < 3.0.0
 croniter >= 1.0.12, < 3.0.0
+exceptiongroup >= 1.0.0
 fastapi >= 0.111.0, < 1.0.0
 fsspec >= 2022.5.0
 graphviz >= 0.20.1

requirements-dev.txt

@@ -2,7 +2,6 @@ ruff
 cairosvg
 codespell>=2.2.6
 ddtrace
-exceptiongroup
 ipython
 jinja2
 mkdocs

src/integrations/prefect-aws/prefect_aws/batch.py

@@ -3,11 +3,12 @@
 from typing import Any, Dict, Optional
 
 from prefect import get_run_logger, task
-from prefect.utilities.asyncutils import run_sync_in_worker_thread
+from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible
 from prefect_aws.credentials import AwsCredentials
 
 
 @task
+@sync_compatible
 async def batch_submit(
     job_name: str,
     job_queue: str,

src/integrations/prefect-aws/prefect_aws/client_parameters.py

@@ -5,15 +5,10 @@
 
 from botocore import UNSIGNED
 from botocore.client import Config
-from pydantic import VERSION as PYDANTIC_VERSION
+from pydantic import BaseModel, Field, FilePath, field_validator, model_validator
 
 from prefect_aws.utilities import hash_collection
 
-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):
     """
@@ -84,7 +79,8 @@ def __hash__(self):
             )
         )
 
-    @validator("config", pre=True)
+    @field_validator("config", mode="before")
+    @classmethod
     def instantiate_config(cls, value: Union[Config, Dict[str, Any]]) -> Dict[str, Any]:
         """
         Casts lists to Config instances.
@@ -93,7 +89,8 @@ def instantiate_config(cls, value: Union[Config, Dict[str, Any]]) -> Dict[str, A
             return value.__dict__["_user_provided_options"]
         return value
 
-    @root_validator
+    @model_validator(mode="before")
+    @classmethod
     def deprecated_verify_cert_path(cls, values: Dict[str, Any]) -> Dict[str, Any]:
         """
         If verify is not a bool, raise a warning.
@@ -112,7 +109,8 @@ def deprecated_verify_cert_path(cls, values: Dict[str, Any]) -> Dict[str, Any]:
             )
         return values
 
-    @root_validator
+    @model_validator(mode="before")
+    @classmethod
     def verify_cert_path_and_verify(cls, values: Dict[str, Any]) -> Dict[str, Any]:
         """
         If verify_cert_path is set but verify is False, raise a warning.
@@ -139,7 +137,7 @@ def get_params_override(self) -> Dict[str, Any]:
         Return the dictionary of the parameters to override.
         The parameters to override are the one which are not None.
         """
-        params = self.dict()
+        params = self.model_dump()
         if params.get("verify_cert_path"):
             # to ensure that verify doesn't re-overwrite verify_cert_path
             params.pop("verify")

src/integrations/prefect-aws/prefect_aws/client_waiter.py

@@ -5,11 +5,12 @@
 from botocore.waiter import WaiterModel, create_waiter_with_client
 
 from prefect import get_run_logger, task
-from prefect.utilities.asyncutils import run_sync_in_worker_thread
+from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible
 from prefect_aws.credentials import AwsCredentials
 
 
 @task
+@sync_compatible
 async def client_waiter(
     client: str,
     waiter_name: str,

src/integrations/prefect-aws/prefect_aws/credentials.py

@@ -8,15 +8,9 @@
 import boto3
 from mypy_boto3_s3 import S3Client
 from mypy_boto3_secretsmanager import SecretsManagerClient
-from pydantic import VERSION as PYDANTIC_VERSION
+from pydantic import ConfigDict, Field, SecretStr
 
 from prefect.blocks.abstract import CredentialsBlock
-
-if PYDANTIC_VERSION.startswith("2."):
-    from pydantic.v1 import Field, SecretStr
-else:
-    from pydantic import Field, SecretStr
-
 from prefect_aws.client_parameters import AwsClientParameters
 
 _LOCK = Lock()
@@ -72,6 +66,8 @@ class AwsCredentials(CredentialsBlock):
         ```
     """  # noqa E501
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/d74b16fe84ce626345adf235a47008fea2869a60-225x225.png"  # noqa
     _block_type_name = "AWS Credentials"
     _documentation_url = "https://prefecthq.github.io/prefect-aws/credentials/#prefect_aws.credentials.AwsCredentials"  # noqa
@@ -107,11 +103,6 @@ class AwsCredentials(CredentialsBlock):
         title="AWS Client Parameters",
     )
 
-    class Config:
-        """Config class for pydantic model."""
-
-        arbitrary_types_allowed = True
-
     def __hash__(self):
         field_hashes = (
             hash(self.aws_access_key_id),
@@ -209,6 +200,8 @@ class MinIOCredentials(CredentialsBlock):
         ```
     """  # noqa E501
 
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     _logo_url = "https://cdn.sanity.io/images/3ugk85nk/production/676cb17bcbdff601f97e0a02ff8bcb480e91ff40-250x250.png"  # noqa
     _block_type_name = "MinIO Credentials"
     _description = (
@@ -231,18 +224,13 @@ class MinIOCredentials(CredentialsBlock):
         description="Extra parameters to initialize the Client.",
     )
 
-    class Config:
-        """Config class for pydantic model."""
-
-        arbitrary_types_allowed = True
-
     def __hash__(self):
         return hash(
             (
                 hash(self.minio_root_user),
                 hash(self.minio_root_password),
                 hash(self.region_name),
-                hash(frozenset(self.aws_client_parameters.dict().items())),
+                hash(frozenset(self.aws_client_parameters.model_dump().items())),
             )
         )
 

src/integrations/prefect-aws/prefect_aws/glue_job.py

@@ -6,21 +6,15 @@
 import time
 from typing import Any, Optional
 
-from pydantic import VERSION as PYDANTIC_VERSION
+from pydantic import BaseModel, Field
 
 from prefect.blocks.abstract import JobBlock, JobRun
-
-if PYDANTIC_VERSION.startswith("2."):
-    from pydantic.v1 import BaseModel, Field
-else:
-    from pydantic import BaseModel, Field
-
 from prefect_aws import AwsCredentials
 
 _GlueJobClient = Any
 
 
-class GlueJobRun(JobRun, BaseModel):
+class GlueJobRun(BaseModel, JobRun):
     """Execute a Glue Job"""
 
     job_name: str = Field(

src/integrations/prefect-aws/prefect_aws/lambda_function.py

@@ -49,19 +49,14 @@
     ```
 
 """
+
 import json
 from typing import Literal, Optional
 
-from pydantic import VERSION as PYDANTIC_VERSION
+from pydantic import Field
 
 from prefect.blocks.core import Block
 from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible
-
-if PYDANTIC_VERSION.startswith("2."):
-    from pydantic.v1 import Field
-else:
-    from pydantic import Field
-
 from prefect_aws.credentials import AwsCredentials
 
 
@@ -109,11 +104,6 @@ class LambdaFunction(Block):
         description="The AWS credentials to invoke the Lambda with.",
     )
 
-    class Config:
-        """Lambda's pydantic configuration."""
-
-        smart_union = True
-
     def _get_lambda_client(self):
         """
         Retrieve a boto3 session and Lambda client

src/integrations/prefect-aws/prefect_aws/s3.py

@@ -1,4 +1,5 @@
 """Tasks for interacting with AWS S3"""
+
 import asyncio
 import io
 import os
@@ -9,19 +10,13 @@
 import boto3
 from botocore.paginate import PageIterator
 from botocore.response import StreamingBody
-from pydantic import VERSION as PYDANTIC_VERSION
+from pydantic import Field
 
 from prefect import get_run_logger, task
 from prefect.blocks.abstract import ObjectStorageBlock
 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
-
-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
 
@@ -394,7 +389,6 @@ async def example_s3_list_objects_flow():
 
 
 class S3Bucket(WritableFileSystem, WritableDeploymentStorage, ObjectStorageBlock):
-
     """
     Block used to store data using AWS S3 or S3-compatible object storage like MinIO.
 

src/integrations/prefect-aws/prefect_aws/secrets_manager.py

@@ -3,17 +3,11 @@
 from typing import Any, Dict, List, Optional, Union
 
 from botocore.exceptions import ClientError
-from pydantic import VERSION as PYDANTIC_VERSION
+from pydantic import Field
 
 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
-
-if PYDANTIC_VERSION.startswith("2."):
-    from pydantic.v1 import Field
-else:
-    from pydantic import Field
-
 from prefect_aws import AwsCredentials