(airflow) branch main updated: AIP-44 Optimize Pydantic Class Serialization Mapping (#39943)

[jscheffl]

This is an automated email from the ASF dual-hosted git repository.

jscheffl pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 93ba263732 AIP-44 Optimize Pydantic Class Serialization Mapping 
(#39943)
93ba263732 is described below

commit 93ba263732a30ed51fdae0b2bff7cd959be115b8
Author: Jens Scheffler <95105677+jsche...@users.noreply.github.com>
AuthorDate: Thu May 30 19:56:21 2024 +0200

    AIP-44 Optimize Pydantic Class Serialization Mapping (#39943)
---
 airflow/serialization/serialized_objects.py | 17 ++++-------------
 1 file changed, 4 insertions(+), 13 deletions(-)

diff --git a/airflow/serialization/serialized_objects.py 
b/airflow/serialization/serialized_objects.py
index b9b774ab7b..6e7f50a87c 100644
--- a/airflow/serialization/serialized_objects.py
+++ b/airflow/serialization/serialized_objects.py
@@ -369,13 +369,15 @@ _orm_to_model = {
     DagRun: DagRunPydantic,
     DagModel: DagModelPydantic,
     LogTemplate: LogTemplatePydantic,
+    Dataset: DatasetPydantic,
 }
-_type_to_class = {
+_type_to_class: dict[DAT, list] = {
     DAT.BASE_JOB: [JobPydantic, Job],
     DAT.TASK_INSTANCE: [TaskInstancePydantic, TaskInstance],
     DAT.DAG_RUN: [DagRunPydantic, DagRun],
     DAT.DAG_MODEL: [DagModelPydantic, DagModel],
     DAT.LOG_TEMPLATE: [LogTemplatePydantic, LogTemplate],
+    DAT.DATA_SET: [DatasetPydantic, Dataset],
 }
 _class_to_type = {cls_: type_ for type_, classes in _type_to_class.items() for 
cls_ in classes}
 
@@ -747,18 +749,7 @@ class BaseSerialization:
         elif type_ == DAT.CONNECTION:
             return Connection(**var)
         elif use_pydantic_models and _ENABLE_AIP_44:
-            if type_ == DAT.BASE_JOB:
-                return JobPydantic.model_validate(var)
-            elif type_ == DAT.TASK_INSTANCE:
-                return TaskInstancePydantic.model_validate(var)
-            elif type_ == DAT.DAG_RUN:
-                return DagRunPydantic.model_validate(var)
-            elif type_ == DAT.DAG_MODEL:
-                return DagModelPydantic.model_validate(var)
-            elif type_ == DAT.DATA_SET:
-                return DatasetPydantic.model_validate(var)
-            elif type_ == DAT.LOG_TEMPLATE:
-                return LogTemplatePydantic.model_validate(var)
+            return _type_to_class[type_][0].model_validate(var)
         elif type_ == DAT.ARG_NOT_SET:
             return NOTSET
         else: