Support using DatasetAlias and fix orphaning unreferenced dataset

[tatiana]

Context

Cosmos versions between 1.1 and 1.6 supported automatically emitting Airflow Datasets when using ExecutionMode.LOCAL. Although the datasets generated by these versions of Cosmos could be utilised for dataset-aware scheduling, the implementation had long-standing issues, as described in #522.

The main problems were:

• Orphaning unreferenced dataset
• Not displaying dataset inlets/outlets in the Airflow UI

These issues were caused by Cosmos defining task outlets and inlets during Airflow task execution, a feature only partially supported before Airflow 2.10: apache/airflow#34206.

Solution

Airflow 2.10 has introduced the concept of DatasetAlias, as described in the official docs, so operators can dynamically define inlets and outlets during task execution.

This PR uses Airflow DatasetAlias, when possible (Airflow 2.10 or above), and does two things:

1. Adds a DatasetAlias to every LocalOperator/VirtualenvOperator subclass
2. Dynamically adds Dataset as outlets during the LocalOperator subclasses execution, associating them to the desired Dataset instance.
3. Exposes to users a function to retrieve Cosmos' DatasetAlias names, programatically

Caveats

• Only works for Airflow 2.10 and above

This feature relies on DatasetAlias, only available in Airflow 2.10 and above. If users use previous versions of Airflow, Cosmos behaves like it did before, and the issues described in this task are not solved.

• Unable to leverage DatasetAlias in airflow dags test

Although the feature described in this PR works well when scheduling DAGs, triggering them via the UI, or using airflow dags trigger, it does not work when users attempt to use dags test or dag.test(). When trying to test the DAG, these commands fail with an sqlalchemy.orm. etc.FlushError. This is a known issue from Airflow 2.10.0 and Airflow 2.10.1 when declaring DatasetAliases, as described in apache/airflow#42495.

To mitigate this second problem, we've introduced a new Airflow variable, AIRFLOW__COSMOS__ENABLE_DATASET_ALIAS, that allows users to disable using dataset aliases when running Cosmos. We'd recommend users who face the sqlalchemy.orm. etc.FlushError in their tests to set this configuration to False only for running tests - until the issue is solved in Airflow. When this configuration is set to False, Cosmos behaves as before the DatasetAlias feature was introduced.

How this feature was validated

Manually triggering this DAG bq_profile_example:

from datetime import datetime

from cosmos import DbtDag, ExecutionConfig, ProjectConfig, ProfileConfig, RenderConfig
from cosmos.constants import TestBehavior, LoadMode

from include.constants import jaffle_shop_path, venv_execution_config

project_config = ProjectConfig(
    dbt_project_path=jaffle_shop_path,
    #partial_parse=True,
    #dbt_vars=self.dbt_configuration,
    #env_vars = self.param_handler.env_vars
)

profile_config = ProfileConfig(
    # these map to dbt/jaffle_shop/profiles.yml
    profile_name="airflow_db",
    target_name="bq",
    profiles_yml_filepath=jaffle_shop_path / "profiles.yml",
)

dbt_profile_example = DbtDag(
    project_config=project_config,
    profile_config=profile_config,
    render_config=render_config,
    execution_config=venv_execution_config,
    schedule_interval="@daily",
    start_date=datetime(2023, 1, 1),
    catchup=False,
    dag_id="bq_profile_example",
    tags=["profiles"],
)

And seeing that the dataset-scheduled DAG dataset_triggered_dag dependent on the first was triggered:

from datetime import datetime

from airflow import DAG
from airflow.datasets import Dataset
from airflow.operators.empty import EmptyOperator


with DAG(
    "dataset_triggered_dag",
    description="A DAG that should be triggered via Dataset/Dataset alias",
    start_date=datetime(2024, 9, 1),
    schedule=[Dataset(uri='bigquery/observability-sandbox-344122.moonactive.stg_customers')],
) as dag:
    t1 = EmptyOperator(
            task_id="task_1",
        )
    t2 = EmptyOperator(
            task_id="task_2",
        )
    t3 = EmptyOperator(
            task_id="task_3",
        )
    
    t1 >> t2 >> t3

Checking Airflow UI:

Related tickets

Closes: #522
Closes: #1119

[netlify]

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

Name	Link
🔨 Latest commit	43c9ff0
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/66fae7d4aa6e3d0008a90988
😎 Deploy Preview	https://deploy-preview-1217--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.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 95.82%. Comparing base (e0a9fd3) to head (43c9ff0).
Report is 1 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1217      +/-   ##
==========================================
+ Coverage   95.78%   95.82%   +0.04%     
==========================================
  Files          65       66       +1     
  Lines        3745     3783      +38     
==========================================
+ Hits         3587     3625      +38     
  Misses        158      158              

Flag	Coverage Δ
	95.82% <100.00%> (+0.04%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

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

[pankajkoti]

LGTM, excited for this one! 👏🏽