Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Upgrade to Pydantic v2 #35551

Merged
merged 2 commits into from
Nov 22, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
4 changes: 0 additions & 4 deletions airflow/configuration.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,10 +62,6 @@
warnings.filterwarnings(action="default", category=DeprecationWarning, module="airflow")
warnings.filterwarnings(action="default", category=PendingDeprecationWarning, module="airflow")

# Temporarily suppress warnings from pydantic until we upgrade minimum version of pydantic to v2
# Which should happen in Airflow 2.8.0
warnings.filterwarnings(action="ignore", category=UserWarning, module=r"pydantic._internal._config")

_SQLITE3_VERSION_PATTERN = re2.compile(r"(?P<version>^\d+(?:\.\d+)*)\D?.*$")

ConfigType = Union[str, int, float, bool]
Expand Down
29 changes: 10 additions & 19 deletions airflow/serialization/pydantic/dag.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,13 @@
from typing import Any, List, Optional

from dateutil import relativedelta
from pydantic import BaseModel as BaseModelPydantic, PlainSerializer, PlainValidator, ValidationInfo
from pydantic import (
BaseModel as BaseModelPydantic,
ConfigDict,
PlainSerializer,
PlainValidator,
ValidationInfo,
)
from typing_extensions import Annotated

from airflow import DAG, settings
Expand Down Expand Up @@ -86,12 +92,7 @@ class DagOwnerAttributesPydantic(BaseModelPydantic):
owner: str
link: str

class Config:
"""Make sure it deals automatically with SQLAlchemy ORM classes."""

from_attributes = True
orm_mode = True # Pydantic 1.x compatibility.
arbitrary_types_allowed = True
model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)


class DagTagPydantic(BaseModelPydantic):
Expand All @@ -100,12 +101,7 @@ class DagTagPydantic(BaseModelPydantic):
name: str
dag_id: str

class Config:
"""Make sure it deals automatically with SQLAlchemy ORM classes."""

from_attributes = True
orm_mode = True # Pydantic 1.x compatibility.
arbitrary_types_allowed = True
model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)


class DagModelPydantic(BaseModelPydantic):
Expand Down Expand Up @@ -141,12 +137,7 @@ class DagModelPydantic(BaseModelPydantic):

_processor_dags_folder: Optional[str] = None

class Config:
"""Make sure it deals automatically with SQLAlchemy ORM classes."""

from_attributes = True
orm_mode = True # Pydantic 1.x compatibility.
arbitrary_types_allowed = True
model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)

@property
def relative_fileloc(self) -> pathlib.Path:
Expand Down
9 changes: 2 additions & 7 deletions airflow/serialization/pydantic/dag_run.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from datetime import datetime
from typing import TYPE_CHECKING, Iterable, List, Optional

from pydantic import BaseModel as BaseModelPydantic
from pydantic import BaseModel as BaseModelPydantic, ConfigDict

from airflow.serialization.pydantic.dag import PydanticDag
from airflow.serialization.pydantic.dataset import DatasetEventPydantic
Expand Down Expand Up @@ -56,12 +56,7 @@ class DagRunPydantic(BaseModelPydantic):
dag: Optional[PydanticDag]
consumed_dataset_events: List[DatasetEventPydantic] # noqa

class Config:
"""Make sure it deals automatically with SQLAlchemy ORM classes."""

from_attributes = True
orm_mode = True # Pydantic 1.x compatibility.
arbitrary_types_allowed = True
model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)

@property
def logical_date(self) -> datetime:
Expand Down
27 changes: 5 additions & 22 deletions airflow/serialization/pydantic/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
from datetime import datetime
from typing import List, Optional

from pydantic import BaseModel as BaseModelPydantic
from pydantic import BaseModel as BaseModelPydantic, ConfigDict


class DagScheduleDatasetReferencePydantic(BaseModelPydantic):
Expand All @@ -28,11 +28,7 @@ class DagScheduleDatasetReferencePydantic(BaseModelPydantic):
created_at: datetime
updated_at: datetime

class Config:
"""Make sure it deals automatically with SQLAlchemy ORM classes."""

from_attributes = True
orm_mode = True # Pydantic 1.x compatibility.
model_config = ConfigDict(from_attributes=True)


class TaskOutletDatasetReferencePydantic(BaseModelPydantic):
Expand All @@ -44,11 +40,7 @@ class TaskOutletDatasetReferencePydantic(BaseModelPydantic):
created_at: datetime
updated_at: datetime

class Config:
"""Make sure it deals automatically with SQLAlchemy ORM classes."""

from_attributes = True
orm_mode = True # Pydantic 1.x compatibility.
model_config = ConfigDict(from_attributes=True)


class DatasetPydantic(BaseModelPydantic):
Expand All @@ -64,11 +56,7 @@ class DatasetPydantic(BaseModelPydantic):
consuming_dags: List[DagScheduleDatasetReferencePydantic]
producing_tasks: List[TaskOutletDatasetReferencePydantic]

class Config:
"""Make sure it deals automatically with SQLAlchemy ORM classes."""

from_attributes = True
orm_mode = True # Pydantic 1.x compatibility.
model_config = ConfigDict(from_attributes=True)


class DatasetEventPydantic(BaseModelPydantic):
Expand All @@ -84,9 +72,4 @@ class DatasetEventPydantic(BaseModelPydantic):
timestamp: datetime
dataset: Optional[DatasetPydantic]

class Config:
"""Make sure it deals automatically with SQLAlchemy ORM classes."""

from_attributes = True
orm_mode = True # Pydantic 1.x compatibility.
arbitrary_types_allowed = True
model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)
8 changes: 2 additions & 6 deletions airflow/serialization/pydantic/job.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
from functools import cached_property
from typing import Optional

from pydantic import BaseModel as BaseModelPydantic
from pydantic import BaseModel as BaseModelPydantic, ConfigDict

from airflow.executors.executor_loader import ExecutorLoader
from airflow.jobs.base_job_runner import BaseJobRunner
Expand All @@ -44,11 +44,7 @@ class JobPydantic(BaseModelPydantic):
hostname: Optional[str]
unixname: Optional[str]

class Config:
"""Make sure it deals automatically with SQLAlchemy ORM classes."""

from_attributes = True
orm_mode = True # Pydantic 1.x compatibility.
model_config = ConfigDict(from_attributes=True)

@cached_property
def executor(self):
Expand Down
9 changes: 2 additions & 7 deletions airflow/serialization/pydantic/taskinstance.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from datetime import datetime
from typing import TYPE_CHECKING, Any, Iterable, Optional

from pydantic import BaseModel as BaseModelPydantic, PlainSerializer, PlainValidator
from pydantic import BaseModel as BaseModelPydantic, ConfigDict, PlainSerializer, PlainValidator
from typing_extensions import Annotated

from airflow.models import Operator
Expand Down Expand Up @@ -105,12 +105,7 @@ class TaskInstancePydantic(BaseModelPydantic, LoggingMixin):
dag_run: Optional[DagRunPydantic]
dag_model: Optional[DagModelPydantic]

class Config:
"""Make sure it deals automatically with SQLAlchemy ORM classes."""

from_attributes = True
orm_mode = True # Pydantic 1.x compatibility.
arbitrary_types_allowed = True
model_config = ConfigDict(from_attributes=True, arbitrary_types_allowed=True)

def init_run_context(self, raw: bool = False) -> None:
"""Set the log context."""
Expand Down
9 changes: 1 addition & 8 deletions airflow/serialization/serde.py
Original file line number Diff line number Diff line change
Expand Up @@ -319,14 +319,7 @@ def _is_pydantic(cls: Any) -> bool:
Checking is done by attributes as it is significantly faster than
using isinstance.
"""
return (
hasattr(cls, "__validators__")
and hasattr(cls, "__fields__")
and hasattr(cls, "dict") # Pydantic v1
or hasattr(cls, "model_config")
and hasattr(cls, "model_fields")
and hasattr(cls, "model_fields_set") # Pydantic v2
)
return hasattr(cls, "model_config") and hasattr(cls, "model_fields") and hasattr(cls, "model_fields_set")


def _register():
Expand Down
5 changes: 1 addition & 4 deletions airflow/serialization/serialized_objects.py
Original file line number Diff line number Diff line change
Expand Up @@ -502,10 +502,7 @@ def serialize(
elif use_pydantic_models and _ENABLE_AIP_44:

def _pydantic_model_dump(model_cls: type[BaseModel], var: Any) -> dict[str, Any]:
try:
return model_cls.model_validate(var).model_dump(mode="json") # type: ignore[attr-defined]
except AttributeError: # Pydantic 1.x compatibility.
return model_cls.from_orm(var).dict() # type: ignore[attr-defined]
return model_cls.model_validate(var).model_dump(mode="json") # type: ignore[attr-defined]

if isinstance(var, Job):
return cls._encode(_pydantic_model_dump(JobPydantic, var), type_=DAT.BASE_JOB)
Expand Down
6 changes: 1 addition & 5 deletions setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -133,11 +133,7 @@ install_requires =
pendulum>=2.0,<3.0
pluggy>=1.0
psutil>=4.2.0
# We should bump it to at least pydantic>=2.3.0 when we prepare Airflow 2.8.0 release
# we keep Pydantic < 1 for compatibility with packages that depend on Pydantic 1
# We should also remove then `filterwarning` for pydantic from airflow/configuration.py
# and # Pydantic v1 check in airflow/serialization/serde.py
pydantic>=1.10.0
pydantic>=2.3.0
pygments>=2.0.1
pyjwt>=2.0.0
python-daemon>=3.0.0
Expand Down