Support rendering DbtDag in Airflow 3

[tatiana]

Add necessary changes so we can render Cosmos DbtDag using LoadMode.DBT_LS when using Airflow 3 main branch (commit 9d6bebe9690a3f21a00044c0aeceaded86484a0c).

Previously, we were unable to run the DAG:

import os
from datetime import datetime
from pathlib import Path
from airflow.models import Variable

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

DEFAULT_DBT_ROOT_PATH = Path(__file__).parent.parent.parent.parent / "dev/dags/dbt"
DBT_ROOT_PATH = Path(os.getenv("DBT_ROOT_PATH", DEFAULT_DBT_ROOT_PATH))
DBT_PROJ_DIR = DBT_ROOT_PATH / "jaffle_shop"
DBT_PROFILE_PATH = DBT_PROJ_DIR / "profiles.yml"
DBT_ARTIFACT = DBT_PROJ_DIR / "target"

project_config = ProjectConfig(
    dbt_project_path=DBT_PROJ_DIR,
    project_name="jaffle_shop",
)

profile_config = ProfileConfig(
    profile_name="default",
    target_name="dev",
    profiles_yml_filepath=DBT_PROFILE_PATH,
)

basic_cosmos_dag = DbtDag(
    # dbt/cosmos-specific parameters
    project_config=project_config,
    profile_config=profile_config,
    operator_args={
        "install_deps": True,  # install any necessary dependencies before running any dbt command
        "full_refresh": True,  # used only in dbt commands that support this flag
    },
    # normal dag parameters
    schedule="@daily",
    start_date=datetime(2023, 1, 1),
    catchup=False,
    dag_id="basic_cosmos_dag",
    default_args={

It would raise the following exception when we attempted to run airflow dags reserialize:

i  File "/Users/tati/Code/cosmos-fresh/astronomer-cosmos/cosmos/airflow/graph.py", line 592, in build_airflow_graph
    task_or_group = conversion_function(  # type: ignore
  File "/Users/tati/Code/cosmos-fresh/astronomer-cosmos/cosmos/airflow/graph.py", line 391, in generate_task_or_group
    task >> test_task
  File "/Users/tati/Code/cosmos-fresh/astronomer-cosmos/venv-af3-main/lib/python3.10/site-packages/airflow/sdk/definitions/_internal/mixins.py", line 97, in __rshift__
    self.set_downstream(other)
  File "/Users/tati/Code/cosmos-fresh/astronomer-cosmos/venv-af3-main/lib/python3.10/site-packages/airflow/sdk/definitions/_internal/node.py", line 216, in set_downstream
    self._set_relatives(task_or_task_list, upstream=False, edge_modifier=edge_modifier)
  File "/Users/tati/Code/cosmos-fresh/astronomer-cosmos/venv-af3-main/lib/python3.10/site-packages/airflow/sdk/definitions/_internal/node.py", line 186, in _set_relatives
    raise ValueError(
ValueError: Tried to create relationships between tasks that don't have DAGs yet. Set the DAG for at least one task and try again: [<Task(DbtRunLocalOperator): run>, <Task(DbtTestLocalOperator): test>]

This issue happened due to an interface change in Airflow 3's BaseOperator, that would lead the line:

base_operator_args = set(inspect.signature(BaseOperator.__init__).parameters.keys())

To return zero arguments.

This was solved by changing the import path:

try:  # Airflow 3
    from airflow.sdk.bases.operator import BaseOperator
except ImportError:  # Airflow 2
    from airflow.models import BaseOperator

After these changes, we were able to successfully run:

• airflow dags reserialize
• airflow dags test basic_cosmos_dag

Co-authored-by: Ash Berlin-Taylor ash@astronomer.io

[netlify]

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

Name	Link
🔨 Latest commit	4629518
🔍 Latest deploy log	https://app.netlify.com/sites/sunny-pastelito-5ecb04/deploys/67f68f8011c9270008cb1220

[Copilot]

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

[cloudflare-workers-and-pages]

Deploying astronomer-cosmos with    Cloudflare Pages

Latest commit:	4629518
Status:	✅  Deploy successful!
Preview URL:	https://eb8c0d59.astronomer-cosmos.pages.dev
Branch Preview URL:	https://fix-dag-rendering-af3main.astronomer-cosmos.pages.dev

View logs

[tatiana]

The test coverage for this PR will fail, but it is expected to be solved once the following ticket is implemented:
#1646

[pankajastro]

Nice!

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 97.47%. Comparing base (a0a1cb5) to head (4629518).
Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1657   +/-   ##
=======================================
  Coverage   97.47%   97.47%           
=======================================
  Files          80       80           
  Lines        4982     4991    +9     
=======================================
+ Hits         4856     4865    +9     
  Misses        126      126           

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

[pankajkoti]

so glad, it was this change and we did not have to do more altercations. Thanks a lot @tatiana for getting this to work 👏🏽

[tatiana]

More information on the current BaseOperator interface (only **kwargs):
https://github.com/apache/airflow/blob/00aec39a4c4987496ee365eca212657dec45be5e/airflow-core/src/airflow/models/baseoperator.py#L80-L288C1