Skip to content

Commit a171f9d

Browse files
axellpadillaaxellpadilla
committed
Refactor index and include tests for included columns
1 parent 4f62063 commit a171f9d

File tree

2 files changed

+183
-57
lines changed

2 files changed

+183
-57
lines changed

‎tests/functional/adapter/mssql/test_index_config.py‎

Lines changed: 180 additions & 54 deletions
Original file line numberDiff line numberDiff line change
@@ -6,37 +6,49 @@
66
base_validation = """
77
with base_query AS (
88
select i.[name] as index_name,
9-
substring(column_names, 1, len(column_names)-1) as [columns],
10-
case when i.[type] = 1 then 'Clustered index'
11-
when i.[type] = 2 then 'Nonclustered unique index'
12-
when i.[type] = 3 then 'XML index'
13-
when i.[type] = 4 then 'Spatial index'
14-
when i.[type] = 5 then 'Clustered columnstore index'
15-
when i.[type] = 6 then 'Nonclustered columnstore index'
16-
when i.[type] = 7 then 'Nonclustered hash index'
17-
end as index_type,
18-
case when i.is_unique = 1 then 'Unique'
19-
else 'Not unique' end as [unique],
20-
schema_name(t.schema_id) + '.' + t.[name] as table_view,
21-
case when t.[type] = 'U' then 'Table'
22-
when t.[type] = 'V' then 'View'
23-
end as [object_type],
9+
substring(column_names, 1, len(column_names)-1) as [columns],
10+
substring(included_column_names, 1, len(included_column_names)-1) as included_columns,
11+
case when i.[type] = 1 then 'clustered'
12+
when i.[type] = 2 then 'nonclustered'
13+
when i.[type] = 3 then 'xml'
14+
when i.[type] = 4 then 'spatial'
15+
when i.[type] = 5 then 'clustered columnstore'
16+
when i.[type] = 6 then 'nonclustered columnstore'
17+
when i.[type] = 7 then 'nonclustered hash'
18+
end as index_type,
19+
case when i.is_unique = 1 then 'Unique'
20+
else 'Not unique' end as [unique],
21+
schema_name(t.schema_id) + '.' + t.[name] as table_view,
22+
case when t.[type] = 'U' then 'Table'
23+
when t.[type] = 'V' then 'View'
24+
end as [object_type],
2425
s.name as schema_name
2526
from sys.objects t
2627
inner join sys.schemas s
27-
on
28-
t.schema_id = s.schema_id
29-
inner join sys.indexes i
30-
on t.object_id = i.object_id
31-
cross apply (select col.[name] + ', '
32-
from sys.index_columns ic
33-
inner join sys.columns col
34-
on ic.object_id = col.object_id
35-
and ic.column_id = col.column_id
36-
where ic.object_id = t.object_id
37-
and ic.index_id = i.index_id
38-
order by key_ordinal
39-
for xml path ('') ) D (column_names)
28+
on
29+
t.schema_id = s.schema_id
30+
inner join sys.indexes i
31+
on t.object_id = i.object_id
32+
cross apply (select col.[name] + ', '
33+
from sys.index_columns ic
34+
inner join sys.columns col
35+
on ic.object_id = col.object_id
36+
and ic.column_id = col.column_id
37+
where ic.object_id = t.object_id
38+
and ic.index_id = i.index_id
39+
and ic.is_included_column = 0
40+
order by key_ordinal
41+
for xml path ('') ) D (column_names)
42+
cross apply (select col.[name] + ', '
43+
from sys.index_columns ic
44+
inner join sys.columns col
45+
on ic.object_id = col.object_id
46+
and ic.column_id = col.column_id
47+
where ic.object_id = t.object_id
48+
and ic.index_id = i.index_id
49+
and ic.is_included_column = 1
50+
order by key_ordinal
51+
for xml path ('') ) E (included_column_names)
4052
where t.is_ms_shipped <> 1
4153
and index_id > 0
4254
)
@@ -46,13 +58,13 @@
4658
base_validation
4759
+ """
4860
select
49-
index_type,
61+
index_type + case when [unique] = 'Unique' then ' unique' else '' end as index_type,
5062
count(*) index_count
5163
from
5264
base_query
5365
WHERE
5466
schema_name='{schema_name}'
55-
group by index_type
67+
group by index_type + case when [unique] = 'Unique' then ' unique' else '' end
5668
"""
5769
)
5870

@@ -62,6 +74,7 @@
6274
SELECT
6375
index_name,
6476
[columns],
77+
[included_columns],
6578
index_type,
6679
[unique],
6780
table_view,
@@ -135,7 +148,26 @@
135148
{'columns': ['column_b']},
136149
{'columns': ['column_a', 'column_b']},
137150
{'columns': ['column_b', 'column_a'], 'type': 'clustered', 'unique': True},
138-
{'columns': ['column_a'], 'type': 'nonclustered'}
151+
{'columns': ['column_a','column_c'],
152+
'type': 'nonclustered',
153+
'included_columns': ['column_b']},
154+
]
155+
)
156+
}}
157+
158+
select 1 as column_a, 2 as column_b, 3 as column_c
159+
160+
"""
161+
162+
163+
models__table_included_sql = """
164+
{{
165+
config(
166+
materialized = "table",
167+
as_columnstore = False,
168+
indexes=[
169+
{'columns': ['column_a'], 'included_columns': ['column_b']},
170+
{'columns': ['column_b'], 'type': 'clustered'}
139171
]
140172
)
141173
}}
@@ -248,6 +280,7 @@ def models(self):
248280
"table.sql": models__table_sql,
249281
"incremental.sql": models__incremental_sql,
250282
"columnstore.sql": models__columnstore_sql,
283+
"table_included.sql": models__table_included_sql,
251284
}
252285

253286
@pytest.fixture(scope="class")
@@ -265,8 +298,12 @@ def project_config_update(self):
265298
"seeds": {
266299
"quote_columns": False,
267300
"indexes": [
268-
{"columns": ["country_code"], "unique": False, "type": "nonclustered"},
269-
{"columns": ["country_code", "country_name"], "unique": True},
301+
{"columns": ["country_code"], "unique": False},
302+
{
303+
"columns": ["country_code", "country_name"],
304+
"unique": True,
305+
"type": "clustered",
306+
},
270307
],
271308
},
272309
"vars": {
@@ -279,26 +316,85 @@ def test_table(self, project, unique_schema):
279316
assert len(results) == 1
280317

281318
indexes = self.get_indexes("table", project, unique_schema)
319+
indexes = self.sort_indexes(indexes)
282320
expected = [
283-
{"columns": "column_a", "unique": False, "type": "nonclustered"},
284-
{"columns": "column_b", "unique": False, "type": "nonclustered"},
285-
{"columns": "column_a, column_b", "unique": False, "type": "nonclustered"},
286-
{"columns": "column_b, column_a", "unique": True, "type": "clustered"},
287-
{"columns": "column_a", "unique": False, "type": "nonclustered"},
321+
{
322+
"columns": "column_a",
323+
"unique": False,
324+
"type": "nonclustered",
325+
"included_columns": None,
326+
},
327+
{
328+
"columns": "column_a, column_b",
329+
"unique": False,
330+
"type": "nonclustered",
331+
"included_columns": None,
332+
},
333+
{
334+
"columns": "column_a, column_c",
335+
"unique": False,
336+
"type": "nonclustered",
337+
"included_columns": "column_b",
338+
},
339+
{
340+
"columns": "column_b",
341+
"unique": False,
342+
"type": "nonclustered",
343+
"included_columns": None,
344+
},
345+
{
346+
"columns": "column_b, column_a",
347+
"unique": True,
348+
"type": "clustered",
349+
"included_columns": None,
350+
},
351+
]
352+
assert indexes == expected
353+
354+
def test_table_included(self, project, unique_schema):
355+
results = run_dbt(["run", "--models", "table_included"])
356+
assert len(results) == 1
357+
358+
indexes = self.get_indexes("table_included", project, unique_schema)
359+
indexes = self.sort_indexes(indexes)
360+
expected = [
361+
{
362+
"columns": "column_a",
363+
"unique": False,
364+
"type": "nonclustered",
365+
"included_columns": "column_b",
366+
},
367+
{
368+
"columns": "column_b",
369+
"unique": False,
370+
"type": "clustered",
371+
"included_columns": None,
372+
},
288373
]
289-
assert len(indexes) == len(expected)
374+
assert indexes == expected
290375

291376
def test_incremental(self, project, unique_schema):
292377
for additional_argument in [[], [], ["--full-refresh"]]:
293378
results = run_dbt(["run", "--models", "incremental"] + additional_argument)
294379
assert len(results) == 1
295380

296381
indexes = self.get_indexes("incremental", project, unique_schema)
382+
indexes = self.sort_indexes(indexes)
297383
expected = [
298-
{"columns": "column_a", "unique": False, "type": "nonclustered"},
299-
{"columns": "column_a, column_b", "unique": True, "type": "nonclustered"},
384+
{
385+
"columns": "column_a",
386+
"unique": False,
387+
"type": "nonclustered",
388+
"included_columns": None,
389+
},
390+
{
391+
"columns": "column_a, column_b",
392+
"unique": True,
393+
"type": "nonclustered",
394+
"included_columns": None,
395+
},
300396
]
301-
assert len(indexes) == len(expected)
397+
assert indexes == expected
302398

303399
def test_columnstore(self, project, unique_schema):
304400
for additional_argument in [[], [], ["--full-refresh"]]:
@@ -307,49 +403,79 @@ def test_columnstore(self, project, unique_schema):
307403

308404
indexes = self.get_indexes("columnstore", project, unique_schema)
309405
expected = [
310-
{"columns": "column_a", "unique": False, "type": "columnstore"},
406+
{
407+
"columns": "column_a",
408+
"unique": False,
409+
"type": "columnstore",
410+
"included_columns": None,
411+
},
311412
]
312-
assert len(indexes) == len(expected)
413+
assert len(indexes) == len(
414+
expected
415+
) # Nonclustered columnstore indexes meta is different
313416

314417
def test_seed(self, project, unique_schema):
315418
for additional_argument in [[], [], ["--full-refresh"]]:
316419
results = run_dbt(["seed"] + additional_argument)
317420
assert len(results) == 1
318421

319422
indexes = self.get_indexes("seed", project, unique_schema)
423+
indexes = self.sort_indexes(indexes)
320424
expected = [
321-
{"columns": "country_code", "unique": False, "type": "nonclustered"},
322-
{"columns": "country_code, country_name", "unique": True, "type": "clustered"},
425+
{
426+
"columns": "country_code",
427+
"unique": False,
428+
"type": "nonclustered",
429+
"included_columns": None,
430+
},
431+
{
432+
"columns": "country_code, country_name",
433+
"unique": True,
434+
"type": "clustered",
435+
"included_columns": None,
436+
},
323437
]
324-
assert len(indexes) == len(expected)
438+
assert indexes == expected
325439

326440
def test_snapshot(self, project, unique_schema):
327441
for version in [1, 2]:
328442
results = run_dbt(["snapshot", "--vars", f"version: {version}"])
329443
assert len(results) == 1
330444

331445
indexes = self.get_indexes("colors", project, unique_schema)
446+
indexes = self.sort_indexes(indexes)
332447
expected = [
333-
{"columns": "id", "unique": False, "type": "nonclustered"},
334-
{"columns": "id, color", "unique": True, "type": "clustered"},
448+
{
449+
"columns": "id",
450+
"unique": False,
451+
"type": "nonclustered",
452+
"included_columns": None,
453+
},
454+
{
455+
"columns": "id, color",
456+
"unique": True,
457+
"type": "nonclustered",
458+
"included_columns": None,
459+
},
335460
]
336-
assert len(indexes) == len(expected)
461+
assert indexes == expected
337462

338463
def get_indexes(self, table_name, project, unique_schema):
339464
sql = indexes_def.format(schema_name=unique_schema, table_name=table_name)
340465
results = project.run_sql(sql, fetch="all")
341466
return [self.index_definition_dict(row) for row in results]
342467

343468
def index_definition_dict(self, index_definition):
344-
is_unique = index_definition[3] == "Unique"
469+
is_unique = index_definition[4] == "Unique"
345470
return {
346471
"columns": index_definition[1],
472+
"included_columns": index_definition[2],
347473
"unique": is_unique,
348-
"type": index_definition[2],
474+
"type": index_definition[3],
349475
}
350476

351-
def assertCountEqual(self, a, b):
352-
assert len(a) == len(b)
477+
def sort_indexes(self, indexes):
478+
return sorted(indexes, key=lambda x: (x["columns"], x["type"]))
353479

354480

355481
class TestSQLServerInvalidIndex:

‎tests/functional/adapter/mssql/test_index_macros.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -104,9 +104,9 @@ def test_create_index(self, project):
104104
)
105105
schema_dict = {_[0]: _[1] for _ in table.rows}
106106
expected = {
107-
"Clustered columnstore index": 1,
108-
"Clustered index": 1,
109-
"Nonclustered unique index": 4,
107+
"clustered columnstore": 1,
108+
"clustered unique": 1,
109+
"nonclustered": 4,
110110
}
111111
self.drop_artifacts(project)
112112
assert schema_dict == expected

0 commit comments

Comments
 (0)