Fix: Retain multiple foreign-keys from one table to another after an incremental run

CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Fixes
 
 - Fix `hard_deletes: invalidate` incorrectly invalidating active records in snapshots (thanks @Zurbste!) ([#1281](https://github.com/databricks/dbt-databricks/issues/1281))
+- Fix multiple foreign-keys from one table to another not being retained after incremental run
 
 ### Under the Hood
 

dbt/include/databricks/macros/materializations/incremental/incremental.sql

@@ -177,6 +177,7 @@
         {% set tags = _configuration_changes.changes.get("tags", None) %}
         {% set tblproperties = _configuration_changes.changes.get("tblproperties", None) %}
         {% set liquid_clustering = _configuration_changes.changes.get("liquid_clustering") %}
+        {% set constraints = _configuration_changes.changes.get("constraints") %}
         {% if tags is not none %}
           {% do apply_tags(target_relation, tags.set_tags) %}
         {%- endif -%}
@@ -186,6 +187,9 @@
         {% if liquid_clustering is not none %}
           {% do apply_liquid_clustered_cols(target_relation, liquid_clustering) %}
         {% endif %}
+        {% if constraints %}
+          {{ apply_constraints(target_relation, constraints) }}
+        {% endif %}
       {%- endif -%}
       {% do persist_docs(target_relation, model, for_relation=True) %}
     {%- endif -%}

tests/functional/adapter/incremental/fixtures.py

@@ -1109,3 +1109,64 @@ def model(dbt, spark):
         to_columns: [id]
         warn_unenforced: false
 """
+
+target_of_multiple_fks_sql = """
+{{ config(
+    materialized='table',
+) }}
+
+SELECT 'a' AS str_key;
+"""
+
+source_of_multiple_fks_sql = """
+{{ config(
+    materialized='incremental',
+    unique_key=['fk_1_col'],
+    incremental_strategy='delete+insert',
+    on_schema_change='fail'
+) }}
+
+SELECT
+    'a' AS fk_1_col,
+    'a' AS fk_2_col
+"""
+
+multiple_fks_on_one_target_schema_yml = """
+version: 2
+
+models:
+  - name: target_of_multiple_fks
+    config:
+      contract:
+        enforced: true
+    columns:
+      - name: str_key
+        data_type: string
+        constraints:
+          - type: not_null
+          - type: primary_key
+            name: pk_target_key
+            warn_unenforced: false
+
+  - name: source_of_multiple_fks
+    config:
+      contract:
+        enforced: true
+    columns:
+      - name: fk_1_col
+        data_type: string
+        constraints:
+          - type: foreign_key
+            to: ref('target_of_multiple_fks')
+            to_columns: ["str_key"]
+            name: fk_1
+            warn_unenforced: false
+      - name: fk_2_col
+        data_type: string
+        constraints:
+          - type: foreign_key
+            to: ref('target_of_multiple_fks')
+            to_columns: ["str_key"]
+            name: fk_2
+            warn_unenforced: false
+"""

tests/functional/adapter/incremental/test_incremental_constraints.py

@@ -328,3 +328,94 @@ def test_remove_foreign_key_constraint(self, project):
         util.run_dbt(["run"])
         referential_constraints = project.run_sql(referential_constraint_sql, fetch="all")
         assert len(referential_constraints) == 0
+
+
+@pytest.mark.skip_profile("databricks_cluster")
+class TestIncrementalMultipleFKsToSameTable:
+    """
+    Test that multiple foreign keys to the same parent table persist correctly
+    after incremental runs. This succeeds for `use_materialization_v2` at the time
+    of writing, but there is a bug in the v1 materialization implementation. This
+    test is included to confirm this.
+    """
+
+    @pytest.fixture(scope="class")
+    def project_config_update(self):
+        return {
+            "flags": {"use_materialization_v2": True},
+        }
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "target_of_multiple_fks.sql": fixtures.target_of_multiple_fks_sql,
+            "source_of_multiple_fks.sql": fixtures.source_of_multiple_fks_sql,
+            "schema.yml": fixtures.multiple_fks_on_one_target_schema_yml,
+        }
+
+    def test_multiple_fks_to_same_table_persist_after_incremental(self, project):
+        expected_constraints = {
+            ("fk_1", "pk_target_key"),
+            ("fk_2", "pk_target_key"),
+        }
+
+        # Initial run - create tables with constraints
+        util.run_dbt(["run"])
+
+        referential_constraints = project.run_sql(referential_constraint_sql, fetch="all")
+
+        constraints = {(row[0], row[1]) for row in referential_constraints}
+        assert constraints == expected_constraints
+
+        # Incremental run - this should NOT lose any foreign keys
+        util.run_dbt(["run"])
+
+        referential_constraints_after = project.run_sql(referential_constraint_sql, fetch="all")
+
+        constraint_names_after = {(row[0], row[1]) for row in referential_constraints_after}
+        assert constraint_names_after == expected_constraints
+
+
+@pytest.mark.skip_profile("databricks_cluster")
+class TestIncrementalMultipleFKsToSameTableNoV2:
+    """
+    Test that multiple foreign keys to the same parent table persist correctly
+    after incremental runs. This is a very specific test to reproduce a very
+    specific bug found when not using `use_materialization_v2`.
+    """
+
+    @pytest.fixture(scope="class")
+    def project_config_update(self):
+        return {
+            "flags": {"use_materialization_v2": False},
+        }
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "target_of_multiple_fks.sql": fixtures.target_of_multiple_fks_sql,
+            "source_of_multiple_fks.sql": fixtures.source_of_multiple_fks_sql,
+            "schema.yml": fixtures.multiple_fks_on_one_target_schema_yml,
+        }
+
+    def test_multiple_fks_to_same_table_persist_after_incremental(self, project):
+        expected_constraints = {
+            ("fk_1", "pk_target_key"),
+            ("fk_2", "pk_target_key"),
+        }
+
+        # Initial run - create tables with constraints
+        util.run_dbt(["run"])
+
+        referential_constraints = project.run_sql(referential_constraint_sql, fetch="all")
+
+        constraints = {(row[0], row[1]) for row in referential_constraints}
+        assert constraints == expected_constraints
+
+        # Incremental run - this should NOT lose any foreign keys
+        util.run_dbt(["run"])
+
+        referential_constraints_after = project.run_sql(referential_constraint_sql, fetch="all")
+
+        constraint_names_after = {(row[0], row[1]) for row in referential_constraints_after}
+        assert constraint_names_after == expected_constraints

tests/unit/relation_configs/test_constraint.py

@@ -352,3 +352,47 @@ def test_get_diff__check_constraints_different_formatting(self):
         )
         diff = config.get_diff(other)
         assert diff is None
+
+    def test_get_diff__multiple_fks_to_same_table(self):
+        """Test that multiple FKs to the same parent table are handled correctly."""
+        config = ConstraintsConfig(
+            set_non_nulls=set(),
+            set_constraints={
+                ForeignKeyConstraint(
+                    type=ConstraintType.foreign_key,
+                    name="fk_start_date",
+                    columns=["start_date"],
+                    to="`catalog`.`schema`.`date_dim`",
+                    to_columns=["date_key"],
+                ),
+                ForeignKeyConstraint(
+                    type=ConstraintType.foreign_key,
+                    name="fk_end_date",
+                    columns=["end_date"],
+                    to="`catalog`.`schema`.`date_dim`",
+                    to_columns=["date_key"],
+                ),
+            },
+        )
+        other = ConstraintsConfig(
+            set_non_nulls=set(),
+            set_constraints={
+                ForeignKeyConstraint(
+                    type=ConstraintType.foreign_key,
+                    name="fk_start_date",
+                    columns=["start_date"],
+                    to="`catalog`.`schema`.`date_dim`",
+                    to_columns=["date_key"],
+                ),
+                ForeignKeyConstraint(
+                    type=ConstraintType.foreign_key,
+                    name="fk_end_date",
+                    columns=["end_date"],
+                    to="`catalog`.`schema`.`date_dim`",
+                    to_columns=["date_key"],
+                ),
+            },
+        )
+        # Both configs have same constraints, diff should be None
+        diff = config.get_diff(other)
+        assert diff is None