A Simple MWAA + dbt Implementation

Airflow and dbt. Both part of the modern data stack.

When we were building our data platform, integrating AWS’ Managed Workflows for Apache Airflow (MWAA) with dbt-core posed some challenges. While AWS documentation offers some guidance, it falls short of providing instructions for a successful production setup. We found that the documentation leaves some crucial questions unanswered, such as: “What is the optimal approach to configuring multiple dbt profiles in MWAA?” and “How should dbt dependencies be managed effectively within this setup?”

This post details the encountered issues and provides two possible solutions, ensuring you can avoid the same hurdles and streamline your integration of MWAA and dbt effortlessly.

The problem

MWAA uses Fargate as the underlaying architecture. This means that — based on your configuration — it spins up Airflow worker containers when they’re needed to handle the workload. However, every new container misses some of the dependencies and information it needs, such as the dbt profile settings, and the required dbt deps. What steps can we take to ensure that workers have all the necessary information to execute the dbt CLI commands?

The solution

There are two solutions I found to tackle this problem. The first one perhaps being the most elegant. It’s Cosmos, an open source Python library from our friends over at Astronomer. The second solution is a very lightweight and easy to implement custom Airflow operator. We’ve decided on the latter, but I’ll get to both.

Cosmos

As mentioned, Astronomer has developed an Airflow library they’ve named Cosmos. Cosmos is an elegant (my words) way to run dbt within Airflow. You have to point Cosmos to where your dbt project is located, and that’s almost everything you have to do. Cosmos will parse your dbt project, and generate the whole DAG for you. It’s aware of all your dbt dependencies, and will visualise your individual snapshots, models & tests individually within an Airflow TaskGroup. Having worked with Cosmos, I can say their interface and implementation is great. They have a getting started guide specifically for MWAA, and you can get it working in no time.

An example of a TaskGroup as generated by Cosmos

Cosmos solves the problem we ran into with a startup script — a script that runs every time a container is spun up — that installs the dependency and inherits the right environment variables. Running the startup script on a new container takes a bit of time. On our setup with MWAA, each process took about 90 seconds. Theoretically, this could be run for every individual task, and thus every individual model. TL;DR: it comes with some overhead.

Cosmos would be a perfect solution for running dbt within MWAA, but I wanted a faster, more lightweight approach. Though I like the visualisations a lot, if it comes at the expense of the additional overhead, it is not worth it to me. So I went with an even simpler approach; a custom operator.

A custom operator

If you think about it, all that dbt does is compile some SQL and run it against a database — it should not be a heavy process. In the AWS documentation, you can see a good example of this:

cli_command = BashOperator(
        task_id="bash_command",
        bash_command="cp -R /usr/local/airflow/dags/dbt /tmp;\
cd /tmp/dbt/dbt-starter-project;\
/usr/local/airflow/.local/bin/dbt  run --project-dir /tmp/dbt/dbt-starter-project/ --profiles-dir ..;\
cat /tmp/dbt_logs/dbt.log"
    )

Here, dbt is run through a BashOperator that runs a few commands:

1. It copies the dbt project to the write-accessible /tmp folder

2. It jumps into the dbt project root in the folder it just created

3. it then runs dbt using the absolute path, and provinding the project directory and profile directory flags

4. Lastly, it outputs the contents of the dbt.log file

Now, I had a few problems with this approach. The first one being the way it handles the profiles. This method assumesthe profile is in the parent directory of where the project is located. It does not cover on how it gets there.

The second problem I have with this approach, is that it includes a lot of bash code within the operator. It’s not readable, and it does not comply with the DRY principle. The last (minor) problem is that it dumps the content of the dbt.log file, while Airflow is perfectly capable of doing that already.

Lets fix all these with a simple solution: a custom dbt operator and a bit of additional configuration. Lets start with the latter.

Configuration

To be able to work with dbt, we need to configure a connection profile. A profile contains secrets, such as the username and password of a database connection. We don’t want to expose these secrets to the world, but we want to automate deployments using CI/CD. The way I fixed this, is by putting a profile.yml file within the dbt project, that contains environment variable placeholders, like so:

