diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 04062bddec77f..267b2e83ed490 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -425,6 +425,9 @@ 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/api_fastapi/core_api/routes/public/test_variables.py$| + ^airflow-core/tests/unit/cli/commands/test_task_command.py$| + ^airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py$| ^airflow-core/tests/unit/api_fastapi/execution_api/versions/head/test_task_instances.py$| ^airflow-core/tests/unit/models/test_renderedtifields.py$| ^airflow-core/tests/unit/models/test_timestamp.py$| diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py index a12a05b45b94a..0617fa184cca4 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py @@ -22,6 +22,7 @@ from unittest.mock import ANY import pytest +from sqlalchemy import select from airflow.models.team import Team from airflow.models.variable import Variable @@ -65,7 +66,7 @@ def create_file_upload(content: dict) -> BytesIO: @provide_session def _create_variables(session) -> None: - team = session.query(Team).where(Team.name == "test").one() + team = session.scalars(select(Team).where(Team.name == "test")).one() Variable.set( key=TEST_VARIABLE_KEY, @@ -130,13 +131,13 @@ def create_variables(self): class TestDeleteVariable(TestVariableEndpoint): def test_delete_should_respond_204(self, test_client, session): self.create_variables() - variables = session.query(Variable).all() + variables = session.scalars(select(Variable)).all() assert len(variables) == 5 response = test_client.delete(f"/variables/{TEST_VARIABLE_KEY}") assert response.status_code == 204 response = test_client.delete(f"/variables/{TEST_VARIABLE_KEY4}") assert response.status_code == 204 - variables = session.query(Variable).all() + variables = session.scalars(select(Variable)).all() assert len(variables) == 3 check_last_log(session, dag_id=None, event="delete_variable", logical_date=None) 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 diff --git a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py index 216f94da5781b..74a89d54823f5 100644 --- a/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py +++ b/airflow-core/tests/unit/dag_processing/bundles/test_dag_bundle_manager.py @@ -23,6 +23,7 @@ from unittest.mock import patch import pytest +from sqlalchemy import func, select from airflow.dag_processing.bundles.base import BaseDagBundle from airflow.dag_processing.bundles.manager import DagBundlesManager @@ -155,7 +156,9 @@ def clear_db(): @conf_vars({("core", "LOAD_EXAMPLES"): "False"}) def test_sync_bundles_to_db(clear_db, session): def _get_bundle_names_and_active(): - return session.query(DagBundleModel.name, DagBundleModel.active).order_by(DagBundleModel.name).all() + return session.execute( + select(DagBundleModel.name, DagBundleModel.active).order_by(DagBundleModel.name) + ).all() # Initial add with patch.dict( @@ -182,7 +185,7 @@ def _get_bundle_names_and_active(): ("my-test-bundle", False), ] # Since my-test-bundle is inactive, the associated import errors should be deleted - assert session.query(ParseImportError).count() == 0 + assert session.scalar(select(func.count(ParseImportError.id))) == 0 # Re-enable one that reappears in config with patch.dict( @@ -250,7 +253,7 @@ def test_sync_bundles_to_db_with_template(clear_db, session): manager.sync_bundles_to_db() # Check that the template and parameters were stored - bundle_model = session.query(DagBundleModel).filter_by(name="template-bundle").first() + bundle_model = session.scalars(select(DagBundleModel).filter_by(name="template-bundle").limit(1)).first() session.merge(bundle_model) @@ -269,7 +272,9 @@ def test_bundle_model_render_url(clear_db, session): ): manager = DagBundlesManager() manager.sync_bundles_to_db() - bundle_model = session.query(DagBundleModel).filter_by(name="template-bundle").first() + bundle_model = session.scalars( + select(DagBundleModel).filter_by(name="template-bundle").limit(1) + ).first() session.merge(bundle_model) assert bundle_model is not None @@ -291,7 +296,7 @@ def test_template_params_update_on_sync(clear_db, session): manager.sync_bundles_to_db() # Verify initial template and parameters - bundle_model = session.query(DagBundleModel).filter_by(name="template-bundle").first() + bundle_model = session.scalars(select(DagBundleModel).filter_by(name="template-bundle").limit(1)).first() url = bundle_model._unsign_url() assert url == "https://github.com/example/repo/tree/{version}/{subdir}" assert bundle_model.template_params == {"subdir": "dags"} @@ -316,7 +321,7 @@ def test_template_params_update_on_sync(clear_db, session): manager.sync_bundles_to_db() # Verify the template and parameters were updated - bundle_model = session.query(DagBundleModel).filter_by(name="template-bundle").first() + bundle_model = session.scalars(select(DagBundleModel).filter_by(name="template-bundle").limit(1)).first() url = bundle_model._unsign_url() assert url == "https://gitlab.com/example/repo/-/tree/{version}/{subdir}" assert bundle_model.template_params == {"subdir": "workflows"} @@ -335,7 +340,7 @@ def test_template_update_on_sync(clear_db, session): manager.sync_bundles_to_db() # Verify initial template - bundle_model = session.query(DagBundleModel).filter_by(name="template-bundle").first() + bundle_model = session.scalars(select(DagBundleModel).filter_by(name="template-bundle").limit(1)).first() url = bundle_model._unsign_url() assert url == "https://github.com/example/repo/tree/{version}/{subdir}" assert bundle_model.render_url(version="v1") == "https://github.com/example/repo/tree/v1/dags" @@ -360,7 +365,7 @@ def test_template_update_on_sync(clear_db, session): manager.sync_bundles_to_db() # Verify the template was updated - bundle_model = session.query(DagBundleModel).filter_by(name="template-bundle").first() + bundle_model = session.scalars(select(DagBundleModel).filter_by(name="template-bundle").limit(1)).first() url = bundle_model._unsign_url() assert url == "https://gitlab.com/example/repo/-/tree/{version}/{subdir}" assert bundle_model.render_url("v1") == "https://gitlab.com/example/repo/-/tree/v1/dags" @@ -438,7 +443,7 @@ def test_multiple_bundles_one_fails(clear_db, session): assert isinstance(bundles[0], BasicBundle) manager.sync_bundles_to_db() - bundle_names = {b.name for b in session.query(DagBundleModel).all()} + bundle_names = {b.name for b in session.scalars(select(DagBundleModel)).all()} assert bundle_names == {"my-test-bundle"}