Orchestrate dbt Core projects with Airflow and Cosmos

dbt Core is an open-source library for analytics engineering that helps users build interdependent SQL models for in-warehouse data transformation, using ephemeral compute of data warehouses. Cosmos is an open-source package developed by Astronomer to run dbt models that are part of a dbt Core project within Airflow.

dbt on Airflow with Cosmos and the Astro CLI

The open-source provider package Cosmos allows you to integrate dbt jobs into Airflow by automatically creating Airflow tasks from dbt models. You can turn your dbt Core projects into an Airflow DAG or task group with just a few lines of code.

Why use Airflow with dbt Core?

dbt Core offers the possibility to build modular, reusable SQL components with built-in dependency management and incremental builds.

With Cosmos, you can integrate dbt jobs into your open-source Airflow orchestration environment as standalone DAGs or as task groups within DAGs.

The benefits of using Airflow with dbt Core include:

• Use Airflow’s data-aware scheduling and Airflow sensors to run models depending on other events in your data ecosystem.
• Turn each dbt model into a task, complete with Airflow features like retries and error notifications, as well as full observability into past runs directly in the Airflow UI.
• Run dbt test on tables created by individual models immediately after a model has completed. Catch issues before moving downstream and integrate additional data quality checks with your preferred tool to run alongside dbt tests.
• Run dbt projects using Airflow connections instead of dbt profiles. You can store all your connections in one place, directly within Airflow or by using a secrets backend.
• Leverage native support for installing and running dbt in a virtual environment to avoid dependency conflicts with Airflow.

With Astro, you get all the above benefits without having to modify your Airflow environment or create connections. You can deploy directly from a dbt project using the Astro CLI. For more information, see Deploy dbt projects to Astro.

Time to complete

This tutorial takes approximately 30 minutes to complete.

Assumed knowledge

To get the most out of this tutorial, make sure you have an understanding of:

• The basics of dbt Core. See What is dbt?.
• Airflow fundamentals, such as writing DAGs and defining tasks. See Get started with Apache Airflow.
• How Airflow and dbt concepts relate to each other. See Similar dbt & Airflow concepts.
• Airflow operators. See Operators 101.
• Airflow task groups. See Airflow task groups.
• Airflow connections. See Manage connections in Apache Airflow.

Prerequisites

• The Astro CLI.
• Access to a data warehouse supported by dbt Core. See dbt documentation for all supported warehouses. This tutorial uses a Postgres database.

You do not need to have dbt Core installed locally in order to complete this tutorial.

Step 1: Configure your Astro project

To use dbt Core with Airflow install dbt Core in a virtual environment and Cosmos in a new Astro project.

1. Create a new Astro project:

   1
   $ mkdir astro-dbt-core-tutorial && cd astro-dbt-core-tutorial
   2
   $ astro dev init

2. Open the Dockerfile and add the following lines to the end of the file:

   # replace dbt-postgres with another supported adapter if you're using a different warehouse type
   RUN python -m venv dbt_venv && source dbt_venv/bin/activate && \
       pip install --no-cache-dir dbt-postgres && deactivate

   This code runs a bash command when the Docker image is built that creates a virtual environment called dbt_venv inside of the Astro CLI scheduler container. The dbt-postgres package, which also contains dbt-core, is installed in the virtual environment. If you are using a different data warehouse, replace dbt-postgres with the adapter package for your data warehouse.

3. Add Cosmos and the Postgres provider to your Astro project requirements.txt file. If you are using a different data warehouse, replace apache-airflow-providers-postgres with the provider package for your data warehouse. You can find information on all provider packages on the Astronomer registry.

   astronomer-cosmos==1
   apache-airflow-providers-postgres==6
   apache-airflow-providers-common-sql==1

Step 2: Prepare your dbt project

To integrate your dbt project with Airflow, you need to add the project folder to your Airflow environment. For this step you can either add your own project in a new dbt folder in your dags directory, or follow the steps below to create a simple project using two models.

1. Create a folder called dbt in your dags folder.

2. In the dbt folder, create a folder called my_simple_dbt_project.

