Conditional imports to support operating with pydantic>2 installed (#159)

* Conditional imports to support pydantic>2

* Using the new naming convention

* Skipping integration tests if the required environment variables aren't set

Author: chrisguidry

.github/workflows/tests.yml

@@ -31,7 +31,7 @@ jobs:
 
       - name: Run unit tests
         env:
-          PREFECT_ORION_DATABASE_CONNECTION_URL: "sqlite+aiosqlite:///./orion-tests.db"
+          PREFECT_API_DATABASE_CONNECTION_URL: "sqlite+aiosqlite:///./orion-tests.db"
         run: |
           coverage run --branch -m pytest tests -vv -m "not integration"
           coverage report
@@ -60,10 +60,12 @@ jobs:
 
       - name: Run integration tests
         env:
-          PREFECT_ORION_DATABASE_CONNECTION_URL: "sqlite+aiosqlite:///./orion-tests.db"
-          DBT_CLOUD_JOB_ID: ${{ secrets.DBT_CLOUD_JOB_ID }}
-          DBT_CLOUD_ACCOUNT_ID: ${{ secrets.DBT_CLOUD_ACCOUNT_ID }}
-          DBT_CLOUD_API_KEY: ${{ secrets.DBT_CLOUD_API_KEY }}
+          PREFECT_API_DATABASE_CONNECTION_URL: "sqlite+aiosqlite:///./orion-tests.db"
+          # Note: our dbt Cloud test account has been cleared out, so we will need to
+          # recreate that before we can run the integration tests against the real thing
+          # DBT_CLOUD_JOB_ID: ${{ secrets.DBT_CLOUD_JOB_ID }}
+          # DBT_CLOUD_ACCOUNT_ID: ${{ secrets.DBT_CLOUD_ACCOUNT_ID }}
+          # DBT_CLOUD_API_KEY: ${{ secrets.DBT_CLOUD_API_KEY }}
         run: |
           pytest tests -vv -m integration
 

prefect_dbt/cli/commands.py

@@ -7,7 +7,12 @@
 from prefect import get_run_logger, task
 from prefect.utilities.filesystem import relative_path_to_current_platform
 from prefect_shell.commands import ShellOperation, shell_run_command
-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_dbt.cli.credentials import DbtCliProfile
 

prefect_dbt/cli/configs/base.py

@@ -5,7 +5,12 @@
 from typing import Any, Dict, Optional
 
 from prefect.blocks.core import Block
-from pydantic import BaseModel, Field, SecretField
+from pydantic import VERSION as PYDANTIC_VERSION
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import BaseModel, Field, SecretField
+else:
+    from pydantic import BaseModel, Field, SecretField
 
 
 class DbtConfigs(Block, abc.ABC):

prefect_dbt/cli/configs/bigquery.py

@@ -8,7 +8,12 @@
 except ImportError:
     from typing_extensions import Literal
 
-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_dbt.cli.configs.base import BaseTargetConfigs, MissingExtrasRequireError
 

prefect_dbt/cli/configs/postgres.py

@@ -2,7 +2,13 @@
 import warnings
 from typing import Any, Dict, Union
 
-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 typing_extensions import Literal
 
 from prefect_dbt.cli.configs.base import BaseTargetConfigs, MissingExtrasRequireError

prefect_dbt/cli/configs/snowflake.py

@@ -6,7 +6,12 @@
 except ImportError:
     from typing_extensions import Literal
 
-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_dbt.cli.configs.base import BaseTargetConfigs, MissingExtrasRequireError
 

prefect_dbt/cli/credentials.py

@@ -2,7 +2,12 @@
 from typing import Any, Dict, Optional, Union
 
 from prefect.blocks.core import Block
-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_dbt.cli.configs import GlobalConfigs, TargetConfigs
 

prefect_dbt/cloud/credentials.py

@@ -2,7 +2,13 @@
 from typing import Union
 
 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 typing_extensions import Literal
 
 from prefect_dbt.cloud.clients import (

prefect_dbt/cloud/jobs.py

@@ -10,7 +10,13 @@
 from prefect.blocks.abstract import JobBlock, JobRun
 from prefect.context import FlowRunContext
 from prefect.utilities.asyncutils import 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 typing_extensions import Literal
 
 from prefect_dbt.cloud.credentials import DbtCloudCredentials

prefect_dbt/cloud/models.py

@@ -2,7 +2,12 @@
 from typing import List, Optional
 
 from prefect.context import FlowRunContext, TaskRunContext, get_run_context
-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
 
 
 def default_cause_factory():