From 46a1bc4d08502668feb56ef381740a17f2ed23dd Mon Sep 17 00:00:00 2001 From: Jarek Potiuk Date: Mon, 10 Apr 2023 11:47:22 +0200 Subject: [PATCH] Move Pydantic class serialization under AIP-44 feature flag The Pydantic representation of the ORM models is only used in AIP-44 in-progress feature, and we are moving to a new seialization implementation (more modular) in a near feature so in order to not unecessarily extend features in old serialization, but allow to test AIP-44, we are moving the use_pydantic_models parameter and it's implementation under _ENABLE_AIP_44 feature flag, so that it is not used accidentally. We will eventually remove it and add Pydantic serialization to the new serialization implementation. --- airflow/serialization/serialized_objects.py | 60 ++++++++++++------- .../serialization/test_serialized_objects.py | 2 + 2 files changed, 41 insertions(+), 21 deletions(-) diff --git a/airflow/serialization/serialized_objects.py b/airflow/serialization/serialized_objects.py index 45ef83f927b03..3fc85d93b3006 100644 --- a/airflow/serialization/serialized_objects.py +++ b/airflow/serialization/serialized_objects.py @@ -58,7 +58,7 @@ from airflow.serialization.pydantic.dag_run import DagRunPydantic from airflow.serialization.pydantic.dataset import DatasetPydantic from airflow.serialization.pydantic.taskinstance import TaskInstancePydantic -from airflow.settings import DAGS_FOLDER, json +from airflow.settings import _ENABLE_AIP_44, DAGS_FOLDER, json from airflow.timetables.base import Timetable from airflow.utils.code_utils import get_python_source from airflow.utils.docs import get_docs_url @@ -404,6 +404,11 @@ def serialize( :meta private: """ + if use_pydantic_models and not _ENABLE_AIP_44: + raise RuntimeError( + "Setting use_pydantic_models = True requires AIP-44 (in progress) feature flag to be true. " + "This parameter will be removed eventually when new serialization is used by AIP-44" + ) if cls._is_primitive(var): # enum.IntEnum is an int instance, it causes json dumps error so we use its value. if isinstance(var, enum.Enum): @@ -473,19 +478,26 @@ def serialize( cls.serialize(var.__dict__, strict=strict, use_pydantic_models=use_pydantic_models), type_=DAT.SIMPLE_TASK_INSTANCE, ) - elif use_pydantic_models and isinstance(var, BaseJob): - return cls._encode(BaseJobPydantic.from_orm(var).dict(), type_=DAT.BASE_JOB) - elif use_pydantic_models and isinstance(var, TaskInstance): - return cls._encode(TaskInstancePydantic.from_orm(var).dict(), type_=DAT.TASK_INSTANCE) - elif use_pydantic_models and isinstance(var, DagRun): - return cls._encode(DagRunPydantic.from_orm(var).dict(), type_=DAT.DAG_RUN) - elif use_pydantic_models and isinstance(var, Dataset): - return cls._encode(DatasetPydantic.from_orm(var).dict(), type_=DAT.DATA_SET) + elif use_pydantic_models and _ENABLE_AIP_44: + if isinstance(var, BaseJob): + return cls._encode(BaseJobPydantic.from_orm(var).dict(), type_=DAT.BASE_JOB) + elif isinstance(var, TaskInstance): + return cls._encode(TaskInstancePydantic.from_orm(var).dict(), type_=DAT.TASK_INSTANCE) + elif isinstance(var, DagRun): + return cls._encode(DagRunPydantic.from_orm(var).dict(), type_=DAT.DAG_RUN) + elif isinstance(var, Dataset): + return cls._encode(DatasetPydantic.from_orm(var).dict(), type_=DAT.DATA_SET) + else: + return cls.default_serialization(strict, var) else: - log.debug("Cast type %s to str in serialization.", type(var)) - if strict: - raise SerializationError("Encountered unexpected type") - return str(var) + return cls.default_serialization(strict, var) + + @classmethod + def default_serialization(cls, strict, var) -> str: + log.debug("Cast type %s to str in serialization.", type(var)) + if strict: + raise SerializationError("Encountered unexpected type") + return str(var) @classmethod def deserialize(cls, encoded_var: Any, use_pydantic_models=False) -> Any: @@ -494,6 +506,11 @@ def deserialize(cls, encoded_var: Any, use_pydantic_models=False) -> Any: :meta private: """ # JSON primitives (except for dict) are not encoded. + if use_pydantic_models and not _ENABLE_AIP_44: + raise RuntimeError( + "Setting use_pydantic_models = True requires AIP-44 (in progress) feature flag to be true. " + "This parameter will be removed eventually when new serialization is used by AIP-44" + ) if cls._is_primitive(encoded_var): return encoded_var elif isinstance(encoded_var, list): @@ -535,14 +552,15 @@ def deserialize(cls, encoded_var: Any, use_pydantic_models=False) -> Any: return Dataset(**var) elif type_ == DAT.SIMPLE_TASK_INSTANCE: return SimpleTaskInstance(**cls.deserialize(var)) - elif use_pydantic_models and type_ == DAT.BASE_JOB: - return BaseJobPydantic.parse_obj(var) - elif use_pydantic_models and type_ == DAT.TASK_INSTANCE: - return TaskInstancePydantic.parse_obj(var) - elif use_pydantic_models and type_ == DAT.DAG_RUN: - return DagRunPydantic.parse_obj(var) - elif use_pydantic_models and type_ == DAT.DATA_SET: - return DatasetPydantic.parse_obj(var) + elif use_pydantic_models and _ENABLE_AIP_44: + if type_ == DAT.BASE_JOB: + return BaseJobPydantic.parse_obj(var) + elif type_ == DAT.TASK_INSTANCE: + return TaskInstancePydantic.parse_obj(var) + elif type_ == DAT.DAG_RUN: + return DagRunPydantic.parse_obj(var) + elif type_ == DAT.DATA_SET: + return DatasetPydantic.parse_obj(var) else: raise TypeError(f"Invalid type {type_!s} in deserialization.") diff --git a/tests/serialization/test_serialized_objects.py b/tests/serialization/test_serialized_objects.py index 8981cafcc44c9..1eb42147836a9 100644 --- a/tests/serialization/test_serialized_objects.py +++ b/tests/serialization/test_serialized_objects.py @@ -23,6 +23,7 @@ from airflow.models.taskinstance import TaskInstance from airflow.operators.empty import EmptyOperator from airflow.serialization.pydantic.taskinstance import TaskInstancePydantic +from airflow.settings import _ENABLE_AIP_44 from airflow.utils.state import State from tests import REPO_ROOT @@ -82,6 +83,7 @@ class Test: BaseSerialization.serialize(obj, strict=True) # now raises +@pytest.mark.skipif(not _ENABLE_AIP_44, reason="AIP-44 is disabled") def test_use_pydantic_models(): """If use_pydantic_models=True the TaskInstance object should be serialized to TaskInstancePydantic."""