[MAINTENANCE] Pydantic 2.0 support

.github/workflows/ci.yml

@@ -57,13 +57,12 @@ jobs:
         run: pip install -r requirements-types.txt -r reqs/requirements-dev-contrib.txt
       - run: invoke lint --no-fmt
       - run: invoke fmt --check
-      - name: Marker-coverage-check
-        run: |
-          invoke marker-coverage
       - name: Type-check
         run: |
           invoke type-check --ci --pretty
           invoke type-check --ci --pretty  --check-stub-sources
+      - name: Marker-coverage-check
+        run: invoke marker-coverage
 
   docs-changes:
     # check whether docs were changed
@@ -405,6 +404,24 @@ jobs:
       - name: Run the tests
         run: invoke ci-tests -m unit --xdist --slowest=10  --timeout=2.0
 
+  pydantic-v1:
+    needs: [unit-tests, static-analysis]
+    if: github.event.pull_request.draft == false
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.8"
+          cache: "pip"
+          cache-dependency-path: reqs/requirements-dev-test.txt
+      - name: Install dependencies
+        run: pip install . -c ci/constraints-test/pydantic-v1-install.txt -r reqs/requirements-dev-test.txt
+      - name: Run the tests
+        run: invoke ci-tests -m unit --xdist --slowest=10  --timeout=2.0
+
   airflow220-min-versions:
     needs: [unit-tests, static-analysis]
     runs-on: ubuntu-latest
@@ -481,6 +498,7 @@ jobs:
       py39-min-versions,
       py310-min-versions,
       py311-min-versions,
+      pydantic-v1,
       airflow220-min-versions,
       import_gx,
       usage_stats_integration
@@ -528,6 +546,7 @@ jobs:
       py39-min-versions,
       py310-min-versions,
       py311-min-versions,
+      pydantic-v1,
       airflow220-min-versions,
       import_gx,
       usage_stats_integration,

ci/constraints-test/pydantic-v1-install.txt

@@ -0,0 +1,3 @@
+# this is higher than our min version but is the earliest version that actually works
+# TODO: update our min version or fix gx for our current min version
+pydantic==1.10.8

contrib/cli/requirements.txt

@@ -2,7 +2,7 @@ black[jupyter]==23.3.0 # Linting / code style
 Click>=7.1.2         # CLI tooling
 cookiecutter==2.1.1  # Project templating
 mypy==1.5.1          # Type checker
-pydantic>=1.0,<2.0    # Needed for mypy plugin
+pydantic>=1.0        # Needed for mypy plugin
 pytest>=5.3.5        # Test framework
 ruff==0.0.284        # Linting / code style
 twine==3.7.1         # Packaging

great_expectations/agent/actions/agent_action.py

@@ -1,9 +1,8 @@
 from abc import abstractmethod
 from typing import Generic, Sequence, TypeVar
 
-from pydantic import BaseModel
-
 from great_expectations.agent.models import CreatedResource, Event
+from great_expectations.compatibility.pydantic import BaseModel
 from great_expectations.data_context import CloudDataContext
 
 

great_expectations/agent/actions/draft_datasource_config_action.py

@@ -1,10 +1,9 @@
 from uuid import UUID
 
-import pydantic
-
 from great_expectations.agent.actions import ActionResult, AgentAction
 from great_expectations.agent.config import GxAgentEnvVars
 from great_expectations.agent.models import DraftDatasourceConfigEvent
+from great_expectations.compatibility import pydantic
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.core.http import create_session
 
@@ -33,7 +32,7 @@ def run(self, event: DraftDatasourceConfigEvent, id: str) -> ActionResult:
 
     def get_draft_config(self, config_id: UUID) -> dict:
         try:
