Add pre-commit To Prevent Usage of session.query #47275

Prab-27 · 2025-03-02T13:40:32Z

closes : #45461
Introduce a pre-commit hook to prevent the use of session.query in the core Airflow code.
This is limited to the source code and excludes the tests/ package.
Removed session.query() and updated a new style

^ Add meaningful description above
Read the Pull Request Guidelines for more information.
In case of fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
In case of a new dependency, check compliance with the ASF 3rd Party License Policy.
In case of backwards incompatible changes please leave a note in a newsfragment file, named {pr_number}.significant.rst or {issue_number}.significant.rst, in newsfragments.

Prab-27 · 2025-03-02T13:58:50Z

I will raise a PR soon for removing session.query() from airflow-core and task-sdk as we discussed in #45714

ephraimbuddy · 2025-03-04T11:29:27Z

.pre-commit-config.yaml

+        language: python
+        additional_dependencies: ['rich>=12.4.4']
+        files: ^airflow.*\.py$|^task_sdk.*\.py
+        exclude: ^tests/.*\.py$|^task_sdk/tests/.*\.py$


Suggested change

exclude: ^tests/.*\.py$|^task_sdk/tests/.*\.py$

Since we have given it specific files to check, do we still need to exclude others?

Ya you are right ! I'll remove this line

ephraimbuddy · 2025-03-04T11:34:09Z

I will raise a PR soon for removing session.query() from airflow-core and task-sdk as we discussed in #45714

You can do it in this PR. This PR should raise and point to where there are usages of session.query

bugraoz93 · 2025-03-04T18:54:05Z

scripts/ci/pre_commit/usage_session_query.py

+                    f"\nSQLAlchemy 2.0 deprecates the `Query` object"
+                    f"use the `select()` construct instead."
+                )
+                errors += 1


Thanks for the changes! We don't need this integer since occurrence isn't important in return and response text. We could make a return bool and the main can still return one similar to below. This can also do an early exit and reduce scans since a single occurrence should trigger the hook. What do you think?

return 1 if check_session_query(ast_module) else 0

ya ! you are right !!
I'll update it soon

Prab-27 · 2025-03-14T14:15:48Z

airflow/models/xcom.py


-        for xcom in query:
+        for xcom in session.execute(query).all():


I am not sure if this works fine for session.delete(xcom).

Could you please clarify when you get some free time?

Prab-27 · 2025-03-18T16:43:17Z

I have been trying to update the code but couldn't
query = session.query(base_table).with_entities(text(f"{base_table_alias}.*")) from here

could you please help me upgrade this ?

github-actions · 2025-05-03T00:16:46Z

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 5 days if no further activity occurs. Thank you for your contributions.

ephraimbuddy · 2025-05-19T17:04:27Z

airflow-core/src/airflow/api_fastapi/execution_api/routes/xcoms.py

