|
| 1 | +import pytest |
| 2 | +from dbt.tests.util import check_table_does_exist, run_dbt |
| 3 | + |
| 4 | +_SEED_CSV = """ |
| 5 | +id,first_name,last_name,email,product_id |
| 6 | + |
| 7 | +2,Kathryn,Walker,[email protected],1 |
| 8 | + |
| 9 | + |
| 10 | +5,Kathryn,Walker,[email protected],5 |
| 11 | + |
| 12 | +""".lstrip() |
| 13 | + |
| 14 | + |
| 15 | +class TestClusterBy: |
| 16 | + @pytest.fixture(scope="class") |
| 17 | + def seeds(self): |
| 18 | + return {"seed.csv": _SEED_CSV} |
| 19 | + |
| 20 | + @pytest.fixture(scope="class") |
| 21 | + def models(self, dbt_profile_target): |
| 22 | + warehouse_name = dbt_profile_target["warehouse"] |
| 23 | + |
| 24 | + _DYNAMIC_TABLE_1_SQL = f""" |
| 25 | + {{{{ config(materialized='dynamic_table', snowflake_warehouse='{warehouse_name}', target_lag='1 minute') }}}} |
| 26 | + select * from {{{{ ref('seed') }}}} |
| 27 | + """.lstrip() |
| 28 | + |
| 29 | + _DYNAMIC_TABLE_2_SQL = f""" |
| 30 | + {{{{ config(materialized='dynamic_table', cluster_by=['last_name'], snowflake_warehouse='{warehouse_name}', target_lag='1 minute') }}}} |
| 31 | + select * from {{{{ ref('dynamic_table_1') }}}} |
| 32 | + """.lstrip() |
| 33 | + |
| 34 | + _DYNAMIC_TABLE_3_SQL = f""" |
| 35 | + {{{{ config(materialized='dynamic_table', cluster_by=['last_name', 'first_name'], snowflake_warehouse='{warehouse_name}', target_lag='1 minute') }}}} |
| 36 | + select |
| 37 | + last_name, |
| 38 | + first_name, |
| 39 | + count(*) as count |
| 40 | + from {{{{ ref('seed') }}}} |
| 41 | + group by 1, 2 |
| 42 | + """.lstrip() |
| 43 | + |
| 44 | + _DYNAMIC_TABLE_4_SQL = f""" |
| 45 | + {{{{ config(materialized='dynamic_table', cluster_by=['last_name', 'product_id % 3'], snowflake_warehouse='{warehouse_name}', target_lag='1 minute') }}}} |
| 46 | + select |
| 47 | + last_name, |
| 48 | + first_name, |
| 49 | + product_id, |
| 50 | + count(*) as count |
| 51 | + from {{{{ ref('seed') }}}} |
| 52 | + group by 1, 2, 3 |
| 53 | + """.lstrip() |
| 54 | + |
| 55 | + return { |
| 56 | + "dynamic_table_1.sql": _DYNAMIC_TABLE_1_SQL, |
| 57 | + "dynamic_table_2.sql": _DYNAMIC_TABLE_2_SQL, |
| 58 | + "dynamic_table_3.sql": _DYNAMIC_TABLE_3_SQL, |
| 59 | + "dynamic_table_4.sql": _DYNAMIC_TABLE_4_SQL, |
| 60 | + } |
| 61 | + |
| 62 | + def test_snowflake_dynamic_table_cluster_by(self, project): |
| 63 | + |
| 64 | + run_dbt(["seed"]) |
| 65 | + |
| 66 | + db_with_schema = f"{project.database}.{project.test_schema}" |
| 67 | + |
| 68 | + check_table_does_exist( |
| 69 | + project.adapter, f"{db_with_schema}.{self._available_models_in_setup()['seed_table']}" |
| 70 | + ) |
| 71 | + |
| 72 | + run_dbt() |
| 73 | + |
| 74 | + # Check that all dynamic tables exist |
| 75 | + check_table_does_exist( |
| 76 | + project.adapter, |
| 77 | + f"{db_with_schema}.{self._available_models_in_setup()['dynamic_table_1']}", |
| 78 | + ) |
| 79 | + check_table_does_exist( |
| 80 | + project.adapter, |
| 81 | + f"{db_with_schema}.{self._available_models_in_setup()['dynamic_table_2']}", |
| 82 | + ) |
| 83 | + check_table_does_exist( |
| 84 | + project.adapter, |
| 85 | + f"{db_with_schema}.{self._available_models_in_setup()['dynamic_table_3']}", |
| 86 | + ) |
| 87 | + check_table_does_exist( |
| 88 | + project.adapter, |
| 89 | + f"{db_with_schema}.{self._available_models_in_setup()['dynamic_table_4']}", |
| 90 | + ) |
| 91 | + |
| 92 | + with project.adapter.connection_named("__test"): |
| 93 | + # Check if cluster_by is applied to dynamic_table_2 (should cluster by last_name) |
| 94 | + cluster_by = self._get_dynamic_table_ddl( |
| 95 | + project, self._available_models_in_setup()["dynamic_table_2"] |
| 96 | + ) |
| 97 | + assert "CLUSTER BY (LAST_NAME)" in cluster_by.upper() |
| 98 | + |
| 99 | + # Check if cluster_by is applied to dynamic_table_3 (should cluster by last_name, first_name) |
| 100 | + cluster_by = self._get_dynamic_table_ddl( |
| 101 | + project, self._available_models_in_setup()["dynamic_table_3"] |
| 102 | + ) |
| 103 | + assert "CLUSTER BY (LAST_NAME, FIRST_NAME)" in cluster_by.upper() |
| 104 | + |
| 105 | + # Check if cluster_by is applied to dynamic_table_4 (should cluster by last_name, product_id % 3) |
| 106 | + cluster_by = self._get_dynamic_table_ddl( |
| 107 | + project, self._available_models_in_setup()["dynamic_table_4"] |
| 108 | + ) |
| 109 | + assert "CLUSTER BY (LAST_NAME, PRODUCT_ID % 3)" in cluster_by.upper() |
| 110 | + |
| 111 | + def _get_dynamic_table_ddl(self, project, table_name: str) -> str: |
| 112 | + ddl_query = f"SELECT GET_DDL('DYNAMIC_TABLE', '{project.database}.{project.test_schema}.{table_name}')" |
| 113 | + ddl = project.run_sql(ddl_query, fetch="one") |
| 114 | + return ddl[0] |
| 115 | + |
| 116 | + def _available_models_in_setup(self) -> dict[str, str]: |
| 117 | + return dict( |
| 118 | + seed_table="SEED", |
| 119 | + dynamic_table_1="DYNAMIC_TABLE_1", |
| 120 | + dynamic_table_2="DYNAMIC_TABLE_2", |
| 121 | + dynamic_table_3="DYNAMIC_TABLE_3", |
| 122 | + dynamic_table_4="DYNAMIC_TABLE_4", |
| 123 | + ) |
0 commit comments