Replace _handle_request elif chains with a handler registry #65827

airflow-core/src/airflow/dag_processing/processor.py

@@ -23,7 +23,7 @@
 import traceback
 from collections.abc import Callable, Sequence
 from pathlib import Path
-from typing import TYPE_CHECKING, Annotated, BinaryIO, ClassVar, Literal
+from typing import TYPE_CHECKING, Annotated, Any, BinaryIO, ClassVar, Literal
 
 import attrs
 from pydantic import BaseModel, Field, TypeAdapter
@@ -71,24 +71,8 @@
     XComSequenceIndexResult,
     XComSequenceSliceResult,
 )
-from airflow.sdk.execution_time.request_handlers import (
-    handle_delete_variable,
-    handle_get_prev_successful_dag_run,
-    handle_get_previous_dag_run,
-    handle_get_previous_ti,
-    handle_get_task_states,
-    handle_get_ti_count,
-    handle_get_variable_keys,
-    handle_get_xcom,
-    handle_get_xcom_count,
-    handle_get_xcom_sequence_item,
-    handle_get_xcom_sequence_slice,
-    handle_mask_secret,
-    handle_put_variable,
-)
 from airflow.sdk.execution_time.supervisor import WatchedSubprocess
 from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance, _send_error_email_notification
-from airflow.sdk.log import mask_secret
 from airflow.serialization.serialized_objects import DagSerialization, LazyDeserializedDAG
 from airflow.utils.dag_version_inflation_checker import check_dag_file_stability
 from airflow.utils.file import iter_airflow_imports
@@ -562,14 +546,12 @@ class DagFileProcessorProcess(WatchedSubprocess, LoggingMixin):
     in core Airflow.
     """
 
+    _msg_union: ClassVar[Any] = ToManager
     logger_filehandle: BinaryIO
     parsing_result: DagFileParsingResult | None = None
     decoder: ClassVar[TypeAdapter[ToManager]] = TypeAdapter[ToManager](ToManager)
     had_callbacks: bool = False  # Track if this process was started with callbacks to prevent stale DAG detection false positives
 
-    client: Client
-    """The HTTP client to use for communication with the API server."""
-
     bundle_name: str
     dag_file_rel_path: str
 
@@ -651,75 +633,11 @@ def _create_log_forwarder(
         )
 
     def _handle_request(self, msg: ToManager, log: FilteringBoundLogger, req_id: int) -> None:
-        from airflow.sdk.api.datamodels._generated import (
-            ConnectionResponse,
-            VariableResponse,
-        )
-
-        resp: BaseModel | None = None
-        dump_opts: dict[str, bool] = {}
         if isinstance(msg, DagFileParsingResult):
             self.parsing_result = msg
-        elif isinstance(msg, GetConnection):
-            conn = self.client.connections.get(msg.conn_id)
-            if isinstance(conn, ConnectionResponse):
-                if conn.password:
-                    mask_secret(conn.password)
-                if conn.extra:
-                    mask_secret(conn.extra)
-                conn_result = ConnectionResult.from_conn_response(conn)
-                resp = conn_result
-                dump_opts = {"exclude_unset": True, "by_alias": True}
-            else:
-                resp = conn
-        elif isinstance(msg, GetVariable):
-            var = self.client.variables.get(msg.key)
-            if isinstance(var, VariableResponse):
-                if var.value:
-                    mask_secret(var.value, var.key)
-                var_result = VariableResult.from_variable_response(var)
-                resp = var_result
-                dump_opts = {"exclude_unset": True}
-            else:
-                resp = var
-        elif isinstance(msg, GetVariableKeys):
-            resp, dump_opts = handle_get_variable_keys(self.client, msg)
-        elif isinstance(msg, PutVariable):
-            resp, dump_opts = handle_put_variable(self.client, msg)
-        elif isinstance(msg, DeleteVariable):
-            resp, dump_opts = handle_delete_variable(self.client, msg)
-        elif isinstance(msg, GetPreviousDagRun):
-            resp, dump_opts = handle_get_previous_dag_run(self.client, msg)
-        elif isinstance(msg, GetPrevSuccessfulDagRun):
-            resp, dump_opts = handle_get_prev_successful_dag_run(self.client, self.id)
-        elif isinstance(msg, GetXCom):
-            resp, dump_opts = handle_get_xcom(self.client, msg)
-        elif isinstance(msg, GetXComCount):
-            resp, dump_opts = handle_get_xcom_count(self.client, msg)
-        elif isinstance(msg, GetXComSequenceItem):
-            resp, dump_opts = handle_get_xcom_sequence_item(self.client, msg)
-        elif isinstance(msg, GetXComSequenceSlice):
-            resp, dump_opts = handle_get_xcom_sequence_slice(self.client, msg)
-        elif isinstance(msg, MaskSecret):
-            handle_mask_secret(msg)
-        elif isinstance(msg, GetTICount):
-            resp, dump_opts = handle_get_ti_count(self.client, msg)
-        elif isinstance(msg, GetTaskStates):
-            resp, dump_opts = handle_get_task_states(self.client, msg)
-        elif isinstance(msg, GetPreviousTI):
-            resp, dump_opts = handle_get_previous_ti(self.client, msg)
-        else:
-            log.error("Unhandled request", msg=msg)
-            self.send_msg(
-                None,
-                request_id=req_id,
-                error=ErrorResponse(
-                    detail={"status_code": 400, "message": "Unhandled request"},
-                ),
-            )
+            self.send_msg(None, request_id=req_id, error=None)
             return
-
-        self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
+        super()._handle_request(msg, log, req_id)
 
     @property
     def is_ready(self) -> bool:

airflow-core/src/airflow/jobs/triggerer_job_runner.py

@@ -17,7 +17,6 @@
 from __future__ import annotations
 
 import asyncio
-import functools
 import logging
 import math
 import os
@@ -101,30 +100,6 @@
     _RequestFrame,
 )
 from airflow.sdk.execution_time.context import AssetStateStoreAccessors
-from airflow.sdk.execution_time.request_handlers import (
-    handle_clear_asset_state_store_by_name,
-    handle_clear_asset_state_store_by_uri,
-    handle_delete_asset_state_store_by_name,
-    handle_delete_asset_state_store_by_uri,
-    handle_delete_variable,
-    handle_delete_xcom,
-    handle_get_asset_state_store_by_name,
-    handle_get_asset_state_store_by_uri,
-    handle_get_connection,
-    handle_get_dag_run_state,
-    handle_get_dr_count,
-    handle_get_previous_ti,
-    handle_get_task_states,
-    handle_get_ti_count,
-    handle_get_variable,
-    handle_get_variable_keys,
-    handle_get_xcom,
-    handle_mask_secret,
-    handle_put_variable,
-    handle_set_asset_state_store_by_name,
-    handle_set_asset_state_store_by_uri,
-    handle_set_xcom,
-)
 from airflow.sdk.execution_time.supervisor import WatchedSubprocess, make_buffered_socket_reader
 from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
 from airflow.serialization.serialized_objects import DagSerialization
@@ -484,6 +459,7 @@ class TriggerRunnerSupervisor(WatchedSubprocess):
     rather than silently zombieing the supervisor.
     """
 