3. In the my_simple_dbt_project folder add your dbt_project.yml. This configuration file needs to contain at least the name of the project. This tutorial additionally shows how to inject a variable called my_name from Airflow into your dbt project.

   1
   name: 'my_simple_dbt_project'
   2
   vars:
   3
       my_name: "No entry"

4. Add your dbt models in a subfolder called models in the my_simple_dbt_project folder. You can add as many models as you want to run. This tutorial uses the following two models:

   model1.sql:

   1
   SELECT '{{ var("my_name") }}' as name

   model2.sql:

   1
   SELECT * FROM {{ ref('model1') }}

   model1.sql selects the variable my_name. model2.sql depends on model1.sql and selects everything from the upstream model.

You should now have the following structure within your Airflow environment:

.
└── dags
    └── dbt
        └── my_simple_dbt_project
           ├── dbt_project.yml
           └── models
               ├── model1.sql
               └── model2.sql

Step 3: Create an Airflow connection to your data warehouse

Cosmos allows you to apply Airflow connections to your dbt project.

1. Start Airflow by running astro dev start.

2. In the Airflow UI, go to Admin -> Connections and click +.

3. Create a new connection named db_conn. Select the connection type and supplied parameters based on the data warehouse you are using. For a Postgres connection, enter the following information:

   • Connection ID: db_conn.
   • Connection Type: Postgres.
   • Host: Your Postgres host address.
   • Schema: Your Postgres database.
   • Login: Your Postgres login username.
   • Password: Your Postgres password.
   • Port: Your Postgres port.

Info

If a connection type for your database isn’t available, you might need to make it available by adding the relevant provider package to requirements.txt and running astro dev restart.

Step 4: Write your Airflow DAG

The DAG you’ll write uses Cosmos to create tasks from existing dbt models and the SQLExecuteQueryOperator to query a table that was created. You can add more upstream and downstream tasks to embed the dbt project within other actions in your data ecosystem.

1. In your dags folder, create a file called my_simple_dbt_dag.py.

2. Copy and paste the following DAG code into the file:

   1
   """
   2
   ### Run a dbt Core project as a task group with Cosmos
   3
   4
   Simple DAG showing how to run a dbt project as a task group, using
   5
   an Airflow connection and injecting a variable into the dbt project.
   6
   """
   7
   8
   from airflow.decorators import dag
   9
   from airflow.providers.postgres.operators.postgres import PostgresOperator
   10
   from cosmos import DbtTaskGroup, ProjectConfig, ProfileConfig, ExecutionConfig
   11
   12
   # adjust for other database types
   13
   from cosmos.profiles import PostgresUserPasswordProfileMapping
   14
   from pendulum import datetime
   15
   import os
   16
   17
   YOUR_NAME = "<your_name>"
   18
   CONNECTION_ID = "db_conn"
   19
   DB_NAME = "<your_db_name>"
   20
   SCHEMA_NAME = "<your_schema_name>"
   21
   MODEL_TO_QUERY = "model2"
   22
   # The path to the dbt project
   23
   DBT_PROJECT_PATH = f"{os.environ['AIRFLOW_HOME']}/dags/dbt/my_simple_dbt_project"
   24
   # The path where Cosmos will find the dbt executable
   25
   # in the virtual environment created in the Dockerfile
   26
   DBT_EXECUTABLE_PATH = f"{os.environ['AIRFLOW_HOME']}/dbt_venv/bin/dbt"
   27
   28
   profile_config = ProfileConfig(
   29
       profile_name="default",
   30
       target_name="dev",
   31
       profile_mapping=PostgresUserPasswordProfileMapping(
   32
           conn_id=CONNECTION_ID,
   33
           profile_args={"schema": SCHEMA_NAME},
   34
       ),
   35
   )
   36
   37
   execution_config = ExecutionConfig(
   38
       dbt_executable_path=DBT_EXECUTABLE_PATH,
   39
   )
   40
   41
   42
   @dag(
   43
       start_date=datetime(2023, 8, 1),
   44
       schedule=None,
   45
       catchup=False,
   46
       params={"my_name": YOUR_NAME},
   47
   )
   48
   def my_simple_dbt_dag():
   49
       transform_data = DbtTaskGroup(
   50
           group_id="transform_data",
   51
           project_config=ProjectConfig(DBT_PROJECT_PATH),
   52
           profile_config=profile_config,
   53
           execution_config=execution_config,
   54
           operator_args={
   55
               "vars": '{"my_name": {{ params.my_name }} }',
   56
           },
   57
           default_args={"retries": 2},
   58
       )
   59
   60
       query_table = PostgresOperator(
   61
           task_id="query_table",
   62
           postgres_conn_id=CONNECTION_ID,
   63
           sql=f"SELECT * FROM {DB_NAME}.{SCHEMA_NAME}.{MODEL_TO_QUERY}",
   64
       )
   65
   66
       transform_data >> query_table
   67
   68
   69
   my_simple_dbt_dag()

   This DAG uses the DbtTaskGroup class from the Cosmos package to create a task group from the models in your dbt project. Dependencies between your dbt models are automatically turned into dependencies between Airflow tasks. Make sure to add your own values for YOUR_NAME, YOUR_DB_NAME, and YOUR_SCHEMA_NAME.

   Using the vars keyword in the dictionary provided to the operator_args parameter, you can inject variables into the dbt project. This DAG injects YOUR_NAME for the my_name variable. If your dbt project contains dbt tests, they will be run directly after a model has completed. Note that it is a best practice to set retries to at least 2 for all tasks that run dbt models.

