Refactor AIRFLOW_ASYNC so that the path in the remote object store is specific per DAG run

[tuantran0910]

Description

Refactor AIRFLOW_ASYNC so that the path in the remote object store is specific per DAG run. The format of remote model path will be:

# test_cosmos/simple_dag_async/run/jaffle_shop/models/example/my_first_dbt_model.sql
remote_model_path = f"{remote_target_path_str}/{dbt_dag_task_group_identifier}/{run_id}/run/{relative_file_path}"

Related Issue(s)

Closes #1613

[netlify]

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

Name	Link
🔨 Latest commit	4553787
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/6817a68b4fad2600085abb3e

[Copilot]

Pull Request Overview

This PR refactors the AIRFLOW_ASYNC remote model path generation to include the run_id, ensuring that paths in the remote object store are specific per DAG run.

• Updates file path construction in local and asynchronous operators to incorporate run_id.
• Adds tests to verify that the run_id is correctly embedded in the constructed paths.
• Adjusts async context initialization in asynchronous operators to include run_id.

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
tests/test_async_remote_target_dir_specific_run_id.py	New tests validate that run_id is correctly used in file path building.
cosmos/operators/local.py	Updates _construct_dest_file_path to inject run_id into the path.
cosmos/operators/_asynchronous/bigquery.py	Modifies remote SQL path construction to include run_id.
cosmos/operators/_asynchronous/init.py	Updates async_context creation to include run_id for async operators.

[netlify]

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

Name	Link
🔨 Latest commit	25367d2
🔍 Latest deploy log	https://app.netlify.com/projects/sunny-pastelito-5ecb04/deploys/682cbae3f14ccb000847ca0f

[pankajastro]

Thanks a lot for the contribution, @tuantran0910! I've triggered the CI for the integration tests.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.72%. Comparing base (15a8d91) to head (25367d2).
Report is 3 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1741   +/-   ##
=======================================
  Coverage   97.71%   97.72%           
=======================================
  Files          84       84           
  Lines        5252     5264   +12     
=======================================
+ Hits         5132     5144   +12     
  Misses        120      120           

☔ 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.

[pankajastro]

Currently, the teardown task deletes SQL files through multiple calls, and the goal was to reduce the risk of deleting unexpected files when the task runs in parallel. Now that we're uploading based on the run-id, would it be worthwhile to delete the entire run-id folder directly, rather than iterating over each file? What are your thoughts?

astronomer-cosmos/cosmos/operators/local.py

Line 338 in 9a334a8

def _delete_sql_files(self, tmp_project_dir: Path, resource_type: str) -> None:

This may also apply to the setup task.

[tuantran0910]

Currently, the teardown task deletes SQL files through multiple calls, and the goal was to reduce the risk of deleting unexpected files when the task runs in parallel. Now that we're uploading based on the run-id, would it be worthwhile to delete the entire run-id folder directly, rather than iterating over each file? What are your thoughts?

astronomer-cosmos/cosmos/operators/local.py

Line 338 in 9a334a8

def _delete_sql_files(self, tmp_project_dir: Path, resource_type: str) -> None:

This may also apply to the setup task.

Hmm, I think it's a good idea to delete the entire folder run-id for optimizing the workflow. I will implement this later.

[tuantran0910]

Hi @pankajastro, can you checkout the logic to delete run-id specific directory at my new commit 2bf7258. Thanks a lot :D

[tatiana]

This looks great, @tuantran0910 ! Thank you very much for improving this part of Cosmos and making it more reliable.

I've given some feedback in-line. Also, it seems there is an integration test failing, could you take a look please.

[tuantran0910]

Sorry @pankajastro, can you enable the integration tests and approve this PR again ? I have just pushed a new commit to ensure that all the changes are coverage.

[tatiana]

@tuantran0910 Thank you for all the patience on getting this to work. There is a minor request regarding code coverage - if you could address this as part of this PR, it would be great. If you can't do this before the release of 1.10.1, please, could you do a follow-up PR?

[tuantran0910]

@tuantran0910 Thank you for all the patience on getting this to work. There is a minor request regarding code coverage - if you could address this as part of this PR, it would be great. If you can't do this before the release of 1.10.1, please, could you do a follow-up PR?

Hey @tatiana, I will try to push a quick commit to fix the coverage within an hour. Thank you for pointing that.

[tuantran0910]

Hi @tatiana, I have added the test in fe92ef0. Thank you so much :D

[pankajkoti]

Thanks for improvising this @tuantran0910 . This is certainly a very valuable contribution 👏🏽

I tested this with an example async DAG running in Airflow UI and it works smoothly as intended for. We will include this fix in the upcoming release 1.10.1.