+    _msg_union: ClassVar[Any] = ToTriggerSupervisor
     job: Job | None = None
     capacity: int
     queues: set[str] | None = None
@@ -547,6 +523,7 @@ def start(  # type: ignore[override]
         proc = super().start(
             id=proc_id,
             job=job,
+            client=cls.make_client(),
             target=cls.run_in_process,
             logger=logger,
             use_exec=supervisor._should_use_exec(),
@@ -558,11 +535,8 @@ def start(  # type: ignore[override]
         proc.send_msg(msg, request_id=0)
         return proc
 
-    @functools.cached_property
-    def client(self) -> Client:
-        return self.make_client()
-
-    def make_client(self) -> Client:
+    @classmethod
+    def make_client(cls) -> Client:
         """
         Build the API client used to talk to the API server.
 
@@ -580,9 +554,6 @@ def make_client(self) -> Client:
         return client
 
     def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger, req_id: int) -> None:
-
-        resp: BaseModel | None = None
-        dump_opts: dict[str, bool] = {}
         self._last_runner_comms = time.monotonic()
 
         if isinstance(msg, messages.TriggerStateChanges):
@@ -608,85 +579,37 @@ def _handle_request(self, msg: ToTriggerSupervisor, log: FilteringBoundLogger, r
             while self.persisted_event_seqs:
                 events_persisted.append(self.persisted_event_seqs.popleft())
 
-            response = messages.TriggerStateSync(
+            sync = messages.TriggerStateSync(
                 to_create=[],
-                to_cancel=self.cancelling_triggers,
+                to_cancel=self.cancelling_triggers.copy(),
                 events_persisted=events_persisted or None,
             )
 
             # Pull out of these dequeues in a thread-safe manner
             while self.creating_triggers:
                 workload = self.creating_triggers.popleft()
-                response.to_create.append(workload)
-            self.running_triggers.update(m.id for m in response.to_create)
-            resp = response
-
-        elif isinstance(msg, GetConnection):
-            resp, dump_opts = handle_get_connection(self.client, msg)
-        elif isinstance(msg, DeleteVariable):
-            resp, dump_opts = handle_delete_variable(self.client, msg)
-        elif isinstance(msg, GetVariable):
-            resp, dump_opts = handle_get_variable(self.client, msg)
-        elif isinstance(msg, GetVariableKeys):
-            resp, dump_opts = handle_get_variable_keys(self.client, msg)
-        elif isinstance(msg, PutVariable):
-            resp, dump_opts = handle_put_variable(self.client, msg)
-        elif isinstance(msg, DeleteXCom):
-            resp, dump_opts = handle_delete_xcom(self.client, msg)
-        elif isinstance(msg, GetXCom):
-            resp, dump_opts = handle_get_xcom(self.client, msg)
-        elif isinstance(msg, SetXCom):
-            resp, dump_opts = handle_set_xcom(self.client, msg)
-        elif isinstance(msg, GetDRCount):
-            resp, dump_opts = handle_get_dr_count(self.client, msg)
-        elif isinstance(msg, GetDagRunState):
-            resp, dump_opts = handle_get_dag_run_state(self.client, msg)
-
-        elif isinstance(msg, GetTICount):
-            resp, dump_opts = handle_get_ti_count(self.client, msg)
-
-        elif isinstance(msg, GetTaskStates):
-            resp, dump_opts = handle_get_task_states(self.client, msg)
-        elif isinstance(msg, GetPreviousTI):
-            resp, dump_opts = handle_get_previous_ti(self.client, msg)
-        elif isinstance(msg, UpdateHITLDetail):
+                sync.to_create.append(workload)
+            self.running_triggers.update(m.id for m in sync.to_create)
+            self.send_msg(sync, request_id=req_id, error=None)
+            return
+
+        if isinstance(msg, UpdateHITLDetail):
             api_resp = self.client.hitl.update_response(
                 ti_id=msg.ti_id,
                 chosen_options=msg.chosen_options,
                 params_input=msg.params_input,
             )
             resp = HITLDetailResponseResult.from_api_response(response=api_resp)
-        elif isinstance(msg, GetHITLDetailResponse):
+            self.send_msg(resp, request_id=req_id, error=None)
+            return
+
+        if isinstance(msg, GetHITLDetailResponse):
             api_resp = self.client.hitl.get_detail_response(ti_id=msg.ti_id)
             resp = HITLDetailResponseResult.from_api_response(response=api_resp)
-        elif isinstance(msg, MaskSecret):
-            handle_mask_secret(msg)
-        elif isinstance(msg, ClearAssetStateStoreByName):
-            handle_clear_asset_state_store_by_name(self.client, msg)
-            resp = OKResponse(ok=True)
-        elif isinstance(msg, ClearAssetStateStoreByUri):
-            handle_clear_asset_state_store_by_uri(self.client, msg)
-            resp = OKResponse(ok=True)
-        elif isinstance(msg, DeleteAssetStateStoreByName):
-            handle_delete_asset_state_store_by_name(self.client, msg)
-            resp = OKResponse(ok=True)
-        elif isinstance(msg, DeleteAssetStateStoreByUri):
-            handle_delete_asset_state_store_by_uri(self.client, msg)
-            resp = OKResponse(ok=True)
-        elif isinstance(msg, GetAssetStateStoreByName):
-            resp, dump_opts = handle_get_asset_state_store_by_name(self.client, msg)
-        elif isinstance(msg, GetAssetStateStoreByUri):
-            resp, dump_opts = handle_get_asset_state_store_by_uri(self.client, msg)
-        elif isinstance(msg, SetAssetStateStoreByName):
-            handle_set_asset_state_store_by_name(self.client, msg)
-            resp = OKResponse(ok=True)
-        elif isinstance(msg, SetAssetStateStoreByUri):
-            handle_set_asset_state_store_by_uri(self.client, msg)
-            resp = OKResponse(ok=True)
-        else:
-            raise ValueError(f"Unknown message type {type(msg)}")
+            self.send_msg(resp, request_id=req_id, error=None)
+            return
 
-        self.send_msg(resp, request_id=req_id, error=None, **dump_opts)
+        super()._handle_request(msg, log, req_id)
 
     def run(self) -> None:
         """Run synchronously and handle all database reads/writes."""

airflow-core/tests/unit/dag_processing/test_processor.py

@@ -2214,7 +2214,7 @@ def test_handle_request_get_connection_masks_password_and_extra(self, proc):
         )
 
         with (
-            patch("airflow.dag_processing.processor.mask_secret") as mock_mask_secret,
+            patch("airflow.sdk.execution_time.request_handlers.mask_secret") as mock_mask_secret,
             patch.object(DagFileProcessorProcess, "send_msg", autospec=True) as mock_send_msg,
         ):
             proc._handle_request(
@@ -2251,7 +2251,7 @@ def test_handle_request_get_variable_masks_value_with_key(self, proc):
         )
 
         with (
-            patch("airflow.dag_processing.processor.mask_secret") as mock_mask_secret,
+            patch("airflow.sdk.execution_time.request_handlers.mask_secret") as mock_mask_secret,
             patch.object(DagFileProcessorProcess, "send_msg", autospec=True) as mock_send_msg,
         ):
             proc._handle_request(