Add asset_key_prefix support for dbt sources with tests and docs

Author: audreydng

docs/docs/integrations/libraries/dbt/reference.md

@@ -309,6 +309,40 @@ sources:
             asset_key: ['snowflake', 'jaffle_shop', 'orders']
 ```
 
+If you want to apply a shared prefix without repeating a full key for each node, you can configure
+`meta.dagster.asset_key_prefix` and Dagster will append the dbt node name (and schema, if configured):
+
+```yaml
+models:
+  jaffle_shop:
+    +meta:
+      dagster:
+        asset_key_prefix: ['snowflake', 'jaffle_shop']
+```
+
+For a model `orders` without a custom schema, this produces the asset key `snowflake/jaffle_shop/orders`.
+For a model `orders` with schema `marketing`, this produces the asset key `snowflake/jaffle_shop/marketing/orders`.
+
+For dbt sources, the source name is preserved between the prefix and the table name:
+
+```yaml
+sources:
+  - name: jaffle_shop
+    tables:
+      - name: orders
+        meta:
+          dagster:
+            asset_key_prefix: ['snowflake']
+```
+
+For a table `orders` in source `jaffle_shop`, this produces the asset key `snowflake/jaffle_shop/orders`.
+
+When both are configured, Dagster uses this precedence:
+
+1. `meta.dagster.asset_key`
+2. `meta.dagster.asset_key_prefix` + `[schema, ]` (if present) + `[node_name]` (models/snapshots) or `[source_name, table_name]` (sources)
+3. Default Dagster dbt key generation
+
 Alternatively, to override the asset key generation for all dbt nodes in your dbt project, you can create a custom <PyObject section="libraries" integration="dbt" module="dagster_dbt" object="DagsterDbtTranslator" /> and implement <PyObject section="libraries" integration="dbt" module="dagster_dbt" object="DagsterDbtTranslator.get_asset_key" />. The following example adds a `snowflake` prefix to the default generated asset key:
 
 <CodeExample

python_modules/libraries/dagster-dbt/dagster_dbt/asset_utils.py

@@ -644,12 +644,17 @@ def default_asset_key_fn(dbt_resource_props: Mapping[str, Any]) -> AssetKey:
         dbt sources: a dbt source's key is the union of its source name and its table name
         dbt models: a dbt model's key is the union of its model name and any schema configured on
             the model itself.
+        If an asset key prefix is configured in its metadata, it is prepended
+            to the default asset key. The prefix supports project or folder level
+            configuration in `dbt_project.yml`.
     """
     dbt_meta = dbt_resource_props.get("config", {}).get("meta", {}) or dbt_resource_props.get(
         "meta", {}
     )
     dagster_metadata = dbt_meta.get("dagster", {})
     asset_key_config = dagster_metadata.get("asset_key", [])
+    asset_key_prefix_config = dagster_metadata.get("asset_key_prefix", [])
+
     if asset_key_config:
         return AssetKey(asset_key_config)
 
@@ -664,6 +669,14 @@ def default_asset_key_fn(dbt_resource_props: Mapping[str, Any]) -> AssetKey:
         else:
             components = [dbt_resource_props["name"]]
 
+    if asset_key_prefix_config:
+        prefix = (
+            [asset_key_prefix_config]
+            if isinstance(asset_key_prefix_config, str)
+            else list(asset_key_prefix_config)
+        )
+        components = [*prefix, *components]
+
     return AssetKey(components)
 
 

python_modules/libraries/dagster-dbt/dagster_dbt_tests/conftest.py

@@ -11,6 +11,7 @@
 from dagster_dbt_tests.dbt_projects import (
     test_asset_checks_path,
     test_asset_key_exceptions_path,
+    test_asset_key_prefix_path,
     test_dagster_dbt_mixed_freshness_path,
     test_dbt_alias_path,
     test_dbt_functions_path,
@@ -115,6 +116,11 @@ def test_asset_key_exceptions_manifest_fixture() -> dict[str, Any]:
     return _create_dbt_invocation(test_asset_key_exceptions_path).get_artifact("manifest.json")
 
 
+@pytest.fixture(name="test_asset_key_prefix_manifest", scope="session")
+def test_asset_key_prefix_manifest_fixture() -> dict[str, Any]:
+    return _create_dbt_invocation(test_asset_key_prefix_path).get_artifact("manifest.json")
+
+
 @pytest.fixture(name="test_dbt_alias_manifest", scope="session")
 def test_dbt_alias_manifest_fixture() -> dict[str, Any]:
     return _create_dbt_invocation(test_dbt_alias_path).get_artifact("manifest.json")

python_modules/libraries/dagster-dbt/dagster_dbt_tests/core/test_asset_key_prefix.py

@@ -0,0 +1,145 @@
+from typing import TYPE_CHECKING, Any
+
+import pytest
+from dagster import AssetKey
+from dagster_dbt import build_dbt_asset_specs, dbt_assets
+from dagster_dbt.asset_utils import default_asset_key_fn
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping
+
+
+def test_default_asset_key_fn_asset_key_takes_precedence_over_prefix() -> None:
+    dbt_resource_props: Mapping[str, Any] = {
+        "name": "orders",
+        "resource_type": "model",
+        "config": {
+            "meta": {
+                "dagster": {
+                    "asset_key": ["warehouse", "jaffle_shop", "orders"],
+                    "asset_key_prefix": ["ignored", "prefix"],
+                }
+            }
+        },
+    }
+
+    assert default_asset_key_fn(dbt_resource_props) == AssetKey(
+        ["warehouse", "jaffle_shop", "orders"]
+    )
+
+
+@pytest.mark.parametrize(
+    "resource_type,name",
+    [
+        ("model", "orders"),
+        ("snapshot", "orders_snapshot"),
+    ],
+)
+def test_default_asset_key_fn_asset_key_prefix_for_supported_resource_types(
+    resource_type: str, name: str
+) -> None:
+    dbt_resource_props = {
+        "name": name,
+        "resource_type": resource_type,
+        "source_name": "jaffle_shop",
+        "config": {"meta": {"dagster": {"asset_key_prefix": ["snowflake", "jaffle_shop"]}}},
+    }
+
+    assert default_asset_key_fn(dbt_resource_props) == AssetKey(["snowflake", "jaffle_shop", name])
+
+
+def test_default_asset_key_fn_asset_key_prefix_for_source() -> None:
+    # For sources, the key is: prefix + source_name + table_name
+    dbt_resource_props = {
+        "name": "raw_orders",
+        "resource_type": "source",
+        "source_name": "jaffle_shop",
+        "config": {"meta": {"dagster": {"asset_key_prefix": ["snowflake"]}}},
+    }
+
+    assert default_asset_key_fn(dbt_resource_props) == AssetKey(
+        ["snowflake", "jaffle_shop", "raw_orders"]
+    )
+
+
+def test_default_asset_key_fn_asset_key_prefix_uses_alias_for_versioned_models() -> None:
+    dbt_resource_props = {
+        "name": "orders",
+        "alias": "orders_v1",
+        "version": 1,
+        "resource_type": "model",
+        "config": {"meta": {"dagster": {"asset_key_prefix": ["snowflake", "jaffle_shop"]}}},
+    }
+
+    assert default_asset_key_fn(dbt_resource_props) == AssetKey(
+        ["snowflake", "jaffle_shop", "orders_v1"]
+    )
+
+
+def test_default_asset_key_fn_asset_key_prefix_from_meta_path() -> None:
+    # Sources read asset_key_prefix from top-level meta, not config.meta
+    dbt_resource_props = {
+        "name": "orders",
+        "resource_type": "source",
+        "source_name": "jaffle_shop",
+        "config": {},
+        "meta": {"dagster": {"asset_key_prefix": ["raw"]}},
+    }
+
+    assert default_asset_key_fn(dbt_resource_props) == AssetKey(["raw", "jaffle_shop", "orders"])
+
+
+def test_default_asset_key_fn_empty_asset_key_prefix_uses_node_name() -> None:
+    dbt_resource_props = {
+        "name": "orders",
+        "resource_type": "model",
+        "config": {"meta": {"dagster": {"asset_key_prefix": []}}},
+    }
+
+    assert default_asset_key_fn(dbt_resource_props) == AssetKey(["orders"])
+
+
+def test_default_asset_key_fn_asset_key_prefix_accepts_string() -> None:
+    dbt_resource_props = {
+        "name": "orders",
+        "resource_type": "model",
+        "config": {"meta": {"dagster": {"asset_key_prefix": "snowflake"}}},
+    }
+
+    assert default_asset_key_fn(dbt_resource_props) == AssetKey(["snowflake", "orders"])
+
+
+def test_default_asset_key_fn_asset_key_prefix_with_schema() -> None:
+    # When both prefix and schema are configured, schema is appended after prefix
+    dbt_resource_props = {
+        "name": "orders",
+        "resource_type": "model",
+        "config": {
+            "schema": "analytics",
+            "meta": {"dagster": {"asset_key_prefix": ["snowflake"]}},
+        },
+    }
+
+    assert default_asset_key_fn(dbt_resource_props) == AssetKey(
+        ["snowflake", "analytics", "orders"]
+    )
+
+
+def test_dbt_meta_asset_key_prefix_with_inherited_config(
+    test_asset_key_prefix_manifest: dict[str, Any],
+) -> None:
+    expected_specs_by_key = {
+        spec.key: spec for spec in build_dbt_asset_specs(manifest=test_asset_key_prefix_manifest)
+    }
+
+    @dbt_assets(manifest=test_asset_key_prefix_manifest)
+    def my_dbt_assets(): ...
+
+    assert AssetKey(["warehouse", "jaffle_shop", "orders_from_source"]) in my_dbt_assets.keys
+    assert AssetKey(["override", "explicit_asset_key"]) in my_dbt_assets.keys
+    assert AssetKey(["warehouse", "jaffle_shop", "explicit_asset_key"]) not in my_dbt_assets.keys
+
+    assert AssetKey(["staging", "raw", "orders"]) in set(my_dbt_assets.keys_by_input_name.values())
+
+    assert AssetKey(["warehouse", "jaffle_shop", "orders_from_source"]) in expected_specs_by_key
+    assert AssetKey(["override", "explicit_asset_key"]) in expected_specs_by_key

python_modules/libraries/dagster-dbt/dagster_dbt_tests/dbt_projects/__init__.py

@@ -31,3 +31,4 @@
 test_dagster_dbt_mixed_freshness_path = projects_path.joinpath("test_dagster_dbt_mixed_freshness")
 test_jaffle_with_profile_vars_path = projects_path.joinpath("test_jaffle_with_profile_vars")
 test_dependencies_path = projects_path.joinpath("test_dagster_dbt_dependencies")
+test_asset_key_prefix_path = projects_path.joinpath("test_dagster_asset_key_prefix")

python_modules/libraries/dagster-dbt/dagster_dbt_tests/dbt_projects/test_dagster_asset_key_prefix/dbt_project.yml

@@ -0,0 +1,32 @@
+name: "test_dagster_asset_key_prefix"
+
+config-version: 2
+version: "0.1"
+
+profile: "jaffle_shop"
+
+model-paths: ["models"]
+seed-paths: ["seeds"]
+test-paths: ["tests"]
+analysis-paths: ["analysis"]
+macro-paths: ["macros"]
+
+target-path: "target"
+clean-targets:
+  - "target"
+  - "dbt_modules"
+  - "logs"
+
+require-dbt-version: [">=1.0.0", "<2.0.0"]
+
+models:
+  test_dagster_asset_key_prefix:
+    +meta:
+      dagster:
+        asset_key_prefix: ["warehouse", "jaffle_shop"]
+
+sources:
+  test_dagster_asset_key_prefix:
+    +meta:
+      dagster:
+        asset_key_prefix: ["staging"]

python_modules/libraries/dagster-dbt/dagster_dbt_tests/dbt_projects/test_dagster_asset_key_prefix/models/explicit_asset_key.sql

@@ -0,0 +1 @@
+select 1 as id

python_modules/libraries/dagster-dbt/dagster_dbt_tests/dbt_projects/test_dagster_asset_key_prefix/models/orders_from_source.sql

@@ -0,0 +1 @@
+select * from {{ source('raw', 'orders') }}

python_modules/libraries/dagster-dbt/dagster_dbt_tests/dbt_projects/test_dagster_asset_key_prefix/models/schema.yml

@@ -0,0 +1,15 @@
+version: 2
+
+sources:
+  - name: raw
+    tables:
+      - name: orders
+
+models:
+  - name: orders_from_source
+
+  - name: explicit_asset_key
+    config:
+      meta:
+        dagster:
+          asset_key: ["override", "explicit_asset_key"]

python_modules/libraries/dagster-dbt/dagster_dbt_tests/dbt_projects/test_dagster_asset_key_prefix/profiles.yml

@@ -0,0 +1,7 @@
+jaffle_shop:
+  target: dev
+  outputs:
+    dev:
+      type: duckdb
+      path: "{{ env_var('DAGSTER_DBT_PYTEST_XDIST_DUCKDB_DBFILE_PATH') }}"
+      threads: 24