dwh:
  target: dev
  outputs:
    dev:
      type: postgres
      host: "{{ env_var('DBT_HOST') }}"
      user: "{{ env_var('DBT_USER') }}"
      password: "{{ env_var('DBT_PASSWORD') }}"
      port: 5432
      dbname: analytics
      schema: "{{ env_var('DBT_SCHEMA', 'dbt') }}"

Since this does not contain any secrets, it’s save to check into your git and deploy with the rest of your project code. If you need help with this, see the excellent AWS documentation. Oh, and the second argument is the optional default value, for when you don’t always want to define the environment variable.

Custom Operator

Now we need to create the custom operator.

import json
from airflow.utils.decorators import apply_defaults
from airflow.models import Variable
from airflow.operators.bash_operator import BashOperator
from airflow.utils.log.secrets_masker import mask_secret

class Dbt(BashOperator):
    """
    This class defines the dbt operator that's able to run dbt CLI against the database.
    This is needed because we need to set the profile & run dbt deps on every worker.
    """

    @apply_defaults
    def __init__(
            self,
            command,
            *args, **kwargs):

        # Set target
        target = 'dev'

        # Get profile vars from environment variables
        dbt_profile = json.loads(Variable.get(f'dbt-profile-{target}', default_var=''))

        # Mask secrets in the UI and in the logs
        mask_secret(dbt_profile['host'])
        mask_secret(dbt_profile['password'])

        # Steps:
        # 1. Copy dbt project to /tmp folder
        # 2. Export required environment variables for profiles.yml
        # 3. Run dbt command
        #
        # Resources:
        # - AWS: https://docs.aws.amazon.com/mwaa/latest/userguide/samples-dbt.html
        super().__init__(bash_command=f"""
            cp -R /usr/local/airflow/dags/dbt /tmp;\
            export DBT_SCHEMA={dbt_profile['schema']};\
            export DBT_HOST={dbt_profile['host']};\
            export DBT_USER={dbt_profile['user']};\
            export DBT_PASSWORD={dbt_profile['password']};\
            /usr/local/airflow/.local/bin/dbt deps --project-dir /tmp/dbt --profiles-dir /tmp/dbt --target {target};\
            /usr/local/airflow/.local/bin/dbt {command} --project-dir /tmp/dbt --profiles-dir /tmp/dbt --target {target};\
        """, *args, **kwargs)
        self.command = command

💡 Tip: save this in a file named dbt.pyin a folder named /operatorsin your Airflow /dags folder.

Using this operator is as simple as:

import os
from datetime import timedelta
import pendulum
from airflow import DAG
from operators.dbt import Dbt

with DAG(
    os.path.basename(__file__).replace(".py", ""),
    description='DAG with all dbt processes',
    schedule_interval="30 6 * * *",
    start_date=pendulum.datetime(2023, 9, 3, tz="Europe/Amsterdam"),
    catchup=False
) as dag:
    run_stg = Dbt(task_id="dbt_run__staging", command='run -m tag:stg')

As you can see, this makes running dbt with MWAA way easier. There are a few things that are good to know:

• I’ve configured AWS KMS as Airflow’s secrets backend. When working with your own connections and variables, it’s wise that you do so too. See the docs. You can configure your secret as a key-value pair, so that you can get all required environment variables at once, which saves you a little money, too!

• I’m using Airflow’s mask_secret() function to hide the host & password form the UI & logs. It will now show up as ***

• The code above is simplified. It imports and runs only the development (dev) profile. For real world application, it would make sense to create multiple profiles for every environment, and import the one you need at runtime.

• You can extend the options of the Dbt operator as you wish. Now, you only control the command that follows the dbt CLI command. If you wish to have a more Pythonic interface, in which you can also control the flags, I highly encourage you to do so!

There you have it! This is how you can easily run dbt within MWAA 🎉 I think it’s still very close to the way that AWS recommends it, while providing an simple, DRY and elegant solution to the problems we encountered.

Hi, I’m Bastiaan 👋🏼 Data Lead at a scale-up. I write about the Modern Data Workflow, where I explore tools & processes to supercharge your data capabilities. Follow me for more!