dbt under Airflow: one task per model, not one per project

The ordinary way to run dbt from Airflow fits on one line: a bash operator, dbt run, and the project sits in a task. It works, and it is paid for on three counts you only notice during an incident. A retry replays the whole project when a single model failed. The dependency graph dbt knows about stays invisible to the orchestrator, which draws one square where there are two hundred nodes. And the logs arrive as one block to be searched.

Cosmos, published by Astronomer, reads the dbt project and builds Airflow tasks from it: one per model, with the dependencies dbt has already worked out. The point is not cosmetic: what the orchestrator can see, it can retry, parallelise and account for.

Three configuration objects, and the DAG follows from the project. `ProjectConfig` says where it is, `ProfileConfig` how to connect, `ExecutionConfig` which dbt binary to call: the last matters more than it looks, since dbt and Airflow almost never agree on their dependencies.

from datetime import datetime

from cosmos import DbtDag, ExecutionConfig, ProfileConfig, ProjectConfig
from cosmos.profiles import AthenaAccessKeyProfileMapping

profile = ProfileConfig(
    profile_name="warehouse",
    target_name="prod",
    # The dbt profile is built from an Airflow connection of type aws:
    # no profiles.yml to deploy, no secret in the repository.
    profile_mapping=AthenaAccessKeyProfileMapping(
        conn_id="aws_warehouse",
        profile_args={"schema": "marts"},
    ),
)

dag = DbtDag(
    dag_id="warehouse_dbt",
    project_config=ProjectConfig("/usr/local/airflow/dags/dbt/warehouse"),
    profile_config=profile,
    # dbt in its own virtual environment: its dependencies and Airflow's do not
    # hold together, and that is the trap that costs you the day.
    execution_config=ExecutionConfig(
        dbt_executable_path="/usr/local/airflow/dbt_venv/bin/dbt",
    ),
    schedule_interval="0 5 * * *",
    start_date=datetime(2023, 11, 1),
    catchup=False,
    default_args={"retries": 2},
)

More often the dbt project is not alone: it follows an ingestion and precedes a publication. `DbtTaskGroup` produces the same graph, set inside an existing DAG, and the dependencies are declared as between any two tasks.

from airflow.decorators import dag, task
from cosmos import DbtTaskGroup, ProjectConfig

@dag(schedule="0 5 * * *", start_date=datetime(2023, 11, 1), catchup=False)
def warehouse():
    @task
    def ingest() -> None:
        ...

    transform = DbtTaskGroup(
        group_id="transform",
        project_config=ProjectConfig("/usr/local/airflow/dags/dbt/warehouse"),
        profile_config=profile,
    )

    @task
    def publish() -> None:
        ...

    ingest() >> transform >> publish()

warehouse()

Athena is not a database in the usual sense: the compute sits with AWS, the data in S3, and a write location is needed for intermediate results. A dbt-athena profile therefore carries more fields than a Postgres one, and two of them, the staging directory and the workgroup, have no equivalent elsewhere.

The Cosmos mapping reads those fields from the `extra` block of an Airflow connection of type `aws`, and takes the credentials through the Amazon provider’s usual mechanism, which means an IAM role works as well as a key, and is worth more.

{
  "region_name": "eu-west-1",
  "database": "awsdatacatalog",
  "schema": "marts",
  "s3_staging_dir": "s3://warehouse-athena/results/",
  "work_group": "dbt"
}

Two warnings, from use rather than from the documentation. First: giving dbt its own Athena workgroup, as above, separates its consumption from the rest and allows a scanned-bytes limit to be set on it. Without which a clumsy join is discovered on the invoice. Second: the adapter installs as `dbt-athena-community`, not `dbt-athena`, which is the first half hour everybody loses.

The graph has to be known when Airflow parses the file, which by default is every thirty seconds. Letting Cosmos call dbt at each parse to discover the models is the configuration that sets itself up, and the worst one: the scheduler slows as the project grows. The remedy is one sentence: produce the manifest when the image is built, and hand it to Cosmos rather than making it recompute it.

The second cost is less obvious: two hundred tasks instead of one means two hundred rows in the metadata database on every run, and a dashboard no longer read at a glance. On a twenty-model project the gain is plain; past a few hundred, grouping by folder rather than by model is the answer, and Cosmos can do that.

I have not run this on Kubernetes, where execution happens in a container per task and the trade-off changes entirely: the detailed graph then costs one container start per model, and the economics of fine-grained retries can invert. Nor do I say anything about dbt tests exposed as separate tasks, which Cosmos allows: I distrust it, because a red test should stop the chain, and making it independent makes that rule negotiable.

One last thing, and it is dated. The Athena mapping is a month old as I write: it landed in version 1.2.0 on 13 October 2023, and five releases followed in six weeks. That is the pace of a project still moving, and the reason the code above should be read as a snapshot rather than as a recipe.

• astronomer-cosmos
• dbt-athena
• Cosmos · profil Athena