fix foreign key constraints

databricks · Sep 13, 2024 · c69ddf4 · c69ddf4
1 parent 67fcb4d
commit c69ddf4
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 7 deletions.
diff --git a/dbt/include/databricks/macros/relations/constraints.sql b/dbt/include/databricks/macros/relations/constraints.sql
@@ -200,7 +200,7 @@
 
       {% set joined_names = quoted_names|join(", ") %}
 
-      {% set parent = constraint.get("parent") %}
+      {% set parent = constraint.get("to") %}
       {% if not parent %}
         {{ exceptions.raise_compiler_error('No parent table defined for foreign key: ' ~ expression) }}
       {% endif %}
@@ -218,7 +218,7 @@
       {% endif %}
 
       {% set stmt = "alter table " ~ relation ~ " add constraint " ~ name ~ " foreign key(" ~ joined_names ~ ") references " ~ parent %}
-      {% set parent_columns = constraint.get("parent_columns") %}
+      {% set parent_columns = constraint.get("to_columns") %}
       {% if parent_columns %}
         {% set stmt = stmt ~ "(" ~ parent_columns|join(", ") ~ ")"%}
       {% endif %}

diff --git a/tests/functional/adapter/constraints/fixtures.py b/tests/functional/adapter/constraints/fixtures.py
@@ -48,3 +48,56 @@
             name: fk_n
             expression: (n) REFERENCES {schema}.raw_numbers
 """
+
+parent_foreign_key = """
+version: 2
+
+models:
+  - name: parent_table
+    config:
+      materialized: table
+      on_schema_change: fail
+      contract:
+        enforced: true
+    columns:
+      - name: id
+        data_type: integer
+        constraints:
+          - type: not_null
+          - type: primary_key
+            name: pk_example__parent_table
+  - name: child_table
+    config:
+      materialized: incremental
+      on_schema_change: fail
+      contract:
+        enforced: true
+    constraints:
+      - type: primary_key
+        name: pk_example__child_table
+        columns: ["id"]
+      - type: not_null
+        columns: ["id", "name", "parent_id"]
+      - type: foreign_key
+        name: fk_example__child_table_1
+        columns: ["parent_id"]
+        to: parent_table
+        to_columns: ["id"]
+    columns:
+      - name: id
+        data_type: integer
+      - name: name
+        data_type: string
+      - name: parent_id
+        data_type: integer
+"""
+
+parent_sql = """
+select 1 as id
+"""
+
+child_sql = """
+ -- depends_on: {{ ref('parent_table') }}
+
+select 2 as id, 'name' as name, 1 as parent_id
+"""
diff --git a/tests/functional/adapter/constraints/test_constraints.py b/tests/functional/adapter/constraints/test_constraints.py
@@ -189,7 +189,7 @@ def models(self):
 
 
 @pytest.mark.skip_profile("databricks_cluster")
-class TestIncrementalForeignKeyConstraint:
+class TestIncrementalForeignKeyExpressionConstraint:
     @pytest.fixture(scope="class")
     def models(self):
         return {
@@ -211,6 +211,20 @@ def test_incremental_foreign_key_constraint(self, project):
         util.run_dbt(["run", "--select", "stg_numbers"])
 
 
+@pytest.mark.skip_profile("databricks_cluster")
+class TestForeignKeyParentConstraint:
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {
+            "schema.yml": override_fixtures.parent_foreign_key,
+            "parent_table.sql": override_fixtures.parent_sql,
+            "child_table.sql": override_fixtures.child_sql,
+        }
+
+    def test_foreign_key_constraint(self, project):
+        util.run_dbt(["build"])
+
+
 class TestConstraintQuotedColumn(BaseConstraintQuotedColumn):
     @pytest.fixture(scope="class")
     def models(self):

diff --git a/tests/unit/macros/relations/test_constraint_macros.py b/tests/unit/macros/relations/test_constraint_macros.py
@@ -362,8 +362,8 @@ def test_macros_get_constraint_sql_foreign_key_parent_column(self, template_bund
             "type": "foreign_key",
             "name": "myconstraint",
             "columns": ["name"],
-            "parent": "parent_table",
-            "parent_columns": ["parent_name"],
+            "to": "parent_table",
+            "to_columns": ["parent_name"],
         }
         r = self.render_constraint_sql(template_bundle, constraint, model)
 
@@ -379,8 +379,8 @@ def test_macros_get_constraint_sql_foreign_key_multiple_columns(self, template_b
             "type": "foreign_key",
             "name": "myconstraint",
             "columns": ["name", "id"],
-            "parent": "parent_table",
-            "parent_columns": ["parent_name", "parent_id"],
+            "to": "parent_table",
+            "to_columns": ["parent_name", "parent_id"],
         }
         r = self.render_constraint_sql(template_bundle, constraint, model)