@@ -270,7 +270,7 @@ def set_xcom(
    if not run_id:
        raise HTTPException(status.HTTP_404_NOT_FOUND, f"Run with ID: `{run_id}` was not found")

-    dag_run_id = session.query(DagRun.id).filter_by(dag_id=dag_id, run_id=run_id).scalar()
+    dag_run_id = session.execute(DagRun.id).where(dag_id=dag_id, run_id=run_id).scalar()


Suggested change

dag_run_id = session.execute(DagRun.id).where(dag_id=dag_id, run_id=run_id).scalar()

dag_run_id = session.scalar(select(DagRun.id).where(dag_id=dag_id, run_id=run_id))

ephraimbuddy · 2025-05-19T17:05:23Z

airflow-core/src/airflow/dag_processing/bundles/manager.py

-        active_bundle_names = set(self._bundle_config.keys())
-        for name in active_bundle_names:
+        stored = {b.name: b for b in session.scalars(select(DagBundleModel)).all()}
+        for name in self._bundle_config.keys():


Looks like you are omitting the active_bundle_names filter?

It was a mistake. Rebased again and resolved

ephraimbuddy · 2025-05-19T17:07:00Z

airflow-core/src/airflow/dag_processing/bundles/manager.py

+            updated_rows = session.execute(
+                update(DagVersion)
+                .where(DagVersion.bundle_name.in_(inactive_bundle_names))
+                .values({DagVersion.bundle_name: new_bundle_name})


Suggested change

.values({DagVersion.bundle_name: new_bundle_name})

.values(bundle_name=new_bundle_name)

This change has been updated so I rebased and resolved it

ephraimbuddy · 2025-05-19T17:07:34Z

airflow-core/src/airflow/dag_processing/collection.py

+                        "filename": filename,
+                        "bundle_name": bundle_name,
+                        "timestamp": utcnow(),
+                        "stacktrace": stacktrace,


values accept keyword args not dictionary

This change has been updated, so I rebased and resolved it

…nnecessary exclude files in pre-commit yaml

uranusjr · 2025-06-23T04:16:33Z

Instead of trying to fix everything in this PR, it might be easier to submit smaller PRs fixing things one by one, and rebase this repeatedly after each PR is merged until CI is green.

Prab-27 · 2025-06-23T04:59:14Z

Instead of trying to fix everything in this PR, it might be easier to submit smaller PRs fixing things one by one, and rebase this repeatedly after each PR is merged until CI is green.

Thanks for the suggestion—I'll proceed with this plan.

potiuk · 2025-06-23T05:17:55Z

Instead of trying to fix everything in this PR, it might be easier to submit smaller PRs fixing things one by one, and rebase this repeatedly after each PR is merged until CI is green.

Yep - it could also be done similar to the #52020

find all files/folders that need fixing
convert them to a list
implement a mechanism (.pre-comit files selection) that includes certain folders in .pre-commit check
ask for help (and do yourself some of that) with PRs that incrementally fix things and add stuff to be checked by .pre-commit so that no session is added to already "cleaned" folders
when completed, switch it to check "everywhere".

Prab-27 · 2025-06-23T06:17:19Z

Instead of trying to fix everything in this PR, it might be easier to submit smaller PRs fixing things one by one, and rebase this repeatedly after each PR is merged until CI is green.

Yep - it could also be done similar to the #52020

find all files/folders that need fixing

convert them to a list

implement a mechanism (.pre-comit files selection) that includes certain folders in .pre-commit check

ask for help (and do yourself some of that) with PRs that incrementally fix things and add stuff to be checked by .pre-commit so that no session is added to already "cleaned" folders

when completed, switch it to check "everywhere".

Appreciate your approach to this; I'm jumping in headfirst

Prab-27 · 2025-06-23T15:24:31Z

I've created a list as you suggested

Total changes needed files : 116
Including -

airflow-core - 10 files
airflow-core tests - 49 files
dev - 1 files
devel-common - 6 files
Providers - including tests - 50 files

My Plan :
Single file checks for all - 66 files
Providers - chnages as per module which has session.query - including tests and Providers both

@potiuk Would love to hear your thoughts on this.

potiuk · 2025-06-23T16:23:56Z

@potiuk Would love to hear your thoughts on this.

Sure. Also you might split per module in "core". I think what works best is if you do one or two yourself - some representative examples of PRs - and provide instructions to others based on thsose PRs. That avoids any ambiguities with description and speeds up people ramping up

Prab-27 requested review from potiuk, ashb and jedcunningham as code owners March 2, 2025 13:40

boring-cyborg bot added the area:dev-tools label Mar 2, 2025

Prab-27 mentioned this pull request Mar 2, 2025

Add pre-commit to prevent usage of session.query in Airflow Core #45461

Open

2 tasks

Prab-27 force-pushed the add-pre-commit-to-prevent-usage-session.query branch from 5d4d4dc to b6775a8 Compare March 2, 2025 14:07

ephraimbuddy reviewed Mar 4, 2025

View reviewed changes

bugraoz93 reviewed Mar 4, 2025

View reviewed changes

Prab-27 force-pushed the add-pre-commit-to-prevent-usage-session.query branch from e9bcf8e to 1c65940 Compare March 14, 2025 14:05

Prab-27 requested review from pierrejeambrun and XD-DENG as code owners March 14, 2025 14:05

Prab-27 commented Mar 14, 2025

View reviewed changes

Prab-27 mentioned this pull request Mar 18, 2025

Update query to use 2.0 sqlalchemy style #47884

Merged

Prab-27 force-pushed the add-pre-commit-to-prevent-usage-session.query branch from 1c65940 to 98be988 Compare March 18, 2025 15:55

github-actions bot added the stale Stale PRs per the .github/workflows/stale.yml policy file label May 3, 2025

Prab-27 force-pushed the add-pre-commit-to-prevent-usage-session.query branch from 98be988 to 376a2dc Compare May 6, 2025 14:28

Prab-27 requested review from gopidesupavan and amoghrajesh as code owners May 6, 2025 14:28

github-actions bot removed the stale Stale PRs per the .github/workflows/stale.yml policy file label May 7, 2025

Prab-27 force-pushed the add-pre-commit-to-prevent-usage-session.query branch from 376a2dc to 7f7cc6c Compare May 15, 2025 10:23

Prab-27 requested a review from kaxil as a code owner May 15, 2025 10:23

ephraimbuddy reviewed May 19, 2025

View reviewed changes

Prab-27 added 4 commits June 21, 2025 22:31

Resolving conflicts part

dbf0e21

update errors variable to bool in usage_session_query.py and remove u…

b1e8674

…nnecessary exclude files in pre-commit yaml

resolving conflicts second

f9d423e

update new style in manager.py and xcom.py

6865227

Prab-27 added 3 commits June 21, 2025 22:46

remove session.query and fix pre-commit checks resolve conflict

4682d7a

fix static checks - resolve conflict

d5b99ff

use scalar() instead of session.execute()

e37f259

Prab-27 force-pushed the add-pre-commit-to-prevent-usage-session.query branch from 7f7cc6c to e37f259 Compare June 22, 2025 11:47

Prab-27 mentioned this pull request Jun 23, 2025

Add Pre-Commit Hook To Prevent Usage of session.query #52083

Merged

	dag_run_id = session.execute(DagRun.id).where(dag_id=dag_id, run_id=run_id).scalar()
	dag_run_id = session.scalar(select(DagRun.id).where(dag_id=dag_id, run_id=run_id))

	.values({DagVersion.bundle_name: new_bundle_name})
	.values(bundle_name=new_bundle_name)

Add pre-commit To Prevent Usage of session.query #47275

Are you sure you want to change the base?

Add pre-commit To Prevent Usage of session.query #47275

Conversation

Prab-27 commented Mar 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Prab-27 commented Mar 2, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Prab-27 Mar 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ephraimbuddy commented Mar 4, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Prab-27 Mar 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Prab-27 commented Mar 18, 2025

Uh oh!

github-actions bot commented May 3, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

uranusjr commented Jun 23, 2025

Uh oh!

Prab-27 commented Jun 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

potiuk commented Jun 23, 2025

Uh oh!

Prab-27 commented Jun 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Prab-27 commented Jun 23, 2025

Uh oh!

potiuk commented Jun 23, 2025

Uh oh!

Uh oh!

Prab-27 commented Mar 2, 2025 •

edited

Loading

Prab-27 commented Mar 2, 2025 •

edited

Loading

Prab-27 Mar 4, 2025 •

edited

Loading

Prab-27 Mar 6, 2025 •

edited

Loading

Prab-27 commented Jun 23, 2025 •

edited

Loading

Prab-27 commented Jun 23, 2025 •

edited

Loading