Add Airflow 3 compatibility for BiqQuery deferrable support (`ExecutionMode.AIRFLOW_ASYNC`) by pankajkoti · Pull Request #1674 · astronomer/astronomer-cosmos

pankajkoti · 2025-04-16T17:19:15Z

Changes

Remove SQLAlchemy session dependency and calls to DB ORM
Fix base class inheritance for proper object initialisation wrt Airflow 3
Add symlink in the airflow3 dags for the simple_dag_async.py example DAG
Set environment variables in airflow3 env to not use setup task for the async DAG until we resolve the Virtualenv execution mode
Add dbt-bigquery and Google provider to requirements.txt in Airflow 3 setup required for the example DAG

closes: #1634

netlify · 2025-04-16T17:19:38Z

✅ Deploy Preview for sunny-pastelito-5ecb04 canceled.

Name	Link
🔨 Latest commit	`143d902`
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/67ffe69617f2a0000893df10

cloudflare-workers-and-pages · 2025-04-16T17:19:38Z

Deploying astronomer-cosmos with Cloudflare Pages

Latest commit:	`27a9681`
Status:	✅ Deploy successful!
Preview URL:	https://d471e8ff.astronomer-cosmos.pages.dev
Branch Preview URL:	https://fix-async-bigquery-af3.astronomer-cosmos.pages.dev

View logs

netlify · 2025-04-16T17:20:16Z

✅ Deploy Preview for sunny-pastelito-5ecb04 canceled.

Name	Link
🔨 Latest commit	`27a9681`
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/6800de3742c0e400084c497b

codecov · 2025-04-16T18:08:34Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.08%. Comparing base (5ae181e) to head (27a9681).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1674      +/-   ##
==========================================
- Coverage   97.09%   97.08%   -0.01%     
==========================================
  Files          80       80              
  Lines        5022     5015       -7     
==========================================
- Hits         4876     4869       -7     
  Misses        146      146

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Copilot

Pull Request Overview

This PR adds Airflow 3 compatibility for BigQuery deferrable support by removing the SQLAlchemy session dependency and updating the operator's inheritance, as well as adjusting the example DAG configuration.

Removed SQLAlchemy session dependency and related DB ORM calls
Updated base class initialization for DbtRunAirflowAsyncBigqueryOperator
Added a symlink and updated settings for the Airflow 3 asynchronous DAG example

Reviewed Changes

Copilot reviewed 4 out of 6 changed files in this pull request and generated 1 comment.

File	Description
tests/operators/_asynchronous/test_bigquery.py	Updated test to patch _override_rtif instead of session-dependent RenderedTaskInstanceFields
scripts/airflow3/dags/simple_dag_async.py	Added a symlink to the example DAG in the dev directory
dev/dags/simple_dag_async.py	Modified path resolution, dataset reference, and schedule parameter for Airflow 3
cosmos/operators/_asynchronous/bigquery.py	Removed session dependency and restructured inheritance for Airflow 3 compatibility

Files not reviewed (2)

scripts/airflow3/env.sh: Language not supported
scripts/airflow3/requirements.txt: Language not supported

Comments suppressed due to low confidence (1)

cosmos/operators/_asynchronous/bigquery.py:110

Reassigning bases here overrides the previous assignment, which may introduce unintended behavior in the operator's inheritance. Consolidate the modifications into a single base class assignment for clarity.

DbtRunAirflowAsyncBigqueryOperator.__bases__ = (BigQueryInsertJobOperator, AbstractDbtLocalBase,)

tatiana

Excellent work, @pankajkoti !

