This is an automated email from the ASF dual-hosted git repository.
ephraimanierobi pushed a commit to branch v2-8-test
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit b331cb2e0e3c47e3e2322decba29cba1d6f1aecf
Author: Sin-Woo Bang <sinwoob...@gmail.com>
AuthorDate: Wed Nov 22 20:56:43 2023 +0900
Upgrade to Pydantic v2 (#35551)
* Replace deprecated Config with ConfigDict
* Drop Pydantic v1 compatibility as bumping it to 2.3.0
---
airflow/configuration.py | 4 ----
airflow/serialization/pydantic/dag.py | 29 +++++++++-----------------
airflow/serialization/pydantic/dag_run.py | 9 ++------
airflow/serialization/pydantic/dataset.py | 27 +++++-------------------
airflow/serialization/pydantic/job.py | 8 ++-----
airflow/serialization/pydantic/taskinstance.py | 9 ++------
airflow/serialization/serde.py | 9 +-------
airflow/serialization/serialized_objects.py | 5 +----
setup.cfg | 6 +-----
9 files changed, 24 insertions(+), 82 deletions(-)
diff --git a/airflow/configuration.py b/airflow/configuration.py
index 6b03759033..1df62c9b99 100644
--- a/airflow/configuration.py
+++ b/airflow/configuration.py
@@ -62,10 +62,6 @@ if not sys.warnoptions:
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]
diff --git a/airflow/serialization/pydantic/dag.py
b/airflow/serialization/pydantic/dag.py
index 6631afdf73..04b2472355 100644
--- a/airflow/serialization/pydantic/dag.py
+++ b/airflow/serialization/pydantic/dag.py
@@ -21,7 +21,13 @@ from datetime import datetime, timedelta
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
@@ -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):
@@ -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):
@@ -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:
diff --git a/airflow/serialization/pydantic/dag_run.py
b/airflow/serialization/pydantic/dag_run.py
index aaa4372a50..cd0886ecaf 100644
--- a/airflow/serialization/pydantic/dag_run.py
+++ b/airflow/serialization/pydantic/dag_run.py
@@ -19,7 +19,7 @@ from __future__ import annotations
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
@@ -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:
diff --git a/airflow/serialization/pydantic/dataset.py
b/airflow/serialization/pydantic/dataset.py
index 14255c8271..0c233a3fd6 100644
--- a/airflow/serialization/pydantic/dataset.py
+++ b/airflow/serialization/pydantic/dataset.py
@@ -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):
@@ -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):
@@ -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):
@@ -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):
@@ -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)
diff --git a/airflow/serialization/pydantic/job.py
b/airflow/serialization/pydantic/job.py
index eb92411090..2db30ab2c8 100644
--- a/airflow/serialization/pydantic/job.py
+++ b/airflow/serialization/pydantic/job.py
@@ -18,7 +18,7 @@ import datetime
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
@@ -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):
diff --git a/airflow/serialization/pydantic/taskinstance.py
b/airflow/serialization/pydantic/taskinstance.py
index 0043bfaef0..106a31186e 100644
--- a/airflow/serialization/pydantic/taskinstance.py
+++ b/airflow/serialization/pydantic/taskinstance.py
@@ -19,7 +19,7 @@ from __future__ import annotations
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
@@ -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."""
diff --git a/airflow/serialization/serde.py b/airflow/serialization/serde.py
index c86f643505..23d67e6162 100644
--- a/airflow/serialization/serde.py
+++ b/airflow/serialization/serde.py
@@ -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():
diff --git a/airflow/serialization/serialized_objects.py
b/airflow/serialization/serialized_objects.py
index 889177f4c1..9d7955bf01 100644
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -502,10 +502,7 @@ class BaseSerialization:
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)
diff --git a/setup.cfg b/setup.cfg
index 1bbf5fb8f1..cfde203f43 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -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