Tip

In some cases, especially in larger dbt projects, you might run into a DagBag import timeout error. This error can be resolved by increasing the value of the Airflow configuration core.dagbag_import_timeout.

3. Run the DAG manually by clicking the play button and view the DAG in the graph view. Double click the task groups in order to expand them and see all tasks.

   

4. Check the XCom returned by the query_table task to see your name in the model2 table.

Info

The DbtTaskGroup class populates an Airflow task group with Airflow tasks created from dbt models inside of a normal DAG. To directly define a full DAG containing only dbt models use the DbtDag class, as shown in the Cosmos documentation.

Congratulations! You’ve run a DAG using Cosmos to automatically create tasks from dbt models. You can learn more about how to configure Cosmos in the Cosmos documentation.

Alternative ways to run dbt Core with Airflow

While using Cosmos is recommended, there are several other ways to run dbt Core with Airflow.

Using the BashOperator

You can use the BashOperator to execute specific dbt commands. It’s recommended to run dbt-core and the dbt adapter for your database in a virtual environment because there often are dependency conflicts between dbt and other packages.

The DAG below uses the BashOperator to activate the virtual environment and execute dbt_run for a dbt project.

1
from pendulum import datetime
2
from airflow.decorators import dag
3
from airflow.operators.bash import BashOperator
4
5
PATH_TO_DBT_PROJECT = "<path to your dbt project>"
6
PATH_TO_DBT_VENV = "<path to your venv activate binary>"
7
8
9
@dag(
10
    start_date=datetime(2023, 3, 23),
11
    schedule="@daily",
12
    catchup=False,
13
)
14
def simple_dbt_dag():
15
    dbt_run = BashOperator(
16
        task_id="dbt_run",
17
        bash_command="source $PATH_TO_DBT_VENV && dbt run --models .",
18
        env={"PATH_TO_DBT_VENV": PATH_TO_DBT_VENV},
19
        cwd=PATH_TO_DBT_PROJECT,
20
    )
21
22
23
simple_dbt_dag()

Using the BashOperator to run dbt run and other dbt commands can be useful during development. However, running dbt at the project level has a couple of issues:

• There is low observability into what execution state the project is in.
• Failures are absolute and require all models in a project to be run again, which can be costly.

Using a manifest file

Using a dbt-generated manifest.json file gives you more visibility into the steps dbt is running in each task. This file is generated in the target directory of your dbt project and contains its full representation. For more information on this file, see the dbt documentation.

Cosmos can parse manifest files, see the Cosmos documentation for more information.