Skip to content

Commit 4f62063

Browse files
axellpadillaaxellpadilla
committed
Added Basic Index configuration using dbt-postgresql reference
1 parent c401e02 commit 4f62063

File tree

5 files changed

+267
-3
lines changed

5 files changed

+267
-3
lines changed

‎dbt/adapters/sqlserver/relation_configs/__init__.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,8 @@
1+
from dbt.adapters.sqlserver.relation_configs.index import (
2+
SQLServerIndexConfig,
3+
SQLServerIndexConfigChange,
4+
SQLServerIndexType,
5+
)
16
from dbt.adapters.sqlserver.relation_configs.policies import (
27
MAX_CHARACTERS_IN_IDENTIFIER,
38
SQLServerIncludePolicy,
@@ -10,4 +15,7 @@
1015
"SQLServerIncludePolicy",
1116
"SQLServerQuotePolicy",
1217
"SQLServerRelationType",
18+
"SQLServerIndexType",
19+
"SQLServerIndexConfig",
20+
"SQLServerIndexConfigChange",
1321
]

‎dbt/adapters/sqlserver/relation_configs/index.py

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
from dataclasses import dataclass, field
2+
from datetime import datetime, timezone
3+
from typing import Optional
4+
5+
import agate
6+
from dbt.adapters.exceptions import IndexConfigError, IndexConfigNotDictError
7+
from dbt.adapters.relation_configs import (
8+
RelationConfigBase,
9+
RelationConfigChange,
10+
RelationConfigChangeAction,
11+
RelationConfigValidationMixin,
12+
RelationConfigValidationRule,
13+
)
14+
from dbt_common.dataclass_schema import StrEnum, ValidationError, dbtClassMixin
15+
from dbt_common.exceptions import DbtRuntimeError
16+
from dbt_common.utils import encoding as dbt_encoding
17+
18+
19+
# ALTERED FROM:
20+
# github.com/dbt-labs/dbt-postgres/blob/main/dbt/adapters/postgres/relation_configs/index.py
21+
class SQLServerIndexType(StrEnum):
22+
# btree = "btree" #All SQL Server common indexes are B-tree indexes
23+
# hash = "hash" #A hash index can exist only on a memory-optimized table.
24+
# TODO Implement memory optimized table materialization.
25+
clustered = "clustered" # Cant't have included columns
26+
nonclustered = "nonclustered"
27+
columnstore = "columnstore" # Cant't have included columns or unique config
28+
29+
@classmethod
30+
def default(cls) -> "SQLServerIndexType":
31+
return cls("nonclustered")
32+
33+
@classmethod
34+
def valid_types(cls):
35+
return list(cls)
36+
37+
38+
@dataclass(frozen=True, eq=True, unsafe_hash=True)
39+
class SQLServerIndexConfig(RelationConfigBase, RelationConfigValidationMixin, dbtClassMixin):
40+
"""
41+
This config follows the specs found here:
42+
43+
https://learn.microsoft.com/en-us/sql/t-sql/statements/create-index-transact-sql
44+
45+
The following parameters are configurable by dbt:
46+
- name: the name of the index in the database, isn't predictable since we apply a timestamp
47+
- unique: checks for duplicate values when the index is created and on data updates
48+
- type: the index type method to be used
49+
- columns: the columns names in the index
50+
- included_columns: the extra included columns names in the index
51+
52+
"""
53+
54+
name: str = field(default="", hash=False, compare=False)
55+
columns: list[str] = field(default_factory=list, hash=True) # Keeping order is important
56+
unique: bool = field(
57+
default=False, hash=True
58+
) # Uniqueness can be a property of both clustered and nonclustered indexes.
59+
type: SQLServerIndexType = field(default=SQLServerIndexType.default(), hash=True)
60+
included_columns: frozenset[str] = field(
61+
default_factory=frozenset, hash=True
62+
) # Keeping order is not important
63+
64+
@property
65+
def validation_rules(self) -> set[RelationConfigValidationRule]:
66+
return {
67+
RelationConfigValidationRule(
68+
validation_check=True if self.columns else False,
69+
validation_error=DbtRuntimeError("'columns' is a required property"),
70+
),
71+
RelationConfigValidationRule(
72+
validation_check=(
73+
True
74+
if not self.included_columns
75+
else self.type == SQLServerIndexType.nonclustered
76+
),
77+
validation_error=DbtRuntimeError(
78+
"Non-clustered indexes are the only index types that can include extra columns"
79+
),
80+
),
81+
RelationConfigValidationRule(
82+
validation_check=(
83+
True
84+
if not self.unique
85+
else self.type
86+
in (SQLServerIndexType.clustered, SQLServerIndexType.nonclustered)
87+
),
88+
validation_error=DbtRuntimeError(
89+
"Clustered and nonclustered indexes are the only types that can be unique"
90+
),
91+
),
92+
RelationConfigValidationRule(
93+
validation_check=True if self.type in SQLServerIndexType.valid_types() else False,
94+
validation_error=DbtRuntimeError(
95+
f"Invalid index type: {self.type}, valid types:"
96+
+ f"{SQLServerIndexType.valid_types()}"
97+
),
98+
),
99+
}
100+
101+
@classmethod
102+
def from_dict(cls, config_dict) -> "SQLServerIndexConfig":
103+
kwargs_dict = {
104+
"name": config_dict.get("name"),
105+
"columns": list(column for column in config_dict.get("columns", list())),
106+
"unique": config_dict.get("unique"),
107+
"type": config_dict.get("type"),
108+
"included_columns": frozenset(
109+
column for column in config_dict.get("included_columns", set())
110+
),
111+
}
112+
index: "SQLServerIndexConfig" = super().from_dict(kwargs_dict) # type: ignore
113+
return index
114+
115+
@classmethod
116+
def parse_model_node(cls, model_node_entry: dict) -> dict:
117+
config_dict = {
118+
"columns": list(model_node_entry.get("columns", list())),
119+
"unique": model_node_entry.get("unique"),
120+
"type": model_node_entry.get("type"),
121+
"included_columns": frozenset(model_node_entry.get("included_columns", set())),
122+
}
123+
return config_dict
124+
125+
@classmethod
126+
def parse_relation_results(cls, relation_results_entry: agate.Row) -> dict:
127+
config_dict = {
128+
"name": relation_results_entry.get("name"),
129+
"columns": list(relation_results_entry.get("columns", "").split(",")),
130+
"unique": relation_results_entry.get("unique"),
131+
"type": relation_results_entry.get("type"),
132+
"included_columns": set(relation_results_entry.get("included_columns", "").split(",")),
133+
}
134+
return config_dict
135+
136+
@property
137+
def as_node_config(self) -> dict:
138+
"""
139+
Returns: a dictionary that can be passed into `get_create_index_sql()`
140+
"""
141+
node_config = {
142+
"columns": list(self.columns),
143+
"unique": self.unique,
144+
"type": self.type.value,
145+
"included_columns": list(self.included_columns),
146+
}
147+
return node_config
148+
149+
def render(self, relation):
150+
# We append the current timestamp to the index name because otherwise
151+
# the index will only be created on every other run. See
152+
# https://github.com/dbt-labs/dbt-core/issues/1945#issuecomment-576714925
153+
# for an explanation.
154+
155+
now = datetime.now(timezone.utc).isoformat()
156+
inputs = self.columns + [relation.render(), str(self.unique), str(self.type), now]
157+
string = "_".join(inputs)
158+
return dbt_encoding.md5(string)
159+
160+
@classmethod
161+
def parse(cls, raw_index) -> Optional["SQLServerIndexConfig"]:
162+
if raw_index is None:
163+
return None
164+
try:
165+
cls.validate(raw_index)
166+
return cls.from_dict(raw_index)
167+
except ValidationError as exc:
168+
raise IndexConfigError(exc)
169+
except TypeError:
170+
raise IndexConfigNotDictError(raw_index)
171+
172+
173+
@dataclass(frozen=True, eq=True, unsafe_hash=True)
174+
class SQLServerIndexConfigChange(RelationConfigChange, RelationConfigValidationMixin):
175+
"""
176+
Example of an index change:
177+
{
178+
"action": "create",
179+
"context": {
180+
"name": "", # we don't know the name since it gets created as a hash at runtime
181+
"columns": ["column_1", "column_3"],
182+
"type": "clustered",
183+
"unique": True
184+
}
185+
},
186+
{
187+
"action": "drop",
188+
"context": {
189+
"name": "index_abc", # we only need this to drop, but we need the rest to compare
190+
"columns": ["column_1"],
191+
"type": "nonclustered",
192+
"unique": True
193+
}
194+
}
195+
"""
196+
197+
# TODO: Implement the change actions on the adapter
198+
context: SQLServerIndexConfig
199+
200+
@property
201+
def requires_full_refresh(self) -> bool:
202+
return False
203+
204+
@property
205+
def validation_rules(self) -> set[RelationConfigValidationRule]:
206+
return {
207+
RelationConfigValidationRule(
208+
validation_check=self.action
209+
in {RelationConfigChangeAction.create, RelationConfigChangeAction.drop},
210+
validation_error=DbtRuntimeError(
211+
"Invalid operation, only `drop` and `create` are supported for indexes."
212+
),
213+
),
214+
RelationConfigValidationRule(
215+
validation_check=not (
216+
self.action == RelationConfigChangeAction.drop and self.context.name is None
217+
),
218+
validation_error=DbtRuntimeError(
219+
"Invalid operation, attempting to drop an index with no name."
220+
),
221+
),
222+
RelationConfigValidationRule(
223+
validation_check=not (
224+
self.action == RelationConfigChangeAction.create
225+
and self.context.columns == set()
226+
),
227+
validation_error=DbtRuntimeError(
228+
"Invalid operations, attempting to create an index with no columns."
229+
),
230+
),
231+
}

