[RFC] Adding support for partition mapping to Dagster's dbt integration

[rexledesma]

Summary

This discussion provides an overview of the API changes to the dagster-dbt integration to support partition mappings for your software-defined assets (SDAs).

We would like to use this document to gather feedback and answer questions about the new APIs. If you have any comments or use cases that are not addressed, please participate in the discussion!

We'll talk about two use cases here:

1. Overriding partition dependency rules for your dbt SDAs
2. Adding self-dependencies for your dbt SDAs

Requirements

• You must use dagster-dbt>=0.21.6, which was released on November 2, 2023.
• You must use @dbt_assets to define your dbt SDAs.

Usage

DagsterDbtTranslator now has a new method get_partition_mapping that can be overridden to provide custom PartitionMapping's for your dbt dependencies.

get_partition_mapping is a function that takes in two dictionary parameters: the first, representing properties of a dbt resource; and the second, representing the properties of a parent dependency to the first dbt resource. The function returns a PartitionMapping specifying the mapping for this dependency.

Example 1: Overriding partition dependency rules for your dbt SDAs

For specific dbt dependencies, you can specify a custom partition mapping. This can be done by overriding DagsterDbtTranslator.get_partition_mapping.

from typing import Any, Mapping, Optional

from dagster import PartitionMapping, TimeWindowPartitionMapping
from dagster_dbt import DagsterDbtTranslator, dbt_assets


class CustomDagsterDbtTranslator(DagsterDbtTranslator):
    @classmethod
    def get_partition_mapping(
        cls,
        dbt_resource_props: Mapping[str, Any],
        dbt_parent_resource_props: Mapping[str, Any],
    ) -> Optional[PartitionMapping]:
        custom_partition_mappings = {
            (
                "model.jaffle_shop.stg_orders",
                "seed.jaffle_shop.raw_orders",
            ): TimeWindowPartitionMapping(start_offset=-1),
            ...
        }

        dbt_unique_id = dbt_resource_props["unique_id"]
        dbt_parent_unique_id = dbt_parent_resource_props["unique_id"]
        dbt_dependency = (dbt_unique_id, dbt_parent_unique_id)

        return custom_partition_mappings.get(dbt_dependency)


@dbt_assets(manifest=..., dagster_dbt_translator=CustomDagsterDbtTranslator())
def my_custom_dbt_assets(...):
    ...

Example 2: Adding self-dependencies for your dbt SDAs

You can also specify a dbt model to have a self-dependency on its previous partitions. Here, you will need to indicate in the dbt metadata that the model has a self-dependency. This can be configured using +meta.dagster.has_self_dependency.

version: 2

models:
  - name: customers
    description: This table has basic information about a customer, as well as some derived facts based on a customer's orders
    meta:
      dagster:
        has_self_dependency: True

Then, you can override DagsterDbtTranslator.get_partition_mapping to provide the partition mapping for the self-dependency.

from typing import Any, Mapping, Optional

from dagster import PartitionMapping, TimeWindowPartitionMapping
from dagster_dbt import DagsterDbtTranslator, dbt_assets


class CustomDagsterDbtTranslator(DagsterDbtTranslator):
    @classmethod
    def get_partition_mapping(
        cls,
        dbt_resource_props: Mapping[str, Any],
        dbt_parent_resource_props: Mapping[str, Any],
    ) -> Optional[PartitionMapping]:
        custom_partition_mappings = {
            (
                "model.jaffle_shop.customers",
                "model.jaffle_shop.customers", 
            ): TimeWindowPartitionMapping(start_offset=-1),
            ...
        }

        dbt_unique_id = dbt_resource_props["unique_id"]
        dbt_parent_unique_id = dbt_parent_resource_props["unique_id"]
        dbt_dependency = (dbt_unique_id, dbt_parent_unique_id)

        return custom_partition_mappings.get(dbt_dependency)


@dbt_assets(manifest=..., dagster_dbt_translator=CustomDagsterDbtTranslator())
def my_custom_dbt_assets():
    ...

References

• Overriding partition dependency rules
• PartitionMapping API docs

[brunodeprez]

Hi, would this work for DynamicPartitionsDefinition as well? I need my different dbt_assets to use different DynamicPartitionsDefinition