Experimental BQ support to run dbt models with `ExecutionMode.AIRFLOW_ASYNC` by pankajastro · Pull Request #1230 · astronomer/astronomer-cosmos

pankajastro · 2024-09-29T18:53:25Z

This work has been inspired by the talk "Airflow at Monzo: Evolving our data platform as the bank scales" by @jonathanrainer @ed-sparkes given at Airflow Summit 2023: https://airflowsummit.org/sessions/2023/airflow-at-monzo-evolving-our-data-platform-as-the-bank-scales/.

Enable BQ users to run dbt models (full_refresh) asynchronously. This releases the Airflow worker node from waiting while the transformation (I/O) happens in the dataware house (some of which can take hours, according to customers), increasing the overall Airflow task throughput (more information: https://airflow.apache.org/docs/apache-airflow/stable/authoring-and-scheduling/deferring.html). As part of this change, we introduce the capability of not using the dbt command to run actual SQL transformations. This also avoids creating subprocesses in the worker node (ExecutionMode.LOCAL with InvocationMode. SUBPROCESS and ExecutionMode.VIRTUALENV) or the overhead of creating a Kubernetes Pod to execute the actual dbt command (ExecutionMode.KUBERNETES). This can avoid issues related to memory and CPU usage.

This PR takes advantage of an already implemented async operator in the Airflow repo by extending it in the Cosmos async operator. It also utilizes the pre-compiled SQL generated as part of the PR #1224. It downloads the generated SQL from a remote location (S3/GCS), which allows us to decouple from dbt during task execution.

Details

Expose get_profile_type on ProfileConfig: This aids in database selection
~~Add async_op_args: A high-level parameter to forward arguments to the upstream operator (Airflow operator). (This may change in this PR itself)~~ The async operator params are process as kwargs in the operator_args parameter
Implement DbtRunAirflowAsyncOperator: This initializes the Airflow Operator, retrieves the SQL query at task runtime from a remote location, modifies the query as needed, and triggers the upstream execute method.

Limitations

This feature only works when using Airflow 2.8 and above
The async execution only works for BigQuery
The async execution only supports running dbt models (other dbt resources, such as seeds, sources, snapshots, tests, are run using the ExecutionMode.LOCAL)
This will work only if the user provides sets full_refresh=True in operator_args (which means tables will be dropped before being populated, as implemented in dbt-core)
Users need to use ProfileMapping in ProfileConfig, since Cosmos relies on having the connection (credentials) to be able to run the transformation in BQ without dbt-core
Users must provide the BQ location in operator_args (this is a limitation from the BigQueryInsertJobOperator that is being used to implement the native Airflow asynchronous support)

Testing

We have added a new dbt project to the repository to facilitate asynchronous task execution. The goal is to accelerate development without disrupting or requiring fixes for the existing tests. Also, we have added DAG for end-to-end testing https://github.com/astronomer/astronomer-cosmos/blob/bd6657a29b111510fc34b2baf0bcc0d65ec0e5b9/dev/dags/simple_dag_async.py

Configuration

Users need to configure the below param to execute deferrable tasks in the Cosmos

Example DAG: https://github.com/astronomer/astronomer-cosmos/blob/bd6657a29b111510fc34b2baf0bcc0d65ec0e5b9/dev/dags/simple_dag_async.py

Installation

You can leverage async operator support by installing an additional dependency

astronomer-cosmos[dbt-bigquery, google]

Documentation

The PR also document the limitations and uses of Airflow async execution in the Cosmos.

Related Issue(s)

Related to: #1120
Closes: #1134

Breaking Change?

No

Notes

This is an experimental feature, and as such, it may undergo breaking changes. We encourage users to share their experiences and feedback to improve it further.

We'd love support and feedback so we can define the next steps.

Checklist

I have made corresponding changes to the documentation (if required)
I have added tests that prove my fix is effective or that my feature works

Credits

This was a result of teamwork and effort:

Co-authored-by: Pankaj Koti pankajkoti699@gmail.com
Co-authored-by: Tatiana Al-Chueyr tatiana.alchueyr@gmail.com

Future Work

Design interface to facilitate the easy addition of new asynchronous databases operators [async] Design ExecutionMode.AIRFLOW_ASYNC interface to incorporate additional database #1238
Improve the test coverage [async] Improve test coverage for BigQuery Async task #1239
Address the limitations (we need to log these issues)

Remove print stmt Fix query Fix query Remove oss execute method code

…L from

This PR is the groundwork for the implementation of `ExecutionMode.AIRFLOW_ASYNC` (#1120), which - once all other epic tasks are completed - will enable asynchronous execution of dbt resources using Apache Airflow’s deferrable operators. As part of this work, this PR introduces a new option to the enum `ExecutionMode` : `AIRFLOW_ASYNC`. When this execution mode is used, Cosmos now creates a setup task that will pre-compile the dbt project SQL and make it available to the remaining dbt tasks. This PR, however, does not yet leverage Airflow's deferrable operators. If users use `ExecutionMode.AIRFLOW_ASYNC` they will actually be running `ExecutionMode.LOCAL` operators with this change. The PR (#1230) has a first experimental version of using deferrable operators for task execution. ## Setup task as the ground work for a new Execution Mode: `ExecutionMode.AIRFLOW_ASYNC`: - Adds a new operator, `DbtCompileAirflowAsyncOperator`, as a root task(analogous to a setup task) in the DAG, running the dbt compile command and uploading the compiled SQL files to a remote storage location for subsequent tasks that fetch these compiled SQL files from the remote storage and run them asynchronously using Airflow's deferrable operators. ## Airflow Configurations: - `remote_target_path`: Introduces a configurable path to store dbt-generated files remotely, supporting any storage scheme that works with Airflow’s Object Store (e.g., S3, GCS, Azure Blob). - `remote_target_path_conn_id`: Allows specifying a custom connection ID for the remote target path, defaulting to the scheme’s associated Airflow connection if not set. ## Example DAG for CI Testing: Introduces an example DAG (`simple_dag_async.py`) demonstrating how to use the new execution mode(The execution like mentioned earlier would still run like Execution Mode LOCAL operators at the moment with this PR alone) This DAG is integrated into the CI pipeline to run integration tests and aims at verifying the functionality of the `ExecutionMode.AIRFLOW_ASYNC` as and when implementation gets added starting with the experimental implementation in #1230 . ## Unit & Integration Tests: - Adds comprehensive unit and integration tests to ensure correct behavior. - Tests include validation for successful uploads, error handling for misconfigured remote paths, and scenarios where `remote_target_path` are not set. ## Documentation: - Adds detailed documentation explaining how to configure and set the `ExecutionMode.AIRFLOW_ASYNC`. ## Scope & Limitations of the feature being introduced: 1. This feature is meant to be released as Experimental and is also marked so in the documentation. 2. Currently, it has been scoped for only dbt models to be executed asynchronously (being worked upon in PR #1230), while other resource types would be run synchronously. 3. `BigQuery` will be the only supported target database for this execution mode ((being worked upon in PR #1230). Thus, this PR enhances Cosmos by providing the ground work for more efficient execution of long-running dbt resources ## Additional Notes: - This feature is planned to be introduced in Cosmos v1.7.0. related: #1134

netlify · 2024-09-30T22:23:06Z

✅ Deploy Preview for sunny-pastelito-5ecb04 ready!

Name	Link
🔨 Latest commit	`78bc069`
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/66fb24b60ca64b00085c5976
😎 Deploy Preview	https://deploy-preview-1230--sunny-pastelito-5ecb04.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

To edit notification comments on pull requests, go to your Netlify site configuration.

netlify · 2024-09-30T22:25:32Z

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

Name	Link
🔨 Latest commit	`a0cb147`
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/66fe73da6f73b00008c98193

tatiana · 2024-10-02T11:47:51Z

On the issue:

There are two issues we have to solve before merging this PR

CI is getting stuck when trying to run
tests/test_example_dags.py::test_example_dag[simple_dag_async]

We identified that Cosmos was using its own dag.test that was buggy when used with Airflow deferrable operators. We changed to use Airflow's one if users are using AF >= 2.5 and it seems to have resolved the issue:
16a87ea (CI jobs were no longer stuck)

tatiana

Incredible work, @pankajastro @pankajkoti ! 🎉

To run dbt transformations without dbt while leveraging Airflow asynchronous processing is far from trivial - and this PR is a significant step in the right direction.

codecov · 2024-10-02T14:09:50Z

Codecov Report

Attention: Patch coverage is 89.74359% with 12 lines in your changes missing coverage. Please review.

Project coverage is 95.73%. Comparing base (225d6e9) to head (a0cb147).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
cosmos/config.py	36.36%	7 Missing ⚠️
cosmos/operators/airflow_async.py	94.04%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1230      +/-   ##
==========================================
- Coverage   95.93%   95.73%   -0.20%     
==========================================
  Files          67       67              
  Lines        3885     3965      +80     
==========================================
+ Hits         3727     3796      +69     
- Misses        158      169      +11

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

jlaneve · 2024-10-02T14:10:38Z

I haven't read through the code but this feature looks so awesome - it's something we've been talking about since day 1 so it's incredible to see it come to life!

pankajastro · 2024-10-03T13:05:45Z

Thank you so much, @pankajkoti and @tatiana, for taking this to completion. I truly appreciate your support and hard work 🙏

tatiana · 2024-10-21T09:30:11Z

I just created a few follow-up tickets:

pankajkoti added 3 commits September 30, 2024 02:29

Draft: dbt compile task

851564f

Put compiled files under dag_id folder & refactor few snippets

9dc2c9c

Add tests & minor refactorings

0ce662e

pankajkoti force-pushed the poc-dbt-compile-task branch from 3f00cc9 to 0ce662e Compare September 29, 2024 20:59

pankajkoti added 2 commits September 30, 2024 02:40

Apply suggestions from code review

1b6f57e

Install deps for the newly added example DAG

cc48161

pankajkoti mentioned this pull request Sep 30, 2024

Introduce ground work for ExecutionMode.AIRFLOW_ASYNC #1224

Merged

Add docs

1068025

pankajastro force-pushed the execute-async-task branch 2 times, most recently from 35e58b6 to 407d311 Compare September 30, 2024 09:32

Add async run operator

faa706d

Remove print stmt Fix query Fix query Remove oss execute method code

pankajastro force-pushed the execute-async-task branch from 407d311 to faa706d Compare September 30, 2024 09:42

pankajastro added 2 commits September 30, 2024 15:51

Fix remote sql path and async args

0e155e4

Fix query

5f1ecaa

pankajastro force-pushed the execute-async-task branch from 7916857 to 5f1ecaa Compare September 30, 2024 14:28

pre-commit-ci Bot and others added 2 commits September 30, 2024 14:28

🎨 [pre-commit.ci] Auto format from pre-commit.com hooks

1278847

Use dbt node's filepath to construct remote path to fetch compiled SQ…

b3d6cf3

…L from

pankajkoti force-pushed the poc-dbt-compile-task branch from 2e56a75 to 77c7c6c Compare September 30, 2024 21:13

Base automatically changed from poc-dbt-compile-task to main September 30, 2024 21:57

Merge branch 'main' into execute-async-task

78bc069

tatiana temporarily deployed to internal September 30, 2024 22:22 — with GitHub Actions Inactive

🎨 [pre-commit.ci] Auto format from pre-commit.com hooks

9ca5e85

pre-commit-ci Bot temporarily deployed to internal September 30, 2024 22:23 Inactive

tatiana added 2 commits October 1, 2024 00:31

Fix unittests

99bf7c0

Improve code

3aaaf9e

tatiana temporarily deployed to internal September 30, 2024 23:32 — with GitHub Actions Inactive

tatiana reviewed Sep 30, 2024

View reviewed changes

Comment thread .github/workflows/test.yml Outdated

pankajkoti reviewed Oct 2, 2024

View reviewed changes

Comment thread cosmos/operators/base.py Outdated

Comment thread cosmos/operators/base.py Outdated

Apply suggestions from code review

ea5816b

pankajkoti temporarily deployed to internal October 2, 2024 11:44 — with GitHub Actions Inactive

tatiana approved these changes Oct 2, 2024

View reviewed changes

dosubot Bot added the lgtm This PR has been approved by a maintainer label Oct 2, 2024

pankajkoti mentioned this pull request Oct 3, 2024

Try upgrading dbt-core & dbt adapters versions to see if the dependencies resolve #1235

Closed

Add install instruction

85f86a4

pankajastro temporarily deployed to internal October 3, 2024 09:01 — with GitHub Actions Inactive

Add min airflow version in limitation

402f823

pankajastro temporarily deployed to internal October 3, 2024 09:08 — with GitHub Actions Inactive

Ignore Async DAG for dbt <=1.5

621a4de

pankajastro temporarily deployed to internal October 3, 2024 09:58 — with GitHub Actions Inactive

Ignore Async DAG for dbt <=1.5

a0cb147

pankajastro temporarily deployed to internal October 3, 2024 10:37 — with GitHub Actions Inactive

This was referenced Oct 3, 2024

Bump dbt adaptor version for sqlite test #1236

Closed

[async] Design ExecutionMode.AIRFLOW_ASYNC interface to incorporate additional database #1238

Closed

[async] Improve test coverage for BigQuery Async task #1239

Closed

tatiana merged commit 111d430 into main Oct 3, 2024

tatiana deleted the execute-async-task branch October 3, 2024 16:13

pankajkoti mentioned this pull request Oct 3, 2024

Refactor docs for async mode execution #1241

Merged

tatiana mentioned this pull request Oct 21, 2024

[bug] Fix ExecutionMode.AIRFLOW_ASYNC query #1260

Closed

2 tasks

tatiana mentioned this pull request Oct 21, 2024

[async] Evaluate the possiblity of using dbt itself to create the full SQL command #1266

Closed

1 task

pankajastro mentioned this pull request Jan 23, 2025

Add structure to support multiple db for async operator execution #1483

Merged

pankajkoti mentioned this pull request Feb 4, 2025

Create and run accurate SQL statements when using ExecutionMode.AIRFLOW_ASYNC #1474

Merged

pankajkoti mentioned this pull request Apr 17, 2025

Remove original_jaffle_shop dbt project #1676

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Experimental BQ support to run dbt models with `ExecutionMode.AIRFLOW_ASYNC`#1230

Experimental BQ support to run dbt models with `ExecutionMode.AIRFLOW_ASYNC`#1230
tatiana merged 51 commits into
mainfrom
execute-async-task

pankajastro commented Sep 29, 2024 •

edited by tatiana

Loading

Uh oh!

netlify Bot commented Sep 30, 2024 •

edited

Loading

Uh oh!

netlify Bot commented Sep 30, 2024 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tatiana commented Oct 2, 2024

Uh oh!

tatiana left a comment

Uh oh!

codecov Bot commented Oct 2, 2024 •

edited

Loading

Uh oh!

jlaneve commented Oct 2, 2024

Uh oh!

pankajastro commented Oct 3, 2024

Uh oh!

tatiana commented Oct 21, 2024 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Conversation

pankajastro commented Sep 29, 2024 • edited by tatiana Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Details

Limitations

Testing

Configuration

Installation

Documentation

Related Issue(s)

Breaking Change?

Notes

Checklist

Credits

Future Work

Uh oh!

netlify Bot commented Sep 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ Deploy Preview for sunny-pastelito-5ecb04 ready!

Uh oh!

netlify Bot commented Sep 30, 2024 • 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!

Uh oh!

Uh oh!

Uh oh!

tatiana commented Oct 2, 2024

Uh oh!

tatiana left a comment

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Oct 2, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

jlaneve commented Oct 2, 2024

Uh oh!

pankajastro commented Oct 3, 2024

Uh oh!

tatiana commented Oct 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

pankajastro commented Sep 29, 2024 •

edited by tatiana

Loading

netlify Bot commented Sep 30, 2024 •

edited

Loading

netlify Bot commented Sep 30, 2024 •

edited

Loading

codecov Bot commented Oct 2, 2024 •

edited

Loading

tatiana commented Oct 21, 2024 •

edited

Loading