From c7310bd0ff2f02f698d7e08fb893cf1b095439a9 Mon Sep 17 00:00:00 2001 From: Henry Chen Date: Mon, 15 Dec 2025 15:05:41 +0800 Subject: [PATCH] Remove Legacy Query Object Usage from test_command.py --- .pre-commit-config.yaml | 1 + .../tests/unit/cli/commands/test_task_command.py | 11 +++++------ 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 81e24f2e8a3b3..cc2f9183e0af4 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -413,6 +413,7 @@ repos: ^airflow-ctl.*\.py$| ^airflow-core/src/airflow/models/.*\.py$| ^airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_assets.py$| + ^airflow-core/tests/unit/cli/commands/test_task_command.py$| ^task_sdk.*\.py$ pass_filenames: true - id: update-supported-versions diff --git a/airflow-core/tests/unit/cli/commands/test_task_command.py b/airflow-core/tests/unit/cli/commands/test_task_command.py index 1c7b227bb6fe0..1592e9d7cdac3 100644 --- a/airflow-core/tests/unit/cli/commands/test_task_command.py +++ b/airflow-core/tests/unit/cli/commands/test_task_command.py @@ -30,6 +30,7 @@ from unittest import mock import pytest +from sqlalchemy import delete from airflow._shared.timezones import timezone from airflow.cli import cli_parser @@ -63,12 +64,10 @@ def reset(dag_id): with create_session() as session: - tis = session.query(TaskInstance).filter_by(dag_id=dag_id) - tis.delete() - runs = session.query(DagRun).filter_by(dag_id=dag_id) - runs.delete() - session.query(DagModel).filter_by(dag_id=dag_id).delete() - session.query(SerializedDagModel).filter_by(dag_id=dag_id).delete() + session.execute(delete(TaskInstance).where(TaskInstance.dag_id == dag_id)) + session.execute(delete(DagRun).where(DagRun.dag_id == dag_id)) + session.execute(delete(DagModel).where(DagModel.dag_id == dag_id)) + session.execute(delete(SerializedDagModel).where(SerializedDagModel.dag_id == dag_id)) @contextmanager