Configure Airflow tasks using dbt model meta

[wornjs]

Description

The various dbt models have unique characteristics, and some may require the use of custom pools, queues, or other specific configurations. To support such cases, this update introduces the ability to add necessary information in the meta section of the dbt model.yaml. This metadata is then passed as kwargs to the corresponding Airflow tasks, enabling model-specific customization and enhanced task configuration.

here is sample

DbtTaskGroup - default_args for all dbt models

    dbt_task_group = DbtTaskGroup(
        project_config=,
        profile_config=ProfileConfig,
        default_args={'pool': dbt_pool}
    )

version: 2

models:
  - name: name
    description: description
    meta:
      owner: '[email protected]'
      cosmos:
        pool: abcd

result

general pool

custom pool

Related Issue(s)

#1325

Breaking Change?

Checklist

• I have made corresponding changes to the documentation (if required)
• I have added tests that prove my fix is effective or that my feature works

[wornjs]

It would be great to add documentation for the features introduced in this PR. However, looking at the current project, it seems there aren’t any markdown files for documentation apart from the CONTRIBUTING file. Where do you think would be the best place to add this documentation?

[pankajastro]

It would be great to add documentation for the features introduced in this PR. However, looking at the current project, it seems there aren’t any markdown files for documentation apart from the CONTRIBUTING file. Where do you think would be the best place to add this documentation?

We use rst. You can find docs at https://github.com/astronomer/astronomer-cosmos/tree/main/docs

[pankajastro]

Could you please rebase this PR

[codecov]

Codecov Report

Attention: Patch coverage is 83.33333% with 2 lines in your changes missing coverage. Please review.

Project coverage is 95.98%. Comparing base (43998d2) to head (26811de).
Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
cosmos/core/airflow.py	50.00%	1 Missing ⚠️
cosmos/dbt/graph.py	85.71%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1339      +/-   ##
==========================================
- Coverage   96.02%   95.98%   -0.04%     
==========================================
  Files          67       67              
  Lines        4025     4036      +11     
==========================================
+ Hits         3865     3874       +9     
- Misses        160      162       +2     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[linchun3]

related issue?
#881

[tatiana]

Hi, @wornjs,

The per-node configuration has been a long-standing request (example: #881 (comment)), and your PR solves this. This is an exciting feature; thanks for contributing to Cosmos!

A question and two requests:

1. What are the advantages/disadvantages of having these properties in meta as opposed to having them in config:

    meta:
      cosmos:
        pool: abcd

versus

    config:
      cosmos:
        pool: abcd

2. Please, could you add tests to cover this feature?

3. We'll also need docs.

If you can address these before 20 December, we can ship Cosmos 1.8. I'm tentatively adding this to that milestone.

[wornjs]

Hi, @wornjs,

The per-node configuration has been a long-standing request (example: #881 (comment)), and your PR solves this. This is an exciting feature; thanks for contributing to Cosmos!

A question and two requests:

1. What are the advantages/disadvantages of having these properties in meta as opposed to having them in config:

    meta:
      cosmos:
        pool: abcd

versus

    config:
      cosmos:
        pool: abcd

2. Please, could you add tests to cover this feature?
3. We'll also need docs.

If you can address these before 20 December, we can ship Cosmos 1.8. I'm tentatively adding this to that milestone.

i'm gonna add test until this week

[dwreeves]

Hi, responding here instead of in #1325 as it seems the discussion has migrated here.

First, this PR is related to #881. This issue was for supporting arbitrary kwargs, not just a single kwarg.

I was a big fan of where we ended up with the API of that proposal, and there are a few differences.

The first difference, as @tatiana brings up, is that it uses config: instead of meta:. I do not believe meta: is a valid attribute of a model specification. https://github.com/dbt-labs/dbt-core/blob/fc6167a2ee5291cf6c562eadcd975b48e6a34d65/core/dbt/artifacts/resources/v1/model.py#L38-L48 Please let me know if I am wrong, but I do believe it should be config: and not meta: based on the dbt-core source code. I am wrong, it is! There is a chain of inheritances that lead to meta. My mistake. With that out of the way, I am not sure what I prefer between config and meta, to be honest. I would ask the same question as Tatiana again then, i.e. what do you see as the difference between one or the other and why should one (in this case meta) be preferred?

The second difference is, as mentioned, it supports arbitrary kwargs, not just pool. There are a lot of reasons why someone may want to inject many different kwargs into their operators, not just pool.

The third difference is that the namespace was cosmos: operator_args: pool:, not just cosmos: pool:. I think this is really important once you move the API from just pool to any kwarg. The issue is, in the future, we may want to add additional options that are not actually kwargs (or alternatively allow users to inject such things themselves), so we would need to keep the kwargs separate of things that are not kwargs.

So my final proposal would be this:

version: 2
models:
  - name: model_a
    config:
      alias: model_a
      cosmos:
        operator_args:
          pool: my-pool-here

(or alternatively, if using meta:)

version: 2
models:
  - name: model_a
    config:
      alias: model_a
    meta:
      cosmos:
        operator_args:
          pool: my-pool-here

^ So the key difference is there is one extra dict, i.e. operator_args.

But really, I think we should just **unpack the whole cosmos.operator_args into the operator so that users can do things like this:

version: 2
models:
  - name: model_a
    config:
      alias: model_a
      cosmos:
        operator_args:
          pool: my-pool-here
          retries: 4
          conn_id: special_conn_id

[dwreeves]

I've confirmed both meta and config have the same merge behavior: both merge the top-level keys, but neither recursively merges, which is annoying because

1. it makes it harder to decide whether cosmos should be in meta: or config:
2. it also means there is no way for users to do something like +config: cosmos: operator_args: x: ... in dbt_project.yml and config(cosmos={"operator_args": "y": ...}) in a sql model and get both x and y 🫤

(The only way to support merging this way would be to use top-level attributes, but I think that's tedious and limiting.)

There is an issue open relating to this in dbt-core where someone suggests recursive merging for meta: dbt-labs/dbt-core#10946