Astronomer Cosmos

import os
from datetime import datetime
from pathlib import Path

from airflow.models.dag import DAG
from cosmos import DbtDag, ProjectConfig, ProfileConfig, ExecutionConfig
from cosmos.profiles import PostgresUserPasswordProfileMapping

# Define path to your dbt project relative to the DAG file
# For this example, assume a dbt project named 'jaffle_shop' exists in 'dags/dbt/'
# Example structure: dags/my_cosmos_dag.py, dags/dbt/jaffle_shop/dbt_project.yml
DEFAULT_DBT_PROJECT_PATH = Path(__file__).parent / "dbt" / "jaffle_shop"

# Configure your dbt profile to use an Airflow connection
# Replace 'your_airflow_postgres_conn_id' with an actual Airflow connection ID
# and adjust profile_args as needed for your database schema.
profile_config = ProfileConfig(
    profile_name="default", # Corresponds to the profile name in your dbt_project.yml
    target_name="dev",      # Corresponds to the target name in your dbt_project.yml
    profile_mapping=PostgresUserPasswordProfileMapping(
        conn_id=os.environ.get('AIRFLOW_DB_CONN_ID', 'airflow_db'), # Airflow connection ID
        profile_args={"schema": "public"}, # Example schema
    ),
)

# Configure dbt execution, e.g., to run dbt in a virtual environment.
# Ensure 'dbt_executable_path' points to your dbt virtual environment's executable.
# If dbt is installed in the same environment as Airflow, this might be omitted or set to 'dbt'.
execution_config = ExecutionConfig(
    dbt_executable_path=os.environ.get('DBT_EXECUTABLE_PATH', '/opt/airflow/dbt_venv/bin/dbt'),
)

with DAG(
    dag_id="cosmos_dbt_example_dag",
    schedule_interval="@daily",
    start_date=datetime(2023, 1, 1),
    catchup=False,
    tags=["dbt", "cosmos", "example"],
    default_args={
        "retries": 2, # Recommended for dbt tasks
    },
) as dag:
    dbt_project = DbtDag(
        project_config=ProjectConfig(
            DEFAULT_DBT_PROJECT_PATH,
        ),
        profile_config=profile_config,
        execution_config=execution_config,
        operator_args={
            "install_deps": True, # Install dbt dependencies before run
        },
    )

    # The dbt_project object automatically creates Airflow tasks for your dbt models,
    # which can be observed in the Airflow UI with their lineage.
    # You can add upstream/downstream Airflow tasks as needed.
    # For example, a simple PythonOperator before or after the dbt_project.