‎dbt/adapters/sqlserver/sqlserver_adapter.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
1-
from typing import Optional
1+
from typing import Any, Optional
22

33
import dbt.exceptions
4-
from dbt.adapters.base.impl import ConstraintSupport
4+
from dbt.adapters.base import ConstraintSupport, available
55
from dbt.adapters.fabric import FabricAdapter
66
from dbt.contracts.graph.nodes import ConstraintType
77

8+
from dbt.adapters.sqlserver.relation_configs import SQLServerIndexConfig
89
from dbt.adapters.sqlserver.sqlserver_column import SQLServerColumn
10+
from dbt.adapters.sqlserver.sqlserver_configs import SQLServerConfigs
911
from dbt.adapters.sqlserver.sqlserver_connections import SQLServerConnectionManager
1012
from dbt.adapters.sqlserver.sqlserver_relation import SQLServerRelation
1113

@@ -18,6 +20,7 @@ class SQLServerAdapter(FabricAdapter):
1820
ConnectionManager = SQLServerConnectionManager
1921
Column = SQLServerColumn
2022
Relation = SQLServerRelation
23+
AdapterSpecificConfigs = SQLServerConfigs
2124

2225
CONSTRAINT_SUPPORT = {
2326
ConstraintType.check: ConstraintSupport.ENFORCED,
@@ -59,3 +62,7 @@ def render_model_constraint(cls, constraint) -> Optional[str]:
5962
@classmethod
6063
def date_function(cls):
6164
return "getdate()"
65+
66+
@available
67+
def parse_index(self, raw_index: Any) -> Optional[SQLServerIndexConfig]:
68+
return SQLServerIndexConfig.parse(raw_index)

‎dbt/adapters/sqlserver/sqlserver_configs.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
from dataclasses import dataclass
2+
from typing import Optional
23

34
from dbt.adapters.fabric import FabricConfigs
45

6+
from dbt.adapters.sqlserver.relation_configs import SQLServerIndexConfig
7+
58

69
@dataclass
710
class SQLServerConfigs(FabricConfigs):
8-
pass
11+
indexes: Optional[list[SQLServerIndexConfig]] = None

‎dbt/include/sqlserver/macros/adapter/indexes.sql

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -168,3 +168,18 @@
168168
{% endif %}
169169
end
170170
{% endmacro %}
171+
172+
173+
{% macro sqlserver__get_create_index_sql(relation, index_dict) -%}
174+
{%- set index_config = adapter.parse_index(index_dict) -%}
175+
{%- set comma_separated_columns = ", ".join(index_config.columns) -%}
176+
{%- set index_name = index_config.render(relation) -%}
177+
178+
{# Validations are made on the adapter class SQLServerIndexConfig to control resulting sql #}
179+
create
180+
{% if index_config.unique -%} unique {% endif %}{{ index_config.type }}
181+
index "{{ index_name }}"
182+
on {{ relation }}
183+
({{ comma_separated_columns }})
184+
185+
{%- endmacro %}

0 commit comments

Comments
 (0)