diff --git a/sqlmesh/core/engine_adapter/redshift.py b/sqlmesh/core/engine_adapter/redshift.py index 39453f0cd2..8f2f098d08 100644 --- a/sqlmesh/core/engine_adapter/redshift.py +++ b/sqlmesh/core/engine_adapter/redshift.py @@ -291,11 +291,18 @@ def _to_identifier_if_string(expression: exp.Expr) -> exp.Expr: distkey = table_properties.get("DISTKEY") if distkey: - properties.append(exp.DistKeyProperty(this=_to_identifier_if_string(distkey))) + properties.append( + exp.DistKeyProperty(this=_to_identifier_if_string(distkey.unnest())) + ) sortkey = table_properties.get("SORTKEY") if sortkey: - sortkey_expressions = sortkey.expressions if sortkey.expressions else [sortkey] + if isinstance(sortkey, (exp.Tuple, exp.Array)): + sortkey_expressions = sortkey.expressions + elif isinstance(sortkey, exp.Paren): + sortkey_expressions = [sortkey.unnest()] + else: + sortkey_expressions = [sortkey] properties.append( exp.SortKeyProperty( this=[ diff --git a/tests/core/engine_adapter/test_redshift.py b/tests/core/engine_adapter/test_redshift.py index ddd2c7c2c8..7211244e91 100644 --- a/tests/core/engine_adapter/test_redshift.py +++ b/tests/core/engine_adapter/test_redshift.py @@ -152,22 +152,33 @@ def test_create_table_physical_properties_with_string_columns( ] +@pytest.mark.parametrize( + "sortkey, expected_sortkey", + [ + ('"batch_time"', 'SORTKEY("batch_time")'), + ("batch_time", 'SORTKEY("batch_time")'), + ('("batch_time")', 'SORTKEY("batch_time")'), + ("(batch_time)", 'SORTKEY("batch_time")'), + ('("batch_time", "id_file")', 'SORTKEY("batch_time", "id_file")'), + ('["batch_time", "id_file"]', 'SORTKEY("batch_time", "id_file")'), + ], +) def test_create_table_physical_properties_from_model_definition( - make_mocked_engine_adapter: t.Callable, + make_mocked_engine_adapter: t.Callable, sortkey: str, expected_sortkey: str ): adapter = make_mocked_engine_adapter(RedshiftEngineAdapter) model: SqlModel = t.cast( SqlModel, load_sql_based_model( d.parse( - """ + f""" MODEL ( name test_schema.test_table, kind full, physical_properties ( diststyle = key, distkey = "id_file", - sortkey = "batch_time" + sortkey = {sortkey} ) ); SELECT id_file::INT, batch_time::TIMESTAMP; @@ -183,7 +194,45 @@ def test_create_table_physical_properties_from_model_definition( ) assert to_sql_calls(adapter) == [ - 'CREATE TABLE IF NOT EXISTS "test_schema"."test_table" ("id_file" INTEGER, "batch_time" TIMESTAMP) DISTSTYLE KEY DISTKEY("id_file") SORTKEY("batch_time")', + f'CREATE TABLE IF NOT EXISTS "test_schema"."test_table" ("id_file" INTEGER, "batch_time" TIMESTAMP) DISTSTYLE KEY DISTKEY("id_file") {expected_sortkey}', + ] + + +@pytest.mark.parametrize( + "distkey", + ['"id_file"', "id_file", "'id_file'", '("id_file")', "(id_file)"], +) +def test_create_table_physical_properties_distkey_from_model_definition( + make_mocked_engine_adapter: t.Callable, distkey: str +): + adapter = make_mocked_engine_adapter(RedshiftEngineAdapter) + model: SqlModel = t.cast( + SqlModel, + load_sql_based_model( + d.parse( + f""" +MODEL ( + name test_schema.test_table, + kind full, + physical_properties ( + diststyle = key, + distkey = {distkey} + ) +); +SELECT id_file::INT; + """ + ) + ), + ) + + adapter.create_table( + model.name, + target_columns_to_types=model.columns_to_types_or_raise, + table_properties=model.physical_properties, + ) + + assert to_sql_calls(adapter) == [ + 'CREATE TABLE IF NOT EXISTS "test_schema"."test_table" ("id_file" INTEGER) DISTSTYLE KEY DISTKEY("id_file")', ]