Skip to content

Commit 1ccba21

Browse files
ericj-dbericj-db
authored
Constraints for incremental materialization (#1013)
2 parents 48624cf + 32dce56 commit 1ccba21

11 files changed

Lines changed: 1218 additions & 1 deletion

File tree

‎CHANGELOG.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
## dbt-databricks 1.10.2 (TBD)
22

3+
### Features
4+
5+
- Support constraint updates on incremental runs (with Materialization V2) ([1013](https://github.com/databricks/dbt-databricks/pull/1013))
6+
37
### Fixes
48

59
- Fix bug with multiple not_null constraints defined on the model level ([1008](https://github.com/databricks/dbt-databricks/pull/1008))

‎dbt/adapters/databricks/constraints.py‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
from abc import ABC, abstractmethod
2+
from dataclasses import dataclass
23
from typing import Any, ClassVar, Optional, TypeVar
34
from uuid import uuid4
45

@@ -27,6 +28,7 @@
2728
T = TypeVar("T", bound="TypedConstraint")
2829

2930

31+
@dataclass
3032
class TypedConstraint(ModelLevelConstraint, ABC):
3133
"""Constraint that enforces type because it has render logic"""
3234

@@ -61,6 +63,19 @@ def _render_error(self, missing: list[list[str]]) -> DbtValidationError:
6163
f"{self.str_type} constraint '{name}' is missing required field(s): {fields}"
6264
)
6365

66+
# Enables set equality checks, especially for convenient unit testing
67+
def __hash__(self) -> int:
68+
# Create a tuple of all the fields that should be used for equality comparison
69+
fields = (
70+
self.type,
71+
self.name,
72+
tuple(self.columns) if self.columns else None,
73+
self.expression,
74+
self.to,
75+
tuple(self.to_columns) if self.to_columns else None,
76+
)
77+
return hash(fields)
78+
6479

6580
class CustomConstraint(TypedConstraint):
6681
str_type = "custom"

‎dbt/adapters/databricks/impl.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -944,6 +944,15 @@ def _describe_relation(
944944

945945
if not relation.is_hive_metastore():
946946
results["information_schema.tags"] = adapter.execute_macro("fetch_tags", kwargs=kwargs)
947+
results["non_null_constraint_columns"] = adapter.execute_macro(
948+
"fetch_non_null_constraint_columns", kwargs=kwargs
949+
)
950+
results["primary_key_constraints"] = adapter.execute_macro(
951+
"fetch_primary_key_constraints", kwargs=kwargs
952+
)
953+
results["foreign_key_constraints"] = adapter.execute_macro(
954+
"fetch_foreign_key_constraints", kwargs=kwargs
955+
)
947956
results["show_tblproperties"] = adapter.execute_macro("fetch_tbl_properties", kwargs=kwargs)
948957

949958
kwargs = {"table_name": relation}

‎dbt/adapters/databricks/relation_configs/constraints.py‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
from dataclasses import asdict
2+
from typing import ClassVar, Optional
3+
4+
from agate import Table
5+
6+
from dbt.adapters.contracts.relation import RelationConfig
7+
from dbt.adapters.databricks.constraints import (
8+
CheckConstraint,
9+
ConstraintType,
10+
ForeignKeyConstraint,
11+
PrimaryKeyConstraint,
12+
TypedConstraint,
13+
parse_constraints,
14+
)
15+
from dbt.adapters.databricks.relation_configs.base import (
16+
DatabricksComponentConfig,
17+
DatabricksComponentProcessor,
18+
)
19+
from dbt.adapters.relation_configs.config_base import RelationResults
20+
21+
22+
class ConstraintsConfig(DatabricksComponentConfig):
23+
"""Component encapsulating the constraints of a relation."""
24+
25+
set_non_nulls: set[str]
26+
unset_non_nulls: set[str] = set()
27+
set_constraints: set[TypedConstraint]
28+
unset_constraints: set[TypedConstraint] = set()
29+
30+
def get_diff(self, other: "ConstraintsConfig") -> Optional["ConstraintsConfig"]:
31+
# Find constraints that need to be unset
32+
constraints_to_unset = other.set_constraints - self.set_constraints
33+
# Find non-nulls that need to be unset
34+
non_nulls_to_unset = other.set_non_nulls - self.set_non_nulls
35+
36+
# Set constraints that exist in self but not in other
37+
set_constraints = self.set_constraints - other.set_constraints
38+
set_non_nulls = self.set_non_nulls - other.set_non_nulls
39+
40+
if set_constraints or set_non_nulls or constraints_to_unset or non_nulls_to_unset:
41+
return ConstraintsConfig(
42+
set_non_nulls=set_non_nulls,
43+
unset_non_nulls=non_nulls_to_unset,
44+
set_constraints=set_constraints,
45+
unset_constraints=constraints_to_unset,
46+
)
47+
return None
48+
49+
50+
class ConstraintsProcessor(DatabricksComponentProcessor[ConstraintsConfig]):
51+
name: ClassVar[str] = "constraints"
52+
53+
@classmethod
54+
def _process_check_constraints(
55+
cls, check_constraints_table: Optional[Table]
56+
) -> set[CheckConstraint]:
57+
check_constraints = set()
58+
if check_constraints_table:
59+
for row in check_constraints_table.rows:
60+
if row[0].startswith("delta.constraints."):
61+
check_constraints.add(
62+
CheckConstraint(
63+
type=ConstraintType.check,
64+
name=row[0].replace("delta.constraints.", ""),
65+
expression=row[1],
66+
)
67+
)
68+
return check_constraints
69+
70+
@classmethod
71+
def _process_primary_key_constraints(
72+
cls, pk_constraints_table: Optional[Table]
73+
) -> set[PrimaryKeyConstraint]:
74+
pk_constraints = set()
75+
if pk_constraints_table:
76+
# Group columns by constraint name
77+
constraint_columns: dict[str, list[str]] = {}
78+
for row in pk_constraints_table.rows:
79+
constraint_name = row[0]
80+
column_name = row[1]
81+
if constraint_name not in constraint_columns:
82+
constraint_columns[constraint_name] = []
83+
constraint_columns[constraint_name].append(column_name)
84+
85+
# Create one PrimaryKeyConstraint per unique constraint name
86+
for constraint_name, columns in constraint_columns.items():
87+
pk_constraints.add(
88+
PrimaryKeyConstraint(
89+
type=ConstraintType.primary_key, name=constraint_name, columns=columns
90+
)
91+
)
92+
return pk_constraints
93+
94+
@classmethod
95+
def _process_foreign_key_constraints(
96+
cls, fk_constraints_table: Optional[Table]
97+
) -> set[ForeignKeyConstraint]:
98+
fk_constraints = set()
99+
if fk_constraints_table:
100+
# Group by constraint_name
101+
fk_data: dict[str, dict] = {}
102+
for row in fk_constraints_table.rows:
103+
constraint_name = row[0]
104+
if constraint_name not in fk_data:
105+
fk_data[constraint_name] = {
106+
"columns": [],
107+
"to": f"`{row[2]}`.`{row[3]}`.`{row[4]}`",
108+
"to_columns": [],
109+
}
110+
fk_data[constraint_name]["columns"].append(row[1])
111+
fk_data[constraint_name]["to_columns"].append(row[5])
112+
113+
# Create ForeignKeyConstraint objects
114+
for constraint_name, data in fk_data.items():
115+
fk_constraints.add(
116+
ForeignKeyConstraint(
117+
type=ConstraintType.foreign_key,
118+
name=constraint_name,
119+
columns=data["columns"],
120+
to=data["to"],
121+
to_columns=data["to_columns"],
122+
)
123+
)
124+
return fk_constraints
125+
126+
@classmethod
127+
def from_relation_results(cls, results: RelationResults) -> ConstraintsConfig:
128+
non_null_constraint_columns = results.get("non_null_constraint_columns")
129+
to_set_non_nulls = (
130+
{row[0] for row in non_null_constraint_columns.rows}
131+
if non_null_constraint_columns
132+
else set()
133+
)
134+
135+
table_properties = results.get("show_tblproperties")
136+
primary_key_constraints = results.get("primary_key_constraints")
137+
foreign_key_constraints = results.get("foreign_key_constraints")
138+
139+
check_constraints = cls._process_check_constraints(table_properties)
140+
pk_constraints = cls._process_primary_key_constraints(primary_key_constraints)
141+
fk_constraints = cls._process_foreign_key_constraints(foreign_key_constraints)
142+
143+
all_constraints: list[TypedConstraint] = []
144+
all_constraints.extend(check_constraints)
145+
all_constraints.extend(pk_constraints)
146+
all_constraints.extend(fk_constraints)
147+
148+
return ConstraintsConfig(
149+
set_non_nulls=to_set_non_nulls,
150+
set_constraints=set(all_constraints),
151+
)
152+
153+
@classmethod
154+
def from_relation_config(cls, relation_config: RelationConfig) -> ConstraintsConfig:
155+
constraints = getattr(relation_config, "constraints", [])
156+
constraints = [
157+
{
158+
"name": (item["name"] if isinstance(item, dict) else item.name),
159+
**(item if isinstance(item, dict) else asdict(item)),
160+
}
161+
for item in constraints
162+
]
163+
164+
columns = getattr(relation_config, "columns", {})
165+
columns = [
166+
{"name": name, **(col if isinstance(col, dict) else asdict(col))}
167+
for name, col in columns.items()
168+
]
169+
170+
non_nulls, other_constraints = parse_constraints(columns, constraints)
171+
172+
return ConstraintsConfig(
173+
set_non_nulls=set(non_nulls),
174+
set_constraints=set(other_constraints),
175+
)

‎dbt/adapters/databricks/relation_configs/incremental.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
)
44
from dbt.adapters.databricks.relation_configs.column_comments import ColumnCommentsProcessor
55
from dbt.adapters.databricks.relation_configs.comment import CommentProcessor
6+
from dbt.adapters.databricks.relation_configs.constraints import ConstraintsProcessor
67
from dbt.adapters.databricks.relation_configs.liquid_clustering import LiquidClusteringProcessor
78
from dbt.adapters.databricks.relation_configs.tags import TagsProcessor
89
from dbt.adapters.databricks.relation_configs.tblproperties import TblPropertiesProcessor
@@ -12,6 +13,7 @@ class IncrementalTableConfig(DatabricksRelationConfigBase):
1213
config_components = [
1314
CommentProcessor,
1415
ColumnCommentsProcessor,
16+
ConstraintsProcessor,
1517
TagsProcessor,
1618
TblPropertiesProcessor,
1719
LiquidClusteringProcessor,

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

Lines changed: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
1+
{% macro fetch_non_null_constraint_columns(relation) -%}
2+
{% if relation.is_hive_metastore() %}
3+
{{ exceptions.raise_compiler_error("Incremental application of constraints is not supported for Hive Metastore") }}
4+
{%- endif %}
5+
{% call statement('list_non_null_constraint_columns', fetch_result=True) -%}
6+
{{ fetch_non_null_constraint_columns_sql(relation) }}
7+
{% endcall %}
8+
{% do return(load_result('list_non_null_constraint_columns').table) %}
9+
{%- endmacro -%}
10+
11+
{% macro fetch_non_null_constraint_columns_sql(relation) -%}
12+
SELECT column_name
13+
FROM `{{ relation.database|lower }}`.`information_schema`.`columns`
14+
WHERE table_catalog = '{{ relation.database|lower }}'
15+
AND table_schema = '{{ relation.schema|lower }}'
16+
AND table_name = '{{ relation.identifier|lower }}'
17+
AND is_nullable = 'NO';
18+
{%- endmacro -%}
19+
20+
{% macro fetch_primary_key_constraints(relation) -%}
21+
{% if relation.is_hive_metastore() %}
22+
{{ exceptions.raise_compiler_error("Incremental application of constraints is not supported for Hive Metastore") }}
23+
{%- endif %}
24+
{% call statement('list_primary_key_constraints', fetch_result=True) -%}
25+
{{ fetch_primary_key_constraints_sql(relation) }}
26+
{% endcall %}
27+
{% do return(load_result('list_primary_key_constraints').table) %}
28+
{%- endmacro -%}
29+
30+
{% macro fetch_primary_key_constraints_sql(relation) -%}
31+
SELECT kcu.constraint_name, kcu.column_name
32+
FROM `{{ relation.database|lower }}`.information_schema.key_column_usage kcu
33+
WHERE kcu.table_catalog = '{{ relation.database|lower }}'
34+
AND kcu.table_schema = '{{ relation.schema|lower }}'
35+
AND kcu.table_name = '{{ relation.identifier|lower }}'
36+
AND kcu.constraint_name = (
37+
SELECT constraint_name
38+
FROM `{{ relation.database|lower }}`.information_schema.table_constraints
39+
WHERE table_catalog = '{{ relation.database|lower }}'
40+
AND table_schema = '{{ relation.schema|lower }}'
41+
AND table_name = '{{ relation.identifier|lower }}'
42+
AND constraint_type = 'PRIMARY KEY'
43+
)
44+
ORDER BY kcu.ordinal_position;
45+
{%- endmacro -%}
46+
47+
{% macro fetch_foreign_key_constraints(relation) -%}
48+
{% if relation.is_hive_metastore() %}
49+
{{ exceptions.raise_compiler_error("Incremental application of constraints is not supported for Hive Metastore") }}
50+
{%- endif %}
51+
{% call statement('list_foreign_key_constraints', fetch_result=True) -%}
52+
{{ fetch_foreign_key_constraints_sql(relation) }}
53+
{% endcall %}
54+
{% do return(load_result('list_foreign_key_constraints').table) %}
55+
{%- endmacro -%}
56+
57+
{% macro fetch_foreign_key_constraints_sql(relation) -%}
58+
SELECT
59+
kcu.constraint_name,
60+
kcu.column_name AS from_column,
61+
ukcu.table_catalog AS to_catalog,
62+
ukcu.table_schema AS to_schema,
63+
ukcu.table_name AS to_table,
64+
ukcu.column_name AS to_column
65+
FROM `{{ relation.database|lower }}`.information_schema.key_column_usage kcu
66+
JOIN `{{ relation.database|lower }}`.information_schema.referential_constraints rc
67+
ON kcu.constraint_name = rc.constraint_name
68+
JOIN `{{ relation.database|lower }}`.information_schema.key_column_usage ukcu
69+
ON rc.unique_constraint_name = ukcu.constraint_name
70+
AND kcu.ordinal_position = ukcu.ordinal_position
71+
WHERE kcu.table_catalog = '{{ relation.database|lower }}'
72+
AND kcu.table_schema = '{{ relation.schema|lower }}'
73+
AND kcu.table_name = '{{ relation.identifier|lower }}'
74+
AND kcu.constraint_name = (
75+
SELECT constraint_name
76+
FROM `{{ relation.database|lower }}`.information_schema.table_constraints
77+
WHERE table_catalog = '{{ relation.database|lower }}'
78+
AND table_schema = '{{ relation.schema|lower }}'
79+
AND table_name = '{{ relation.identifier|lower }}'
80+
AND constraint_type = 'FOREIGN KEY'
81+
)
82+
ORDER BY kcu.ordinal_position;
83+
{%- endmacro -%}
84+
85+
{% macro apply_constraints(relation, constraints) -%}
86+
{{ log("Applying constraints to relation " ~ constraints) }}
87+
{%- if constraints and relation.is_hive_metastore() -%}
88+
{{ exceptions.raise_compiler_error("Constraints are only supported for Unity Catalog") }}
89+
{%- endif -%}
90+
{# Order matters here because key constraints depend on non-null constraints #}
91+
{%- if constraints.unset_constraints %}
92+
{%- for constraint in constraints.unset_constraints -%}
93+
{%- call statement('main') -%}
94+
{{ alter_unset_constraint(relation, constraint) }}
95+
{%- endcall -%}
96+
{%- endfor -%}
97+
{%- endif %}
98+
{%- if constraints.unset_non_nulls %}
99+
{%- for column in constraints.unset_non_nulls -%}
100+
{%- call statement('main') -%}
101+
{{ alter_unset_non_null_constraint(relation, column) }}
102+
{%- endcall -%}
103+
{%- endfor -%}
104+
{%- endif %}
105+
{%- if constraints.set_non_nulls %}
106+
{%- for column in constraints.set_non_nulls -%}
107+
{%- call statement('main') -%}
108+
{{ alter_set_non_null_constraint(relation, column) }}
109+
{%- endcall -%}
110+
{%- endfor -%}
111+
{%- endif %}
112+
{%- if constraints.set_constraints %}
113+
{%- for constraint in constraints.set_constraints -%}
114+
{%- call statement('main') -%}
115+
{{ alter_set_constraint(relation, constraint) }}
116+
{%- endcall -%}
117+
{%- endfor -%}
118+
{%- endif %}
119+
{%- endmacro -%}
120+
121+
{% macro alter_set_non_null_constraint(relation, column) -%}
122+
ALTER {{ relation.type }} {{ relation.render() }} ALTER COLUMN {{ column }} SET NOT NULL;
123+
{%- endmacro -%}
124+
125+
{% macro alter_unset_non_null_constraint(relation, column) -%}
126+
ALTER {{ relation.type }} {{ relation.render() }} ALTER COLUMN {{ column }} DROP NOT NULL;
127+
{%- endmacro -%}
128+
129+
{% macro alter_set_constraint(relation, constraint) -%}
130+
ALTER {{ relation.type }} {{ relation.render() }} ADD {{ constraint.render() }};
131+
{%- endmacro -%}
132+
133+
{% macro alter_unset_constraint(relation, constraint) -%}
134+
ALTER {{ relation.type }} {{ relation.render() }} DROP CONSTRAINT {{ constraint.name }};
135+
{%- endmacro -%}

0 commit comments

Comments
 (0)