Skip to content

45461 remove remaining session query usages#62758

Open
kunaljubce wants to merge 8 commits intoapache:mainfrom
kunaljubce:45461-remove-remaining-session-query-usages
Open

45461 remove remaining session query usages#62758
kunaljubce wants to merge 8 commits intoapache:mainfrom
kunaljubce:45461-remove-remaining-session-query-usages

Conversation

@kunaljubce
Copy link
Contributor

@kunaljubce kunaljubce commented Mar 3, 2026

relates to: #45461

Details
As part of 45461, as far as I understand, only a last couple of session.query instances were left behind:

→ grep -R "session.query"
Binary file ./generated/_doctrees/docs/apache-airflow-providers-airbyte/environment.pickle matches
./providers/teradata/docs/connections/teradata.rst:QueryBand can be specified using extra connection configuration ... (Multiple instances of 'mock_session.query_tag' matches like the one below)
./providers/snowflake/tests/unit/snowflake/decorators/test_snowpark.py:            return session.query_tag
./providers/standard/tests/unit/standard/utils/test_skipmixin.py:            assert not session.query.called # CHANGE 1
./airflow-core/tests/unit/models/test_deadline.py:            mock_session.query.assert_not_called()         # CHANGE 2
./scripts/ci/prek/prevent_deprecated_sqlalchemy_usage.py:def check_session_query(mod: ast.Module, file_path: str) -> bool:
./scripts/ci/prek/prevent_deprecated_sqlalchemy_usage.py:        errors = check_session_query(ast_module, file_path)

These instances have been replaced with session.execute and session.scalars, dependending on their corresponding usages in the source code. Also fixed a potential dangling import issue I ran into with the prek pre-push hook:

Run mypy for providers...........................................................Failed
- hook id: mypy-providers
- exit code: 1

  Docker container engine detected.

  You are skipping the image upgrade check, but the image needs an upgrade. This 
  might lead to out-dated results of the check!
  Consider running `breeze ci-image build --python 3.10` at earliest convenience!


  Using 'uv' to install Airflow


  Using airflow version from current sources

  providers/standard/tests/unit/standard/utils/test_skipmixin.py:29: error:
  Module "airflow.utils" has no attribute "timezone"  [attr-defined]
      from airflow.utils import timezone
      ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  Found 1 error in 1 file (checked 1 source file)

  Error 1 returned
  If you see strange stacktraces above, and can't reproduce it, please run this command and try again:

  breeze ci-image build --python 3.10

  You can also run `breeze down --cleanup-mypy-cache` to clean up the cache used.


Was generative AI tooling used to co-author this PR?
  • Yes (please specify the tool below)

  • Read the Pull Request Guidelines for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
  • For fundamental code changes, an Airflow Improvement Proposal (AIP) is needed.
  • When adding dependency, check compliance with the ASF 3rd Party License Policy.
  • For significant user-facing changes create newsfragment: {pr_number}.significant.rst or {issue_number}.significant.rst, in airflow-core/newsfragments.

mock_session.delete.assert_called_once_with(mock_deadline)
else:
mock_session.query.assert_not_called()
mock_session.execute.assert_not_called()
Copy link
Contributor Author

@kunaljubce kunaljubce Mar 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

session = Mock()
assert SkipMixin().skip(dag_run=None, execution_date=None, tasks=[]) is None
assert not session.query.called
assert not session.scalars.called
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Prab-27
Copy link
Contributor

Prab-27 commented Mar 4, 2026

Thanks @kunaljubce !! Good catch !!

Will work on this prek hooks again to see what we can do to catch this types of patterns !

Comment on lines +30 to +33
try:
from airflow.sdk import timezone
except ImportError: # Fallback for Airflow < 3.1
from airflow.utils import timezone # type: ignore[attr-defined,no-redef]
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Import from from airflow.providers.common.compat.sdk instead

Copy link
Contributor Author

@kunaljubce kunaljubce Mar 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@uranusjr Done! Out of curiosity, is this change only to conform to a particular standard for imports? If that's the case, should we also do the same change in other places, e.g. https://github.com/kunaljubce/airflow/blob/main/airflow-core/tests/unit/models/test_deadline.py#L33

# Module map: module_name -> module_path(s)
# For entire modules that have been moved (e.g., timezone)
# Usage: from airflow.providers.common.compat.lazy_compat import timezone
# Usage: from airflow.providers.common.compat.sdk import timezone
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

>>> from airflow.providers.common.compat.lazy_compat import timezone
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ModuleNotFoundError: No module named 'airflow.providers.common.compat.lazy_compat'
>>> from airflow.providers.common.compat.sdk import timezone
>>> 

@kunaljubce kunaljubce marked this pull request as ready for review March 4, 2026 08:48
@kunaljubce
Copy link
Contributor Author

Thanks @kunaljubce !! Good catch !!

Will work on this prek hooks again to see what we can do to catch this types of patterns !

Sure @Prab-27, please lemme know if you're occupied with other tasks. I can pick it up as well.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants