From 0669c8606aa3b45ffee5a7eb0ccddd264b368045 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 18:52:35 +0000 Subject: [PATCH 01/12] Carry task handler bindings and probed answers on the parse result The manager has to persist which artifact runs each stub task, so a worker no longer searches for it (ADR-0013), and the answers the parse probed, so no other parse probes those artifacts again. None tells a parse that did not evaluate task handlers apart from one that found none, so a failed or skipped evaluation cannot wipe the recorded bindings. The answers travel apart from the bindings because they are worth keeping when a Dag fails validation and has none, and a binding names only its artifact, whose fingerprint and answer belong to the artifact rather than to the stub task. DagFileParsingResult already existed at 2026-06-16, so each new field needs a VersionChange: a runtime pinned there is upgraded with the bindings unset, which means "leave them alone", and no probed artifacts, which records nothing. --- .../src/airflow/dag_processing/processor.py | 68 +++++++--- .../unit/dag_processing/test_processor.py | 68 ++++++++++ .../sdk/execution_time/schema/schema.json | 52 ++++++++ .../schema/versions/__init__.py | 4 + .../schema/versions/v2026_10_30.py | 20 +++ .../execution_time/schema/test_migrator.py | 123 ++++++++++++++++++ 6 files changed, 316 insertions(+), 19 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/processor.py b/airflow-core/src/airflow/dag_processing/processor.py index 408d108869437..6cda1e3b75bf3 100644 --- a/airflow-core/src/airflow/dag_processing/processor.py +++ b/airflow-core/src/airflow/dag_processing/processor.py @@ -97,25 +97,6 @@ from airflow.typing_compat import Self -class DagFileParsingResult(BaseModel): - """ - Result of DAG File Parsing. - - This is the result of a successful DAG parse, in this class, we gather all serialized DAGs, - import errors and warnings to send back to the scheduler to store in the DB. - """ - - fileloc: str - serialized_dags: list[LazyDeserializedDAG] - warnings: list | None = None - import_errors: dict[str, str] | None = None - parsed_definitions: list[str] = Field(default_factory=list) - """Bundle-relative locations of the Dag definitions imported from ``fileloc``.""" - dag_source_codes: dict[str, DagSourceCode] = Field(default_factory=dict) - """Source code of the parsed Dags, keyed by Dag fileloc.""" - type: Literal["DagFileParsingResult"] = "DagFileParsingResult" - - TaskHandlerBindingMode = Literal["positional", "named"] @@ -201,6 +182,55 @@ class DagFileParseRequest(BaseModel): type: Literal["DagFileParseRequest"] = "DagFileParseRequest" +class TaskHandlerBinding(BaseModel): + """A stub task resolved to the Lang-SDK artifact that runs it.""" + + dag_id: str + task_id: str + + artifact_bundle_name: str + """The bundle the artifact was found in.""" + + artifact_rel_path: str = Field(max_length=2000) + """Path of the artifact within its bundle.""" + + +class DagFileParsingResult(BaseModel): + """ + Result of DAG File Parsing. + + This is the result of a successful DAG parse, in this class, we gather all serialized DAGs, + import errors and warnings to send back to the scheduler to store in the DB. + """ + + fileloc: str + serialized_dags: list[LazyDeserializedDAG] + warnings: list | None = None + import_errors: dict[str, str] | None = None + parsed_definitions: list[str] = Field(default_factory=list) + """Bundle-relative locations of the Dag definitions imported from ``fileloc``.""" + dag_source_codes: dict[str, DagSourceCode] = Field(default_factory=dict) + """Source code of the parsed Dags, keyed by Dag fileloc.""" + + task_handler_bindings: list[TaskHandlerBinding] | None = None + """ + The stub-task bindings of every Dag in ``serialized_dags``. + + ``None`` when task handlers were not evaluated, and the recorded bindings are left as they are. A list + replaces the recorded bindings of each Dag in ``serialized_dags``, so a Dag with no entry has none. + """ + + probed_artifacts: list[TaskHandlerArtifact] = Field(default_factory=list) + """ + Every artifact this parse probed successfully, with its answer. + + Recorded even when ``task_handler_bindings`` is ``None``, so a Dag that fails validation is not probed + again on every parse. + """ + + type: Literal["DagFileParsingResult"] = "DagFileParsingResult" + + class TaskHandlerParseRequest(BaseModel): """ Request for Task Handler Parsing. diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 2de2b89aa313b..8dff0a31c3a6e 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -59,6 +59,7 @@ DagFileParsingResult, DagFileProcessorProcess, TaskHandlerArtifact, + TaskHandlerBinding, TaskHandlerDeclaration, TaskHandlerParam, TaskHandlerParseRequest, @@ -2489,6 +2490,73 @@ def test_rejects_invalid_binding(self, declaration): TaskHandlerDeclaration.model_validate(declaration) +def _make_binding_body(**overrides) -> dict: + return { + "dag_id": "etl", + "task_id": "extract", + "artifact_bundle_name": "java-task-handlers", + "artifact_rel_path": "etl.jar", + **overrides, + } + + +class TestTaskHandlerBinding: + def test_parsing_result_decodes_bindings(self): + result = TypeAdapter(ToManager).validate_python( + { + "type": "DagFileParsingResult", + "fileloc": "/files/dags/etl.py", + "serialized_dags": [], + "task_handler_bindings": [_make_binding_body()], + } + ) + + assert isinstance(result, DagFileParsingResult) + assert result.task_handler_bindings == [ + TaskHandlerBinding( + dag_id="etl", + task_id="extract", + artifact_bundle_name="java-task-handlers", + artifact_rel_path="etl.jar", + ) + ] + assert result.probed_artifacts == [] + + def test_parsing_result_decodes_probed_artifacts(self): + result = TypeAdapter(ToManager).validate_python( + { + "type": "DagFileParsingResult", + "fileloc": "/files/dags/etl.py", + "serialized_dags": [], + "probed_artifacts": [_make_known_artifact_body()], + } + ) + + assert isinstance(result, DagFileParsingResult) + assert result.task_handler_bindings is None + assert result.probed_artifacts == [ + TaskHandlerArtifact( + bundle_name="java-task-handlers", + relative_fileloc="etl.jar", + size_bytes=1024, + cache_digest="ab12", + task_handlers={ + "etl": [ + TaskHandlerDeclaration( + task_id="extract", + binding="positional", + params=[TaskHandlerParam(name=None, required=True)], + ) + ] + }, + ) + ] + + def test_rejects_a_path_wider_than_its_column(self): + with pytest.raises(ValidationError, match="artifact_rel_path"): + TaskHandlerBinding.model_validate(_make_binding_body(artifact_rel_path="a" * 2001)) + + class TestDagFileProcessorProcess: def test_registered_message_types(self): expected = set(typing.get_args(typing.get_args(ToManager)[0])) - {TaskHandlerParsingResult} diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json index e7a512de7b574..a3d1a8560a986 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/schema.json +++ b/task-sdk/src/airflow/sdk/execution_time/schema/schema.json @@ -1036,6 +1036,28 @@ "title": "Dag Source Codes", "type": "object" }, + "task_handler_bindings": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/TaskHandlerBinding" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Task Handler Bindings" + }, + "probed_artifacts": { + "items": { + "$ref": "#/$defs/TaskHandlerArtifact" + }, + "title": "Probed Artifacts", + "type": "array" + }, "type": { "const": "DagFileParsingResult", "default": "DagFileParsingResult", @@ -4212,6 +4234,36 @@ "title": "TaskHandlerArtifact", "type": "object" }, + "TaskHandlerBinding": { + "description": "A stub task resolved to the Lang-SDK artifact that runs it.", + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "artifact_bundle_name": { + "title": "Artifact Bundle Name", + "type": "string" + }, + "artifact_rel_path": { + "maxLength": 2000, + "title": "Artifact Rel Path", + "type": "string" + } + }, + "required": [ + "dag_id", + "task_id", + "artifact_bundle_name", + "artifact_rel_path" + ], + "title": "TaskHandlerBinding", + "type": "object" + }, "TaskHandlerDeclaration": { "description": "A task handler that a Lang-SDK artifact registers for one task.", "properties": { diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py index aec4e0874b728..834f5cc984e62 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/__init__.py @@ -41,7 +41,9 @@ def get_bundle() -> VersionBundle: AddArgBindingsToSupervisorTIRunContext, AddDagDefinitionsToDagFileParsingResult, AddKnownArtifactsToDagFileParseRequest, + AddProbedArtifactsToDagFileParsingResult, AddRetryReasonToTaskState, + AddTaskHandlerBindingsToDagFileParsingResult, ) return VersionBundle( @@ -52,6 +54,8 @@ def get_bundle() -> VersionBundle: AddRetryReasonToTaskState, AddDagDefinitionsToDagFileParsingResult, AddKnownArtifactsToDagFileParseRequest, + AddTaskHandlerBindingsToDagFileParsingResult, + AddProbedArtifactsToDagFileParsingResult, ), Version("2026-06-16"), ) diff --git a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py index 3b810cbaaa254..247f97d56c1cb 100644 --- a/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py +++ b/task-sdk/src/airflow/sdk/execution_time/schema/versions/v2026_10_30.py @@ -74,3 +74,23 @@ class AddKnownArtifactsToDagFileParseRequest(VersionChange): instructions_to_migrate_to_previous_version = ( schema(DagFileParseRequest).field("known_artifacts").didnt_exist, ) + + +class AddTaskHandlerBindingsToDagFileParsingResult(VersionChange): + """Add ``task_handler_bindings``, the resolved stub-task bindings, to ``DagFileParsingResult``.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(DagFileParsingResult).field("task_handler_bindings").didnt_exist, + ) + + +class AddProbedArtifactsToDagFileParsingResult(VersionChange): + """Add ``probed_artifacts``, the task-handler artifacts a parse probed and their answers, to ``DagFileParsingResult``.""" + + description = __doc__ + + instructions_to_migrate_to_previous_version = ( + schema(DagFileParsingResult).field("probed_artifacts").didnt_exist, + ) diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index d837181c6f6d4..4163c3a93201f 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -51,7 +51,9 @@ from airflow.dag_processing.processor import ( DagFileParseRequest, + DagFileParsingResult, TaskHandlerArtifact, + TaskHandlerBinding, TaskHandlerDeclaration, TaskHandlerParam, TaskHandlerParseRequest, @@ -578,3 +580,124 @@ def test_downgrade_keeps_known_artifacts_at_head(self, real_migrator, request_wi }, } ] + + +class TestRealBundleTaskHandlerBindings: + """ + Drive the *real* supervisor bundle through the ``task_handler_bindings`` migration. + + ``DagFileParsingResult`` flows runtime -> supervisor, so ``upgrade`` is the direction a pinned + runtime travels. Only ``downgrade`` re-validates against the versioned class, so that is the + direction that fails if ``AddTaskHandlerBindingsToDagFileParsingResult`` is dropped. + """ + + @pytest.fixture + def result_with_bindings(self) -> DagFileParsingResult: + return DagFileParsingResult( + fileloc="/files/dags/etl.py", + serialized_dags=[], + task_handler_bindings=[ + TaskHandlerBinding( + dag_id="etl", + task_id="extract", + artifact_bundle_name="java-task-handlers", + artifact_rel_path="etl.jar", + ) + ], + ) + + @pytest.fixture + def real_migrator(self) -> SchemaVersionMigrator: + return get_schema_version_migrator() + + def test_downgrade_strips_task_handler_bindings_for_previous_version( + self, real_migrator, result_with_bindings + ): + out = real_migrator.downgrade(result_with_bindings, "2026-06-16").model_dump() + assert "task_handler_bindings" not in out + + def test_downgrade_keeps_task_handler_bindings_at_head(self, real_migrator, result_with_bindings): + out = real_migrator.downgrade(result_with_bindings, "2026-10-30").model_dump() + assert out["task_handler_bindings"] == [ + { + "dag_id": "etl", + "task_id": "extract", + "artifact_bundle_name": "java-task-handlers", + "artifact_rel_path": "etl.jar", + } + ] + + def test_upgrade_leaves_missing_task_handler_bindings_unset(self, real_migrator): + body = {"type": "DagFileParsingResult", "fileloc": "/files/dags/etl.py", "serialized_dags": []} + out = real_migrator.upgrade(body, DagFileParsingResult, "2026-06-16") + assert out["task_handler_bindings"] is None + + +class TestRealBundleProbedArtifacts: + """ + Drive the *real* supervisor bundle through the ``probed_artifacts`` migration. + + ``DagFileParsingResult`` flows runtime -> supervisor, so ``upgrade`` is the direction a pinned + runtime travels, and a result without the field records nothing. + """ + + @pytest.fixture + def result_with_probed_artifacts(self) -> DagFileParsingResult: + return DagFileParsingResult( + fileloc="/files/dags/etl.py", + serialized_dags=[], + probed_artifacts=[ + TaskHandlerArtifact( + bundle_name="java-task-handlers", + relative_fileloc="etl.jar", + size_bytes=1024, + cache_digest=None, + task_handlers={ + "etl": [ + TaskHandlerDeclaration( + task_id="extract", + binding="positional", + params=[TaskHandlerParam(name=None, required=True)], + ) + ] + }, + ) + ], + ) + + @pytest.fixture + def real_migrator(self) -> SchemaVersionMigrator: + return get_schema_version_migrator() + + def test_downgrade_strips_probed_artifacts_for_previous_version( + self, real_migrator, result_with_probed_artifacts + ): + out = real_migrator.downgrade(result_with_probed_artifacts, "2026-06-16").model_dump() + assert "probed_artifacts" not in out + + def test_downgrade_keeps_probed_artifacts_at_head(self, real_migrator, result_with_probed_artifacts): + out = real_migrator.downgrade(result_with_probed_artifacts, "2026-10-30").model_dump() + assert out["probed_artifacts"] == [ + { + "bundle_name": "java-task-handlers", + "relative_fileloc": "etl.jar", + "size_bytes": 1024, + "cache_digest": None, + "task_handlers": { + "etl": [ + { + "task_id": "extract", + "binding": "positional", + "params": [ + {"name": None, "value_schema": None, "required": True, "exact_name": False} + ], + } + ] + }, + } + ] + + def test_upgrade_leaves_missing_probed_artifacts_empty(self, real_migrator): + body = {"type": "DagFileParsingResult", "fileloc": "/files/dags/etl.py", "serialized_dags": []} + out = real_migrator.upgrade(body, DagFileParsingResult, "2026-06-16") + assert out["probed_artifacts"] == [] From a7df3be0b732ed76dd96be4a833ebfb08baf3be5 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 18:53:16 +0000 Subject: [PATCH 02/12] Regenerate TS SDK supervisor types for task handler bindings check-ts-sdk-supervisor-schema fails whenever supervisor.ts lags the snapshot, and CI runs it over all files on any ts-sdk change. The shifted numeric suffixes are the generator's own renumbering; no hand-written TS code refers to them. --- ts-sdk/src/generated/supervisor.ts | 128 +++++++++++++++++------------ 1 file changed, 74 insertions(+), 54 deletions(-) diff --git a/ts-sdk/src/generated/supervisor.ts b/ts-sdk/src/generated/supervisor.ts index af0c907a5ed05..5e480c9216ffd 100644 --- a/ts-sdk/src/generated/supervisor.ts +++ b/ts-sdk/src/generated/supervisor.ts @@ -290,8 +290,14 @@ export type ImportErrors = { export type ParsedDefinitions = string[]; export type SourceCode = string; export type Language = string; -export type Type16 = "DagFileParsingResult"; +export type TaskHandlerBindings = TaskHandlerBinding[] | null; export type DagId4 = string; +export type TaskId3 = string; +export type ArtifactBundleName = string; +export type ArtifactRelPath = string; +export type ProbedArtifacts = TaskHandlerArtifact[]; +export type Type16 = "DagFileParsingResult"; +export type DagId5 = string; export type IsPaused = boolean; export type BundleName5 = string | null; export type BundleVersion3 = string | null; @@ -300,7 +306,7 @@ export type Owners = string | null; export type Tags = string[]; export type NextDagrun = string | null; export type Type17 = "DagResult"; -export type DagId5 = string; +export type DagId6 = string; export type RunId4 = string; export type LogicalDate2 = string | null; export type DataIntervalStart2 = string | null; @@ -348,9 +354,9 @@ export type Type23 = "DeleteTaskStateStore"; export type Key4 = string; export type Type24 = "DeleteVariable"; export type Key5 = string; -export type DagId6 = string; +export type DagId7 = string; export type RunId5 = string; -export type TaskId3 = string; +export type TaskId4 = string; export type MapIndex1 = number | null; export type Type25 = "DeleteXCom"; /** @@ -425,33 +431,33 @@ export type AliasName1 = string; export type Type33 = "GetAssetsByAlias"; export type ConnId2 = string; export type Type34 = "GetConnection"; -export type DagId7 = string; +export type DagId8 = string; export type LogicalDates = string[] | null; export type RunIds = string[] | null; export type States = string[] | null; export type Type35 = "GetDRCount"; -export type DagId8 = string; -export type Type36 = "GetDag"; export type DagId9 = string; +export type Type36 = "GetDag"; +export type DagId10 = string; export type RunId6 = string; export type Type37 = "GetDagRun"; -export type DagId10 = string; +export type DagId11 = string; export type RunId7 = string; export type Type38 = "GetDagRunState"; export type TiId3 = string; export type Type39 = "GetHITLDetailResponse"; export type TiId4 = string; export type Type40 = "GetPrevSuccessfulDagRun"; -export type DagId11 = string; +export type DagId12 = string; export type LogicalDate3 = string; export type State3 = string | null; export type Type41 = "GetPreviousDagRun"; -export type DagId12 = string; -export type TaskId4 = string; +export type DagId13 = string; +export type TaskId5 = string; export type LogicalDate4 = string | null; export type MapIndex2 = number; export type Type42 = "GetPreviousTI"; -export type DagId13 = string; +export type DagId14 = string; export type MapIndex3 = number | null; export type TaskIds = string[] | null; export type TaskGroupId = string | null; @@ -459,7 +465,7 @@ export type LogicalDates1 = string[] | null; export type RunIds1 = string[] | null; export type States1 = string[] | null; export type Type43 = "GetTICount"; -export type DagId14 = string; +export type DagId15 = string; export type RunId8 = string; export type Type44 = "GetTaskBreadcrumbs"; export type TiId5 = string; @@ -468,7 +474,7 @@ export type Type45 = "GetTaskRescheduleStartDate"; export type TiId6 = string; export type Key8 = string; export type Type46 = "GetTaskStateStore"; -export type DagId15 = string; +export type DagId16 = string; export type MapIndex4 = number | null; export type TaskIds1 = string[] | null; export type TaskGroupId1 = string | null; @@ -482,27 +488,27 @@ export type Limit2 = number; export type Offset = number; export type Type49 = "GetVariableKeys"; export type Key10 = string; -export type DagId16 = string; +export type DagId17 = string; export type RunId9 = string; -export type TaskId5 = string; +export type TaskId6 = string; export type MapIndex5 = number | null; export type IncludePriorDates = boolean; export type Type50 = "GetXCom"; export type Key11 = string; -export type DagId17 = string; +export type DagId18 = string; export type RunId10 = string; -export type TaskId6 = string; +export type TaskId7 = string; export type Type51 = "GetXComCount"; export type Key12 = string; -export type DagId18 = string; +export type DagId19 = string; export type RunId11 = string; -export type TaskId7 = string; +export type TaskId8 = string; export type Offset1 = number; export type Type52 = "GetXComSequenceItem"; export type Key13 = string; -export type DagId19 = string; +export type DagId20 = string; export type RunId12 = string; -export type TaskId8 = string; +export type TaskId9 = string; export type Start = number | null; export type Stop = number | null; export type Step = number | null; @@ -534,8 +540,8 @@ export type StartDate4 = string | null; export type EndDate3 = string | null; export type Type58 = "PrevSuccessfulDagRunResult"; export type Type59 = "PreviousDagRunResult"; -export type TaskId9 = string; -export type DagId20 = string; +export type TaskId10 = string; +export type DagId21 = string; export type RunId13 = string; export type LogicalDate5 = string | null; export type StartDate5 = string | null; @@ -576,9 +582,9 @@ export type Key17 = string; export type ExpiresAt = string | null; export type Type70 = "SetTaskStateStore"; export type Key18 = string; -export type DagId21 = string; +export type DagId22 = string; export type RunId14 = string; -export type TaskId10 = string; +export type TaskId11 = string; export type MapIndex7 = number | null; export type DagResult1 = boolean; export type MappedLength = number | null; @@ -632,7 +638,7 @@ export type Conf2 = { export type ResetDagRun = boolean | null; export type PartitionKey7 = string | null; export type Note2 = string | null; -export type DagId22 = string; +export type DagId23 = string; export type DagRunId = string; export type Type83 = "TriggerDagRun"; export type TiId9 = string; @@ -1176,6 +1182,8 @@ export interface DagFileParsingResult { import_errors?: ImportErrors; parsed_definitions?: ParsedDefinitions; dag_source_codes?: DagSourceCodes; + task_handler_bindings?: TaskHandlerBindings; + probed_artifacts?: ProbedArtifacts; type?: Type16; } /** @@ -1207,12 +1215,24 @@ export interface DagSourceCode { source_code: SourceCode; language: Language; } +/** + * A stub task resolved to the Lang-SDK artifact that runs it. + * + * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema + * via the `definition` "TaskHandlerBinding". + */ +export interface TaskHandlerBinding { + dag_id: DagId4; + task_id: TaskId3; + artifact_bundle_name: ArtifactBundleName; + artifact_rel_path: ArtifactRelPath; +} /** * This interface was referenced by `SupervisorWireSchema`'s JSON-Schema * via the `definition` "DagResult". */ export interface DagResult { - dag_id: DagId4; + dag_id: DagId5; is_paused: IsPaused; bundle_name: BundleName5; bundle_version: BundleVersion3; @@ -1227,7 +1247,7 @@ export interface DagResult { * via the `definition` "DagRunResult". */ export interface DagRunResult { - dag_id: DagId5; + dag_id: DagId6; run_id: RunId4; logical_date: LogicalDate2; data_interval_start: DataIntervalStart2; @@ -1313,9 +1333,9 @@ export interface DeleteVariable { */ export interface DeleteXCom { key: Key5; - dag_id: DagId6; + dag_id: DagId7; run_id: RunId5; - task_id: TaskId3; + task_id: TaskId4; map_index?: MapIndex1; type?: Type25; } @@ -1414,7 +1434,7 @@ export interface GetConnection { * via the `definition` "GetDRCount". */ export interface GetDRCount { - dag_id: DagId7; + dag_id: DagId8; logical_dates?: LogicalDates; run_ids?: RunIds; states?: States; @@ -1425,7 +1445,7 @@ export interface GetDRCount { * via the `definition` "GetDag". */ export interface GetDag { - dag_id: DagId8; + dag_id: DagId9; type?: Type36; } /** @@ -1433,7 +1453,7 @@ export interface GetDag { * via the `definition` "GetDagRun". */ export interface GetDagRun { - dag_id: DagId9; + dag_id: DagId10; run_id: RunId6; type?: Type37; } @@ -1442,7 +1462,7 @@ export interface GetDagRun { * via the `definition` "GetDagRunState". */ export interface GetDagRunState { - dag_id: DagId10; + dag_id: DagId11; run_id: RunId7; type?: Type38; } @@ -1469,7 +1489,7 @@ export interface GetPrevSuccessfulDagRun { * via the `definition` "GetPreviousDagRun". */ export interface GetPreviousDagRun { - dag_id: DagId11; + dag_id: DagId12; logical_date: LogicalDate3; state?: State3; type?: Type41; @@ -1481,8 +1501,8 @@ export interface GetPreviousDagRun { * via the `definition` "GetPreviousTI". */ export interface GetPreviousTI { - dag_id: DagId12; - task_id: TaskId4; + dag_id: DagId13; + task_id: TaskId5; logical_date?: LogicalDate4; map_index?: MapIndex2; state?: TaskInstanceState | null; @@ -1493,7 +1513,7 @@ export interface GetPreviousTI { * via the `definition` "GetTICount". */ export interface GetTICount { - dag_id: DagId13; + dag_id: DagId14; map_index?: MapIndex3; task_ids?: TaskIds; task_group_id?: TaskGroupId; @@ -1507,7 +1527,7 @@ export interface GetTICount { * via the `definition` "GetTaskBreadcrumbs". */ export interface GetTaskBreadcrumbs { - dag_id: DagId14; + dag_id: DagId15; run_id: RunId8; type?: Type44; } @@ -1534,7 +1554,7 @@ export interface GetTaskStateStore { * via the `definition` "GetTaskStates". */ export interface GetTaskStates { - dag_id: DagId15; + dag_id: DagId16; map_index?: MapIndex4; task_ids?: TaskIds1; task_group_id?: TaskGroupId1; @@ -1566,9 +1586,9 @@ export interface GetVariableKeys { */ export interface GetXCom { key: Key10; - dag_id: DagId16; + dag_id: DagId17; run_id: RunId9; - task_id: TaskId5; + task_id: TaskId6; map_index?: MapIndex5; include_prior_dates?: IncludePriorDates; type?: Type50; @@ -1581,9 +1601,9 @@ export interface GetXCom { */ export interface GetXComCount { key: Key11; - dag_id: DagId17; + dag_id: DagId18; run_id: RunId10; - task_id: TaskId6; + task_id: TaskId7; type?: Type51; } /** @@ -1592,9 +1612,9 @@ export interface GetXComCount { */ export interface GetXComSequenceItem { key: Key12; - dag_id: DagId18; + dag_id: DagId19; run_id: RunId11; - task_id: TaskId7; + task_id: TaskId8; offset: Offset1; type?: Type52; } @@ -1604,9 +1624,9 @@ export interface GetXComSequenceItem { */ export interface GetXComSequenceSlice { key: Key13; - dag_id: DagId19; + dag_id: DagId20; run_id: RunId12; - task_id: TaskId8; + task_id: TaskId9; start: Start; stop: Stop; step: Step; @@ -1687,8 +1707,8 @@ export interface PreviousDagRunResult { * via the `definition` "PreviousTIResponse". */ export interface PreviousTIResponse { - task_id: TaskId9; - dag_id: DagId20; + task_id: TaskId10; + dag_id: DagId21; run_id: RunId13; logical_date?: LogicalDate5; start_date?: StartDate5; @@ -1820,9 +1840,9 @@ export interface SetTaskStateStore { export interface SetXCom { key: Key18; value: JsonValue; - dag_id: DagId21; + dag_id: DagId22; run_id: RunId14; - task_id: TaskId10; + task_id: TaskId11; map_index?: MapIndex7; dag_result?: DagResult1; mapped_length?: MappedLength; @@ -1976,7 +1996,7 @@ export interface TriggerDagRun { reset_dag_run?: ResetDagRun; partition_key?: PartitionKey7; note?: Note2; - dag_id: DagId22; + dag_id: DagId23; run_id: DagRunId; type?: Type83; } From fb12a7612524454c17fee51ec1034cc55c27f121 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 18:53:32 +0000 Subject: [PATCH 03/12] Sync the Java SDK supervisor schema for task handler bindings The vendored copy tracks the monorepo snapshot while both declare the same api_version, and sync-java-sdk-supervisor-schema rewrites it on the next Java SDK change otherwise. --- java-sdk/sdk/schema/schema.json | 52 +++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/java-sdk/sdk/schema/schema.json b/java-sdk/sdk/schema/schema.json index e7a512de7b574..a3d1a8560a986 100644 --- a/java-sdk/sdk/schema/schema.json +++ b/java-sdk/sdk/schema/schema.json @@ -1036,6 +1036,28 @@ "title": "Dag Source Codes", "type": "object" }, + "task_handler_bindings": { + "anyOf": [ + { + "items": { + "$ref": "#/$defs/TaskHandlerBinding" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "default": null, + "title": "Task Handler Bindings" + }, + "probed_artifacts": { + "items": { + "$ref": "#/$defs/TaskHandlerArtifact" + }, + "title": "Probed Artifacts", + "type": "array" + }, "type": { "const": "DagFileParsingResult", "default": "DagFileParsingResult", @@ -4212,6 +4234,36 @@ "title": "TaskHandlerArtifact", "type": "object" }, + "TaskHandlerBinding": { + "description": "A stub task resolved to the Lang-SDK artifact that runs it.", + "properties": { + "dag_id": { + "title": "Dag Id", + "type": "string" + }, + "task_id": { + "title": "Task Id", + "type": "string" + }, + "artifact_bundle_name": { + "title": "Artifact Bundle Name", + "type": "string" + }, + "artifact_rel_path": { + "maxLength": 2000, + "title": "Artifact Rel Path", + "type": "string" + } + }, + "required": [ + "dag_id", + "task_id", + "artifact_bundle_name", + "artifact_rel_path" + ], + "title": "TaskHandlerBinding", + "type": "object" + }, "TaskHandlerDeclaration": { "description": "A task handler that a Lang-SDK artifact registers for one task.", "properties": { From c6c70072a4c4428711dc4fca196049f5fdf323d5 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Fri, 2 Oct 2026 02:22:31 +0000 Subject: [PATCH 04/12] Regenerate Go SDK models for task handler bindings check-go-sdk-generated-drift regenerates the Go SDK's coordinator-protocol models from the supervisor schema snapshot and fails on any difference, so DagFileParsingResult must carry task_handler_bindings and probed_artifacts on the Go side too. go-jsonschema moves some type blocks, so the diff is larger than the new fields. --- go-sdk/pkg/execution/genmodels/models.gen.go | 170 +++++++++++-------- 1 file changed, 97 insertions(+), 73 deletions(-) diff --git a/go-sdk/pkg/execution/genmodels/models.gen.go b/go-sdk/pkg/execution/genmodels/models.gen.go index a6b4542757690..1a92971f649f0 100644 --- a/go-sdk/pkg/execution/genmodels/models.gen.go +++ b/go-sdk/pkg/execution/genmodels/models.gen.go @@ -422,9 +422,16 @@ type DagFileParsingResult struct { // ParsedDefinitions corresponds to the JSON schema field "parsed_definitions". ParsedDefinitions []string `msgpack:"parsed_definitions,omitempty"` + // ProbedArtifacts corresponds to the JSON schema field "probed_artifacts". + ProbedArtifacts []TaskHandlerArtifact `msgpack:"probed_artifacts,omitempty"` + // SerializedDags corresponds to the JSON schema field "serialized_dags". SerializedDags []LazyDeserializedDAG `msgpack:"serialized_dags"` + // TaskHandlerBindings corresponds to the JSON schema field + // "task_handler_bindings". + TaskHandlerBindings *TaskHandlerBindings `msgpack:"task_handler_bindings,omitempty"` + // Type corresponds to the JSON schema field "type". Type string `msgpack:"type,omitempty"` @@ -1717,10 +1724,6 @@ type TaskCallbackRequest struct { VersionData *VersionData `msgpack:"version_data,omitempty"` } -type Warnings []interface{} - -type TriggerKwargs map[string]interface{} - // A Lang-SDK artifact and every task handler it registers. type TaskHandlerArtifact struct { // BundleName corresponds to the JSON schema field "bundle_name". @@ -1739,7 +1742,22 @@ type TaskHandlerArtifact struct { TaskHandlers TaskHandlers `msgpack:"task_handlers"` } -type TaskHandlers map[string][]TaskHandlerDeclaration +// A stub task resolved to the Lang-SDK artifact that runs it. +type TaskHandlerBinding struct { + // ArtifactBundleName corresponds to the JSON schema field "artifact_bundle_name". + ArtifactBundleName string `msgpack:"artifact_bundle_name"` + + // ArtifactRelPath corresponds to the JSON schema field "artifact_rel_path". + ArtifactRelPath string `msgpack:"artifact_rel_path"` + + // DagID corresponds to the JSON schema field "dag_id". + DagID string `msgpack:"dag_id"` + + // TaskID corresponds to the JSON schema field "task_id". + TaskID string `msgpack:"task_id"` +} + +type TaskHandlerBindings []TaskHandlerBinding // A task handler that a Lang-SDK artifact registers for one task. type TaskHandlerDeclaration struct { @@ -1753,7 +1771,10 @@ type TaskHandlerDeclaration struct { TaskID string `msgpack:"task_id"` } -type TaskHandlerParams []TaskHandlerParam +type TaskHandlerDeclarationBinding string + +const TaskHandlerDeclarationBindingNamed TaskHandlerDeclarationBinding = "named" +const TaskHandlerDeclarationBindingPositional TaskHandlerDeclarationBinding = "positional" // One parameter of a task handler. type TaskHandlerParam struct { @@ -1767,39 +1788,50 @@ type TaskHandlerParam struct { ValueSchema *ArgValueSchema `msgpack:"value_schema,omitempty"` } -type TaskIds []string +type TaskHandlerParams []TaskHandlerParam -const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" -const TaskInstanceStateDeferred TaskInstanceState = "deferred" -const TaskInstanceStateSkipped TaskInstanceState = "skipped" -const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" -const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" -const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" -const TaskInstanceStateFailed TaskInstanceState = "failed" -const TaskInstanceStateRestarting TaskInstanceState = "restarting" -const TaskInstanceStateSuccess TaskInstanceState = "success" -const TaskInstanceStateRunning TaskInstanceState = "running" -const TaskInstanceStateQueued TaskInstanceState = "queued" -const TaskInstanceStateScheduled TaskInstanceState = "scheduled" -const TaskInstanceStateRemoved TaskInstanceState = "removed" +// Request for Task Handler Parsing. +// +// Asks a Lang-SDK runtime for every task handler an artifact registers. +type TaskHandlerParseRequest struct { + // BundleName corresponds to the JSON schema field "bundle_name". + BundleName string `msgpack:"bundle_name"` -type TaskInstanceState string + // BundlePath corresponds to the JSON schema field "bundle_path". + BundlePath string `msgpack:"bundle_path"` -// Variable schema for responses with fields that are needed for Runtime. -type VariableResponse struct { - // Key corresponds to the JSON schema field "key". - Key string `msgpack:"key"` + // File corresponds to the JSON schema field "file". + File string `msgpack:"file"` - // Value corresponds to the JSON schema field "value". - Value interface{} `msgpack:"value"` + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` } -const TaskHandlerDeclarationBindingNamed TaskHandlerDeclarationBinding = "named" -const TaskHandlerDeclarationBindingPositional TaskHandlerDeclarationBinding = "positional" +// Result of Task Handler Parsing. +// +// Every task handler a Lang-SDK artifact registers, keyed by Dag id. +// +// The answer depends only on the artifact, never on the request. +type TaskHandlerParsingResult struct { + // Fileloc corresponds to the JSON schema field "fileloc". + Fileloc string `msgpack:"fileloc"` -type TaskHandlerDeclarationBinding string + // ImportErrors corresponds to the JSON schema field "import_errors". + ImportErrors *ImportErrors `msgpack:"import_errors,omitempty"` -type TaskOutlets []AssetProfile + // TaskHandlers corresponds to the JSON schema field "task_handlers". + TaskHandlers TaskHandlers `msgpack:"task_handlers"` + + // Type corresponds to the JSON schema field "type". + Type string `msgpack:"type,omitempty"` + + // Warnings corresponds to the JSON schema field "warnings". + Warnings *Warnings `msgpack:"warnings,omitempty"` +} + +type TaskHandlers map[string][]TaskHandlerDeclaration + +type TaskIds []string // Schema for TaskInstance model with minimal required fields needed for Runtime. type TaskInstance struct { @@ -1834,44 +1866,23 @@ type TaskInstance struct { TryNumber int `msgpack:"try_number"` } -// Request for Task Handler Parsing. -// -// Asks a Lang-SDK runtime for every task handler an artifact registers. -type TaskHandlerParseRequest struct { - // BundleName corresponds to the JSON schema field "bundle_name". - BundleName string `msgpack:"bundle_name"` - - // BundlePath corresponds to the JSON schema field "bundle_path". - BundlePath string `msgpack:"bundle_path"` - - // File corresponds to the JSON schema field "file". - File string `msgpack:"file"` - - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` -} - -// Result of Task Handler Parsing. -// -// Every task handler a Lang-SDK artifact registers, keyed by Dag id. -// -// The answer depends only on the artifact, never on the request. -type TaskHandlerParsingResult struct { - // Fileloc corresponds to the JSON schema field "fileloc". - Fileloc string `msgpack:"fileloc"` - - // ImportErrors corresponds to the JSON schema field "import_errors". - ImportErrors *ImportErrors `msgpack:"import_errors,omitempty"` - - // TaskHandlers corresponds to the JSON schema field "task_handlers". - TaskHandlers TaskHandlers `msgpack:"task_handlers"` +type TaskInstanceState string - // Type corresponds to the JSON schema field "type". - Type string `msgpack:"type,omitempty"` +const TaskInstanceStateAwaitingInput TaskInstanceState = "awaiting_input" +const TaskInstanceStateDeferred TaskInstanceState = "deferred" +const TaskInstanceStateFailed TaskInstanceState = "failed" +const TaskInstanceStateQueued TaskInstanceState = "queued" +const TaskInstanceStateRemoved TaskInstanceState = "removed" +const TaskInstanceStateRestarting TaskInstanceState = "restarting" +const TaskInstanceStateRunning TaskInstanceState = "running" +const TaskInstanceStateScheduled TaskInstanceState = "scheduled" +const TaskInstanceStateSkipped TaskInstanceState = "skipped" +const TaskInstanceStateSuccess TaskInstanceState = "success" +const TaskInstanceStateUpForReschedule TaskInstanceState = "up_for_reschedule" +const TaskInstanceStateUpForRetry TaskInstanceState = "up_for_retry" +const TaskInstanceStateUpstreamFailed TaskInstanceState = "upstream_failed" - // Warnings corresponds to the JSON schema field "warnings". - Warnings *Warnings `msgpack:"warnings,omitempty"` -} +type TaskOutlets []AssetProfile // Response containing the first reschedule date for a task instance. type TaskRescheduleStartDate struct { @@ -1882,12 +1893,6 @@ type TaskRescheduleStartDate struct { Type string `msgpack:"type,omitempty"` } -type TaskStateState string - -const TaskStateStateFailed TaskStateState = "failed" -const TaskStateStateSkipped TaskStateState = "skipped" -const TaskStateStateRemoved TaskStateState = "removed" - // Update a task's state. // // If a process exits without sending one of these the state will be derived from @@ -1911,6 +1916,12 @@ type TaskState struct { Type string `msgpack:"type,omitempty"` } +type TaskStateState string + +const TaskStateStateFailed TaskStateState = "failed" +const TaskStateStateRemoved TaskStateState = "removed" +const TaskStateStateSkipped TaskStateState = "skipped" + // Response to GetTaskStateStore; wraps the generated API response for supervisor // to worker comms. type TaskStateStoreResult struct { @@ -1960,6 +1971,19 @@ type TriggerDagRun struct { Type string `msgpack:"type,omitempty"` } +type TriggerKwargs map[string]interface{} + +type Warnings []interface{} + +// Variable schema for responses with fields that are needed for Runtime. +type VariableResponse struct { + // Key corresponds to the JSON schema field "key". + Key string `msgpack:"key"` + + // Value corresponds to the JSON schema field "value". + Value interface{} `msgpack:"value"` +} + type UpdateDagRunNote struct { // Note corresponds to the JSON schema field "note". Note interface{} `msgpack:"note"` From 1665a67b89b569b7b12a708d575da53ab9addf0a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 19:00:27 +0000 Subject: [PATCH 05/12] Record probed artifacts and reconcile task handler bindings A handler row belongs to a Dag, not to the file that last wrote it, so the rows are reconciled by the result's Dag ids: a Dag that moved files keeps its rows, and a parse that yields fewer Dags cannot delete rows it did not evaluate. The Dag rows are already locked for the persist, so handler rows need no upsert. Artifact rows are shared across files and Dag processors and hold answers that other files trust, so only a probe writes them, only within the Dag file's scope, and in a transaction of its own before the persist: its exclusive locks are released before the reconcile takes shared ones, so two Dag processors cannot deadlock on crossed artifacts, and an answer survives a failed persist. The reconcile only looks artifacts up, under a shared lock that keeps the orphan sweep from deleting one before its handler rows are in; a Dag whose artifact has no row or lies outside the scope keeps its rows, and the next parse probes again. --- .../src/airflow/dag_processing/collection.py | 257 +++++++++- .../src/airflow/dag_processing/manager.py | 43 +- .../unit/dag_processing/test_collection.py | 473 +++++++++++++++++- .../tests/unit/dag_processing/test_manager.py | 134 ++++- 4 files changed, 901 insertions(+), 6 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/collection.py b/airflow-core/src/airflow/dag_processing/collection.py index ec8889db31845..896dec91ed489 100644 --- a/airflow-core/src/airflow/dag_processing/collection.py +++ b/airflow-core/src/airflow/dag_processing/collection.py @@ -28,11 +28,12 @@ from __future__ import annotations import traceback +from collections import defaultdict from typing import TYPE_CHECKING, Any, NamedTuple, TypeVar import structlog from sqlalchemy import delete, false, func, insert, or_, select, tuple_, update -from sqlalchemy.exc import OperationalError +from sqlalchemy.exc import DataError, IntegrityError, OperationalError from sqlalchemy.orm import joinedload, load_only from airflow._shared.timezones.timezone import utcnow @@ -54,6 +55,11 @@ from airflow.models.dagrun import DagRun from airflow.models.dagwarning import DagWarning, DagWarningType from airflow.models.errors import ParseImportError +from airflow.models.lang_sdk_task_handler import ( + LangSDKTaskHandler, + LangSDKTaskHandlerArtifact, + compute_fileloc_hash, +) from airflow.models.serialized_dag import SerializedDagModel from airflow.models.trigger import Trigger from airflow.plugins_manager import get_scheduling_class_teams @@ -68,15 +74,17 @@ from airflow.serialization.serialized_objects import BaseSerialization, LazyDeserializedDAG from airflow.triggers.base import BaseEventTrigger from airflow.utils.retries import MAX_DB_RETRIES, run_with_db_retries -from airflow.utils.sqlalchemy import get_dialect_name, with_row_locks +from airflow.utils.sqlalchemy import build_upsert_stmt, get_dialect_name, with_row_locks from airflow.utils.types import DagRunType if TYPE_CHECKING: from collections.abc import Collection, Iterable, Iterator + from uuid import UUID from sqlalchemy.orm import Session from sqlalchemy.sql import Select + from airflow.dag_processing.processor import TaskHandlerArtifact, TaskHandlerBinding from airflow.models.serialized_dag import DagWriteMetadata from airflow.sdk.importers import DagSourceCode # noqa: SDK001 from airflow.typing_compat import Self, Unpack @@ -581,6 +589,209 @@ def _reject_other_teams_plugin_classes( return accepted +_ArtifactKey = tuple[str, str] +"""``(bundle_name, relative_fileloc_hash)``, the unique key of ``lang_sdk_task_handler_artifact``.""" + + +def record_probed_task_handler_artifacts( + artifacts: Collection[TaskHandlerArtifact], *, bundle_names: Collection[str], session: Session +) -> None: + """Upsert each probed artifact's fingerprint and answer; those outside *bundle_names* are dropped with a warning.""" + if outside := sorted( + { + f"{artifact.bundle_name}/{artifact.relative_fileloc}" + for artifact in artifacts + if artifact.bundle_name not in bundle_names + } + ): + log.warning("Ignoring probed task handler artifacts outside the Dag file's scope", artifacts=outside) + by_key: dict[_ArtifactKey, TaskHandlerArtifact] = {} + for artifact in artifacts: + if artifact.bundle_name in bundle_names: + by_key.setdefault( + (artifact.bundle_name, compute_fileloc_hash(artifact.relative_fileloc)), artifact + ) + if not by_key: + return + dialect = get_dialect_name(session) + now = utcnow() + # In key order, so two transactions recording the same artifacts lock them in the same order. + for key in sorted(by_key): + artifact = by_key[key] + probe = { + "size_bytes": artifact.size_bytes, + "cache_digest": artifact.cache_digest, + "task_handlers": { + dag_id: [declaration.model_dump(mode="json") for declaration in declarations] + for dag_id, declarations in artifact.task_handlers.items() + }, + "last_probed_at": now, + } + session.execute( + build_upsert_stmt( + dialect, + LangSDKTaskHandlerArtifact, + conflict_cols=["bundle_name", "relative_fileloc_hash"], + values={ + "bundle_name": artifact.bundle_name, + "relative_fileloc": artifact.relative_fileloc, + # Upserts skip the model's @validates hook, which keeps the hash in sync on ORM writes. + "relative_fileloc_hash": key[1], + **probe, + }, + update_fields=probe, + ) + ) + + +def _get_artifact_key(binding: TaskHandlerBinding) -> _ArtifactKey: + return binding.artifact_bundle_name, compute_fileloc_hash(binding.artifact_rel_path) + + +def _find_task_handler_artifacts( + keys: Collection[_ArtifactKey], *, session: Session +) -> dict[_ArtifactKey, UUID]: + """Return the ids of the recorded artifacts among *keys*, share-locked until commit.""" + # The orphan sweep skips locked rows, so it cannot delete an artifact between this read and the + # insert of the handler rows that reference it. size_bytes is not in the unique index: reading it + # makes MySQL lock the row itself, which the sweep checks, and not only the index entry. + query = with_row_locks( + select( + LangSDKTaskHandlerArtifact.bundle_name, + LangSDKTaskHandlerArtifact.relative_fileloc_hash, + LangSDKTaskHandlerArtifact.id, + LangSDKTaskHandlerArtifact.size_bytes, + ) + .where( + tuple_( + LangSDKTaskHandlerArtifact.bundle_name, LangSDKTaskHandlerArtifact.relative_fileloc_hash + ).in_(list(keys)) + ) + .order_by(LangSDKTaskHandlerArtifact.bundle_name, LangSDKTaskHandlerArtifact.relative_fileloc_hash), + session, + read=True, + key_share=True, + ) + return { + (bundle_name, fileloc_hash): artifact_id + for bundle_name, fileloc_hash, artifact_id, _ in session.execute(query) + } + + +def _sync_task_handlers( + bindings: dict[tuple[str, str], TaskHandlerBinding], + artifact_ids: dict[_ArtifactKey, UUID], + *, + dag_ids: Collection[str], + dag_bundle_name: str, + dag_relative_fileloc: str, + session: Session, +) -> None: + desired = { + key: { + "dag_id": binding.dag_id, + "task_id": binding.task_id, + "artifact_id": artifact_ids[_get_artifact_key(binding)], + "dag_bundle_name": dag_bundle_name, + "dag_relative_fileloc": dag_relative_fileloc, + # Bulk writes skip the model's @validates hook, which keeps the hash in sync on ORM writes. + "dag_relative_fileloc_hash": compute_fileloc_hash(dag_relative_fileloc), + } + for key, binding in bindings.items() + } + recorded = { + (row.dag_id, row.task_id): row._asdict() + for row in session.execute( + select( + LangSDKTaskHandler.dag_id, + LangSDKTaskHandler.task_id, + LangSDKTaskHandler.artifact_id, + LangSDKTaskHandler.dag_bundle_name, + LangSDKTaskHandler.dag_relative_fileloc, + LangSDKTaskHandler.dag_relative_fileloc_hash, + ).where(LangSDKTaskHandler.dag_id.in_(dag_ids)) + ) + } + if gone := sorted(recorded.keys() - desired.keys()): + session.execute( + delete(LangSDKTaskHandler) + .where(tuple_(LangSDKTaskHandler.dag_id, LangSDKTaskHandler.task_id).in_(gone)) + .execution_options(synchronize_session=False) + ) + if added := [desired[key] for key in sorted(desired.keys() - recorded.keys())]: + session.execute(insert(LangSDKTaskHandler), added) + if changed := [ + desired[key] for key in sorted(desired.keys() & recorded.keys()) if desired[key] != recorded[key] + ]: + session.execute(update(LangSDKTaskHandler), changed) + + +def _reconcile_task_handler_bindings( + bindings: Iterable[TaskHandlerBinding], + *, + dag_ids: Collection[str], + artifact_bundle_names: Collection[str], + dag_bundle_name: str, + dag_relative_fileloc: str, + session: Session, +) -> None: + """ + Make the recorded task handler bindings of each Dag in *dag_ids* match *bindings*. + + Rows are keyed by Dag id, not by file. Rows of other Dags are left alone, including those of a Dag + that moved to another file, which that file's parse reconciles. A Dag bound to an artifact outside + *artifact_bundle_names*, or to one with no recorded row, keeps its rows as they are. Artifact rows are + only read: the probed-artifact write records them, and the Dag processor's orphan sweep deletes them. + """ + by_task: dict[tuple[str, str], list[TaskHandlerBinding]] = defaultdict(list) + ignored_dag_ids: set[str] = set() + for binding in bindings: + if binding.dag_id in dag_ids: + by_task[(binding.dag_id, binding.task_id)].append(binding) + else: + ignored_dag_ids.add(binding.dag_id) + if ignored_dag_ids: + log.warning( + "Ignoring task handler bindings of Dags this parse does not persist", + dag_ids=sorted(ignored_dag_ids), + ) + # The parse reports a task bound twice as an import error; this only keeps the rows as they are. + if conflicting := {dag_id for (dag_id, _), group in by_task.items() if len(group) > 1}: + log.warning( + "Ignoring task handler bindings of Dags that bind a task twice", dag_ids=sorted(conflicting) + ) + resolved = {key: group[0] for key, group in by_task.items() if key[0] not in conflicting} + if out_of_scope := { + binding.dag_id + for binding in resolved.values() + if binding.artifact_bundle_name not in artifact_bundle_names + }: + log.warning( + "Ignoring task handler bindings of Dags bound to an artifact outside their scope", + dag_ids=sorted(out_of_scope), + ) + in_scope = [binding for key, binding in resolved.items() if key[0] not in out_of_scope] + keys = {_get_artifact_key(binding) for binding in in_scope} + artifact_ids = _find_task_handler_artifacts(keys, session=session) if keys else {} + # A concurrent orphan sweep or a failed probed-artifact write leaves no row; the next parse probes again. + if unrecorded := { + binding.dag_id for binding in in_scope if _get_artifact_key(binding) not in artifact_ids + }: + log.warning( + "Ignoring task handler bindings of Dags bound to an unrecorded artifact", + dag_ids=sorted(unrecorded), + ) + unchanged = conflicting | out_of_scope | unrecorded + _sync_task_handlers( + {key: binding for key, binding in resolved.items() if key[0] not in unchanged}, + artifact_ids, + dag_ids=[dag_id for dag_id in dag_ids if dag_id not in unchanged], + dag_bundle_name=dag_bundle_name, + dag_relative_fileloc=dag_relative_fileloc, + session=session, + ) + + def update_dag_parsing_results_in_db( bundle_name: str, bundle_version: str | None, @@ -598,6 +809,9 @@ def update_dag_parsing_results_in_db( ), files_parsed: set[tuple[str, str]] | None = None, dag_source_codes: dict[str, DagSourceCode] | None = None, + relative_fileloc: str | None = None, + task_handler_bindings: list[TaskHandlerBinding] | None = None, + task_handler_artifact_bundles: Collection[str] | None = None, ): """ Update everything to do with DAG parsing in the DB. @@ -609,6 +823,7 @@ def update_dag_parsing_results_in_db( - ParseImportError (including with any errors as a result of serialization, not just parsing) - DagWarning - DAG Permissions + - LangSDKTaskHandler, when ``task_handler_bindings`` is given This function will not remove any rows for dags not passed in. It will remove parse errors and warnings from dags/dag files that are passed in. In order words, if a DAG is passed in with a fileloc of `a.py` @@ -621,7 +836,17 @@ def update_dag_parsing_results_in_db( import errors are cleared for files that were parsed but no longer contain DAGs. :param dag_source_codes: Source code read by the Dag importers, keyed by Dag fileloc. Dags without an entry have their source read from ``fileloc``. + :param relative_fileloc: The parsed file, relative to its bundle. Required with ``task_handler_bindings``. + :param task_handler_bindings: The stub-task bindings of every Dag in ``dags``. ``None`` leaves the + recorded bindings as they are; a list replaces the recorded bindings of each Dag in ``dags``. If + the database rejects that write, the recorded bindings stay and the rest is still written. + :param task_handler_artifact_bundles: The bundles whose artifacts the bindings may reference. Required + with ``task_handler_bindings``; a Dag bound to an artifact outside them keeps its recorded bindings. """ + if task_handler_bindings is not None and relative_fileloc is None: + raise ValueError("relative_fileloc is required with task_handler_bindings") + if task_handler_bindings is not None and task_handler_artifact_bundles is None: + raise ValueError("task_handler_artifact_bundles is required with task_handler_bindings") accepted = _reject_other_teams_plugin_classes(bundle_name, dags, import_errors, session=session) if len(accepted) != len(dags): # A rejected Dag may have no ``dag`` row yet, and dag_warning has a foreign key to it. @@ -653,6 +878,34 @@ def update_dag_parsing_results_in_db( SerializedDAG.bulk_write_to_db( bundle_name, bundle_version, dags, parse_duration, session=session ) + # After the Dag rows, which the handler rows reference, and before the serialized Dags. + if ( + task_handler_bindings is not None + and relative_fileloc is not None + and task_handler_artifact_bundles is not None + and dags + ): + # A savepoint, so bindings the database rejects cannot discard the rest of the result. + savepoint = session.begin_nested() + try: + _reconcile_task_handler_bindings( + task_handler_bindings, + dag_ids={dag.dag_id for dag in dags}, + artifact_bundle_names=task_handler_artifact_bundles, + dag_bundle_name=bundle_name, + dag_relative_fileloc=relative_fileloc, + session=session, + ) + except (IntegrityError, DataError): + savepoint.rollback() + log.warning( + "Failed to record task handler bindings; the rest of the parse result is kept", + bundle_name=bundle_name, + relative_fileloc=relative_fileloc, + exc_info=True, + ) + else: + savepoint.commit() # Bulk prefetch metadata for all DAGs to avoid the standard per-DAG # metadata lookups in write_dag. This replaces the update-interval, # hash, and version queries with 2 bulk queries total; DAGs with diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index fa8705493215c..7b678ae0881c7 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -55,7 +55,10 @@ unpack_bundle_version, ) from airflow.dag_processing.bundles.manager import DagBundlesManager -from airflow.dag_processing.collection import update_dag_parsing_results_in_db +from airflow.dag_processing.collection import ( + record_probed_task_handler_artifacts, + update_dag_parsing_results_in_db, +) from airflow.dag_processing.processor import ( DagFileParsingResult, DagFileProcessorProcess, @@ -81,7 +84,7 @@ from airflow.utils.process_utils import ( kill_child_processes_by_pids, ) -from airflow.utils.retries import retry_db_transaction +from airflow.utils.retries import retry_db_transaction, run_with_db_retries from airflow.utils.session import NEW_SESSION, create_session, provide_session from airflow.utils.sqlalchemy import ( is_lock_not_available_error, @@ -1318,6 +1321,19 @@ def handle_parsing_result( ) if proc.parsing_result is not None: + if proc.parsing_result.probed_artifacts: + try: + self.persist_probed_task_handler_artifacts( + bundle_name=file.bundle_name, artifacts=proc.parsing_result.probed_artifacts + ) + except Exception: + # The parse result is still persisted; a binding whose artifact has no row leaves its + # Dag unchanged, and the next parse probes the artifact again. + self.log.exception( + "Failed to record probed task handler artifacts", + bundle_name=file.bundle_name, + relative_fileloc=str(file.rel_path), + ) try: self.persist_parsing_result( bundle_name=file.bundle_name, @@ -1390,8 +1406,31 @@ def persist_parsing_result( session=session, files_parsed=files_parsed, dag_source_codes=parsing_result.dag_source_codes, + relative_fileloc=relative_fileloc, + task_handler_bindings=parsing_result.task_handler_bindings, + task_handler_artifact_bundles=( + None + if parsing_result.task_handler_bindings is None + else self._get_task_handler_artifact_bundle_names(bundle_name) + ), ) + def persist_probed_task_handler_artifacts( + self, *, bundle_name: str, artifacts: Sequence[TaskHandlerArtifact] + ) -> None: + """ + Record the artifacts a Dag file's parse probed, with their answers, in a transaction of their own. + + *bundle_name* is the Dag file's bundle. Default implementation writes to the metadata DB; override + to do it through an API. + """ + artifact_bundle_names = self._get_task_handler_artifact_bundle_names(bundle_name) + for attempt in run_with_db_retries(logger=self.log): + with attempt, create_session() as session: + record_probed_task_handler_artifacts( + artifacts, bundle_names=artifact_bundle_names, session=session + ) + def _collect_results(self): finished = [] for file, proc in self._processors.items(): diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py b/airflow-core/tests/unit/dag_processing/test_collection.py index d7e5213fad5f1..7bba53b577997 100644 --- a/airflow-core/tests/unit/dag_processing/test_collection.py +++ b/airflow-core/tests/unit/dag_processing/test_collection.py @@ -18,11 +18,13 @@ from __future__ import annotations +import contextlib import importlib import logging import os import sys import textwrap +import uuid import warnings from collections.abc import Generator from datetime import timedelta @@ -32,7 +34,7 @@ import pytest from sqlalchemy import delete, event, func, inspect as sa_inspect, select -from sqlalchemy.exc import OperationalError, SAWarning +from sqlalchemy.exc import DataError, OperationalError, SAWarning import airflow.dag_processing.collection from airflow import plugins_manager @@ -46,8 +48,15 @@ _get_latest_runs_stmt_partitioned, _update_dag_tags, _update_import_errors, + record_probed_task_handler_artifacts, update_dag_parsing_results_in_db, ) +from airflow.dag_processing.processor import ( + TaskHandlerArtifact, + TaskHandlerBinding, + TaskHandlerDeclaration, + TaskHandlerParam, +) from airflow.example_dags.plugins.business_day_window import BusinessDayWindow from airflow.example_dags.plugins.custom_partition_mapper import PrefixStripMapper from airflow.example_dags.plugins.workday import AfterWorkdayTimetable @@ -64,6 +73,11 @@ from airflow.models.dagcode import DagCode from airflow.models.dagwarning import DagWarning, DagWarningType from airflow.models.errors import ParseImportError +from airflow.models.lang_sdk_task_handler import ( + LangSDKTaskHandler, + LangSDKTaskHandlerArtifact, + compute_fileloc_hash, +) from airflow.models.serialized_dag import SerializedDagModel from airflow.models.trigger import Trigger from airflow.partition_mappers.base import RollupMapper @@ -91,6 +105,7 @@ from airflow.timetables.simple import PartitionedAtRuntime from airflow.timetables.trigger import CronTriggerTimetable from airflow.triggers.base import BaseEventTrigger +from airflow.utils.session import create_session from airflow.utils.types import DagRunType from tests_common.test_utils.config import conf_vars @@ -2107,3 +2122,459 @@ class WorkdayPlugin(AirflowPlugin): assert stored == set() assert "belonging to other_team" in errors[(bundle_name, "team_dag.py")] + + +ARTIFACT_BUNDLE = "java-task-handlers" +OTHER_TEAM_BUNDLE = "go-task-handlers" +TASK_HANDLERS = { + "etl": [ + TaskHandlerDeclaration( + task_id="extract", binding="positional", params=[TaskHandlerParam(name=None, required=True)] + ) + ] +} + + +def _make_artifact( + rel_path: str = "etl.jar", + *, + bundle_name: str = ARTIFACT_BUNDLE, + size_bytes: int = 1024, + cache_digest: str | None = "a" * 64, + task_handlers: dict[str, list[TaskHandlerDeclaration]] | None = None, +) -> TaskHandlerArtifact: + return TaskHandlerArtifact( + bundle_name=bundle_name, + relative_fileloc=rel_path, + size_bytes=size_bytes, + cache_digest=cache_digest, + task_handlers=TASK_HANDLERS if task_handlers is None else task_handlers, + ) + + +def _record(session, *rel_paths: str, bundle_name: str = ARTIFACT_BUNDLE) -> None: + record_probed_task_handler_artifacts( + [_make_artifact(rel_path, bundle_name=bundle_name) for rel_path in rel_paths], + bundle_names={bundle_name}, + session=session, + ) + + +def _make_binding( + dag_id: str = "etl", + task_id: str = "extract", + *, + rel_path: str = "etl.jar", + bundle_name: str = ARTIFACT_BUNDLE, +) -> TaskHandlerBinding: + return TaskHandlerBinding( + dag_id=dag_id, task_id=task_id, artifact_bundle_name=bundle_name, artifact_rel_path=rel_path + ) + + +def _persist(session, *dag_ids: str, relative_fileloc: str = "etl.py", bindings) -> None: + update_dag_parsing_results_in_db( + bundle_name="testing", + bundle_version=None, + dags=[LazyDeserializedDAG.from_dag(DAG(dag_id=dag_id)) for dag_id in dag_ids], + import_errors={}, + parse_duration=None, + warnings=set(), + session=session, + relative_fileloc=relative_fileloc, + task_handler_bindings=bindings, + task_handler_artifact_bundles={ARTIFACT_BUNDLE}, + ) + + +def _get_recorded_handlers(session) -> set[tuple[str, str, str]]: + """Return ``(dag_id, task_id, artifact path)`` for every recorded handler.""" + rows = session.execute( + select( + LangSDKTaskHandler.dag_id, + LangSDKTaskHandler.task_id, + LangSDKTaskHandlerArtifact.relative_fileloc, + ).join(LangSDKTaskHandlerArtifact, LangSDKTaskHandlerArtifact.id == LangSDKTaskHandler.artifact_id) + ) + return {tuple(row) for row in rows} + + +def _get_recorded_artifacts(session) -> list[tuple]: + return session.execute( + select( + LangSDKTaskHandlerArtifact.relative_fileloc, + LangSDKTaskHandlerArtifact.size_bytes, + LangSDKTaskHandlerArtifact.cache_digest, + LangSDKTaskHandlerArtifact.task_handlers, + LangSDKTaskHandlerArtifact.last_probed_at, + ).order_by(LangSDKTaskHandlerArtifact.relative_fileloc) + ).all() + + +@contextlib.contextmanager +def _capture_task_handler_writes(session) -> Generator[list[str], None, None]: + writes: list[str] = [] + + def _capture(conn, cursor, statement, parameters, context, executemany): + if "lang_sdk_task_handler" in statement and not statement.lstrip().upper().startswith("SELECT"): + writes.append(statement) + + bind = session.get_bind() + event.listen(bind, "before_cursor_execute", _capture) + try: + yield writes + finally: + event.remove(bind, "before_cursor_execute", _capture) + + +def _sweep_after_the_lookup(find): + """Wrap ``_find_task_handler_artifacts`` so ``etl-v2.jar`` gets an id no row has, as if swept since.""" + + def _find(keys, *, session): + artifact_ids = find(keys, session=session) + artifact_ids[(ARTIFACT_BUNDLE, compute_fileloc_hash("etl-v2.jar"))] = uuid.uuid4() + return artifact_ids + + return _find + + +def _fail_after_writing(sync): + """Wrap ``_sync_task_handlers`` so it raises a ``DataError`` after its writes.""" + + def _sync(*args, **kwargs): + sync(*args, **kwargs) + raise DataError("INSERT INTO lang_sdk_task_handler", {}, Exception("value too long")) + + return _sync + + +@pytest.mark.db_test +class TestRecordProbedTaskHandlerArtifacts: + @pytest.fixture(autouse=True) + def _clear_committed_rows(self): + # Autouse, so it tears down after the ``session`` fixture has rolled back; the race test + # commits rows from a second session. + yield + with create_session() as session: + session.execute(delete(LangSDKTaskHandlerArtifact)) + + def test_records_fingerprint_answer_and_probe_time(self, session, time_machine): + probed_at = tz.datetime(2026, 9, 30, 1) + time_machine.move_to(probed_at, tick=False) + + _record(session, "etl.jar") + + assert _get_recorded_artifacts(session) == [ + ("etl.jar", 1024, "a" * 64, {"etl": [TASK_HANDLERS["etl"][0].model_dump(mode="json")]}, probed_at) + ] + assert session.scalar( + select(LangSDKTaskHandlerArtifact.relative_fileloc_hash) + ) == compute_fileloc_hash("etl.jar") + + def test_reprobe_updates_fingerprint_answer_and_probe_time(self, session, time_machine): + time_machine.move_to(tz.datetime(2026, 9, 30, 1), tick=False) + _record(session, "etl.jar") + artifact_id = session.scalar(select(LangSDKTaskHandlerArtifact.id)) + + reprobed_at = tz.datetime(2026, 9, 30, 2) + time_machine.move_to(reprobed_at, tick=False) + record_probed_task_handler_artifacts( + [_make_artifact(size_bytes=2048, cache_digest="b" * 64, task_handlers={})], + bundle_names={ARTIFACT_BUNDLE}, + session=session, + ) + + assert _get_recorded_artifacts(session) == [("etl.jar", 2048, "b" * 64, {}, reprobed_at)] + assert session.scalar(select(LangSDKTaskHandlerArtifact.id)) == artifact_id + + def test_records_an_artifact_without_cache_digest(self, session): + record_probed_task_handler_artifacts( + [_make_artifact(cache_digest=None)], bundle_names={ARTIFACT_BUNDLE}, session=session + ) + + assert session.scalar(select(LangSDKTaskHandlerArtifact.cache_digest)) is None + + def test_artifact_outside_the_scope_is_dropped(self, session, caplog): + record_probed_task_handler_artifacts( + [_make_artifact("etl.jar"), _make_artifact("etl", bundle_name=OTHER_TEAM_BUNDLE)], + bundle_names={ARTIFACT_BUNDLE}, + session=session, + ) + + assert [row[0] for row in _get_recorded_artifacts(session)] == ["etl.jar"] + assert { + "event": "Ignoring probed task handler artifacts outside the Dag file's scope", + "artifacts": [f"{OTHER_TEAM_BUNDLE}/etl"], + } in caplog + + def test_duplicate_artifact_is_written_once(self, session): + with _capture_task_handler_writes(session) as writes: + record_probed_task_handler_artifacts( + [_make_artifact(), _make_artifact()], bundle_names={ARTIFACT_BUNDLE}, session=session + ) + + assert len(writes) == 1 + assert [row[0] for row in _get_recorded_artifacts(session)] == ["etl.jar"] + + def test_artifact_recorded_by_a_concurrent_parse(self, session): + """Two Dag processors probe one new artifact; the second upsert must update, not fail.""" + with create_session() as other: + _record(other, "etl.jar") + other_id = other.scalar(select(LangSDKTaskHandlerArtifact.id)) + + record_probed_task_handler_artifacts( + [_make_artifact(size_bytes=2048)], bundle_names={ARTIFACT_BUNDLE}, session=session + ) + + assert session.execute( + select(LangSDKTaskHandlerArtifact.id, LangSDKTaskHandlerArtifact.size_bytes) + ).all() == [(other_id, 2048)] + + +@pytest.mark.db_test +@pytest.mark.usefixtures("testing_dag_bundle") +class TestTaskHandlerBindingReconcile: + @pytest.fixture + def recorded(self, testing_dag_bundle, session): + _record(session, "etl.jar", "etl-v2.jar", "other.jar") + _persist( + session, + "etl", + bindings=[_make_binding(task_id="extract"), _make_binding(task_id="transform")], + ) + _persist( + session, + "other", + relative_fileloc="other.py", + bindings=[_make_binding("other", "load", rel_path="other.jar")], + ) + + @pytest.mark.parametrize( + ("bindings", "expected_handlers"), + [ + pytest.param( + None, + { + ("etl", "extract", "etl.jar"), + ("etl", "transform", "etl.jar"), + ("other", "load", "other.jar"), + }, + id="none-leaves-rows", + ), + pytest.param( + [], + {("other", "load", "other.jar")}, + id="empty-deletes-this-results-dag-rows", + ), + pytest.param( + [_make_binding(task_id="extract", rel_path="etl-v2.jar"), _make_binding(task_id="publish")], + { + ("etl", "extract", "etl-v2.jar"), + ("etl", "publish", "etl.jar"), + ("other", "load", "other.jar"), + }, + id="list-deletes-updates-and-inserts", + ), + ], + ) + @pytest.mark.usefixtures("recorded") + def test_task_handler_bindings_reconcile_by_value(self, session, bindings, expected_handlers): + _persist(session, "etl", bindings=bindings) + + assert _get_recorded_handlers(session) == expected_handlers + + @pytest.mark.usefixtures("recorded") + def test_task_handler_bindings_of_a_result_without_dags_change_nothing(self, session): + before = _get_recorded_handlers(session) + + _persist(session, bindings=[]) + + assert _get_recorded_handlers(session) == before + + def test_task_handler_rows_follow_a_dag_to_its_new_file(self, session): + _record(session, "etl.jar") + _persist(session, "etl", relative_fileloc="a.py", bindings=[_make_binding()]) + + _persist(session, "etl", relative_fileloc="dags/b.py", bindings=[_make_binding()]) + + assert session.execute( + select(LangSDKTaskHandler.dag_relative_fileloc, LangSDKTaskHandler.dag_relative_fileloc_hash) + ).all() == [("dags/b.py", compute_fileloc_hash("dags/b.py"))] + + def test_task_handler_artifact_is_shared_across_files(self, session): + _record(session, "etl.jar") + _persist(session, "etl", relative_fileloc="etl.py", bindings=[_make_binding("etl")]) + _persist(session, "other", relative_fileloc="other.py", bindings=[_make_binding("other")]) + + artifact_ids = session.scalars(select(LangSDKTaskHandlerArtifact.id)).all() + assert len(artifact_ids) == 1 + assert set(session.scalars(select(LangSDKTaskHandler.artifact_id))) == set(artifact_ids) + + _persist(session, "etl", relative_fileloc="etl.py", bindings=[]) + + assert session.scalars(select(LangSDKTaskHandlerArtifact.id)).all() == artifact_ids + assert _get_recorded_handlers(session) == {("other", "extract", "etl.jar")} + + def test_bindings_never_write_the_artifact_table(self, session): + _record(session, "etl.jar") + + with _capture_task_handler_writes(session) as writes: + _persist(session, "etl", bindings=[_make_binding()]) + assert writes + assert not [statement for statement in writes if "lang_sdk_task_handler_artifact" in statement] + + with _capture_task_handler_writes(session) as writes: + _persist(session, "etl", bindings=[_make_binding()]) + assert writes == [] + + @pytest.mark.usefixtures("recorded") + def test_binding_to_a_missing_artifact_leaves_its_dag_unchanged(self, session, caplog): + before = _get_recorded_handlers(session) + + _persist( + session, + "etl", + "other", + bindings=[ + _make_binding(task_id="extract", rel_path="swept.jar"), + _make_binding("other", "export"), + ], + ) + + assert _get_recorded_handlers(session) == {row for row in before if row[0] == "etl"} | { + ("other", "export", "etl.jar") + } + assert { + "event": "Ignoring task handler bindings of Dags bound to an unrecorded artifact", + "dag_ids": ["etl"], + } in caplog + + @pytest.mark.usefixtures("recorded") + def test_binding_outside_the_scope_leaves_its_dag_unchanged(self, session, caplog): + _record(session, "etl.jar", bundle_name=OTHER_TEAM_BUNDLE) + before = _get_recorded_handlers(session) + + _persist( + session, + "etl", + "other", + bindings=[ + _make_binding(task_id="extract", bundle_name=OTHER_TEAM_BUNDLE), + _make_binding("other", "export"), + ], + ) + + assert _get_recorded_handlers(session) == {row for row in before if row[0] == "etl"} | { + ("other", "export", "etl.jar") + } + assert { + "event": "Ignoring task handler bindings of Dags bound to an artifact outside their scope", + "dag_ids": ["etl"], + } in caplog + + def test_task_handler_bindings_for_dags_outside_the_result_are_dropped(self, session, caplog): + _record(session, "etl.jar") + _persist(session, "etl", bindings=[_make_binding("etl"), _make_binding("ghost")]) + + assert _get_recorded_handlers(session) == {("etl", "extract", "etl.jar")} + assert { + "event": "Ignoring task handler bindings of Dags this parse does not persist", + "dag_ids": ["ghost"], + } in caplog + + @pytest.mark.usefixtures("recorded") + def test_task_handler_bindings_with_a_duplicate_task_leave_the_dag_untouched(self, session, caplog): + before = _get_recorded_handlers(session) + + _persist( + session, + "etl", + "other", + bindings=[ + _make_binding(task_id="extract", rel_path="etl.jar"), + _make_binding(task_id="extract", rel_path="etl-v2.jar"), + _make_binding("other", "load", rel_path="etl-v2.jar"), + ], + ) + + assert _get_recorded_handlers(session) == {row for row in before if row[0] == "etl"} | { + ("other", "load", "etl-v2.jar") + } + assert { + "event": "Ignoring task handler bindings of Dags that bind a task twice", + "dag_ids": ["etl"], + } in caplog + + @pytest.mark.parametrize( + ("target", "wrap"), + [ + pytest.param("_find_task_handler_artifacts", _sweep_after_the_lookup, id="integrity-error"), + pytest.param("_sync_task_handlers", _fail_after_writing, id="data-error"), + ], + ) + @pytest.mark.usefixtures("recorded") + def test_bindings_the_database_rejects_leave_the_rest_of_the_result(self, session, caplog, target, wrap): + before = _get_recorded_handlers(session) + + with mock.patch.object( + airflow.dag_processing.collection, + target, + side_effect=wrap(getattr(airflow.dag_processing.collection, target)), + ): + update_dag_parsing_results_in_db( + bundle_name="testing", + bundle_version=None, + dags=[LazyDeserializedDAG.from_dag(DAG(dag_id=dag_id)) for dag_id in ("etl", "fresh")], + import_errors={("testing", "broken.py"): "SyntaxError"}, + parse_duration=None, + warnings=set(), + session=session, + files_parsed={("testing", "etl.py"), ("testing", "broken.py")}, + relative_fileloc="etl.py", + # Deletes the transform row before inserting the publish row. + task_handler_bindings=[ + _make_binding(task_id="extract"), + _make_binding(task_id="publish", rel_path="etl-v2.jar"), + ], + task_handler_artifact_bundles={ARTIFACT_BUNDLE}, + ) + + assert _get_recorded_handlers(session) == before + assert session.scalars( + select(SerializedDagModel.dag_id).where(SerializedDagModel.dag_id == "fresh") + ).all() == ["fresh"] + assert session.scalars(select(ParseImportError.filename)).all() == ["broken.py"] + assert { + "event": "Failed to record task handler bindings; the rest of the parse result is kept", + "bundle_name": "testing", + "relative_fileloc": "etl.py", + "log_level": "warning", + } in caplog + + def test_task_handler_bindings_require_relative_fileloc(self, session): + with pytest.raises(ValueError, match="relative_fileloc"): + update_dag_parsing_results_in_db( + bundle_name="testing", + bundle_version=None, + dags=[LazyDeserializedDAG.from_dag(DAG(dag_id="etl"))], + import_errors={}, + parse_duration=None, + warnings=set(), + session=session, + task_handler_bindings=[], + task_handler_artifact_bundles={ARTIFACT_BUNDLE}, + ) + + def test_task_handler_bindings_require_a_scope(self, session): + with pytest.raises(ValueError, match="task_handler_artifact_bundles"): + update_dag_parsing_results_in_db( + bundle_name="testing", + bundle_version=None, + dags=[LazyDeserializedDAG.from_dag(DAG(dag_id="etl"))], + import_errors={}, + parse_duration=None, + warnings=set(), + session=session, + relative_fileloc="etl.py", + task_handler_bindings=[], + ) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 1f0a66cbf303d..d3b73ef20ab6b 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -41,7 +41,7 @@ import pytest import structlog import time_machine -from sqlalchemy import event, func, select +from sqlalchemy import delete, event, func, select from sqlalchemy.exc import OperationalError from uuid6 import uuid7 @@ -62,6 +62,7 @@ DagFileParsingResult, DagFileProcessorProcess, TaskHandlerArtifact, + TaskHandlerBinding, TaskHandlerDeclaration, TaskHandlerParam, _parse_file, @@ -4496,6 +4497,137 @@ def test_get_known_task_handler_artifacts_from_db_reads_only_the_given_bundles(s } +class TestTaskHandlerBindings: + @staticmethod + def _finished_parse(**result) -> tuple[DagFileInfo, MagicMock]: + file = DagFileInfo(bundle_name="dags-a", rel_path=Path("etl.py"), bundle_path=TEST_DAGS_FOLDER) + proc = MagicMock( + had_callbacks=False, + start_time=time.monotonic(), + parsing_result=DagFileParsingResult(fileloc="etl.py", serialized_dags=[], **result), + ) + return file, proc + + @staticmethod + def _manager() -> DagFileProcessorManager: + manager = DagFileProcessorManager(max_runs=1) + manager._bundle_versions["dags-a"] = None + return manager + + @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) + @mock.patch.object(DagFileProcessorManager, "persist_probed_task_handler_artifacts", autospec=True) + def test_handle_parsing_result_records_probed_artifacts_first(self, persist_probed, persist_result): + calls: list[str] = [] + persist_probed.side_effect = lambda *args, **kwargs: calls.append("probed") + persist_result.side_effect = lambda *args, **kwargs: calls.append("result") + file, proc = self._finished_parse(probed_artifacts=[_known_artifact(JAVA_TASK_HANDLERS)]) + manager = self._manager() + + manager.handle_parsing_result(file, proc, session=mock.MagicMock()) + + persist_probed.assert_called_once_with( + manager, bundle_name="dags-a", artifacts=[_known_artifact(JAVA_TASK_HANDLERS)] + ) + assert calls == ["probed", "result"] + + @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) + @mock.patch.object(DagFileProcessorManager, "persist_probed_task_handler_artifacts", autospec=True) + def test_handle_parsing_result_skips_the_artifact_write_without_probes( + self, persist_probed, persist_result + ): + file, proc = self._finished_parse() + + self._manager().handle_parsing_result(file, proc, session=mock.MagicMock()) + + persist_probed.assert_not_called() + persist_result.assert_called_once() + + @mock.patch.object(DagFileProcessorManager, "persist_parsing_result", autospec=True) + @mock.patch.object(DagFileProcessorManager, "persist_probed_task_handler_artifacts", autospec=True) + def test_handle_parsing_result_persists_when_the_artifact_write_fails( + self, persist_probed, persist_result, caplog + ): + persist_probed.side_effect = OperationalError("INSERT", {}, Exception("locked")) + file, proc = self._finished_parse(probed_artifacts=[_known_artifact(JAVA_TASK_HANDLERS)]) + + self._manager().handle_parsing_result(file, proc, session=mock.MagicMock()) + + persist_result.assert_called_once() + assert { + "event": "Failed to record probed task handler artifacts", + "bundle_name": "dags-a", + "relative_fileloc": "etl.py", + "log_level": "error", + } in caplog + + @pytest.mark.parametrize( + ("bindings", "expected_scope"), + [ + pytest.param(None, None, id="not-evaluated"), + pytest.param( + [ + TaskHandlerBinding( + dag_id="etl", + task_id="extract", + artifact_bundle_name=JAVA_TASK_HANDLERS, + artifact_rel_path="etl.jar", + ) + ], + frozenset({JAVA_TASK_HANDLERS}), + id="bindings", + ), + ], + ) + @mock.patch.object( + DagFileProcessorManager, + "_get_task_handler_artifact_bundle_names", + autospec=True, + return_value=frozenset({JAVA_TASK_HANDLERS}), + ) + @mock.patch("airflow.dag_processing.manager.update_dag_parsing_results_in_db", autospec=True) + def test_persist_parsing_result_forwards_task_handler_bindings( + self, update_in_db, get_scope, bindings, expected_scope + ): + DagFileProcessorManager(max_runs=1).persist_parsing_result( + bundle_name="dags-a", + bundle_version=None, + version_data=None, + parsing_result=DagFileParsingResult( + fileloc="/files/dags/etl.py", serialized_dags=[], task_handler_bindings=bindings + ), + run_duration=0.1, + relative_fileloc="etl.py", + session=mock.MagicMock(), + ) + + assert update_in_db.call_args.kwargs["relative_fileloc"] == "etl.py" + assert update_in_db.call_args.kwargs["task_handler_bindings"] == bindings + assert update_in_db.call_args.kwargs["task_handler_artifact_bundles"] == expected_scope + assert get_scope.call_count == (bindings is not None) + + @pytest.mark.db_test + @conf_vars(_routed_coordinators({"java": _coordinator(JAVA_TASK_HANDLERS)})) + def test_persist_probed_task_handler_artifacts_commits_the_scoped_artifacts( + self, configure_dag_bundles, tmp_path + ): + reset_coordinator_manager() + try: + with configure_dag_bundles({"dags-a": tmp_path, JAVA_TASK_HANDLERS: tmp_path}): + DagFileProcessorManager(max_runs=1).persist_probed_task_handler_artifacts( + bundle_name="dags-a", + artifacts=[_known_artifact(JAVA_TASK_HANDLERS), _known_artifact("dags-b")], + ) + + with create_session(scoped=False) as other: + assert other.scalars(select(LangSDKTaskHandlerArtifact.bundle_name)).all() == [ + JAVA_TASK_HANDLERS + ] + finally: + reset_coordinator_manager() + with create_session() as session: + session.execute(delete(LangSDKTaskHandlerArtifact)) + + class TestMultiTeamMetrics: """Tests for team_name tag on dag processing metrics in multi-team mode.""" From f7fd2bacc9cd675d1a25fe100bf559c701d15e67 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 19:02:37 +0000 Subject: [PATCH 06/12] Index task handler artifacts by last_probed_at The orphan sweep deletes unreferenced artifact rows only once last_probed_at is past a grace period, and a batched cleanup must filter on indexed columns. It is a separate migration because the tables' migration already ships in an earlier, separately reviewed PR. --- airflow-core/docs/migrations-ref.rst | 4 +- ...index_to_lang_sdk_task_handler_artifact.py | 51 +++++++++++++++++++ .../airflow/models/lang_sdk_task_handler.py | 1 + airflow-core/src/airflow/utils/db.py | 2 +- 4 files changed, 56 insertions(+), 2 deletions(-) create mode 100644 airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_last_probed_at_index_to_lang_sdk_task_handler_artifact.py diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index 70847aee7d4d3..c25f4acf8889f 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,9 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``f7ed13533d23`` (head) | ``90e4d18ccadf`` | ``3.4.0`` | Add lang_sdk_task_handler_artifact and lang_sdk_task_handler | +| ``37d645374a9c`` (head) | ``f7ed13533d23`` | ``3.4.0`` | Add last_probed_at index to lang_sdk_task_handler_artifact. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``f7ed13533d23`` | ``90e4d18ccadf`` | ``3.4.0`` | Add lang_sdk_task_handler_artifact and lang_sdk_task_handler | | | | | tables. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``90e4d18ccadf`` | ``e5a91c7f42b3`` | ``3.4.0`` | Add timetable_asset_gated to DagModel. | diff --git a/airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_last_probed_at_index_to_lang_sdk_task_handler_artifact.py b/airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_last_probed_at_index_to_lang_sdk_task_handler_artifact.py new file mode 100644 index 0000000000000..8eb9c65d2b596 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0143_3_4_0_add_last_probed_at_index_to_lang_sdk_task_handler_artifact.py @@ -0,0 +1,51 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +""" +Add last_probed_at index to lang_sdk_task_handler_artifact. + +Revision ID: 37d645374a9c +Revises: f7ed13533d23 +Create Date: 2026-09-30 19:00:42.253263 + +""" + +from __future__ import annotations + +from alembic import op + +# revision identifiers, used by Alembic. +revision = "37d645374a9c" +down_revision = "f7ed13533d23" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Add last_probed_at index to lang_sdk_task_handler_artifact.""" + with op.batch_alter_table("lang_sdk_task_handler_artifact", schema=None) as batch_op: + batch_op.create_index( + "idx_lang_sdk_task_handler_artifact_last_probed_at", ["last_probed_at"], unique=False + ) + + +def downgrade(): + """Remove last_probed_at index from lang_sdk_task_handler_artifact.""" + with op.batch_alter_table("lang_sdk_task_handler_artifact", schema=None) as batch_op: + batch_op.drop_index("idx_lang_sdk_task_handler_artifact_last_probed_at") diff --git a/airflow-core/src/airflow/models/lang_sdk_task_handler.py b/airflow-core/src/airflow/models/lang_sdk_task_handler.py index 86cb020f3f1d7..ad8ad1ef59203 100644 --- a/airflow-core/src/airflow/models/lang_sdk_task_handler.py +++ b/airflow-core/src/airflow/models/lang_sdk_task_handler.py @@ -83,6 +83,7 @@ class LangSDKTaskHandlerArtifact(Base): relative_fileloc_hash, name="lang_sdk_task_handler_artifact_bundle_fileloc_uq", ), + Index("idx_lang_sdk_task_handler_artifact_last_probed_at", last_probed_at), ) @validates("relative_fileloc") diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index 689d57ded759f..3a489bfe4bf76 100644 --- a/airflow-core/src/airflow/utils/db.py +++ b/airflow-core/src/airflow/utils/db.py @@ -117,7 +117,7 @@ class MappedClassProtocol(Protocol): "3.1.8": "509b94a1042d", "3.2.0": "1d6611b6ab7c", "3.3.0": "d2f4e1b3c5a7", - "3.4.0": "f7ed13533d23", + "3.4.0": "37d645374a9c", } # Prefix used to identify tables holding data moved during migration. From b1fb803adeb520a34323e2f02e8f9f94d17733f4 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 19:05:44 +0000 Subject: [PATCH 07/12] Delete task handler artifacts no stub task references A parse never evicts artifact rows, because one artifact serves Dags in many files, so the Dag processor reclaims them itself; db clean leaves the table alone. The grace period keeps rows for probed artifacts that bind nothing, which later parses reuse instead of probing again. Batches commit on their own so the sweep never holds locks across the whole table, and SKIP LOCKED keeps it from waiting on, or deleting under, a parse that is binding an artifact. --- .../src/airflow/config_templates/config.yml | 3 +- .../src/airflow/dag_processing/manager.py | 69 ++++++- .../tests/unit/dag_processing/test_manager.py | 183 +++++++++++++++++- 3 files changed, 250 insertions(+), 5 deletions(-) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index eafaecdda739b..595c38f0ccb1e 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2727,7 +2727,8 @@ scheduler: description: | How often (in seconds) to check for stale DAGs (DAGs which are no longer present in the expected files) which should be deactivated, as well as assets that are no longer - referenced and should be marked as orphaned. + referenced and should be marked as orphaned, and Lang-SDK task handler artifacts that no + stub task uses any more, which are deleted. version_added: 2.5.0 type: integer example: ~ diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index 7b678ae0881c7..fb64d1781fdf7 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -40,7 +40,7 @@ import attrs import structlog from pydantic import ValidationError -from sqlalchemy import select, update +from sqlalchemy import delete, exists, select, update from sqlalchemy.exc import OperationalError from sqlalchemy.orm import load_only from tabulate import tabulate @@ -71,7 +71,7 @@ from airflow.models.dagwarning import DagWarning from airflow.models.db_callback_request import DbCallbackRequest from airflow.models.errors import ParseImportError -from airflow.models.lang_sdk_task_handler import LangSDKTaskHandlerArtifact +from airflow.models.lang_sdk_task_handler import LangSDKTaskHandler, LangSDKTaskHandlerArtifact from airflow.observability.metrics import stats_utils from airflow.sdk import SecretCache from airflow.sdk.execution_time.coordinator import get_coordinator_manager # noqa: SDK001 @@ -171,6 +171,14 @@ class _TaskHandlerBundles(NamedTuple): """Whether a coordinator without ``task_handler_bundle_name`` reads the task's own Dag bundle.""" +_TASK_HANDLER_ARTIFACT_SWEEP_BATCH_SIZE = 1000 + +# A probed artifact that serves no stub task can have a row too, so that parses do not probe it again, +# and such a row is unreferenced from the start. An unreferenced row is therefore kept for this many +# sweep intervals after its last probe, long enough for the Dag files parsed next to reuse it. +_TASK_HANDLER_ARTIFACT_SWEEP_GRACE_INTERVALS = 5 + + def _config_int_factory(section: str, key: str): return functools.partial(conf.getint, section, key) @@ -261,6 +269,7 @@ class DagFileProcessorManager(LoggingMixin): ) _last_deactivate_stale_dags_time: float = attrs.field(default=0, init=False) + _last_task_handler_artifact_sweep_time: float = attrs.field(default=0, init=False) _last_stale_bundle_cleanup_time: float = attrs.field(default=0, init=False) print_stats_interval: float = attrs.field( factory=_config_int_factory("dag_processor", "print_stats_interval") @@ -451,6 +460,61 @@ def _scan_stale_dags(self): self.deactivate_stale_dags(last_parsed=last_parsed) self._last_deactivate_stale_dags_time = time.monotonic() + def _sweep_task_handler_artifacts(self) -> None: + now = time.monotonic() + if now - self._last_task_handler_artifact_sweep_time <= self.parsing_cleanup_interval: + return + try: + self.delete_unreferenced_task_handler_artifacts() + except Exception: + self.log.exception("Error deleting unreferenced task handler artifacts") + finally: + self._last_task_handler_artifact_sweep_time = now + + def delete_unreferenced_task_handler_artifacts(self) -> int: + """ + Delete the task handler artifact rows that no stub task references and no recent parse probed. + + Deletes in batches, each committed on its own, and returns how many rows it deleted. Default + implementation writes to the metadata DB; override to do it through an API. + """ + cutoff = timezone.utcnow() - timedelta( + seconds=self.parsing_cleanup_interval * _TASK_HANDLER_ARTIFACT_SWEEP_GRACE_INTERVALS + ) + sweepable = ( + LangSDKTaskHandlerArtifact.last_probed_at < cutoff, + ~exists().where(LangSDKTaskHandler.artifact_id == LangSDKTaskHandlerArtifact.id), + ) + deleted = 0 + while True: + with create_session() as session: + # SKIP LOCKED passes over the artifacts a concurrent parse holds while it binds them. + ids = session.scalars( + with_row_locks( + select(LangSDKTaskHandlerArtifact.id) + .where(*sweepable) + .order_by(LangSDKTaskHandlerArtifact.last_probed_at) + .limit(_TASK_HANDLER_ARTIFACT_SWEEP_BATCH_SIZE), + session, + of=LangSDKTaskHandlerArtifact, + skip_locked=True, + key_share=False, + ) + ).all() + if ids: + # Checked again for handler rows a parse committed since the select. + result = session.execute( + delete(LangSDKTaskHandlerArtifact) + .where(LangSDKTaskHandlerArtifact.id.in_(ids), *sweepable) + .execution_options(synchronize_session=False) + ) + deleted += getattr(result, "rowcount", 0) + if len(ids) < _TASK_HANDLER_ARTIFACT_SWEEP_BATCH_SIZE: + break + if deleted: + self.log.info("Deleted %i unreferenced task handler artifacts.", deleted) + return deleted + def _cleanup_stale_bundle_versions(self): if self.stale_bundle_cleanup_interval <= 0: return @@ -604,6 +668,7 @@ def _run_parsing_loop(self): for callback in self.fetch_callbacks(): self._add_callback_to_queue(callback) self._scan_stale_dags() + self._sweep_task_handler_artifacts() self._cleanup_stale_bundle_versions() self.purge_inactive_dag_warnings() diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index d3b73ef20ab6b..07df6b113e1c4 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -49,7 +49,11 @@ from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion from airflow.dag_processing.bundles.manager import DagBundlesManager -from airflow.dag_processing.collection import update_dag_parsing_results_in_db +from airflow.dag_processing.collection import ( + _find_task_handler_artifacts, + record_probed_task_handler_artifacts, + update_dag_parsing_results_in_db, +) from airflow.dag_processing.dagbag import DagBag from airflow.dag_processing.manager import ( BundleState, @@ -72,7 +76,11 @@ from airflow.models.dag_version import DagVersion from airflow.models.dagbundle import DagBundleModel from airflow.models.dagcode import DagCode -from airflow.models.lang_sdk_task_handler import LangSDKTaskHandlerArtifact +from airflow.models.lang_sdk_task_handler import ( + LangSDKTaskHandler, + LangSDKTaskHandlerArtifact, + compute_fileloc_hash, +) from airflow.models.serialized_dag import SerializedDagModel from airflow.models.team import Team from airflow.providers.standard.operators.empty import EmptyOperator @@ -4628,6 +4636,177 @@ def test_persist_probed_task_handler_artifacts_commits_the_scoped_artifacts( session.execute(delete(LangSDKTaskHandlerArtifact)) +@pytest.mark.db_test +@pytest.mark.usefixtures("testing_dag_bundle") +class TestTaskHandlerArtifactSweep: + @pytest.fixture(autouse=True) + def _clear_task_handler_rows(self): + yield + with create_session() as session: + session.execute(delete(LangSDKTaskHandler)) + session.execute(delete(LangSDKTaskHandlerArtifact)) + clear_db_serialized_dags() + clear_db_dags() + + @staticmethod + def _add_artifacts(*paths: str, probed_at: datetime, referenced: tuple[str, ...] = ()) -> None: + """Commit one artifact per path, and a handler row for each path in *referenced*.""" + with create_session() as session: + artifacts = { + path: LangSDKTaskHandlerArtifact( + bundle_name=JAVA_TASK_HANDLERS, + relative_fileloc=path, + size_bytes=1024, + cache_digest="ab12", + task_handlers={}, + last_probed_at=probed_at, + ) + for path in paths + } + session.add_all(artifacts.values()) + if referenced: + session.add(DagModel(dag_id="etl", bundle_name="testing")) + session.flush() + session.add_all( + LangSDKTaskHandler( + dag_id="etl", + task_id=path, + artifact_id=artifacts[path].id, + dag_bundle_name="testing", + dag_relative_fileloc="etl.py", + ) + for path in referenced + ) + + @staticmethod + def _get_artifact_paths() -> set[str]: + with create_session() as session: + return set(session.scalars(select(LangSDKTaskHandlerArtifact.relative_fileloc))) + + @staticmethod + def _sweep() -> int: + manager = DagFileProcessorManager(max_runs=1) + manager.parsing_cleanup_interval = 60 + return manager.delete_unreferenced_task_handler_artifacts() + + @mock.patch.object(DagFileProcessorManager, "delete_unreferenced_task_handler_artifacts", autospec=True) + def test_task_handler_artifact_sweep_runs_once_per_cleanup_interval(self, delete_unreferenced): + manager = DagFileProcessorManager(max_runs=1) + manager.parsing_cleanup_interval = 10 + manager._last_task_handler_artifact_sweep_time = time.monotonic() - 20 + + manager._sweep_task_handler_artifacts() + manager._sweep_task_handler_artifacts() + + delete_unreferenced.assert_called_once_with(manager) + + @mock.patch.object(DagFileProcessorManager, "delete_unreferenced_task_handler_artifacts", autospec=True) + def test_task_handler_artifact_sweep_logs_errors(self, delete_unreferenced, caplog): + delete_unreferenced.side_effect = OperationalError("DELETE", {}, Exception("locked")) + manager = DagFileProcessorManager(max_runs=1) + + manager._sweep_task_handler_artifacts() + + assert {"event": "Error deleting unreferenced task handler artifacts", "log_level": "error"} in caplog + assert manager._last_task_handler_artifact_sweep_time > 0 + + @mock.patch("airflow.dag_processing.manager._TASK_HANDLER_ARTIFACT_SWEEP_BATCH_SIZE", 2) + def test_delete_unreferenced_task_handler_artifacts_commits_each_batch(self, session): + self._add_artifacts( + "a.jar", + "b.jar", + "c.jar", + "d.jar", + "e.jar", + "used.jar", + probed_at=timezone.utcnow() - timedelta(days=1), + referenced=("used.jar",), + ) + events: list[str] = [] + + def on_execute(conn, cursor, statement, parameters, context, executemany): + if statement.lstrip().upper().startswith("DELETE FROM LANG_SDK_TASK_HANDLER_ARTIFACT"): + events.append("DELETE") + + def on_commit(conn): + events.append("COMMIT") + + bind = session.get_bind() + event.listen(bind, "before_cursor_execute", on_execute) + event.listen(bind, "commit", on_commit) + try: + deleted = self._sweep() + finally: + event.remove(bind, "before_cursor_execute", on_execute) + event.remove(bind, "commit", on_commit) + + assert deleted == 5 + assert events == ["DELETE", "COMMIT"] * 3 + assert self._get_artifact_paths() == {"used.jar"} + + @pytest.mark.parametrize( + ("age", "referenced", "expected"), + [ + pytest.param(timedelta(seconds=301), (), set(), id="unreferenced-past-grace"), + pytest.param(timedelta(seconds=299), (), {"etl.jar"}, id="unreferenced-within-grace"), + pytest.param(timedelta(days=30), ("etl.jar",), {"etl.jar"}, id="referenced"), + ], + ) + @time_machine.travel(datetime(2026, 9, 30, 12, tzinfo=timezone.utc), tick=False) + def test_task_handler_artifact_sweep_grace_period(self, age, referenced, expected): + self._add_artifacts("etl.jar", probed_at=timezone.utcnow() - age, referenced=referenced) + + self._sweep() + + assert self._get_artifact_paths() == expected + + def test_task_handler_artifact_of_a_deleted_dag_is_swept(self, time_machine): + time_machine.move_to(datetime(2026, 9, 30, 12, tzinfo=timezone.utc), tick=False) + with create_session() as session: + record_probed_task_handler_artifacts( + [_known_artifact(JAVA_TASK_HANDLERS)], bundle_names={JAVA_TASK_HANDLERS}, session=session + ) + update_dag_parsing_results_in_db( + "testing", + None, + [LazyDeserializedDAG.from_dag(SdkDAG(dag_id="etl"))], + {}, + None, + set(), + session, + relative_fileloc="etl.py", + task_handler_bindings=[ + TaskHandlerBinding( + dag_id="etl", + task_id="extract", + artifact_bundle_name=JAVA_TASK_HANDLERS, + artifact_rel_path=f"{JAVA_TASK_HANDLERS}.jar", + ) + ], + task_handler_artifact_bundles={JAVA_TASK_HANDLERS}, + ) + with create_session() as session: + session.execute(delete(DagModel).where(DagModel.dag_id == "etl")) + time_machine.shift(timedelta(hours=1)) + + self._sweep() + + assert self._get_artifact_paths() == set() + + @pytest.mark.backend("postgres", "mysql") + def test_task_handler_artifact_sweep_skips_rows_a_parse_holds(self): + self._add_artifacts("held.jar", "free.jar", probed_at=timezone.utcnow() - timedelta(days=1)) + + with create_session() as parse: + _find_task_handler_artifacts( + [(JAVA_TASK_HANDLERS, compute_fileloc_hash("held.jar"))], session=parse + ) + self._sweep() + parse.rollback() + + assert self._get_artifact_paths() == {"held.jar"} + + class TestMultiTeamMetrics: """Tests for team_name tag on dag processing metrics in multi-team mode.""" From f2c1b9efa019f62ba320320dffb76c5777124834 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 19:07:00 +0000 Subject: [PATCH 08/12] Describe the binding reconcile and artifact sweep in ADR-0013 A binding names only its artifact, and the probed answers travel in a field of their own, because an answer belongs to the artifact and is worth keeping even when a Dag fails validation. Writing them in their own transaction before the persist takes the upsert out of the reconcile, which then only looks artifacts up, so the ADR has to say what happens to a Dag whose artifact is missing or outside the file's scope. An empty list deletes only the rows of the Dags in the result, since the reconcile is keyed by Dag id rather than by file. db clean excludes the artifact table, so the orphan sweep is the only reclaim path. --- .../0013-persisted-task-handler-bindings.md | 51 ++++++++++--------- 1 file changed, 28 insertions(+), 23 deletions(-) diff --git a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md index e38c4c3c3cb58..bba078b7b1b1f 100644 --- a/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md +++ b/airflow-core/adr/lang-sdk/0013-persisted-task-handler-bindings.md @@ -271,17 +271,14 @@ none. The answer must not depend on the request, because the manager records it every Dag file that resolves against the artifact. **Dag-parsing child → manager.** [`DagFileParsingResult`](https://github.com/apache/airflow/blob/79991cd4db0c9346a28b23c453377f6df0c6b4ed/airflow-core/src/airflow/dag_processing/processor.py#L133-L145) gains the resolved -bindings. +bindings and the probed artifacts. ```python -class SDKTaskHandlerBinding(BaseModel): +class TaskHandlerBinding(BaseModel): dag_id: str task_id: str artifact_bundle_name: str artifact_rel_path: str - artifact_size_bytes: int - artifact_cache_digest: str - handler_params: list[TaskHandlerParam] class DagFileParsingResult(BaseModel): @@ -289,15 +286,18 @@ class DagFileParsingResult(BaseModel): serialized_dags: list[LazyDeserializedDAG] warnings: list | None = None import_errors: dict[str, str] | None = None - task_handler_bindings: list[SDKTaskHandlerBinding] | None = None # new + task_handler_bindings: list[TaskHandlerBinding] | None = None # new + probed_artifacts: list[TaskHandlerArtifact] = [] # new ``` -`None` and `[]` mean different things, and the difference is load-bearing: +`probed_artifacts` holds every artifact the parse probed, with its answer; it is recorded even when the bindings are `None`, and it is the only way an answer gets written. + +For `task_handler_bindings`, `None` and `[]` mean different things, and the difference is load-bearing: | value | meaning | manager does | |--------|------------------------------------------|-----------------------| | `None` | handlers were not evaluated in this parse | **nothing** — no reconcile | -| `[]` | evaluated, this file has no stub handlers | delete this file's rows | +| `[]` | evaluated, this file has no stub handlers | delete the rows of this result's Dags | | `[…]` | evaluated, these are the bindings | reconcile to this set | `None` covers the stability-check early return ([`_parse_file`](https://github.com/apache/airflow/blob/79991cd4db0c9346a28b23c453377f6df0c6b4ed/airflow-core/src/airflow/dag_processing/processor.py#L245-L251)), callback-only runs @@ -413,29 +413,34 @@ DagFileProcessorProcess(etl.py) [no DB — client con on mismatch → import_errors[etl.py]=… AND bindings=None │ ├── DagFileParsingResult(serialized_dags=[…], import_errors={…}, - ▼ task_handler_bindings=[…] | [] | None) -DagProcessorManager.persist_parsing_result [writes DB] - └── update_dag_parsing_results_in_db — one transaction, in this order: + ▼ task_handler_bindings=[…] | [] | None, + probed_artifacts=[…]) +DagProcessorManager.handle_parsing_result [writes DB] + ├── record the probed artifacts, a transaction of their own, skipped when there are none ← new + │ lang_sdk_task_handler_artifact UPSERT by (bundle_name, relative_fileloc): + │ fingerprint, answer, last_probed_at + │ + └── persist_parsing_result → update_dag_parsing_results_in_db — one transaction, in this order: 1. add_dags / update_dags → DagModel rows exist 2. asset reference tables (existing) - 3. lang_sdk_task_handler_artifact UPSERT by (bundle_name, relative_fileloc) ← new - 4. lang_sdk_task_handler reconcile by dag_id ← new - 5. SerializedDagModel / DagVersion / DagCode - 6. ParseImportError / DagWarning + 3. lang_sdk_task_handler reconcile by dag_id, artifact ids looked up by key ← new + 4. SerializedDagModel / DagVersion / DagCode + 5. ParseImportError / DagWarning ``` -Step 3 must upsert: two parse children can discover the same artifact in the same loop and race on -the unique key. +The probed-artifact write must upsert: two parse children can probe the same artifact in the same loop and race on the unique key. It commits before the persist starts, so its exclusive row locks are gone before step 3 takes shared ones, and an answer is kept even when the persist fails. -Step 4 reconciles **by the `dag_id`s in the result**, not by file path: delete rows for those +Step 3 reconciles **by the `dag_id`s in the result**, not by file path: delete rows for those `dag_id`s whose `task_id` is absent from the returned set, then insert or update the rest. Path-keyed eviction breaks when a Dag moves between files — the old rows stay keyed to a path nothing parses any more, and the primary key then blocks the new insert. With `dag_id` as the key a move simply updates -`dag_relative_fileloc`, and a Dag that disappears entirely is reclaimed by the `ON DELETE CASCADE`. +`dag_relative_fileloc`. A Dag removed from its file is only marked stale, so its rows stay until its `dag` row is deleted, and the `ON DELETE CASCADE` then removes them. Until then they keep their artifacts from the orphan sweep. + +Step 3 writes nothing to the artifact table. A Dag bound to an artifact outside the Dag file's scope, or to one with no row because a concurrent orphan sweep deleted it, keeps its rows as they are, with a warning; in the second case the next parse finds no recorded answer and probes again. Artifact rows are **never** evicted from one file's result. One artifact backs handlers owned by many Python files, so this file seeing fewer candidates says nothing about another file's. They are a -cache; they are reclaimed by orphan sweep or by `db clean`, never by a per-file reconcile. +cache; they are reclaimed by orphan sweep, never by a per-file reconcile. ### Flow 2) Scheduling: the read path @@ -557,8 +562,8 @@ argument error instead of catching it as an import error. ### Failure handling -**Validation fails.** Record the import error; leave every existing row untouched. `bindings=None` -already expresses "do not reconcile", so this needs no additional mechanism. +**Validation fails.** Record the import error; leave every binding row untouched. `bindings=None` +already expresses "do not reconcile", so this needs no additional mechanism. The answers probed for it are still recorded, so the next parse validates against them without probing. **A stub task has no binding at all.** The scheduler fails it with the reason rather than queueing a workload that would die on the worker at `ValueError("dag_path is required")`, far from the cause. @@ -578,7 +583,7 @@ definition whose author can act) naming both artifact paths, since the fix is in that owns rows in it; `lang_sdk_task_handler_artifact` is a cache with no per-file eviction. - The artifact bundle is a second bundle on the execution path. `ExecuteTask` and `StartupDetails` each grow one optional `SDKTaskHandlerRef`, and the worker performs a second `initialize()`. Workload payloads grow by roughly one `BundleInfo`. -- `DagFileParseRequest` and `DagFileParsingResult` each gain a field, and `ToSDKTaskHandlerProcessor` +- `DagFileParseRequest` gains a field and `DagFileParsingResult` two, and `ToSDKTaskHandlerProcessor` becomes a fifth union the supervisor-schema registry introspects. Both messages already appear in the generated schemas of all three SDKs, so the snapshot is regenerated and the two prek hooks guarding it run. - Every Lang SDK runtime must answer `SDKTaskHandlerParseRequest`. From 36e5a7ac1a3b630875af00d3514f4f70db6749a1 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 19:12:22 +0000 Subject: [PATCH 09/12] Give the task handler race tests their own sessions create_session() hands back the thread's scoped session, so the session standing in for a second Dag processor shared the test's transaction: on Postgres the sweep then deleted a row the "parse" had locked, because the lock was its own. A non-scoped session is a separate transaction, as a second processor would be. --- airflow-core/tests/unit/dag_processing/test_collection.py | 2 +- airflow-core/tests/unit/dag_processing/test_manager.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py b/airflow-core/tests/unit/dag_processing/test_collection.py index 7bba53b577997..75e552693db4d 100644 --- a/airflow-core/tests/unit/dag_processing/test_collection.py +++ b/airflow-core/tests/unit/dag_processing/test_collection.py @@ -2318,7 +2318,7 @@ def test_duplicate_artifact_is_written_once(self, session): def test_artifact_recorded_by_a_concurrent_parse(self, session): """Two Dag processors probe one new artifact; the second upsert must update, not fail.""" - with create_session() as other: + with create_session(scoped=False) as other: _record(other, "etl.jar") other_id = other.scalar(select(LangSDKTaskHandlerArtifact.id)) diff --git a/airflow-core/tests/unit/dag_processing/test_manager.py b/airflow-core/tests/unit/dag_processing/test_manager.py index 07df6b113e1c4..4095cc217db44 100644 --- a/airflow-core/tests/unit/dag_processing/test_manager.py +++ b/airflow-core/tests/unit/dag_processing/test_manager.py @@ -45,6 +45,7 @@ from sqlalchemy.exc import OperationalError from uuid6 import uuid7 +from airflow import settings from airflow._shared.timezones import timezone from airflow.callbacks.callback_requests import DagCallbackRequest from airflow.dag_processing.bundles.base import BaseDagBundle, BundleVersion @@ -4711,7 +4712,7 @@ def test_task_handler_artifact_sweep_logs_errors(self, delete_unreferenced, capl assert manager._last_task_handler_artifact_sweep_time > 0 @mock.patch("airflow.dag_processing.manager._TASK_HANDLER_ARTIFACT_SWEEP_BATCH_SIZE", 2) - def test_delete_unreferenced_task_handler_artifacts_commits_each_batch(self, session): + def test_delete_unreferenced_task_handler_artifacts_commits_each_batch(self): self._add_artifacts( "a.jar", "b.jar", @@ -4731,7 +4732,7 @@ def on_execute(conn, cursor, statement, parameters, context, executemany): def on_commit(conn): events.append("COMMIT") - bind = session.get_bind() + bind = settings.engine event.listen(bind, "before_cursor_execute", on_execute) event.listen(bind, "commit", on_commit) try: @@ -4797,7 +4798,7 @@ def test_task_handler_artifact_of_a_deleted_dag_is_swept(self, time_machine): def test_task_handler_artifact_sweep_skips_rows_a_parse_holds(self): self._add_artifacts("held.jar", "free.jar", probed_at=timezone.utcnow() - timedelta(days=1)) - with create_session() as parse: + with create_session(scoped=False) as parse: _find_task_handler_artifacts( [(JAVA_TASK_HANDLERS, compute_fileloc_hash("held.jar"))], session=parse ) From 3cccb23d8e489052304928b720348153eff6fcdc Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 19:39:18 +0000 Subject: [PATCH 10/12] Warn only about bindings of Dags missing from the parse result A Dag rejected for another team's plugin class already gets an import error on every parse, so warning about its dropped bindings as well was noise that buried the case worth a warning: a binding for a Dag the parse never produced, which points at a bug in the parse. --- .../src/airflow/dag_processing/collection.py | 16 ++++++++----- .../unit/dag_processing/test_collection.py | 23 ++++++++++++++++++- 2 files changed, 32 insertions(+), 7 deletions(-) diff --git a/airflow-core/src/airflow/dag_processing/collection.py b/airflow-core/src/airflow/dag_processing/collection.py index 896dec91ed489..d9bb2c3954e43 100644 --- a/airflow-core/src/airflow/dag_processing/collection.py +++ b/airflow-core/src/airflow/dag_processing/collection.py @@ -730,6 +730,7 @@ def _reconcile_task_handler_bindings( bindings: Iterable[TaskHandlerBinding], *, dag_ids: Collection[str], + rejected_dag_ids: Collection[str], artifact_bundle_names: Collection[str], dag_bundle_name: str, dag_relative_fileloc: str, @@ -742,18 +743,19 @@ def _reconcile_task_handler_bindings( that moved to another file, which that file's parse reconciles. A Dag bound to an artifact outside *artifact_bundle_names*, or to one with no recorded row, keeps its rows as they are. Artifact rows are only read: the probed-artifact write records them, and the Dag processor's orphan sweep deletes them. + Bindings of *rejected_dag_ids*, Dags the parse produced but that are not persisted, are dropped + quietly: their import error explains why. """ by_task: dict[tuple[str, str], list[TaskHandlerBinding]] = defaultdict(list) - ignored_dag_ids: set[str] = set() + unknown_dag_ids: set[str] = set() for binding in bindings: if binding.dag_id in dag_ids: by_task[(binding.dag_id, binding.task_id)].append(binding) - else: - ignored_dag_ids.add(binding.dag_id) - if ignored_dag_ids: + elif binding.dag_id not in rejected_dag_ids: + unknown_dag_ids.add(binding.dag_id) + if unknown_dag_ids: log.warning( - "Ignoring task handler bindings of Dags this parse does not persist", - dag_ids=sorted(ignored_dag_ids), + "Ignoring task handler bindings of Dags not in the parse result", dag_ids=sorted(unknown_dag_ids) ) # The parse reports a task bound twice as an import error; this only keeps the rows as they are. if conflicting := {dag_id for (dag_id, _), group in by_task.items() if len(group) > 1}: @@ -848,6 +850,7 @@ def update_dag_parsing_results_in_db( if task_handler_bindings is not None and task_handler_artifact_bundles is None: raise ValueError("task_handler_artifact_bundles is required with task_handler_bindings") accepted = _reject_other_teams_plugin_classes(bundle_name, dags, import_errors, session=session) + rejected_ids: set[str] = set() if len(accepted) != len(dags): # A rejected Dag may have no ``dag`` row yet, and dag_warning has a foreign key to it. rejected_ids = {dag.dag_id for dag in dags} - {dag.dag_id for dag in accepted} @@ -891,6 +894,7 @@ def update_dag_parsing_results_in_db( _reconcile_task_handler_bindings( task_handler_bindings, dag_ids={dag.dag_id for dag in dags}, + rejected_dag_ids=rejected_ids, artifact_bundle_names=task_handler_artifact_bundles, dag_bundle_name=bundle_name, dag_relative_fileloc=relative_fileloc, diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py b/airflow-core/tests/unit/dag_processing/test_collection.py index 75e552693db4d..126b1364e93fa 100644 --- a/airflow-core/tests/unit/dag_processing/test_collection.py +++ b/airflow-core/tests/unit/dag_processing/test_collection.py @@ -1961,6 +1961,27 @@ def test_team_class_is_only_stored_for_its_team( assert list(errors) == [(bundle_name, "team_dag.py")] assert f"belonging to {plugin_team}" in errors[(bundle_name, "team_dag.py")] + @conf_vars({("core", "multi_team"): "True"}) + def test_task_handler_bindings_of_a_rejected_dag_are_dropped_quietly(self, bundle, session, caplog): + bundle_name = bundle(True) + _record(session, "etl.jar") + with mock_plugin_manager(plugins=[self._plugin("other_team", "timetables", AfterWorkdayTimetable)]): + update_dag_parsing_results_in_db( + bundle_name=bundle_name, + bundle_version=None, + dags=[self._serialized(schedule=AfterWorkdayTimetable()), self._serialized("plain_dag")], + import_errors={}, + parse_duration=None, + warnings=set(), + session=session, + relative_fileloc="team_dag.py", + task_handler_bindings=[_make_binding("team_dag"), _make_binding("plain_dag")], + task_handler_artifact_bundles={ARTIFACT_BUNDLE}, + ) + + assert session.scalars(select(LangSDKTaskHandler.dag_id)).all() == ["plain_dag"] + assert not any("task handler bindings" in entry["event"] for entry in caplog.entries) + @conf_vars({("core", "multi_team"): "True"}) @pytest.mark.parametrize( "weight_rule", @@ -2478,7 +2499,7 @@ def test_task_handler_bindings_for_dags_outside_the_result_are_dropped(self, ses assert _get_recorded_handlers(session) == {("etl", "extract", "etl.jar")} assert { - "event": "Ignoring task handler bindings of Dags this parse does not persist", + "event": "Ignoring task handler bindings of Dags not in the parse result", "dag_ids": ["ghost"], } in caplog From 459da041c11acbdb25eb7fbde6d53ccd4b15da1a Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Wed, 30 Sep 2026 19:40:40 +0000 Subject: [PATCH 11/12] Align the artifact sweep with the other cleanup checks The sweep's interval check now reads like the bundle-version cleanup next to it, and its deleted-row count no longer hides a missing rowcount behind a default that can never apply. The parsing_cleanup_interval description said the artifact deletion in a clause that did not parse cleanly. --- airflow-core/src/airflow/config_templates/config.yml | 4 ++-- airflow-core/src/airflow/dag_processing/manager.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/config_templates/config.yml b/airflow-core/src/airflow/config_templates/config.yml index 595c38f0ccb1e..d11111a7cb8a5 100644 --- a/airflow-core/src/airflow/config_templates/config.yml +++ b/airflow-core/src/airflow/config_templates/config.yml @@ -2727,8 +2727,8 @@ scheduler: description: | How often (in seconds) to check for stale DAGs (DAGs which are no longer present in the expected files) which should be deactivated, as well as assets that are no longer - referenced and should be marked as orphaned, and Lang-SDK task handler artifacts that no - stub task uses any more, which are deleted. + referenced and should be marked as orphaned. The same check deletes the Lang-SDK task + handler artifacts that no stub task uses any more. version_added: 2.5.0 type: integer example: ~ diff --git a/airflow-core/src/airflow/dag_processing/manager.py b/airflow-core/src/airflow/dag_processing/manager.py index fb64d1781fdf7..1d8dfa6417bac 100644 --- a/airflow-core/src/airflow/dag_processing/manager.py +++ b/airflow-core/src/airflow/dag_processing/manager.py @@ -97,6 +97,7 @@ from collections.abc import Callable, Collection, Iterable, Iterator, Sequence from socket import socket + from sqlalchemy.engine import CursorResult from sqlalchemy.orm import Session from sqlalchemy.sql import Select @@ -462,7 +463,7 @@ def _scan_stale_dags(self): def _sweep_task_handler_artifacts(self) -> None: now = time.monotonic() - if now - self._last_task_handler_artifact_sweep_time <= self.parsing_cleanup_interval: + if now - self._last_task_handler_artifact_sweep_time < self.parsing_cleanup_interval: return try: self.delete_unreferenced_task_handler_artifacts() @@ -508,7 +509,7 @@ def delete_unreferenced_task_handler_artifacts(self) -> int: .where(LangSDKTaskHandlerArtifact.id.in_(ids), *sweepable) .execution_options(synchronize_session=False) ) - deleted += getattr(result, "rowcount", 0) + deleted += cast("CursorResult", result).rowcount if len(ids) < _TASK_HANDLER_ARTIFACT_SWEEP_BATCH_SIZE: break if deleted: From aa3f77152c53d202c8b0e66f2c4c4c38b66a9fc2 Mon Sep 17 00:00:00 2001 From: LIU ZHE YOU Date: Thu, 1 Oct 2026 17:50:29 +0000 Subject: [PATCH 12/12] Build the binding and probed-artifact test params without required TaskHandlerParam no longer has required. A model dump no longer carries it, so the expected result bodies must drop it, and the constructors would otherwise pass a keyword pydantic only ignores. --- airflow-core/tests/unit/dag_processing/test_collection.py | 4 +--- airflow-core/tests/unit/dag_processing/test_processor.py | 2 +- .../tests/task_sdk/execution_time/schema/test_migrator.py | 6 ++---- 3 files changed, 4 insertions(+), 8 deletions(-) diff --git a/airflow-core/tests/unit/dag_processing/test_collection.py b/airflow-core/tests/unit/dag_processing/test_collection.py index 126b1364e93fa..1be65b46f6dcb 100644 --- a/airflow-core/tests/unit/dag_processing/test_collection.py +++ b/airflow-core/tests/unit/dag_processing/test_collection.py @@ -2149,9 +2149,7 @@ class WorkdayPlugin(AirflowPlugin): OTHER_TEAM_BUNDLE = "go-task-handlers" TASK_HANDLERS = { "etl": [ - TaskHandlerDeclaration( - task_id="extract", binding="positional", params=[TaskHandlerParam(name=None, required=True)] - ) + TaskHandlerDeclaration(task_id="extract", binding="positional", params=[TaskHandlerParam(name=None)]) ] } diff --git a/airflow-core/tests/unit/dag_processing/test_processor.py b/airflow-core/tests/unit/dag_processing/test_processor.py index 8dff0a31c3a6e..bd02675b3ae30 100644 --- a/airflow-core/tests/unit/dag_processing/test_processor.py +++ b/airflow-core/tests/unit/dag_processing/test_processor.py @@ -2545,7 +2545,7 @@ def test_parsing_result_decodes_probed_artifacts(self): TaskHandlerDeclaration( task_id="extract", binding="positional", - params=[TaskHandlerParam(name=None, required=True)], + params=[TaskHandlerParam(name=None)], ) ] }, diff --git a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py index 4163c3a93201f..c694fde8ad7e4 100644 --- a/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py +++ b/task-sdk/tests/task_sdk/execution_time/schema/test_migrator.py @@ -657,7 +657,7 @@ def result_with_probed_artifacts(self) -> DagFileParsingResult: TaskHandlerDeclaration( task_id="extract", binding="positional", - params=[TaskHandlerParam(name=None, required=True)], + params=[TaskHandlerParam(name=None)], ) ] }, @@ -688,9 +688,7 @@ def test_downgrade_keeps_probed_artifacts_at_head(self, real_migrator, result_wi { "task_id": "extract", "binding": "positional", - "params": [ - {"name": None, "value_schema": None, "required": True, "exact_name": False} - ], + "params": [{"name": None, "value_schema": None, "exact_name": False}], } ] },