kaxil commented on code in PR #76806: URL: https://github.com/apache/airflow/pull/76806#discussion_r4539991572
########## airflow-core/docs/core-concepts/resumable-tasks.rst: ########## @@ -159,13 +159,12 @@ To resume from the checkpoint instead, set ``keep_task_state`` when clearing, or corresponding box in the clear dialog. That is the right choice when nothing about the inputs or the code changed and you only want the task to carry on where it stopped. -This applies to clearing individual task instances through the REST API, the UI, or -``airflowctl dags clear`` (which clears every task instance in the matched Dag run(s) through this -same endpoint). Clearing an entire Dag run through the Clear Run dialog/API, marking a task as -failed or success (which clears downstream tasks as a side effect), and the core CLI's -``airflow tasks clear`` / ``airflow dags clear`` (which call ``clear_task_instances()`` directly and -never reach the discard endpoint) all still keep task state unconditionally today; see -`#72929 <https://github.com/apache/airflow/issues/72929>`_. +This applies to every way of clearing a task: clearing task instances or whole Dag runs through Review Comment: Provider docs still describe the old behaviour, and after this they're wrong for every clear path. `providers/common/ai/docs/frameworks/strands.rst` says clearing the whole Dag run "keeps task state" and that a later Dag run clear resumes from the last pause, and the Glue, Redshift Data, Databricks (`run_now`, `submit_run`), Snowflake and Spark operator pages all say "clearing does not delete the stored run id". The `AgentOperator` docstring in common.ai makes the same claim. Someone following those pages expects a cleared run to reuse a job that already succeeded, and pays for a second one instead. Fine as a follow-up PR since providers release separately, but can it land before this ships? ########## airflow-core/newsfragments/76806.significant.rst: ########## @@ -0,0 +1,28 @@ +Clearing a Dag run now discards task state store entries by default + +Airflow 3.4.0 made clearing a task instance discard its ``task_state_store`` entries. Every other Review Comment: Should this say the Dag run clear endpoints now delete `task_state_store` rows? They only check Dag Run edit, while the task state store API guards the same rows with Task Instance delete. #72100 had that note for the TI clear. `dag.test()` on an existing logical date also discards now, since it calls `clear_dags`. ########## airflow-core/src/airflow/models/taskinstance.py: ########## @@ -145,6 +147,9 @@ log = logging.getLogger(__name__) +# Five bind parameters per key; stays well under the bind parameter limits of every supported DB. +_TASK_STATE_DISCARD_BATCH_SIZE = 500 Review Comment: 500 keys x 5 columns is 2,500 bind params per `DELETE`. SQLite before 3.32 caps that at 999 by default, and we still accept 3.15+ (`_validate_sqlite3_version`). With the limit set to 999 I get `too many SQL variables` from 200 keys up, so clearing a run with 200+ tasks on an older system libsqlite3 now fails, even when there are no state rows. Could the batch be ~150 (or sized from the dialect), and the comment corrected? ########## airflow-core/src/airflow/ui/src/pages/DagRuns/BulkClearDagRunsButton.tsx: ########## @@ -42,13 +43,16 @@ const BulkClearDagRunsButton = ({ deselectKeys, selectedDagRuns }: Props) => { const { onClose, onOpen, open } = useDisclosure(); const [selectedOptions, setSelectedOptions] = useState<Array<string>>(["existingTasks"]); const [note, setNote] = useState<string | null>(null); + const [keepTaskStateDefault] = useClearKeepTaskStateDefault(); + const [keepTaskState, setKeepTaskState] = useState(keepTaskStateDefault); Review Comment: This is the fourth copy of the keep-task-state state and checkbox (with `ClearTaskInstanceDialog`, `ClearGroupTaskInstanceDialog` and `ClearRunDialog`), and the copies already reset differently: `ClearRunDialog` resets when it opens, this one only in `handleClose`, which the success path skips (`onSuccessConfirm: onClose`). Could they share one hook? ########## airflow-core/src/airflow/ui/src/components/Clear/Run/ClearRunDialog.test.tsx: ########## @@ -0,0 +1,81 @@ +/*! + * 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. + */ +import "@testing-library/jest-dom/vitest"; +import { fireEvent, render, screen, waitFor } from "@testing-library/react"; +import { http, HttpResponse } from "msw"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +import type { DAGRunResponse } from "openapi/requests/types.gen"; + +import { server } from "src/mocks/server"; +import { Wrapper } from "src/utils/Wrapper"; + +import ClearRunDialog from "./ClearRunDialog"; + +const DAG_ID = "test_dag"; +const DAG_RUN_ID = "run_1"; + +const dagRun = { + bundle_version: null, + dag_id: DAG_ID, + dag_run_id: DAG_RUN_ID, + dag_versions: [], + note: null, +} as unknown as DAGRunResponse; + +vi.mock("src/queries/useClearDagRunDryRun", () => ({ + useClearDagRunDryRun: () => ({ data: { task_instances: [{ task_id: "t" }], total_entries: 1 } }), + useClearDagRunDryRunKey: "clearDagRunDryRun", +})); + +let clearBody: Record<string, unknown> | undefined; + +beforeEach(() => { + server.use( + http.get(`/api/v2/dags/${DAG_ID}/details`, () => HttpResponse.json({})), + http.post(`/api/v2/dags/${DAG_ID}/dagRuns/${DAG_RUN_ID}/clear`, async ({ request }) => { + clearBody = (await request.json()) as Record<string, unknown>; + + return HttpResponse.json({ dag_id: DAG_ID, dag_run_id: DAG_RUN_ID }); + }), + ); +}); +afterEach(() => { + clearBody = undefined; +}); + +describe("ClearRunDialog", () => { + it.each([ + { expected: undefined, tick: false }, + { expected: true, tick: true }, + ])("sends keep_task_state=$expected when the box is ticked=$tick", async ({ expected, tick }) => { + render(<ClearRunDialog dagRun={dagRun} onClose={vi.fn()} open />, { wrapper: Wrapper }); + + const checkbox = screen.getByRole("checkbox", { name: /keepTaskState/iu }); + + expect(checkbox).not.toBeChecked(); Review Comment: Both new UI tests start from an empty `localStorage`, so `useState(false)` in place of `useState(keepTaskStateDefault)` still passes them. Could one case seed the per-browser default like `ClearTaskInstanceDialog.test.tsx` does and expect the box pre-ticked? ########## airflow-core/docs/administration-and-deployment/task-and-asset-state-store.rst: ########## @@ -61,8 +61,8 @@ Number of days after which task state store rows expire. When a key is written w ~~~~~~~~~~~~~~~~~~~~ Retention and ``clear_on_success`` are not the only ways task state store entries get removed: -clearing a task instance through the REST API, the UI, or ``airflowctl`` also discards them by -default unless ``keep_task_state`` is set. See :doc:`/core-concepts/resumable-tasks`. +clearing a task instance or a Dag run (REST API, UI, ``airflowctl`` or the ``airflow`` CLI) also Review Comment: `airflowctl dags clear` has no way to set `keep_task_state` (it builds `ClearTaskInstancesBody` without it), so "unless `keep_task_state` is set" isn't true there. #72100's newsfragment called that gap out. Now that `airflow dags clear` gets `--keep-task-state`, should its airflowctl replacement get one too, or should the docs keep the caveat? ########## airflow-core/tests/unit/cli/commands/test_dag_command.py: ########## @@ -1522,6 +1522,17 @@ def test_rejects_multiple_selectors(self, parser, extra_args): with pytest.raises(SystemExit, match="mutually exclusive"): dag_command.dag_clear(args) + @pytest.mark.usefixtures("seeded_partitioned_runs") + @pytest.mark.parametrize(("extra_args", "expected"), [([], False), (["--keep-task-state"], True)]) + @mock.patch("airflow.cli.commands.dag_command._bulk_clear_runs", autospec=True, return_value=1) + def test_keep_task_state(self, mock_bulk_clear_runs, parser, extra_args, expected): Review Comment: This and `test_task_clear_keep_task_state` mock one level above the forward (`_bulk_clear_runs`, `clear_dags`), so dropping `keep_task_state` in `_bulk_clear_runs`, `clear_dags` or the loop clear helper keeps CI green. A `keep_task_state` parametrize on `TestWholeRunClear` with seeded rows would pin the real path. Not blocking. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
