From 566a581953604ab06195c781a5c505513922eb6b Mon Sep 17 00:00:00 2001 From: Pedrinhonitz Date: Mon, 31 Aug 2026 17:32:12 -0300 Subject: [PATCH 1/8] fix(core): test connection api ui solved #72318 --- .../core_api/routes/public/connections.py | 37 +++++++++-- .../routes/public/test_connections.py | 62 +++++++++++++++++++ 2 files changed, 95 insertions(+), 4 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index 19378dd8aeaaf..1bc83edac628e 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -104,6 +104,35 @@ def _ensure_executor_is_configured(executor: str | None) -> None: ) +_MASKED_CREDENTIAL_SENTINEL = "***" + + +def _same_endpoint(requested: str | int | None, stored: str | int | None) -> bool: + """Return True when request and stored host/port refer to the same destination. + + The UI sends empty string for hidden unused host/port fields; the ORM stores + those as NULL. Treat blank as unset so connection types that do not use + host/port still reuse stored credentials. + """ + + def _norm(value: str | int | None) -> str | int | None: + return None if value is None or value == "" else value + + return _norm(requested) == _norm(stored) + + +def _supplies_own_credentials(test_body: ConnectionBody) -> bool: + """Return True when the request includes a real (non-masked) password. + + The UI always posts the masked sentinel for unchanged secrets. That is not + a caller-supplied credential and must not skip restoring stored extras. + """ + if "password" not in test_body.model_fields_set: + return False + password = test_body.password + return bool(password) and password != _MASKED_CREDENTIAL_SENTINEL + + @connections_router.delete( "/{connection_id}", status_code=status.HTTP_204_NO_CONTENT, @@ -353,10 +382,10 @@ def test_connection( # Stored credentials are only reused to test the connection's own # host/port; testing a different destination must supply its own. fields_set = test_body.model_fields_set - if ("host" in fields_set and test_body.host != existing_conn.host) or ( - "port" in fields_set and test_body.port != existing_conn.port - ): - if "password" not in fields_set: + host_changed = "host" in fields_set and not _same_endpoint(test_body.host, existing_conn.host) + port_changed = "port" in fields_set and not _same_endpoint(test_body.port, existing_conn.port) + if host_changed or port_changed: + if not _supplies_own_credentials(test_body): raise HTTPException( status.HTTP_400_BAD_REQUEST, "The host or port to test differs from the stored connection. " diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py index 231a7ca7b22ef..a9822e5afaff9 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_connections.py @@ -1413,6 +1413,68 @@ def test_stored_secret_reused_only_for_same_target( tested_connection = mock_test.call_args.args[0] assert tested_connection.password == expected_password + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_should_reuse_stored_extra_when_host_and_port_are_blank(self, test_client, session): + """Hidden unused host/port (None vs "") must not skip restoring masked extra.""" + stored_path = "/real.pem" + session.add( + Connection( + conn_id=TEST_CONN_ID, + conn_type="snowflake", + host=None, + port=None, + extra=json.dumps({"private_key_file": stored_path, "account": "acct"}), + ) + ) + session.commit() + + captured = {} + + def mock_test_connection(self): + captured["extra"] = self.extra + return True, "mocked" + + body = { + "connection_id": TEST_CONN_ID, + "conn_type": "snowflake", + "host": "", + "password": "***", + "extra": json.dumps({"private_key_file": "***", "account": "acct"}), + } + + with mock.patch.object(Connection, "test_connection", mock_test_connection): + response = test_client.post("/connections/test", json=body) + + assert response.status_code == 200 + assert json.loads(captured["extra"])["private_key_file"] == stored_path + + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) + def test_should_reject_overridden_target_when_password_is_masked(self, test_client, session): + """A masked password is not caller-supplied credentials for a new destination.""" + session.add( + Connection( + conn_id=TEST_CONN_ID, + conn_type="sqlite", + host="stored_host", + port=1234, + password="existing_password", + ) + ) + session.commit() + + body = { + "connection_id": TEST_CONN_ID, + "conn_type": "sqlite", + "host": "other_host", + "password": "***", + } + with mock.patch.object(Connection, "test_connection", autospec=True) as mock_test: + mock_test.return_value = (True, "mocked") + response = test_client.post("/connections/test", json=body) + + assert response.status_code == 400 + mock_test.assert_not_called() + @mock.patch.dict(os.environ, {"AIRFLOW__CORE__TEST_CONNECTION": "Enabled"}) def test_should_test_new_connection_without_existing(self, test_client): body = { From 27fef7b25686c4d7900dca0fdd4e8fc651bb7679 Mon Sep 17 00:00:00 2001 From: Pedrinhonitz Date: Mon, 31 Aug 2026 19:35:48 -0300 Subject: [PATCH 2/8] fix(ruff): format docstring --- .../api_fastapi/core_api/routes/public/connections.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index 1bc83edac628e..e828b932dea8e 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -108,7 +108,8 @@ def _ensure_executor_is_configured(executor: str | None) -> None: def _same_endpoint(requested: str | int | None, stored: str | int | None) -> bool: - """Return True when request and stored host/port refer to the same destination. + """ + Return True when request and stored host/port refer to the same destination. The UI sends empty string for hidden unused host/port fields; the ORM stores those as NULL. Treat blank as unset so connection types that do not use @@ -122,7 +123,8 @@ def _norm(value: str | int | None) -> str | int | None: def _supplies_own_credentials(test_body: ConnectionBody) -> bool: - """Return True when the request includes a real (non-masked) password. + """ + Return True when the request includes a real (non-masked) password. The UI always posts the masked sentinel for unchanged secrets. That is not a caller-supplied credential and must not skip restoring stored extras. From 453055601a5cae3ed3f3f2552af72bccc83a2d24 Mon Sep 17 00:00:00 2001 From: Pedrinhonitz Date: Mon, 31 Aug 2026 19:35:48 -0300 Subject: [PATCH 3/8] fix connection test with null host and port --- .../api_fastapi/core_api/routes/public/connections.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py index 1bc83edac628e..e828b932dea8e 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/connections.py @@ -108,7 +108,8 @@ def _ensure_executor_is_configured(executor: str | None) -> None: def _same_endpoint(requested: str | int | None, stored: str | int | None) -> bool: - """Return True when request and stored host/port refer to the same destination. + """ + Return True when request and stored host/port refer to the same destination. The UI sends empty string for hidden unused host/port fields; the ORM stores those as NULL. Treat blank as unset so connection types that do not use @@ -122,7 +123,8 @@ def _norm(value: str | int | None) -> str | int | None: def _supplies_own_credentials(test_body: ConnectionBody) -> bool: - """Return True when the request includes a real (non-masked) password. + """ + Return True when the request includes a real (non-masked) password. The UI always posts the masked sentinel for unchanged secrets. That is not a caller-supplied credential and must not skip restoring stored extras. From 31635d43d786875e2dcd72c4d57644808465522a Mon Sep 17 00:00:00 2001 From: Pedrinhonitz Date: Fri, 18 Sep 2026 10:52:08 -0300 Subject: [PATCH 4/8] Ignore stale executor success with per-invocation workload_run_id --- .../src/airflow/executors/base_executor.py | 80 +++++++++++++++-- .../src/airflow/executors/local_executor.py | 12 +-- .../src/airflow/executors/workloads/task.py | 1 + .../src/airflow/executors/workloads/types.py | 3 +- .../src/airflow/jobs/scheduler_job_runner.py | 64 +++++++++----- ..._0_add_workload_run_id_to_task_instance.py | 56 ++++++++++++ .../src/airflow/models/taskinstance.py | 5 ++ .../src/airflow/models/taskinstancehistory.py | 1 + .../unit/executors/test_base_executor.py | 24 +++-- .../tests/unit/jobs/test_scheduler_job.py | 87 +++++++++++++++++++ .../tests_common/test_utils/mock_executor.py | 8 +- .../celery/executors/celery_executor.py | 19 +++- .../executors/kubernetes_executor.py | 10 ++- 13 files changed, 317 insertions(+), 53 deletions(-) create mode 100644 airflow-core/src/airflow/migrations/versions/0135_3_4_0_add_workload_run_id_to_task_instance.py diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index ef5fafded8b39..49e762cd1245b 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -85,8 +85,8 @@ def get_execution_api_server_url(conf_source: AirflowConfigParser | ExecutorConf from airflow.models.taskinstance import TaskInstance # Event_buffer dict value type - # Tuple of: state, info - EventBufferValueType = tuple[str | None, Any] + # Tuple of: state, info, workload_run_id (optional; None when unknown/legacy) + EventBufferValueType = tuple[Any, Any, str | None] log = logging.getLogger(__name__) @@ -227,6 +227,10 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None) self.queued_callbacks: dict[CallbackKey, workloads.ExecuteCallback] = {} self.queued_connection_tests: dict[ConnectionTestKey, workloads.TestConnection] = {} self.running: set[WorkloadKey] = set() + # FIFO of workload_run_id values per TaskInstanceKey, appended on queue and + # consumed when a terminal SUCCESS/FAILED event is recorded. Survives a + # re-enqueue of the same key so a stale defer-exit SUCCESS keeps its old id. + self._workload_run_ids: dict[TaskInstanceKey, deque[str]] = defaultdict(deque) self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {} self._task_event_logs: deque[Log] = deque() self.conf = ExecutorConf(team_name) @@ -266,6 +270,8 @@ def queue_workload(self, workload: ExecutorWorkload, session: Session) -> None: if isinstance(workload, workloads.ExecuteTask): ti = workload.ti self.queued_tasks[ti.key] = workload + if ti.workload_run_id: + self._workload_run_ids[ti.key].append(ti.workload_run_id) elif isinstance(workload, workloads.ExecuteCallback): if not self.supports_callbacks: raise NotImplementedError( @@ -288,6 +294,53 @@ def queue_workload(self, workload: ExecutorWorkload, session: Session) -> None: f"Workload must be one of: ExecuteTask, ExecuteCallback, TestConnection." ) + def _pop_workload_run_id(self, key: WorkloadKey) -> str | None: + """Consume the oldest queued workload_run_id for a task key, if any.""" + if not isinstance(key, TaskInstanceKey): + return None + pending = self._workload_run_ids.get(key) + if not pending: + return None + run_id = pending.popleft() + if not pending: + del self._workload_run_ids[key] + return run_id + + @staticmethod + def unpack_event(value: EventBufferValueType | tuple[Any, ...]) -> tuple[Any, Any, str | None]: + """Normalize event buffer values to ``(state, info, workload_run_id)``.""" + if len(value) >= 3: + return value[0], value[1], value[2] + return value[0], value[1], None + + def record_event( + self, + key: WorkloadKey, + state: WorkloadState, + info=None, + *, + workload_run_id: str | None = None, + consume_run_id: bool | None = None, + ) -> None: + """ + Write an executor event, optionally attaching a workload_run_id. + + Terminal task SUCCESS/FAILED events consume the next queued run id when + one was not supplied explicitly. Intermediate states (QUEUED/RUNNING) + must not consume it — those events are drained separately. + """ + from airflow.utils.state import TaskInstanceState + + if consume_run_id is None: + consume_run_id = isinstance(key, TaskInstanceKey) and state in ( + TaskInstanceState.SUCCESS, + TaskInstanceState.FAILED, + ) + popped_run_id = self._pop_workload_run_id(key) if consume_run_id else None + if workload_run_id is None: + workload_run_id = popped_run_id + self.event_buffer[key] = (state, info, workload_run_id) + def _get_workloads_to_schedule(self, open_slots: int) -> list[tuple[WorkloadKey, ExecutorWorkload]]: """ Select and return the next batch of workloads to schedule, respecting priority policy. @@ -478,7 +531,14 @@ def trigger_tasks(self, open_slots: int) -> None: # TODO: This should not be using `TaskInstanceState` here, this is just "did the process complete, or did # it die". It is possible for the task itself to finish with success, but the state of the task to be set # to FAILED. By using TaskInstanceState enum here it confuses matters! - def change_state(self, key: WorkloadKey, state: WorkloadState, info=None, remove_running=True) -> None: + def change_state( + self, + key: WorkloadKey, + state: WorkloadState, + info=None, + remove_running=True, + workload_run_id: str | None = None, + ) -> None: """ Change state of the task. @@ -486,6 +546,8 @@ def change_state(self, key: WorkloadKey, state: WorkloadState, info=None, remove :param state: State to set for the task. :param info: Executor information for the task instance :param remove_running: Whether or not to remove the TI key from running set + :param workload_run_id: Invocation id for the finished worker; when omitted for + terminal task events, the next id queued for this key is consumed. """ self.log.debug("Changing state: %s", key) if remove_running: @@ -493,25 +555,27 @@ def change_state(self, key: WorkloadKey, state: WorkloadState, info=None, remove self.running.remove(key) except KeyError: self.log.debug("Could not find key: %s", key) - self.event_buffer[key] = state, info + self.record_event(key, state, info, workload_run_id=workload_run_id) - def fail(self, key: WorkloadKey, info=None) -> None: + def fail(self, key: WorkloadKey, info=None, workload_run_id: str | None = None) -> None: """ Set fail state for the event. :param info: Executor information for the task instance :param key: Unique key for the task instance + :param workload_run_id: Invocation id for the finished worker, if known """ - self.change_state(key, state_class_for_key(key).FAILED, info) + self.change_state(key, state_class_for_key(key).FAILED, info, workload_run_id=workload_run_id) - def success(self, key: WorkloadKey, info=None) -> None: + def success(self, key: WorkloadKey, info=None, workload_run_id: str | None = None) -> None: """ Set success state for the event. :param info: Executor information for the task instance :param key: Unique key for the task instance + :param workload_run_id: Invocation id for the finished worker, if known """ - self.change_state(key, state_class_for_key(key).SUCCESS, info) + self.change_state(key, state_class_for_key(key).SUCCESS, info, workload_run_id=workload_run_id) def queued(self, key: WorkloadKey, info=None) -> None: """ diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 24ab737f5ab20..9790a0e98fc2c 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -97,7 +97,7 @@ def _run_worker( unread_messages.value -= 1 if workload.running_state is not None: - output.put((workload.key, workload.running_state, None)) + output.put((workload.key, workload.running_state, None, None)) try: BaseExecutor.run_workload( @@ -106,10 +106,12 @@ def _run_worker( proctitle=f"{_get_executor_process_title_prefix(team_conf.team_name)} {workload.display_name}", subprocess_logs_to_stdout=True, ) - output.put((workload.key, workload.success_state, None)) + run_id = getattr(getattr(workload, "ti", None), "workload_run_id", None) + output.put((workload.key, workload.success_state, None, run_id)) except Exception as e: log.exception("Workload execution failed.", workload_type=type(workload).__name__) - output.put((workload.key, workload.failure_state, e)) + run_id = getattr(getattr(workload, "ti", None), "workload_run_id", None) + output.put((workload.key, workload.failure_state, e, run_id)) class LocalExecutor(BaseExecutor): @@ -246,8 +248,8 @@ def sync(self) -> None: def _read_results(self): try: while not self.result_queue.empty(): - key, state, exc = self.result_queue.get() - self.change_state(key, state) + key, state, exc, workload_run_id = self.result_queue.get() + self.change_state(key, state, workload_run_id=workload_run_id) except (OSError, EOFError): self.log.exception("Error reading from result queue") diff --git a/airflow-core/src/airflow/executors/workloads/task.py b/airflow-core/src/airflow/executors/workloads/task.py index 3099fe1d77485..51968a499d353 100644 --- a/airflow-core/src/airflow/executors/workloads/task.py +++ b/airflow-core/src/airflow/executors/workloads/task.py @@ -46,6 +46,7 @@ class TaskInstanceDTO(TaskInstance): priority_weight: int external_executor_id: str | None = Field(default=None, exclude=True) + workload_run_id: str | None = Field(default=None, exclude=True) executor_config: dict | None = Field(default=None, exclude=True) # TODO: Task-SDK: Can we replace TaskInstanceKey with just the uuid across the codebase? diff --git a/airflow-core/src/airflow/executors/workloads/types.py b/airflow-core/src/airflow/executors/workloads/types.py index 3e5d0b06f10ac..c37f8feb9f43a 100644 --- a/airflow-core/src/airflow/executors/workloads/types.py +++ b/airflow-core/src/airflow/executors/workloads/types.py @@ -32,7 +32,8 @@ WorkloadState: TypeAlias = TaskInstanceState | CallbackState | ConnectionTestState # Type alias for executor workload results (used by executor implementations) - WorkloadResultType: TypeAlias = tuple[WorkloadKey, WorkloadState, Exception | None] + # key, state, exception, optional workload_run_id from the finished invocation + WorkloadResultType: TypeAlias = tuple[WorkloadKey, WorkloadState, Exception | None, str | None] # Type alias for scheduler workloads (ORM models that can be routed to executors) # Must be outside TYPE_CHECKING for use in function signatures diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 90aefb30be99d..a5f2267593870 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1081,6 +1081,9 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - "state": TaskInstanceState.QUEUED, "queued_dttm": timezone.utcnow(), "queued_by_job_id": self.job.id, + # Per-invocation id so a stale executor SUCCESS from a previous worker + # (e.g. defer exit) cannot be matched to a later enqueue of the same key. + "workload_run_id": random_db_uuid(), } # Pre-assign external_executor_id atomically with the QUEUED state so it @@ -1116,25 +1119,26 @@ def _executable_task_instances_to_queued(self, max_tis: int, session: Session) - .execution_options(synchronize_session=False) ) + # Always read workload_run_id (and external_executor_id when pre-assigned) + # back onto in-memory objects so ExecuteTask.make carries them through + # make_transient. Use RETURNING on PostgreSQL; SELECT elsewhere. + returning_cols = [TI.id, TI.workload_run_id] if pre_assign_executors: - # Read the DB-generated UUIDs back onto the in-memory objects so the - # workload DTO carries them through to send_workload_to_executor (the - # objects are about to be detached by make_transient). Use RETURNING - # where supported (PostgreSQL); fall back to a SELECT for MySQL and - # SQLite (RETURNING requires SQLite 3.35+ which isn't guaranteed). - if get_dialect_name(session) == "postgresql": - result = session.execute(queued_update.returning(TI.id, TI.external_executor_id)) - id_map = {row[0]: row[1] for row in result} - else: - session.execute(queued_update) - id_rows = session.execute( - select(TI.id, TI.external_executor_id).where(filter_for_tis) - ).all() - id_map = {row[0]: row[1] for row in id_rows} - for ti in executable_tis: - ti.external_executor_id = id_map.get(ti.id) + returning_cols.append(TI.external_executor_id) + + if get_dialect_name(session) == "postgresql": + result = session.execute(queued_update.returning(*returning_cols)) + rows = list(result) else: session.execute(queued_update) + rows = list(session.execute(select(*returning_cols).where(filter_for_tis)).all()) + + workload_run_id_map = {row[0]: row[1] for row in rows} + external_id_map = {row[0]: row[2] for row in rows} if pre_assign_executors else {} + for ti in executable_tis: + ti.workload_run_id = workload_run_id_map.get(ti.id) + if pre_assign_executors: + ti.external_executor_id = external_id_map.get(ti.id) for ti in executable_tis: ti.emit_state_change_metric(TaskInstanceState.QUEUED) @@ -1412,8 +1416,14 @@ def process_executor_events( tis_with_right_state: list[TaskInstanceKey] = [] callback_keys_with_events: list[CallbackKey] = [] + def _unpack_event(value: tuple) -> tuple[Any, Any, str | None]: + if len(value) >= 3: + return value[0], value[1], value[2] + return value[0], value[1], None + # Report execution - handle both task and callback events - for key, (state, _) in event_buffer.items(): + for key, event_value in event_buffer.items(): + state, _, _ = _unpack_event(event_value) if isinstance(key, TaskInstanceKey): existing_try = ti_primary_key_to_try_number_map.get(key.primary) if existing_try is not None and existing_try != key.try_number: @@ -1447,7 +1457,7 @@ def process_executor_events( # Handle callback state events for callback_id in callback_keys_with_events: - state, info = event_buffer.pop(callback_id) + state, info, _ = _unpack_event(event_buffer.pop(callback_id)) callback = session.get(Callback, UUID(str(callback_id))) if not callback: # This should not normally happen - we just received an event for this callback. @@ -1510,7 +1520,7 @@ def process_executor_events( ti.state, job_id, ) - state, info = event_buffer.pop(buffer_key) + state, info, event_workload_run_id = _unpack_event(event_buffer.pop(buffer_key)) if state in (TaskInstanceState.QUEUED, TaskInstanceState.RUNNING): ti.external_executor_id = info @@ -1559,9 +1569,8 @@ def process_executor_events( # from the worker exit after defer() has not been processed yet - should not fail it. # 4) the trigger already put the TI back to queued (resume after defer) but the executor success # from the worker exit after defer() has not been processed yet - should not fail it. - - # All of this could also happen if the state is "running", - # but that is handled by the scheduler detecting task instances without heartbeats. + # 5) the resumed attempt is already RUNNING when the stale defer-exit SUCCESS arrives — + # workload_run_id mismatch identifies this without relying on next_method/state heuristics. ti_queued = ti.try_number == buffer_key.try_number and ti.state in ( TaskInstanceState.SCHEDULED, @@ -1569,12 +1578,18 @@ def process_executor_events( TaskInstanceState.RUNNING, TaskInstanceState.RESTARTING, ) + stale_workload_run = ( + event_workload_run_id is not None + and ti.workload_run_id is not None + and event_workload_run_id != ti.workload_run_id + ) ti_requeued = ( ti.queued_by_job_id != job_id # Another scheduler has queued this task again or executor.has_task(ti) # This scheduler has this task already + or stale_workload_run or ( - # Resume-after-defer: trigger moved TI to scheduled or queued (next_method set) - # before we saw the executor success from the defer exit for the same try_number. + # Defense in depth for older events without workload_run_id: resume-after-defer + # while next_method is still set (SCHEDULED/QUEUED variants from #66431/#68741). ti.state in (TaskInstanceState.SCHEDULED, TaskInstanceState.QUEUED) and state == TaskInstanceState.SUCCESS and ti.next_method is not None @@ -3544,6 +3559,7 @@ def adopt_or_reset_orphaned_tasks(self, *, session: Session = NEW_SESSION) -> in ti.state = None ti.queued_by_job_id = None ti.external_executor_id = None + ti.workload_run_id = None ti.clear_next_method_args() for ti in set(tis_to_adopt_or_reset) - set(to_reset): diff --git a/airflow-core/src/airflow/migrations/versions/0135_3_4_0_add_workload_run_id_to_task_instance.py b/airflow-core/src/airflow/migrations/versions/0135_3_4_0_add_workload_run_id_to_task_instance.py new file mode 100644 index 0000000000000..ac6bb3ec9b246 --- /dev/null +++ b/airflow-core/src/airflow/migrations/versions/0135_3_4_0_add_workload_run_id_to_task_instance.py @@ -0,0 +1,56 @@ +# +# 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 workload_run_id to task_instance and task_instance_history. + +Per-invocation UUID generated on every scheduler enqueue so executor +completion events can be matched to the attempt that produced them. +This prevents stale SUCCESS from a defer-exit worker from failing a +resumed attempt that shares the same TaskInstanceKey. + +Revision ID: c7d4e8f1a203 +Revises: b6a9c2e7d410 +Create Date: 2026-09-18 00:00:00.000000 + +""" + +from __future__ import annotations + +import sqlalchemy as sa +from alembic import op + +revision = "c7d4e8f1a203" +down_revision = "b6a9c2e7d410" +branch_labels = None +depends_on = None +airflow_version = "3.4.0" + + +def upgrade(): + """Add workload_run_id to task_instance and task_instance_history.""" + for table in ("task_instance", "task_instance_history"): + with op.batch_alter_table(table, schema=None) as batch_op: + batch_op.add_column(sa.Column("workload_run_id", sa.Text(), nullable=True)) + + +def downgrade(): + """Remove workload_run_id from task_instance and task_instance_history.""" + for table in ("task_instance", "task_instance_history"): + with op.batch_alter_table(table, schema=None) as batch_op: + batch_op.drop_column("workload_run_id") diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 5c8559d11f4b1..f7996442df34a 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -443,6 +443,7 @@ def clear_task_instances( ti.max_tries = max(ti.max_tries, ti.try_number) ti.state = None ti.external_executor_id = None + ti.workload_run_id = None ti.clear_next_method_args() # Match DagVersion to latest serialized DAG when running on the latest version. if use_latest_version: @@ -660,6 +661,10 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload): external_executor_id: Mapped[str | None] = mapped_column(Text(), nullable=True) + # Per-invocation id generated on every enqueue. Used to ignore stale executor + # SUCCESS events from a previous worker exit (e.g. defer) after the TI was resumed. + workload_run_id: Mapped[str | None] = mapped_column(Text(), nullable=True) + # The trigger to resume on if we are in state DEFERRED trigger_id: Mapped[int | None] = mapped_column(Integer, nullable=True) diff --git a/airflow-core/src/airflow/models/taskinstancehistory.py b/airflow-core/src/airflow/models/taskinstancehistory.py index 1f741507bd5aa..0763a1c807b1c 100644 --- a/airflow-core/src/airflow/models/taskinstancehistory.py +++ b/airflow-core/src/airflow/models/taskinstancehistory.py @@ -104,6 +104,7 @@ class TaskInstanceHistory(Base): context_carrier: Mapped[dict | None] = mapped_column(MutableDict.as_mutable(ExtendedJSON), nullable=True) external_executor_id: Mapped[str | None] = mapped_column(Text(), nullable=True) + workload_run_id: Mapped[str | None] = mapped_column(Text(), nullable=True) trigger_id: Mapped[int | None] = mapped_column(Integer, nullable=True) trigger_timeout: Mapped[DateTime | None] = mapped_column(DateTime, nullable=True) next_method: Mapped[str | None] = mapped_column(String(1000), nullable=True) diff --git a/airflow-core/tests/unit/executors/test_base_executor.py b/airflow-core/tests/unit/executors/test_base_executor.py index 1d8a2dfa32936..1550c133a3e94 100644 --- a/airflow-core/tests/unit/executors/test_base_executor.py +++ b/airflow-core/tests/unit/executors/test_base_executor.py @@ -160,7 +160,7 @@ def test_state_methods_pick_callback_state_for_callback_key(method_name, expecte getattr(executor, method_name)(callback_key) - assert executor.event_buffer[callback_key] == (expected_state, None) + assert executor.event_buffer[callback_key] == (expected_state, None, None) def test_fail_and_success(): @@ -448,6 +448,20 @@ def test_running_retry_attempt_type(loop_duration, total_tries): assert a.tries_after_min == 1 +def test_success_consumes_queued_workload_run_id_fifo(): + """Re-enqueue of the same key keeps earlier run ids for stale events.""" + executor = BaseExecutor() + key = TaskInstanceKey("my_dag1", "my_task1", "run1", 1) + executor._workload_run_ids[key].append("run-a") + executor._workload_run_ids[key].append("run-b") + + executor.success(key) + assert executor.event_buffer[key] == (TaskInstanceState.SUCCESS, None, "run-a") + + executor.success(key) + assert executor.event_buffer[key] == (TaskInstanceState.SUCCESS, None, "run-b") + + def test_state_fail(): executor = BaseExecutor() key = TaskInstanceKey("my_dag1", "my_task1", timezone.utcnow(), 1) @@ -455,7 +469,7 @@ def test_state_fail(): info = "info" executor.fail(key, info=info) assert not executor.running - assert executor.event_buffer[key] == (TaskInstanceState.FAILED, info) + assert executor.event_buffer[key] == (TaskInstanceState.FAILED, info, None) def test_state_success(): @@ -465,7 +479,7 @@ def test_state_success(): info = "info" executor.success(key, info=info) assert not executor.running - assert executor.event_buffer[key] == (TaskInstanceState.SUCCESS, info) + assert executor.event_buffer[key] == (TaskInstanceState.SUCCESS, info, None) def test_state_queued(): @@ -475,7 +489,7 @@ def test_state_queued(): info = "info" executor.queued(key, info=info) assert not executor.running - assert executor.event_buffer[key] == (TaskInstanceState.QUEUED, info) + assert executor.event_buffer[key] == (TaskInstanceState.QUEUED, info, None) def test_state_running(): @@ -486,7 +500,7 @@ def test_state_running(): executor.running_state(key, info=info) # Running state should not remove a command as running assert executor.running - assert executor.event_buffer[key] == (TaskInstanceState.RUNNING, info) + assert executor.event_buffer[key] == (TaskInstanceState.RUNNING, info, None) def test_repr(): diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 4d1a540fedacb..1eb6fdfbf52c1 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -1027,6 +1027,67 @@ def test_process_executor_events_stale_success_when_queued_after_defer( tags={"dag_id": dag_id, "task_id": ti1.task_id}, ) + @pytest.mark.parametrize( + "ti_state", + [State.SCHEDULED, State.QUEUED, State.RUNNING], + ) + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") + @mock.patch("airflow._shared.observability.metrics.stats._get_backend") + def test_process_executor_events_stale_success_mismatched_workload_run_id( + self, mock_get_backend, mock_task_callback, dag_maker, ti_state + ): + """ + Stale defer-exit SUCCESS carries an older workload_run_id than the resumed attempt. + + Must not treat as state mismatch for SCHEDULED, QUEUED, or RUNNING — even when + next_method is already cleared (the RUNNING variant of #72716). + """ + mock_stats = mock.MagicMock(spec=StatsLogger) + mock_get_backend.return_value = mock_stats + dag_id = f"test_stale_success_workload_run_id_{ti_state}" + task_id_1 = "dummy_task" + + session = settings.Session() + with dag_maker(dag_id=dag_id, fileloc="/test_path1/"): + task1 = EmptyOperator(task_id=task_id_1) + ti1 = dag_maker.create_dagrun().get_task_instance(task1.task_id) + + executor = MockExecutor(do_update=False) + mock_task_callback.return_value = mock.MagicMock() + scheduler_job = Job() + session.add(scheduler_job) + session.flush() + self.job_runner = SchedulerJobRunner(scheduler_job, executors=[executor]) + + ti1.state = ti_state + ti1.next_method = None + ti1.queued_by_job_id = scheduler_job.id + ti1.try_number = 1 + ti1.workload_run_id = "current-run-id" + session.merge(ti1) + session.commit() + + # Event from the previous (defer-exit) invocation + executor.event_buffer[ti1.key] = State.SUCCESS, None, "stale-defer-run-id" + executor.has_task = mock.MagicMock(return_value=False) + mock_stats.incr.reset_mock() + + self.job_runner._process_executor_events(executor=executor, session=session) + ti1.refresh_from_db(session=session) + assert ti1.state == ti_state + self.job_runner.executor.callback_sink.send.assert_not_called() + mock_stats.incr.assert_called_once_with("scheduler.executor_events.processed", count=1) + + # Matching run id with no other requeue signal is still an external kill. + executor.event_buffer[ti1.key] = State.SUCCESS, None, "current-run-id" + mock_stats.incr.reset_mock() + + self.job_runner._process_executor_events(executor=executor, session=session) + mock_stats.incr.assert_any_call( + "scheduler.tasks.killed_externally", + tags={"dag_id": dag_id, "task_id": ti1.task_id}, + ) + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") @mock.patch("airflow._shared.observability.metrics.stats._get_backend") def test_process_executor_events_multiple_try_numbers_warns( @@ -3225,6 +3286,32 @@ class PreAssigningExecutor(MockExecutor): session.rollback() + def test_executable_task_instances_to_queued_sets_workload_run_id(self, dag_maker, session): + """workload_run_id is written for every TI on enqueue, for all executors.""" + dag_id = "SchedulerJobTest.test_executable_sets_workload_run_id" + session = settings.Session() + with dag_maker(dag_id=dag_id, start_date=DEFAULT_DATE, session=session): + EmptyOperator(task_id="task_a") + + self.job_runner = SchedulerJobRunner(job=Job(), executors=[MockExecutor()]) + + dr = dag_maker.create_dagrun() + ti = dr.get_task_instance("task_a", session=session) + ti.state = State.SCHEDULED + session.flush() + + returned_tis = self.job_runner._executable_task_instances_to_queued(max_tis=32, session=session) + assert len(returned_tis) == 1 + assert returned_tis[0].workload_run_id is not None + assert UUID(returned_tis[0].workload_run_id), "is valid uuid" + + db_value = session.scalar( + select(TaskInstance.workload_run_id).where(TaskInstance.id == ti.id) + ) + assert db_value == returned_tis[0].workload_run_id + + session.rollback() + @pytest.mark.parametrize("state", [State.FAILED, State.SUCCESS]) def test_enqueue_task_instances_sets_ti_state_to_None_if_dagrun_in_finish_state(self, state, dag_maker): """This tests that task instances whose dagrun is in finished state are not queued""" diff --git a/devel-common/src/tests_common/test_utils/mock_executor.py b/devel-common/src/tests_common/test_utils/mock_executor.py index c7a2f26315234..7d770287d6e73 100644 --- a/devel-common/src/tests_common/test_utils/mock_executor.py +++ b/devel-common/src/tests_common/test_utils/mock_executor.py @@ -113,11 +113,13 @@ def terminate(self): def end(self): self.sync() - def change_state(self, key, state, info=None, remove_running=False): - super().change_state(key, state, info=info, remove_running=remove_running) + def change_state(self, key, state, info=None, remove_running=False, workload_run_id=None): + super().change_state( + key, state, info=info, remove_running=remove_running, workload_run_id=workload_run_id + ) # The normal event buffer is cleared after reading, we want to keep # a list of all events for testing - self.sorted_tasks.append((key, (state, info))) + self.sorted_tasks.append((key, (state, info, workload_run_id))) def mock_task_fail(self, dag_id, task_id, run_id: str, try_number=1): """ diff --git a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py index ba3af24415a0d..ca7f6c826275e 100644 --- a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py +++ b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py @@ -227,7 +227,7 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): self.workload_publish_retries.pop(key, None) if isinstance(result, ExceptionWithTraceback): self.log.error("%s: %s\n%s\n", CELERY_SEND_ERR_MSG_HEADER, result.exception, result.traceback) - self.event_buffer[key] = (TaskInstanceState.FAILED, None) + self.record_event(key, TaskInstanceState.FAILED, None) elif result is not None: result.backend = cached_celery_backend self.running.add(key) @@ -236,7 +236,9 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): # Store the Celery task_id (workload execution ID) in the event buffer. This will get "overwritten" if the task # has another event, but that is fine, because the only other events are success/failed at # which point we don't need the ID anymore anyway. - self.event_buffer[key] = (TaskInstanceState.QUEUED, result.task_id) + self.record_event( + key, TaskInstanceState.QUEUED, result.task_id, consume_run_id=False + ) def _send_workloads_to_celery(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): from airflow.providers.celery.executors.celery_executor_utils import send_workload_to_executor @@ -284,8 +286,17 @@ def update_all_workload_states(self) -> None: if state: self.update_task_state(cast("TaskInstanceKey", key), state, info) - def change_state(self, key: WorkloadKey, state: WorkloadState, info=None, remove_running=True) -> None: - super().change_state(key, state, info, remove_running=remove_running) + def change_state( + self, + key: WorkloadKey, + state: WorkloadState, + info=None, + remove_running=True, + workload_run_id: str | None = None, + ) -> None: + super().change_state( + key, state, info, remove_running=remove_running, workload_run_id=workload_run_id + ) self.workloads.pop(key, None) def update_task_state(self, key: TaskInstanceKey, state: str, info: Any) -> None: diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index c9f597e13ad2c..771aa18b29a12 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -367,7 +367,9 @@ def execute_async( queue, ) - self.event_buffer[key] = (TaskInstanceState.QUEUED, self.scheduler_job_id) + self.record_event( + key, TaskInstanceState.QUEUED, self.scheduler_job_id, consume_run_id=False + ) job = KubernetesJob(key, command, kube_executor_config, pod_template_file, coordinator_kube_image) self.pod_launch_attempts[key] = _PodLaunchAttempt(job=job) self.task_queue.put(job) @@ -379,6 +381,8 @@ def queue_workload(self, workload: workloads.All, session: Session | None) -> No raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}") ti = workload.ti self.queued_tasks[ti.key] = workload + if ti.workload_run_id: + self._workload_run_ids[ti.key].append(ti.workload_run_id) def _process_workloads(self, workloads: Sequence[workloads.All]) -> None: from airflow.executors.workloads import ExecuteTask @@ -735,7 +739,7 @@ def _change_state( return if state == TaskInstanceState.RUNNING: - self.event_buffer[key] = state, None + self.record_event(key, state, None, consume_run_id=False) return if self.kube_config.delete_worker_pods: @@ -809,7 +813,7 @@ def _change_state( if state is None: state = self._get_task_instance_state(key, session=session) - self.event_buffer[key] = state, termination_reason + self.record_event(key, state, termination_reason) def _get_task_instance_state(self, key: TaskInstanceKey, *, session: Session) -> TaskInstanceState | None: """Look up the current task instance state from the metadata database.""" From 00d22d8a00d8fedb1435de4d525a9cf059f25a20 Mon Sep 17 00:00:00 2001 From: Pedrinhonitz Date: Mon, 21 Sep 2026 11:22:20 -0300 Subject: [PATCH 5/8] Fix actions failed in executor --- .../src/airflow/executors/base_executor.py | 2 +- .../src/airflow/executors/local_executor.py | 9 ++++++- .../unit/executors/test_local_executor.py | 3 ++- .../tests_common/test_utils/mock_executor.py | 9 ++++--- .../celery/executors/celery_executor.py | 26 +++++++++++++++---- .../executors/kubernetes_executor.py | 22 ++++++++++++---- 6 files changed, 55 insertions(+), 16 deletions(-) diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index 49e762cd1245b..ea11740ee1518 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -270,7 +270,7 @@ def queue_workload(self, workload: ExecutorWorkload, session: Session) -> None: if isinstance(workload, workloads.ExecuteTask): ti = workload.ti self.queued_tasks[ti.key] = workload - if ti.workload_run_id: + if getattr(ti, "workload_run_id", None): self._workload_run_ids[ti.key].append(ti.workload_run_id) elif isinstance(workload, workloads.ExecuteCallback): if not self.supports_callbacks: diff --git a/airflow-core/src/airflow/executors/local_executor.py b/airflow-core/src/airflow/executors/local_executor.py index 9790a0e98fc2c..36a528d3c2ffe 100644 --- a/airflow-core/src/airflow/executors/local_executor.py +++ b/airflow-core/src/airflow/executors/local_executor.py @@ -248,7 +248,14 @@ def sync(self) -> None: def _read_results(self): try: while not self.result_queue.empty(): - key, state, exc, workload_run_id = self.result_queue.get() + result = self.result_queue.get() + # Support legacy 3-tuples (key, state, exc) and current 4-tuples + # (key, state, exc, workload_run_id) from workers / tests. + if len(result) == 3: + key, state, _exc = result + workload_run_id = None + else: + key, state, _exc, workload_run_id = result self.change_state(key, state, workload_run_id=workload_run_id) except (OSError, EOFError): self.log.exception("Error reading from result queue") diff --git a/airflow-core/tests/unit/executors/test_local_executor.py b/airflow-core/tests/unit/executors/test_local_executor.py index 79987f9ab1cb4..8d57266d42771 100644 --- a/airflow-core/tests/unit/executors/test_local_executor.py +++ b/airflow-core/tests/unit/executors/test_local_executor.py @@ -95,7 +95,8 @@ def _write_large_results_to_queue(result_queue, result_count, payload_size): payload = RuntimeError("x" * payload_size) for index in range(result_count): key = TaskInstanceKey("test_dag", f"test_task_{index}", "test_run") - result_queue.put((key, State.SUCCESS, payload)) + # (key, state, exc, workload_run_id) — matches LocalExecutor worker output + result_queue.put((key, State.SUCCESS, payload, None)) class TestLocalExecutor: diff --git a/devel-common/src/tests_common/test_utils/mock_executor.py b/devel-common/src/tests_common/test_utils/mock_executor.py index 7d770287d6e73..e518b9ca5fd6a 100644 --- a/devel-common/src/tests_common/test_utils/mock_executor.py +++ b/devel-common/src/tests_common/test_utils/mock_executor.py @@ -114,9 +114,12 @@ def end(self): self.sync() def change_state(self, key, state, info=None, remove_running=False, workload_run_id=None): - super().change_state( - key, state, info=info, remove_running=remove_running, workload_run_id=workload_run_id - ) + try: + super().change_state( + key, state, info=info, remove_running=remove_running, workload_run_id=workload_run_id + ) + except TypeError: + super().change_state(key, state, info=info, remove_running=remove_running) # The normal event buffer is cleared after reading, we want to keep # a list of all events for testing self.sorted_tasks.append((key, (state, info, workload_run_id))) diff --git a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py index ca7f6c826275e..ecbc9a844b782 100644 --- a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py +++ b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py @@ -227,7 +227,7 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): self.workload_publish_retries.pop(key, None) if isinstance(result, ExceptionWithTraceback): self.log.error("%s: %s\n%s\n", CELERY_SEND_ERR_MSG_HEADER, result.exception, result.traceback) - self.record_event(key, TaskInstanceState.FAILED, None) + self._emit_task_event(key, TaskInstanceState.FAILED, None) elif result is not None: result.backend = cached_celery_backend self.running.add(key) @@ -236,10 +236,22 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): # Store the Celery task_id (workload execution ID) in the event buffer. This will get "overwritten" if the task # has another event, but that is fine, because the only other events are success/failed at # which point we don't need the ID anymore anyway. - self.record_event( + self._emit_task_event( key, TaskInstanceState.QUEUED, result.task_id, consume_run_id=False ) + def _emit_task_event(self, key, state, info=None, *, consume_run_id: bool | None = None) -> None: + """Write an executor event; compatible with cores that lack ``record_event``.""" + record_event = getattr(self, "record_event", None) + if callable(record_event): + if consume_run_id is None: + record_event(key, state, info) + else: + record_event(key, state, info, consume_run_id=consume_run_id) + return + # Older BaseExecutor: event buffer is a plain (state, info) tuple. + self.event_buffer[key] = (state, info) + def _send_workloads_to_celery(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): from airflow.providers.celery.executors.celery_executor_utils import send_workload_to_executor @@ -294,9 +306,13 @@ def change_state( remove_running=True, workload_run_id: str | None = None, ) -> None: - super().change_state( - key, state, info, remove_running=remove_running, workload_run_id=workload_run_id - ) + try: + super().change_state( + key, state, info, remove_running=remove_running, workload_run_id=workload_run_id + ) + except TypeError: + # Older BaseExecutor.change_state has no workload_run_id kwarg (provider compat matrix). + super().change_state(key, state, info, remove_running=remove_running) self.workloads.pop(key, None) def update_task_state(self, key: TaskInstanceKey, state: str, info: Any) -> None: diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index 771aa18b29a12..2f16177326487 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -367,13 +367,24 @@ def execute_async( queue, ) - self.record_event( + self._emit_task_event( key, TaskInstanceState.QUEUED, self.scheduler_job_id, consume_run_id=False ) job = KubernetesJob(key, command, kube_executor_config, pod_template_file, coordinator_kube_image) self.pod_launch_attempts[key] = _PodLaunchAttempt(job=job) self.task_queue.put(job) + def _emit_task_event(self, key, state, info=None, *, consume_run_id: bool | None = None) -> None: + """Write an executor event; compatible with cores that lack ``record_event``.""" + record_event = getattr(self, "record_event", None) + if callable(record_event): + if consume_run_id is None: + record_event(key, state, info) + else: + record_event(key, state, info, consume_run_id=consume_run_id) + return + self.event_buffer[key] = (state, info) + def queue_workload(self, workload: workloads.All, session: Session | None) -> None: from airflow.executors import workloads @@ -381,8 +392,9 @@ def queue_workload(self, workload: workloads.All, session: Session | None) -> No raise RuntimeError(f"{type(self)} cannot handle workloads of type {type(workload)}") ti = workload.ti self.queued_tasks[ti.key] = workload - if ti.workload_run_id: - self._workload_run_ids[ti.key].append(ti.workload_run_id) + workload_run_id = getattr(ti, "workload_run_id", None) + if workload_run_id and hasattr(self, "_workload_run_ids"): + self._workload_run_ids[ti.key].append(workload_run_id) def _process_workloads(self, workloads: Sequence[workloads.All]) -> None: from airflow.executors.workloads import ExecuteTask @@ -739,7 +751,7 @@ def _change_state( return if state == TaskInstanceState.RUNNING: - self.record_event(key, state, None, consume_run_id=False) + self._emit_task_event(key, state, None, consume_run_id=False) return if self.kube_config.delete_worker_pods: @@ -813,7 +825,7 @@ def _change_state( if state is None: state = self._get_task_instance_state(key, session=session) - self.record_event(key, state, termination_reason) + self._emit_task_event(key, state, termination_reason) def _get_task_instance_state(self, key: TaskInstanceKey, *, session: Session) -> TaskInstanceState | None: """Look up the current task instance state from the metadata database.""" From 002954c0e614396f6cce0061b942dd61825e231d Mon Sep 17 00:00:00 2001 From: Pedrinhonitz Date: Mon, 21 Sep 2026 14:35:19 -0300 Subject: [PATCH 6/8] Fix providers actions run --- airflow-core/docs/migrations-ref.rst | 5 ++++- airflow-core/src/airflow/executors/base_executor.py | 5 +++-- airflow-core/src/airflow/utils/db.py | 2 +- airflow-core/tests/unit/executors/test_workloads.py | 2 ++ airflow-core/tests/unit/models/test_taskinstance.py | 1 + .../airflow/providers/celery/executors/celery_executor.py | 8 ++++---- .../providers/celery/executors/celery_executor_utils.py | 4 ++-- .../cncf/kubernetes/executors/kubernetes_executor.py | 4 ++-- 8 files changed, 19 insertions(+), 12 deletions(-) diff --git a/airflow-core/docs/migrations-ref.rst b/airflow-core/docs/migrations-ref.rst index b2eec9366d6ae..900ab447c50af 100644 --- a/airflow-core/docs/migrations-ref.rst +++ b/airflow-core/docs/migrations-ref.rst @@ -39,7 +39,10 @@ Here's the list of all the Database Migrations that are executed via when you ru +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | Revision ID | Revises ID | Airflow Version | Description | +=========================+==================+===================+==============================================================+ -| ``b6a9c2e7d410`` (head) | ``f8c2a1d94e03`` | ``3.4.0`` | Add draining state to DagModel. | +| ``c7d4e8f1a203`` (head) | ``b6a9c2e7d410`` | ``3.4.0`` | Add workload_run_id to task_instance and | +| | | | task_instance_history. | ++-------------------------+------------------+-------------------+--------------------------------------------------------------+ +| ``b6a9c2e7d410`` | ``f8c2a1d94e03`` | ``3.4.0`` | Add draining state to DagModel. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ | ``f8c2a1d94e03`` | ``8d3f1a6b2c47`` | ``3.4.0`` | Add team_name and bundle_names scope columns to job table. | +-------------------------+------------------+-------------------+--------------------------------------------------------------+ diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index ea11740ee1518..e88b209a70d18 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -270,8 +270,9 @@ def queue_workload(self, workload: ExecutorWorkload, session: Session) -> None: if isinstance(workload, workloads.ExecuteTask): ti = workload.ti self.queued_tasks[ti.key] = workload - if getattr(ti, "workload_run_id", None): - self._workload_run_ids[ti.key].append(ti.workload_run_id) + run_id = getattr(ti, "workload_run_id", None) + if isinstance(run_id, str): + self._workload_run_ids[ti.key].append(run_id) elif isinstance(workload, workloads.ExecuteCallback): if not self.supports_callbacks: raise NotImplementedError( diff --git a/airflow-core/src/airflow/utils/db.py b/airflow-core/src/airflow/utils/db.py index c29af4a971c47..bbcb70223ebba 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": "b6a9c2e7d410", + "3.4.0": "c7d4e8f1a203", } # Prefix used to identify tables holding data moved during migration. diff --git a/airflow-core/tests/unit/executors/test_workloads.py b/airflow-core/tests/unit/executors/test_workloads.py index 7c639c43594cd..794d718fdf9fc 100644 --- a/airflow-core/tests/unit/executors/test_workloads.py +++ b/airflow-core/tests/unit/executors/test_workloads.py @@ -179,6 +179,7 @@ def test_workload_ti_round_trips_through_sdk_generated_model(): dumped = ti.model_dump(mode="json") assert "external_executor_id" not in dumped + assert "workload_run_id" not in dumped assert "executor_config" not in dumped # Executor-side scheduling fields stay on the workload wire (older workers # deserialize the workload with a model that requires them) but are not @@ -237,6 +238,7 @@ def _make_mock_ti( ti.context_carrier = None ti.hostname = None ti.external_executor_id = None + ti.workload_run_id = None ti.dag_model.bundle_name = "test-bundle" ti.dag_model.relative_fileloc = "dags/test_dag.py" diff --git a/airflow-core/tests/unit/models/test_taskinstance.py b/airflow-core/tests/unit/models/test_taskinstance.py index 7401c7781f3cc..00d1ba0de2b5d 100644 --- a/airflow-core/tests/unit/models/test_taskinstance.py +++ b/airflow-core/tests/unit/models/test_taskinstance.py @@ -2580,6 +2580,7 @@ def test_refresh_from_db(self, create_task_instance): "executor": "some_executor", "executor_config": {"Some": {"extra": "information"}}, "external_executor_id": "some_executor_id", + "workload_run_id": "some-run-id", "trigger_timeout": None, "trigger_id": None, "next_kwargs": None, diff --git a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py index ecbc9a844b782..d404481e264f1 100644 --- a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py +++ b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py @@ -249,8 +249,8 @@ def _emit_task_event(self, key, state, info=None, *, consume_run_id: bool | None else: record_event(key, state, info, consume_run_id=consume_run_id) return - # Older BaseExecutor: event buffer is a plain (state, info) tuple. - self.event_buffer[key] = (state, info) + # Older BaseExecutor: event buffer is a plain (state, info[, workload_run_id]) tuple. + self.event_buffer[key] = (state, info, None) def _send_workloads_to_celery(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): from airflow.providers.celery.executors.celery_executor_utils import send_workload_to_executor @@ -294,7 +294,7 @@ def update_all_workload_states(self) -> None: self.log.debug("Inquiries completed.") for key, async_result in list(self.workloads.items()): - state, info = state_and_info_by_celery_task_id.get(async_result.task_id) + state, info, *_ = state_and_info_by_celery_task_id.get(async_result.task_id) if state: self.update_task_state(cast("TaskInstanceKey", key), state, info) @@ -378,7 +378,7 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task adopted = [] cached_celery_backend = next(iter(celery_tasks.values()))[0].backend - for celery_task_id, (state, info) in states_by_celery_task_id.items(): + for celery_task_id, (state, info, *_) in states_by_celery_task_id.items(): result, ti = celery_tasks[celery_task_id] result.backend = cached_celery_backend if isinstance(result.result, BaseException): diff --git a/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py b/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py index 39d52deb1dd13..ff4359661f709 100644 --- a/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py +++ b/providers/celery/src/airflow/providers/celery/executors/celery_executor_utils.py @@ -556,7 +556,7 @@ def _prepare_state_and_info_by_task_dict( else: state = celery_states.PENDING info = None - state_info[task_id] = state, info + state_info[task_id] = state, info, None return state_info def _get_many_using_multiprocessing( @@ -581,5 +581,5 @@ def _get_many_using_multiprocessing( state_or_exception.traceback, ) else: - states_and_info_by_task_id[task_id] = state_or_exception, info + states_and_info_by_task_id[task_id] = state_or_exception, info, None return states_and_info_by_task_id diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index 2f16177326487..a2a899764443a 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -383,7 +383,7 @@ def _emit_task_event(self, key, state, info=None, *, consume_run_id: bool | None else: record_event(key, state, info, consume_run_id=consume_run_id) return - self.event_buffer[key] = (state, info) + self.event_buffer[key] = (state, info, None) def queue_workload(self, workload: workloads.All, session: Session | None) -> None: from airflow.executors import workloads @@ -393,7 +393,7 @@ def queue_workload(self, workload: workloads.All, session: Session | None) -> No ti = workload.ti self.queued_tasks[ti.key] = workload workload_run_id = getattr(ti, "workload_run_id", None) - if workload_run_id and hasattr(self, "_workload_run_ids"): + if isinstance(workload_run_id, str) and hasattr(self, "_workload_run_ids"): self._workload_run_ids[ti.key].append(workload_run_id) def _process_workloads(self, workloads: Sequence[workloads.All]) -> None: From d7703b95f08a5eaeddfb03e544c33c2cc57b649c Mon Sep 17 00:00:00 2001 From: Pedrinhonitz Date: Mon, 21 Sep 2026 16:43:27 -0300 Subject: [PATCH 7/8] Updated asserts celery --- .../tests/integration/celery/test_celery_executor.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/providers/celery/tests/integration/celery/test_celery_executor.py b/providers/celery/tests/integration/celery/test_celery_executor.py index f068f93375cf7..5b5322f82c909 100644 --- a/providers/celery/tests/integration/celery/test_celery_executor.py +++ b/providers/celery/tests/integration/celery/test_celery_executor.py @@ -413,7 +413,7 @@ def test_should_support_kv_backend(self, mock_mget, caplog): assert set(mget_args[0]) == {b"celery-task-meta-456", b"celery-task-meta-123"} mock_mget.assert_called_once_with(mock.ANY) - assert result == {"123": ("SUCCESS", None), "456": ("PENDING", None)} + assert result == {"123": ("SUCCESS", None, None), "456": ("PENDING", None, None)} assert caplog.messages == ["Fetched 2 state(s) for 2 task(s)"] @mock.patch("celery.backends.database.DatabaseBackend.ResultSession") @@ -440,7 +440,7 @@ def test_should_support_db_backend(self, mock_session, caplog): ] ) - assert result == {"123": ("SUCCESS", None), "456": ("PENDING", None)} + assert result == {"123": ("SUCCESS", None, None), "456": ("PENDING", None, None)} assert caplog.messages == ["Fetched 2 state(s) for 2 task(s)"] @mock.patch("celery.backends.database.DatabaseBackend.ResultSession") @@ -474,7 +474,7 @@ def test_should_retry_db_backend(self, mock_session, caplog): ] ) assert mock_retry_db_result.call_count == 2 - assert result == {"123": ("SUCCESS", None), "456": ("PENDING", None)} + assert result == {"123": ("SUCCESS", None, None), "456": ("PENDING", None, None)} assert caplog.messages == [ "Failed operation _query_task_cls_from_db_backend. Retrying 2 more times.", "Fetched 2 state(s) for 2 task(s)", @@ -499,5 +499,5 @@ def test_should_support_base_backend(self, caplog): ] ) - assert result == {"123": ("SUCCESS", None), "456": ("PENDING", None)} + assert result == {"123": ("SUCCESS", None, None), "456": ("PENDING", None, None)} assert caplog.messages == ["Fetched 2 state(s) for 2 task(s)"] From 7a18167feb5ab200610c5ffe89a6bb092c9fc431 Mon Sep 17 00:00:00 2001 From: Pedrinhonitz Date: Mon, 21 Sep 2026 20:11:56 -0300 Subject: [PATCH 8/8] fix prek and ruff --- airflow-core/tests/unit/jobs/test_scheduler_job.py | 4 +--- .../src/airflow/providers/celery/executors/celery_executor.py | 4 +--- .../cncf/kubernetes/executors/kubernetes_executor.py | 4 +--- 3 files changed, 3 insertions(+), 9 deletions(-) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 1eb6fdfbf52c1..6039483ad4f31 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -3305,9 +3305,7 @@ def test_executable_task_instances_to_queued_sets_workload_run_id(self, dag_make assert returned_tis[0].workload_run_id is not None assert UUID(returned_tis[0].workload_run_id), "is valid uuid" - db_value = session.scalar( - select(TaskInstance.workload_run_id).where(TaskInstance.id == ti.id) - ) + db_value = session.scalar(select(TaskInstance.workload_run_id).where(TaskInstance.id == ti.id)) assert db_value == returned_tis[0].workload_run_id session.rollback() diff --git a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py index d404481e264f1..736f8f35b5af1 100644 --- a/providers/celery/src/airflow/providers/celery/executors/celery_executor.py +++ b/providers/celery/src/airflow/providers/celery/executors/celery_executor.py @@ -236,9 +236,7 @@ def _send_workloads(self, workload_tuples_to_send: Sequence[WorkloadInCelery]): # Store the Celery task_id (workload execution ID) in the event buffer. This will get "overwritten" if the task # has another event, but that is fine, because the only other events are success/failed at # which point we don't need the ID anymore anyway. - self._emit_task_event( - key, TaskInstanceState.QUEUED, result.task_id, consume_run_id=False - ) + self._emit_task_event(key, TaskInstanceState.QUEUED, result.task_id, consume_run_id=False) def _emit_task_event(self, key, state, info=None, *, consume_run_id: bool | None = None) -> None: """Write an executor event; compatible with cores that lack ``record_event``.""" diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index a2a899764443a..cf3e0f12aff24 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -367,9 +367,7 @@ def execute_async( queue, ) - self._emit_task_event( - key, TaskInstanceState.QUEUED, self.scheduler_job_id, consume_run_id=False - ) + self._emit_task_event(key, TaskInstanceState.QUEUED, self.scheduler_job_id, consume_run_id=False) job = KubernetesJob(key, command, kube_executor_config, pod_template_file, coordinator_kube_image) self.pod_launch_attempts[key] = _PodLaunchAttempt(job=job) self.task_queue.put(job)