Add generic rule: models should implement uniqueness test for their PK (#90)

For models materialized as `table` or `incremental`:
- Loop over their columns to extract PK
- Loop over their tests to find a uniqueness test matching the PK
columns

Consider both table- and column-level constraints and tests.

Author: matthieucan

CHANGELOG.md

@@ -10,6 +10,11 @@ and this project adheres to
 
 - Add debug mode to help writing new rules. (#91)
 - Fix tests without metadata. (#88)
+- Add new rule to enforce presence of uniqueness test. (#90)
+- Add new rule to enforce single-column PK to be defined at column level. (#90)
+- Add new rule to enforce 1 column uniqueness test to be defined at column
+  level. (#90)
+- Add `constraints` to the model schema. (#90)
 
 ## [0.9.0] - 2024-12-19
 

docs/rules/filters.md

@@ -0,0 +1,3 @@
+# Rule filters
+
+::: dbt_score.rules.filters

mkdocs.yml

@@ -28,6 +28,7 @@ nav:
   - Programmatic invocations: programmatic_invocations.md
   - Rules:
       - rules/generic.md
+      - rules/filters.md
   - Reference:
       - reference/cli.md
       - reference/config.md

src/dbt_score/models.py

@@ -15,18 +15,20 @@
 
 @dataclass
 class Constraint:
-    """Constraint for a column.
+    """Constraint for a model or a column.
 
     Attributes:
         type: The type of the constraint, e.g. `foreign_key`.
         name: The name of the constraint.
         expression: The expression of the constraint, e.g. `schema.other_table`.
+        columns: The columns for the constraint (only for model-level constraints).
         _raw_values: The raw values of the constraint in the manifest.
     """
 
     type: str
     name: str | None = None
     expression: str | None = None
+    columns: list[str] | None = None
     _raw_values: dict[str, Any] = field(default_factory=dict)
 
     @classmethod
@@ -36,6 +38,7 @@ def from_raw_values(cls, raw_values: dict[str, Any]) -> "Constraint":
             type=raw_values["type"],
             name=raw_values["name"],
             expression=raw_values["expression"],
+            columns=raw_values.get("columns"),
             _raw_values=raw_values,
         )
 
@@ -198,6 +201,7 @@ class Model(HasColumnsMixin):
     tags: list[str] = field(default_factory=list)
     tests: list[Test] = field(default_factory=list)
     depends_on: dict[str, list[str]] = field(default_factory=dict)
+    constraints: list[Constraint] = field(default_factory=list)
     _raw_values: dict[str, Any] = field(default_factory=dict)
     _raw_test_values: list[dict[str, Any]] = field(default_factory=list)
 
@@ -232,6 +236,10 @@ def from_node(
                 .get("column_name")
             ],
             depends_on=node_values["depends_on"],
+            constraints=[
+                Constraint.from_raw_values(constraint)
+                for constraint in node_values["constraints"]
+            ],
             _raw_values=node_values,
             _raw_test_values=test_values,
         )

src/dbt_score/rules/filters.py

@@ -0,0 +1,10 @@
+"""Rule filters."""
+
+from dbt_score import rule_filter
+from dbt_score.models import Model
+
+
+@rule_filter
+def is_table(model: Model) -> bool:
+    """Models that are tables."""
+    return model.config.get("materialized") in {"table", "incremental"}

src/dbt_score/rules/generic.py

@@ -1,6 +1,7 @@
 """All generic rules."""
 
 from dbt_score import Model, RuleViolation, Severity, rule
+from dbt_score.rules.filters import is_table
 
 
 @rule
@@ -47,7 +48,65 @@ def sql_has_reasonable_number_of_lines(
 def has_example_sql(model: Model) -> RuleViolation | None:
     """The documentation of a model should have an example query."""
     if model.language == "sql":
-        if "```sql" not in model.description:
+        if "```sql" not in (model.description or ""):
             return RuleViolation(
                 "The model description does not include an example SQL query."
             )
+
+
+@rule(rule_filters={is_table()})
+def single_pk_defined_at_column_level(model: Model) -> RuleViolation | None:
+    """Single-column PK must be defined as a column constraint."""
+    for constraint in model.constraints:
+        if constraint.type == "primary_key":
+            if constraint.columns is not None and len(constraint.columns) == 1:
+                return RuleViolation(
+                    f"Single-column PK {constraint.columns[0]} must be defined as a "
+                    f"column constraint."
+                )
+
+
+@rule(rule_filters={is_table()})
+def single_column_uniqueness_at_column_level(model: Model) -> RuleViolation | None:
+    """Single-column uniqueness test must be defined as a column test."""
+    for data_test in model.tests:
+        if data_test.type == "unique_combination_of_columns":
+            if len(data_test.kwargs.get("combination_of_columns", [])) == 1:
+                return RuleViolation(
+                    "Single-column uniqueness test must be defined as a column test."
+                )
+
+
+@rule(rule_filters={is_table()})
+def has_uniqueness_test(model: Model) -> RuleViolation | None:
+    """Model has uniqueness test for primary key."""
+    # ruff: noqa: C901 [too-complex]
+
+    # Single-column PK
+    for column in model.columns:
+        for column_constraint in column.constraints:
+            if column_constraint.type == "primary_key":
+                for data_test in column.tests:
+                    if data_test.type == "unique":
+                        return None
+                return RuleViolation(
+                    f"No unique constraint defined on PK column {column.name}."
+                )
+
+    # Composite PK
+    pk_columns: list[str] = []
+    for model_constraint in model.constraints:
+        if model_constraint.type == "primary_key":
+            pk_columns = model_constraint.columns or []
+            break
+
+    if not pk_columns: # No PK, no need for uniqueness test
+        return None
+
+    for data_test in model.tests:
+        if data_test.type == "unique_combination_of_columns":
+            if set(data_test.kwargs.get("combination_of_columns")) == set(pk_columns): # type: ignore
+                return None
+    return RuleViolation(
+        f"No uniqueness test defined and matching PK {','.join(pk_columns)}."
+    )

tests/resources/manifest.json

@@ -26,6 +26,7 @@
           "tags": []
         }
       },
+      "constraints": [],
       "package_name": "package",
       "database": "db",
       "schema": "schema",
@@ -56,6 +57,7 @@
           "tags": []
         }
       },
+      "constraints": [],
       "package_name": "package",
       "database": "db",
       "schema": "schema",
@@ -86,6 +88,7 @@
           "tags": []
         }
       },
+      "constraints": [],
       "package_name": "package2",
       "database": "db",
       "schema": "schema",