…nMode.AIRFLOW_ASYNC) (#1674) - Remove SQLAlchemy session dependency and calls to DB ORM - Fix base class inheritance for proper object initialisation wrt Airflow 3 - Add symlink in the airflow3 dags for the simple_dag_async.py example DAG - Set environment variables in airflow3 env to not use setup task for the async DAG until we resolve the Virtualenv execution mode - Add dbt-bigquery and Google provider to requirements.txt in Airflow 3 setup required for the example DAG Closes: #1634

PR #1230 introduced the `original_jaffle_shop` project for `ExecutionMode.AIRFLOW_ASYNC`, likely to support early customisations for async execution. Since Cosmos 1.9, we’ve shifted to relying on dbt to generate the SQL, making those customisations unnecessary. As this project isn’t used in any other example DAG, it now seems redundant. I’ve updated the test and example DAG to use `jaffle_shop` instead, which is consistently used across most of our examples. This is a follow-up PR from the discussion in #1674 (comment) related: #1677

@tatiana

Main is currently failing for AF3 after merging PR #1674, with the following error: ``` File "<frozen importlib._bootstrap>", line 241, in _call_with_frames_removed File "/Users/tati/Code/cosmos-fresh/astronomer-cosmos/cosmos/operators/airflow_async.py", line 8, in <module> from cosmos.operators._asynchronous.base import DbtRunAirflowAsyncFactoryOperator File "/Users/tati/Code/cosmos-fresh/astronomer-cosmos/cosmos/operators/_asynchronous/__init__.py", line 12, in <module> from cosmos.operators.virtualenv import DbtRunVirtualenvOperator File "/Users/tati/Code/cosmos-fresh/astronomer-cosmos/cosmos/operators/virtualenv.py", line 11, in <module> from airflow.utils.python_virtualenv import prepare_virtualenv ModuleNotFoundError: No module named 'airflow.utils.python_virtualenv' ``` (as reported by @tatiana) It looks like we need to update the import path for the python_virtualenv module. I had the necessary fix locally but didn’t commit it, as I didn’t realize it would break after removing the setup task. However, since it’s a top-level import, it seems we do need to include the fix after all.

pankajkoti temporarily deployed to internal April 16, 2025 17:19 — with GitHub Actions Inactive

pankajkoti force-pushed the fix-async-bigquery-af3 branch from 143d902 to 6e6d083 Compare April 16, 2025 17:19

pankajkoti temporarily deployed to internal April 16, 2025 17:20 — with GitHub Actions Inactive

pankajkoti force-pushed the fix-async-bigquery-af3 branch from 6e6d083 to fe592f5 Compare April 16, 2025 17:35

pankajkoti temporarily deployed to internal April 16, 2025 17:36 — with GitHub Actions Inactive

add example DAG symlink and remove DB ORM access

4cd2f6a

pankajkoti force-pushed the fix-async-bigquery-af3 branch from fe592f5 to 4cd2f6a Compare April 16, 2025 17:46

pankajkoti temporarily deployed to internal April 16, 2025 17:46 — with GitHub Actions Inactive

pankajkoti marked this pull request as ready for review April 17, 2025 08:20

Copilot AI review requested due to automatic review settings April 17, 2025 08:20

dosubot Bot added the size:M This PR changes 30-99 lines, ignoring generated files. label Apr 17, 2025

pankajkoti requested review from pankajastro and tatiana April 17, 2025 08:20

dosubot Bot added area:dependencies Related to dependencies, like Python packages, library versions, etc area:execution Related to the execution environment/mode, like Docker, Kubernetes, Local, VirtualEnv, etc labels Apr 17, 2025

Copilot AI reviewed Apr 17, 2025

View reviewed changes

Comment thread cosmos/operators/_asynchronous/bigquery.py

tatiana approved these changes Apr 17, 2025

View reviewed changes

Comment thread cosmos/operators/_asynchronous/bigquery.py

Comment thread scripts/airflow3/env.sh

Comment thread cosmos/operators/_asynchronous/bigquery.py

dosubot Bot added the lgtm This PR has been approved by a maintainer label Apr 17, 2025

Add comment on why we alter the base classes twice

fee4f95

pankajkoti temporarily deployed to internal April 17, 2025 10:17 — with GitHub Actions Inactive

pankajastro reviewed Apr 17, 2025

View reviewed changes

Comment thread dev/dags/simple_dag_async.py Outdated

pankajastro approved these changes Apr 17, 2025

View reviewed changes

pankajkoti mentioned this pull request Apr 17, 2025

Validate Cosmos ExecutionMode.VIRTUALENV operators with Airflow 3 #1633

Closed

Merge branch 'main' into fix-async-bigquery-af3

3ff7de7

pankajkoti temporarily deployed to internal April 17, 2025 10:22 — with GitHub Actions Inactive

pankajastro approved these changes Apr 17, 2025

View reviewed changes

pankajkoti commented Apr 17, 2025

View reviewed changes

Comment thread dev/dags/simple_dag_async.py Outdated

Update dev/dags/simple_dag_async.py

27a9681

pankajkoti temporarily deployed to internal April 17, 2025 10:56 — with GitHub Actions Inactive

tatiana merged commit 30b6222 into main Apr 17, 2025

tatiana deleted the fix-async-bigquery-af3 branch April 17, 2025 10:58

tatiana mentioned this pull request Apr 17, 2025

Release 1.10.0 #1672

Merged

tatiana changed the title ~~Add Airflow 3 compatibility for BiqQuery deferrable support (ExecutionMode.AIRFLOW_ASYNC)~~ Add Airflow 3 compatibility for BiqQuery deferrable support (ExecutionMode.AIRFLOW_ASYNC) Apr 17, 2025

This was referenced Apr 17, 2025

Remove original_jaffle_shop dbt project #1676

Merged

Consolidate existing dbt projects into a single jaffle shop project #1677

Open

tatiana added this to the Cosmos 1.10.0 milestone Apr 17, 2025

pankajkoti mentioned this pull request Apr 17, 2025

Fix python_virtualenv.prepare_env top-level import for AF3 #1678

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add Airflow 3 compatibility for BiqQuery deferrable support (`ExecutionMode.AIRFLOW_ASYNC`)#1674

Add Airflow 3 compatibility for BiqQuery deferrable support (`ExecutionMode.AIRFLOW_ASYNC`)#1674
tatiana merged 4 commits into
mainfrom
fix-async-bigquery-af3

pankajkoti commented Apr 16, 2025

Uh oh!

netlify Bot commented Apr 16, 2025 •

edited

Loading

Uh oh!

cloudflare-workers-and-pages Bot commented Apr 16, 2025 •

edited

Loading

Uh oh!

netlify Bot commented Apr 16, 2025 •

edited

Loading

Uh oh!

codecov Bot commented Apr 16, 2025 •

edited

Loading

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

tatiana left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

pankajkoti commented Apr 16, 2025

Changes

Uh oh!

netlify Bot commented Apr 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ Deploy Preview for sunny-pastelito-5ecb04 canceled.

Uh oh!

cloudflare-workers-and-pages Bot commented Apr 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Deploying astronomer-cosmos with Cloudflare Pages

Uh oh!

netlify Bot commented Apr 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ Deploy Preview for sunny-pastelito-5ecb04 canceled.

Uh oh!

codecov Bot commented Apr 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

tatiana left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

netlify Bot commented Apr 16, 2025 •

edited

Loading

cloudflare-workers-and-pages Bot commented Apr 16, 2025 •

edited

Loading

netlify Bot commented Apr 16, 2025 •

edited

Loading

codecov Bot commented Apr 16, 2025 •

edited

Loading