move to pydantic v2 models

great_expectations_cloud/agent/actions/agent_action.py

@@ -5,7 +5,7 @@
 from typing import TYPE_CHECKING, Generic, Optional, Sequence, TypeVar
 from uuid import UUID
 
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from great_expectations_cloud.agent.models import CreatedResource, Event
 

great_expectations_cloud/agent/agent.py

@@ -19,8 +19,7 @@
 from great_expectations.data_context.data_context.context_factory import get_context
 from packaging.version import Version
 from pika.exceptions import AuthenticationError, ProbableAuthenticationError
-from pydantic import v1 as pydantic_v1
-from pydantic.v1 import AmqpDsn, AnyUrl
+from pydantic import AmqpDsn, AnyUrl, ValidationError
 from tenacity import after_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential
 
 from great_expectations_cloud.agent.config import (
@@ -370,7 +369,7 @@ def _get_config(cls) -> GXAgentConfig:
 
         try:
             env_vars = GxAgentEnvVars()
-        except pydantic_v1.ValidationError as validation_err:
+        except ValidationError as validation_err:
             raise GXAgentConfigError(
                 generate_config_validation_error_text(validation_err)
             ) from validation_err
@@ -402,7 +401,7 @@ def _get_config(cls) -> GXAgentConfig:
                 gx_cloud_organization_id=env_vars.gx_cloud_organization_id,
                 gx_cloud_access_token=env_vars.gx_cloud_access_token,
             )
-        except pydantic_v1.ValidationError as validation_err:
+        except ValidationError as validation_err:
             raise GXAgentConfigError(
                 generate_config_validation_error_text(validation_err)
             ) from validation_err

great_expectations_cloud/agent/config.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
 from great_expectations.data_context.cloud_constants import CLOUD_DEFAULT_BASE_URL
-from pydantic.v1 import AnyUrl, BaseSettings, ValidationError
+from pydantic import AnyUrl, ValidationError
+from pydantic_settings import BaseSettings
 
 
 class GxAgentEnvVars(BaseSettings):

great_expectations_cloud/agent/event_handler.py

@@ -10,12 +10,13 @@
 import great_expectations as gx
 from packaging.version import Version
 from packaging.version import parse as parse_version
-from pydantic import v1 as pydantic_v1
+from pydantic import ValidationError
 
 from great_expectations_cloud.agent.actions.unknown import UnknownEventAction
 from great_expectations_cloud.agent.exceptions import GXAgentError
 from great_expectations_cloud.agent.models import (
     Event,
+    EventTA,
     UnknownEvent,
 )
 
@@ -103,8 +104,8 @@ def handle_event(  # Refactor opportunity
     @classmethod
     def parse_event_from(cls, msg_body: bytes) -> Event:
         try:
-            event: Event = pydantic_v1.parse_raw_as(Event, msg_body)  # type: ignore[arg-type] # FIXME
-        except (pydantic_v1.ValidationError, JSONDecodeError):
+            event: Event = EventTA.validate_json(msg_body)
+        except (ValidationError, JSONDecodeError):
             # Log as bytes
             LOGGER.exception("Unable to parse event type", extra={"msg_body": f"{msg_body!r}"})
             return UnknownEvent()
@@ -114,8 +115,8 @@ def parse_event_from(cls, msg_body: bytes) -> Event:
     @classmethod
     def parse_event_from_dict(cls, msg_body: dict[str, Any]) -> Event:
         try:
-            event: Event = pydantic_v1.parse_obj_as(Event, msg_body)  # type: ignore[arg-type] # FIXME
-        except pydantic_v1.ValidationError:
+            event: Event = EventTA.validate_python(msg_body)
+        except ValidationError:
             # Log as dict
             LOGGER.exception("Unable to parse event type", extra={"msg_body": msg_body})
             return UnknownEvent()

great_expectations_cloud/agent/models.py

@@ -5,22 +5,20 @@
 from uuid import UUID
 
 from great_expectations.experimental.metric_repository.metrics import MetricTypes
-from pydantic.v1 import BaseModel, Extra, Field
+from pydantic import BaseModel, Field, TypeAdapter
 from typing_extensions import Annotated
 
 from great_expectations_cloud.agent.exceptions import GXCoreError
 
 
 class AgentBaseExtraForbid(BaseModel):
-    class Config:
-        # 2024-03-04: ZEL-501 Strictly enforce models for handling outdated APIs
-        extra: str = Extra.forbid
+    # 2024-03-04: ZEL-501 Strictly enforce models for handling outdated APIs
+    model_config = {"extra": "forbid"}
 
 
 class AgentBaseExtraIgnore(BaseModel):
-    class Config:
-        # Extra fields on Events are not strictly enforced
-        extra: str = Extra.ignore
+    # Extra fields on Events are not strictly enforced
+    model_config = {"extra": "ignore"}
 
 
 class EventBase(AgentBaseExtraIgnore):
@@ -109,6 +107,8 @@ class UnknownEvent(AgentBaseExtraForbid):
     Field(discriminator="type"),
 ]
 
+EventTA: TypeAdapter[Event] = TypeAdapter(Event)
+
 
 class EventMessage(AgentBaseExtraForbid):
     event: Event

pyproject.toml

@@ -177,6 +177,7 @@ runtime-evaluated-base-classes = [
     # NOTE: ruff is unable to detect that these are subclasses of pydantic.BaseModel
     "pydantic.BaseModel",
     "pydantic.v1.BaseModel",
+    "pydantic_settings.BaseSettings",
     "great_expectations_cloud.agent.models.AgentBaseExtraForbid",
     "great_expectations_cloud.agent.models.AgentBaseExtraIgnore",
     "great_expectations_cloud.agent.models.EventBase",