From c5afcf921d947b099bf20dcb7e87f28dc7a16d5b Mon Sep 17 00:00:00 2001 From: Gabriel Gore Date: Thu, 22 Aug 2024 14:54:59 -0400 Subject: [PATCH] move to pydantic v2 models --- .../agent/actions/agent_action.py | 2 +- great_expectations_cloud/agent/agent.py | 7 +++---- great_expectations_cloud/agent/config.py | 3 ++- great_expectations_cloud/agent/event_handler.py | 11 ++++++----- great_expectations_cloud/agent/models.py | 14 +++++++------- pyproject.toml | 1 + 6 files changed, 20 insertions(+), 18 deletions(-) diff --git a/great_expectations_cloud/agent/actions/agent_action.py b/great_expectations_cloud/agent/actions/agent_action.py index eb09141b..0fb7a54b 100644 --- a/great_expectations_cloud/agent/actions/agent_action.py +++ b/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 diff --git a/great_expectations_cloud/agent/agent.py b/great_expectations_cloud/agent/agent.py index 5d73006d..3acf7f0a 100644 --- a/great_expectations_cloud/agent/agent.py +++ b/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 diff --git a/great_expectations_cloud/agent/config.py b/great_expectations_cloud/agent/config.py index e4d43464..410533cb 100644 --- a/great_expectations_cloud/agent/config.py +++ b/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): diff --git a/great_expectations_cloud/agent/event_handler.py b/great_expectations_cloud/agent/event_handler.py index d3b8cf4c..fcd1af9f 100644 --- a/great_expectations_cloud/agent/event_handler.py +++ b/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() diff --git a/great_expectations_cloud/agent/models.py b/great_expectations_cloud/agent/models.py index e54139bd..3881d87c 100644 --- a/great_expectations_cloud/agent/models.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index c1d7f7df..9bc4e570 100644 --- a/pyproject.toml +++ b/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",