-            config = GxAgentEnvVars()
+            config = GxAgentEnvVars()  # type: ignore[call-arg] # args pulled from env vars
         except pydantic.ValidationError as validation_err:
             raise RuntimeError(
                 f"Missing or badly formed environment variable\n{validation_err.errors()}"

great_expectations/agent/actions/list_table_names.py

@@ -1,7 +1,5 @@
 from typing import TYPE_CHECKING, List
 
-import pydantic
-
 from great_expectations.agent.actions.agent_action import (
     ActionResult,
     AgentAction,
@@ -10,6 +8,7 @@
 from great_expectations.agent.models import (
     ListTableNamesEvent,
 )
+from great_expectations.compatibility import pydantic
 from great_expectations.compatibility.sqlalchemy import inspect
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.core.http import create_session
@@ -47,7 +46,7 @@ def _add_or_update_table_names_list(
         self, datasource_id: str, table_names: List[str]
     ) -> None:
         try:
-            cloud_config = GxAgentEnvVars()
+            cloud_config = GxAgentEnvVars()  # type: ignore[call-arg] # args pulled from env vars
         except pydantic.ValidationError as validation_err:
             raise RuntimeError(
                 f"Missing or badly formed environment variable\n{validation_err.errors()}"

great_expectations/agent/agent.py

@@ -6,9 +6,6 @@
 from functools import partial
 from typing import TYPE_CHECKING, Dict, Optional
 
-import pydantic
-from pydantic import AmqpDsn, AnyUrl
-
 from great_expectations import get_context
 from great_expectations.agent.actions.agent_action import ActionResult
 from great_expectations.agent.config import GxAgentEnvVars
@@ -32,6 +29,8 @@
     JobStatus,
     UnknownEvent,
 )
+from great_expectations.compatibility import pydantic
+from great_expectations.compatibility.pydantic import AmqpDsn, AnyUrl
 from great_expectations.core.http import create_session
 from great_expectations.data_context.cloud_constants import CLOUD_DEFAULT_BASE_URL
 
@@ -222,7 +221,7 @@ def _get_config(cls) -> GXAgentConfig:
         # ensure we have all required env variables, and provide a useful error if not
 
         try:
-            env_vars = GxAgentEnvVars()
+            env_vars = GxAgentEnvVars()  # type: ignore[call-arg] # args pulled from env vars
         except pydantic.ValidationError as validation_err:
             raise GXAgentError(
                 f"Missing or badly formed environment variable\n{validation_err.errors()}"

great_expectations/agent/config.py

@@ -1,6 +1,5 @@
-import pydantic
-from pydantic import AnyUrl
-
+from great_expectations.compatibility import pydantic
+from great_expectations.compatibility.pydantic import AnyUrl
 from great_expectations.data_context.cloud_constants import CLOUD_DEFAULT_BASE_URL
 
 

great_expectations/agent/message_service/subscriber.py

@@ -5,14 +5,14 @@
 from json import JSONDecodeError
 from typing import Callable, Coroutine, Union
 
-import pydantic
 from pika.exceptions import AMQPError, ChannelError
 
 from great_expectations.agent.message_service.asyncio_rabbit_mq_client import (
     AsyncRabbitMQClient,
     OnMessagePayload,
 )
 from great_expectations.agent.models import Event, UnknownEvent
+from great_expectations.compatibility import pydantic
 
 
 @dataclass(frozen=True)

great_expectations/agent/models.py

@@ -2,9 +2,10 @@
 from typing import Literal, Sequence, Union
 from uuid import UUID
 
-from pydantic import BaseModel, Extra, Field
 from typing_extensions import Annotated
 
+from great_expectations.compatibility.pydantic import BaseModel, Extra, Field
+
 
 class AgentBaseModel(BaseModel):
     class Config:

great_expectations/compatibility/pydantic.py

@@ -0,0 +1,56 @@
+import pydantic
+
+from great_expectations.compatibility.not_imported import (
+    is_version_greater_or_equal,
+)
+
+if is_version_greater_or_equal(version=pydantic.VERSION, compare_version="2.0.0"):
+    # TODO: don't use star imports
+    from pydantic.v1 import *  # noqa: F403
+    from pydantic.v1 import (
+        AnyUrl,
+        UrlError,
+        error_wrappers,
+        errors,
+        fields,
+        generics,
+        json,
+        networks,
+        schema,
+        typing,
+    )
+    from pydantic.v1.generics import GenericModel
+    from pydantic.v1.main import ModelMetaclass
+
+else:
+    # TODO: don't use star imports
+    from pydantic import *  # type: ignore[assignment,no-redef] # noqa: F403
+    from pydantic import (  # type: ignore[no-redef]
+        AnyUrl,
+        UrlError,
+        error_wrappers,
+        errors,
+        fields,
+        generics,
+        json,
+        networks,
+        schema,
+        typing,
+    )
+    from pydantic.generics import GenericModel  # type: ignore[no-redef]
+    from pydantic.main import ModelMetaclass  # type: ignore[no-redef]
+
+__all__ = [
+    "AnyUrl",
+    "error_wrappers",
+    "errors",
+    "fields",
+    "GenericModel",
+    "generics",
+    "json",
+    "ModelMetaclass",
+    "networks",
+    "schema",
+    "typing",
+    "UrlError",
+]

great_expectations/core/util.py

@@ -30,11 +30,10 @@
 import dateutil.parser
 import numpy as np
 import pandas as pd
-import pydantic
 from IPython import get_ipython
 
 from great_expectations import exceptions as gx_exceptions
-from great_expectations.compatibility import pyspark, sqlalchemy
+from great_expectations.compatibility import pydantic, pyspark, sqlalchemy
 from great_expectations.compatibility.sqlalchemy import (
     SQLALCHEMY_NOT_IMPORTED,
     LegacyRow,

great_expectations/datasource/data_connector/batch_filter.py

@@ -4,9 +4,8 @@
 import logging
 from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Sequence, Union
 
-from pydantic import StrictInt, StrictStr
-
 import great_expectations.exceptions as gx_exceptions
+from great_expectations.compatibility.pydantic import StrictInt, StrictStr
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.core.id_dict import IDDict
 

great_expectations/datasource/fluent/batch_request.py

@@ -11,11 +11,14 @@
     Union,
 )
 
-import pydantic
-from pydantic import StrictStr
-from pydantic.json import pydantic_encoder
-from pydantic.schema import default_ref_template
+from great_expectations.compatibility import pydantic
+from great_expectations.compatibility.pydantic import Field, StrictStr
+from great_expectations.compatibility.pydantic import json as pydantic_json
+from great_expectations.compatibility.pydantic import (
+    schema as pydantic_schema,
+)
 
+# default_ref_template
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.core._docs_decorators import public_api
 from great_expectations.datasource.data_connector.batch_filter import (
@@ -59,17 +62,17 @@ class BatchRequest(pydantic.BaseModel):
         BatchRequest
     """
 
-    datasource_name: StrictStr = pydantic.Field(
+    datasource_name: StrictStr = Field(
         ...,
         allow_mutation=False,
         description="The name of the Datasource used to connect to the data.",
     )
-    data_asset_name: StrictStr = pydantic.Field(
+    data_asset_name: StrictStr = Field(
         ...,
         allow_mutation=False,
         description="The name of the Data Asset used to connect to the data.",
     )
-    options: BatchRequestOptions = pydantic.Field(
+    options: BatchRequestOptions = Field(
         default_factory=dict,
         allow_mutation=True,
         description=(
@@ -215,7 +218,7 @@ def schema_json(
         cls,
         *,
         by_alias: bool = True,
-        ref_template: str = default_ref_template,
+        ref_template: str = pydantic_schema.default_ref_template,
         **dumps_kwargs: Any,
     ) -> str:
         # batch_slice is only a property/pydantic setter, so we need to add a field
@@ -230,7 +233,7 @@ def schema_json(
         )
         result = cls.__config__.json_dumps(
             cls.schema(by_alias=by_alias, ref_template=ref_template),
-            default=pydantic_encoder,
+            default=pydantic_json.pydantic_encoder,
             **dumps_kwargs,
         )
         # revert model changes

great_expectations/datasource/fluent/batch_request.pyi

@@ -1,9 +1,9 @@
 from typing import Any, Dict, Optional
 
-import pydantic
-from pydantic import StrictStr
 from typing_extensions import TypeAlias
 
+from great_expectations.compatibility import pydantic
+from great_expectations.compatibility.pydantic import StrictStr
 from great_expectations.datasource.data_connector.batch_filter import BatchSlice
 
 BatchRequestOptions: TypeAlias = Dict[StrictStr, Any]

great_expectations/datasource/fluent/config.py

@@ -21,9 +21,9 @@
     overload,
 )
 
-from pydantic import Extra, Field, validator
 from ruamel.yaml import YAML
 
+from great_expectations.compatibility.pydantic import Extra, Field, validator
 from great_expectations.compatibility.sqlalchemy import TextClause
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.datasource.fluent.constants import (
@@ -40,8 +40,9 @@
 )
 
 if TYPE_CHECKING:
-    from pydantic.error_wrappers import ErrorDict as PydanticErrorDict
-
+    from great_expectations.compatibility.pydantic.error_wrappers import (
+        ErrorDict as PydanticErrorDict,
+    )
     from great_expectations.datasource.fluent.fluent_base_model import (
         AbstractSetIntStr,
         MappingIntStrAny,

great_expectations/datasource/fluent/config_str.py

@@ -4,8 +4,7 @@
 import warnings
 from typing import TYPE_CHECKING, Mapping
 
-from pydantic import SecretStr
-
+from great_expectations.compatibility.pydantic import SecretStr
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.core.config_substitutor import TEMPLATE_STR_REGEX
 

great_expectations/datasource/fluent/data_asset/data_connector/azure_blob_storage_data_connector.py

@@ -4,8 +4,7 @@
 import re
 from typing import TYPE_CHECKING, Callable, ClassVar, List, Optional, Type
 
-import pydantic
-
+from great_expectations.compatibility import pydantic
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.core.batch_spec import AzureBatchSpec, PathBatchSpec
 from great_expectations.datasource.data_connector.util import (

great_expectations/datasource/fluent/data_asset/data_connector/filesystem_data_connector.py

@@ -5,8 +5,7 @@
 import re
 from typing import TYPE_CHECKING, Callable, ClassVar, List, Optional, Type
 
-import pydantic
-
+from great_expectations.compatibility import pydantic
 from great_expectations.compatibility.typing_extensions import override
 from great_expectations.datasource.data_connector.util import (
     get_filesystem_one_level_directory_glob_path_list,