Add custom constraint option (#792)

Signed-off-by: Roy Dobbe [email protected]
Co-authored-by: Roy Dobbe <[email protected]>

CHANGELOG.md

@@ -10,6 +10,8 @@
   - `matched`, `not matched` and `not matched by source` condition clauses;
   - custom aliases for source and target tables can be specified and used in condition clauses;
   - `matched` and `not matched` steps can now be skipped;
+- Allow for the use of custom constraints, using the `custom` constraint type with an `expression` as the constraint (thanks @roydobbe). ([792](https://github.com/databricks/dbt-databricks/pull/792))
+
 
 ### Under the Hood
 

dbt/include/databricks/macros/relations/constraints.sql

@@ -79,8 +79,18 @@
 
 {% macro get_constraints_sql(relation, constraints, model, column={}) %}
   {% set statements = [] %}
-  -- Hack so that not null constraints will be applied before primary key constraints
-  {% for constraint in constraints|sort(attribute='type') %}
+  -- Hack so that not null constraints will be applied before other constraints
+  {% for constraint in constraints|selectattr('type', 'eq', 'not_null') %}
+    {% if constraint %}
+      {% set constraint_statements = get_constraint_sql(relation, constraint, model, column) %}
+      {% for statement in constraint_statements %}
+        {% if statement %}
+          {% do statements.append(statement) %}
+        {% endif %}
+      {% endfor %}
+    {% endif %}
+  {% endfor %}
+  {% for constraint in constraints|rejectattr('type', 'eq', 'not_null') %}
     {% if constraint %}
       {% set constraint_statements = get_constraint_sql(relation, constraint, model, column) %}
       {% for statement in constraint_statements %}
@@ -225,6 +235,24 @@
     {% endif %}
     {% set stmt = stmt ~ ";" %}
     {% do statements.append(stmt) %}
+  {% elif type == 'custom' %}
+    {% set expression = constraint.get("expression", "") %}
+    {% if not expression %}
+      {{ exceptions.raise_compiler_error('Missing custom constraint expression') }}
+    {% endif %}
+
+    {% set name = constraint.get("name") %}
+    {% set expression = constraint.get("expression") %}
+    {% if not name %}
+      {% if local_md5 %}
+        {{ exceptions.warn("Constraint of type " ~ type ~ " with no `name` provided. Generating hash instead for relation " ~ relation.identifier) }}
+        {%- set name = local_md5 (relation.identifier ~ ";" ~ expression ~ ";") -%}
+      {% else %}
+        {{ exceptions.raise_compiler_error("Constraint of type " ~ type ~ " with no `name` provided, and no md5 utility.") }}
+      {% endif %}
+    {% endif %}
+    {% set stmt = "alter table " ~ relation ~ " add constraint " ~ name ~ " " ~ expression ~ ";" %}
+    {% do statements.append(stmt) %}
   {% elif constraint.get('warn_unsupported') %}
     {{ exceptions.warn("unsupported constraint type: " ~ constraint.type)}}
   {% endif %}

tests/unit/macros/relations/test_constraint_macros.py

@@ -409,3 +409,40 @@ def test_macros_get_constraint_sql_foreign_key_columns_supplied_separately(
             "some_schema.parent_table(parent_name);']"
         )
         assert expected in r
+
+    def test_macros_get_constraint_sql_custom(self, template_bundle, model):
+        constraint = {
+            "type": "custom",
+            "name": "myconstraint",
+            "expression": "PRIMARY KEY(valid_at TIMESERIES)",
+        }
+        r = self.render_constraint_sql(template_bundle, constraint, model)
+
+        expected = (
+            "['alter table `some_database`.`some_schema`.`some_table` add constraint "
+            "myconstraint PRIMARY KEY(valid_at TIMESERIES);']"
+        )
+        assert expected in r
+
+    def test_macros_get_constraint_sql_custom_noname_constraint(self, template_bundle, model):
+        constraint = {
+            "type": "custom",
+            "expression": "PRIMARY KEY(valid_at TIMESERIES)",
+        }
+        r = self.render_constraint_sql(template_bundle, constraint, model)
+
+        expected = (
+            "['alter table `some_database`.`some_schema`.`some_table` "
+            "add constraint hash(some_table;PRIMARY KEY(valid_at TIMESERIES);) "
+            "PRIMARY KEY(valid_at TIMESERIES);']"
+        )  # noqa: E501
+        assert expected in r
+
+    def test_macros_get_constraint_sql_custom_missing_expression(self, template_bundle, model):
+        constraint = {
+            "type": "check",
+            "expression": "",
+            "name": "myconstraint",
+        }
+        r = self.render_constraint_sql(template_bundle, constraint, model)
+        assert "raise_compiler_error" in r