Conditional imports to support pydantic>2 (#215)

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
 

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
 

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

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

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

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
 

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

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

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

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

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