diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index 930cdf7cd4..a196114c70 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -18,7 +18,7 @@ from sqlglot import Dialect, exp from sqlglot.errors import ErrorLevel -from sqlglot.helper import ensure_list, seq_get +from sqlglot.helper import ensure_list, find_new_name, seq_get from sqlglot.optimizer.qualify_columns import quote_identifiers from sqlmesh.core.dialect import ( @@ -2073,6 +2073,11 @@ def remove_managed_columns( prefixed_col = exp.column(column).copy() prefixed_col.this.set("this", f"t_{prefixed_col.name}") prefixed_unmanaged_columns.append(prefixed_col) + target_exists_column = find_new_name( + {col.name.lower() for col in prefixed_columns_to_types} + | {col.lower() for col in unmanaged_columns_to_types}, + "t__exists", + ) query = ( exp.Select() # type: ignore .select(*table_columns) @@ -2141,6 +2146,7 @@ def remove_managed_columns( "joined", exp.select( exp.column("_exists", table="source").as_("_exists"), + exp.column("_exists", table="latest").as_(target_exists_column), *( exp.column(col, table="latest").as_(prefixed_columns_to_types[i].this) for i, col in enumerate(target_columns_to_types) @@ -2164,6 +2170,7 @@ def remove_managed_columns( .union( exp.select( exp.column("_exists", table="source").as_("_exists"), + exp.column("_exists", table="latest").as_(target_exists_column), *( exp.column(col, table="latest").as_( prefixed_columns_to_types[i].this @@ -2195,11 +2202,15 @@ def remove_managed_columns( "updated_rows", exp.select( *( - exp.func( - "COALESCE", + exp.Case() + .when( + exp.column(target_exists_column, table="joined") + .is_(exp.Null()) + .not_(), exp.column(prefixed_unmanaged_columns[i].this, table="joined"), - exp.column(col, table="joined"), - ).as_(col) + ) + .else_(exp.column(col, table="joined")) + .as_(col) for i, col in enumerate(unmanaged_columns_to_types) ), valid_from_case_stmt, diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py index 1971ba3bbc..f14afed6b3 100644 --- a/tests/core/engine_adapter/test_base.py +++ b/tests/core/engine_adapter/test_base.py @@ -1520,6 +1520,7 @@ def test_scd_type_2_by_time(make_mocked_engine_adapter: t.Callable): ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -1541,6 +1542,7 @@ def test_scd_type_2_by_time(make_mocked_engine_adapter: t.Callable): UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -1563,10 +1565,26 @@ def test_scd_type_2_by_time(make_mocked_engine_adapter: t.Callable): "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id", "joined"."id") AS "id", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", - COALESCE("joined"."t_test_UPDATED_at", "joined"."test_UPDATED_at") AS "test_UPDATED_at", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id" + ELSE "joined"."id" + END AS "id", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_test_UPDATED_at" + ELSE "joined"."test_UPDATED_at" + END AS "test_UPDATED_at", CASE WHEN "t_test_valid_from" IS NULL AND NOT "latest_deleted"."_exists" IS NULL THEN CASE @@ -1725,6 +1743,7 @@ def test_scd_type_2_by_time_source_columns(make_mocked_engine_adapter: t.Callabl ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -1741,6 +1760,7 @@ def test_scd_type_2_by_time_source_columns(make_mocked_engine_adapter: t.Callabl UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -1758,10 +1778,26 @@ def test_scd_type_2_by_time_source_columns(make_mocked_engine_adapter: t.Callabl "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id", "joined"."id") AS "id", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", - COALESCE("joined"."t_test_UPDATED_at", "joined"."test_UPDATED_at") AS "test_UPDATED_at", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id" + ELSE "joined"."id" + END AS "id", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_test_UPDATED_at" + ELSE "joined"."test_UPDATED_at" + END AS "test_UPDATED_at", CASE WHEN "t_test_valid_from" IS NULL AND NOT "latest_deleted"."_exists" IS NULL THEN CASE @@ -1926,6 +1962,7 @@ def test_scd_type_2_by_time_no_invalidate_hard_deletes(make_mocked_engine_adapte ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -1942,6 +1979,7 @@ def test_scd_type_2_by_time_no_invalidate_hard_deletes(make_mocked_engine_adapte UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -1959,10 +1997,26 @@ def test_scd_type_2_by_time_no_invalidate_hard_deletes(make_mocked_engine_adapte "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id", "joined"."id") AS "id", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", - COALESCE("joined"."t_test_updated_at", "joined"."test_updated_at") AS "test_updated_at", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id" + ELSE "joined"."id" + END AS "id", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_test_updated_at" + ELSE "joined"."test_updated_at" + END AS "test_updated_at", CASE WHEN "t_test_valid_from" IS NULL AND NOT "latest_deleted"."_exists" IS NULL THEN CASE @@ -2124,6 +2178,7 @@ def test_merge_scd_type_2_pandas(make_mocked_engine_adapter: t.Callable): ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id1" AS "t_id1", "latest"."id2" AS "t_id2", "latest"."name" AS "t_name", @@ -2142,6 +2197,7 @@ def test_merge_scd_type_2_pandas(make_mocked_engine_adapter: t.Callable): UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id1" AS "t_id1", "latest"."id2" AS "t_id2", "latest"."name" AS "t_name", @@ -2161,11 +2217,31 @@ def test_merge_scd_type_2_pandas(make_mocked_engine_adapter: t.Callable): "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id1", "joined"."id1") AS "id1", - COALESCE("joined"."t_id2", "joined"."id2") AS "id2", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", - COALESCE("joined"."t_test_updated_at", "joined"."test_updated_at") AS "test_updated_at", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id1" + ELSE "joined"."id1" + END AS "id1", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id2" + ELSE "joined"."id2" + END AS "id2", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_test_updated_at" + ELSE "joined"."test_updated_at" + END AS "test_updated_at", CASE WHEN "t_test_valid_from" IS NULL AND NOT "latest_deleted"."_exists" IS NULL THEN CASE @@ -2291,6 +2367,7 @@ def test_scd_type_2_by_column(make_mocked_engine_adapter: t.Callable): ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -2305,6 +2382,7 @@ def test_scd_type_2_by_column(make_mocked_engine_adapter: t.Callable): UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -2320,9 +2398,21 @@ def test_scd_type_2_by_column(make_mocked_engine_adapter: t.Callable): "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id", "joined"."id") AS "id", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id" + ELSE "joined"."id" + END AS "id", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", COALESCE("t_test_VALID_from", CAST('2020-01-01 00:00:00' AS TIMESTAMP)) AS "test_VALID_from", CASE WHEN "joined"."_exists" IS NULL @@ -2477,6 +2567,7 @@ def test_scd_type_2_by_column_composite_key(make_mocked_engine_adapter: t.Callab ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id_a" AS "t_id_a", "latest"."id_b" AS "t_id_b", "latest"."name" AS "t_name", @@ -2493,6 +2584,7 @@ def test_scd_type_2_by_column_composite_key(make_mocked_engine_adapter: t.Callab UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id_a" AS "t_id_a", "latest"."id_b" AS "t_id_b", "latest"."name" AS "t_name", @@ -2510,10 +2602,26 @@ def test_scd_type_2_by_column_composite_key(make_mocked_engine_adapter: t.Callab "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id_a", "joined"."id_a") AS "id_a", - COALESCE("joined"."t_id_b", "joined"."id_b") AS "id_b", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id_a" + ELSE "joined"."id_a" + END AS "id_a", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id_b" + ELSE "joined"."id_b" + END AS "id_b", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", COALESCE("t_test_VALID_from", CAST('2020-01-01 00:00:00' AS TIMESTAMP)) AS "test_VALID_from", CASE WHEN "joined"."_exists" IS NULL @@ -2666,6 +2774,7 @@ def test_scd_type_2_truncate(make_mocked_engine_adapter: t.Callable): ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -2680,6 +2789,7 @@ def test_scd_type_2_truncate(make_mocked_engine_adapter: t.Callable): UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -2695,9 +2805,21 @@ def test_scd_type_2_truncate(make_mocked_engine_adapter: t.Callable): "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id", "joined"."id") AS "id", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id" + ELSE "joined"."id" + END AS "id", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", COALESCE("t_test_valid_from", CAST('1970-01-01 00:00:00' AS TIMESTAMP)) AS "test_valid_from", CASE WHEN "joined"."_exists" IS NULL @@ -2846,6 +2968,7 @@ def test_scd_type_2_by_column_star_check(make_mocked_engine_adapter: t.Callable) ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -2860,6 +2983,7 @@ def test_scd_type_2_by_column_star_check(make_mocked_engine_adapter: t.Callable) UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -2875,9 +2999,21 @@ def test_scd_type_2_by_column_star_check(make_mocked_engine_adapter: t.Callable) "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id", "joined"."id") AS "id", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id" + ELSE "joined"."id" + END AS "id", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", COALESCE("t_test_valid_from", CAST('2020-01-01 00:00:00' AS TIMESTAMP)) AS "test_valid_from", CASE WHEN "joined"."_exists" IS NULL @@ -3041,6 +3177,7 @@ def test_scd_type_2_by_column_no_invalidate_hard_deletes(make_mocked_engine_adap ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -3055,6 +3192,7 @@ def test_scd_type_2_by_column_no_invalidate_hard_deletes(make_mocked_engine_adap UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -3070,9 +3208,21 @@ def test_scd_type_2_by_column_no_invalidate_hard_deletes(make_mocked_engine_adap "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id", "joined"."id") AS "id", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id" + ELSE "joined"."id" + END AS "id", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", COALESCE("t_test_valid_from", CAST('2020-01-01 00:00:00' AS TIMESTAMP)) AS "test_valid_from", CASE WHEN ( diff --git a/tests/core/engine_adapter/test_clickhouse.py b/tests/core/engine_adapter/test_clickhouse.py index a3dfe0fdda..d14e2fb260 100644 --- a/tests/core/engine_adapter/test_clickhouse.py +++ b/tests/core/engine_adapter/test_clickhouse.py @@ -738,6 +738,7 @@ def test_scd_type_2_by_time( ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -759,6 +760,7 @@ def test_scd_type_2_by_time( UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -781,10 +783,26 @@ def test_scd_type_2_by_time( "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id", "joined"."id") AS "id", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", - COALESCE("joined"."t_test_UPDATED_at", "joined"."test_UPDATED_at") AS "test_UPDATED_at", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id" + ELSE "joined"."id" + END AS "id", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_test_UPDATED_at" + ELSE "joined"."test_UPDATED_at" + END AS "test_UPDATED_at", CASE WHEN "t_test_valid_from" IS NULL AND NOT "latest_deleted"."_exists" IS NULL THEN CASE @@ -938,6 +956,7 @@ def test_scd_type_2_by_column( ), "joined" AS ( SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -952,6 +971,7 @@ def test_scd_type_2_by_column( UNION ALL SELECT "source"."_exists" AS "_exists", + "latest"."_exists" AS "t__exists", "latest"."id" AS "t_id", "latest"."name" AS "t_name", "latest"."price" AS "t_price", @@ -967,9 +987,21 @@ def test_scd_type_2_by_column( "latest"."_exists" IS NULL ), "updated_rows" AS ( SELECT - COALESCE("joined"."t_id", "joined"."id") AS "id", - COALESCE("joined"."t_name", "joined"."name") AS "name", - COALESCE("joined"."t_price", "joined"."price") AS "price", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_id" + ELSE "joined"."id" + END AS "id", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_name" + ELSE "joined"."name" + END AS "name", + CASE + WHEN NOT "joined"."t__exists" IS NULL + THEN "joined"."t_price" + ELSE "joined"."price" + END AS "price", COALESCE("t_test_VALID_from", CAST('2020-01-01 00:00:00' AS Nullable(DateTime64(6)))) AS "test_VALID_from", CASE WHEN "joined"."_exists" IS NULL diff --git a/tests/core/engine_adapter/test_duckdb.py b/tests/core/engine_adapter/test_duckdb.py index 9fd65a6e66..fa449430ef 100644 --- a/tests/core/engine_adapter/test_duckdb.py +++ b/tests/core/engine_adapter/test_duckdb.py @@ -1,4 +1,5 @@ import typing as t +from datetime import date import pandas as pd # noqa: TID253 import pytest @@ -60,6 +61,249 @@ def test_create_table(adapter: EngineAdapter, duck_conn): assert duck_conn.execute("DESCRIBE test_table").fetchall() == expected_columns +@pytest.mark.parametrize("by_time", [True, False]) +@pytest.mark.parametrize("use_dataframe", [True, False]) +def test_scd_type_2_preserves_nulls(adapter: EngineAdapter, by_time: bool, use_dataframe: bool): + columns_to_types = { + name: exp.DataType.build(data_type) + for name, data_type in { + "id": "INT", + "value": "TEXT", + "tracked_value": "TEXT", + "updated_at": "TIMESTAMP", + "valid_from": "TIMESTAMP", + "valid_to": "TIMESTAMP", + }.items() + } + source_columns_to_types = { + name: data_type + for name, data_type in columns_to_types.items() + if name not in {"valid_from", "valid_to"} + } + adapter.create_table("history", columns_to_types) + scd_kwargs: t.Dict[str, t.Any] = { + "unique_key": [exp.column("id")], + "valid_from_col": exp.column("valid_from"), + "valid_to_col": exp.column("valid_to"), + "target_columns_to_types": columns_to_types, + } + loads: t.List[t.Tuple[str, t.List[t.Tuple[t.Any, ...]]]] = [ + ( + "2026-09-29", + [ + (1, None, "initial", "2026-09-29"), + (2, None, "stable", "2026-09-29"), + (3, "C", "initial", "2026-09-29"), + (4, "deleted", "stable", "2026-09-29"), + (None, "null-key", "stable", "2026-09-29"), + ], + ), + ( + "2026-09-30", + [ + (1, "A", "next", "2026-09-30"), + (2, "ignored", "stable", "2026-09-29"), + (3, None, "next", "2026-09-30"), + (5, "new", "stable", "2026-09-30"), + ], + ), + ( + "2026-10-01", + [ + (1, "B", "final", "2026-10-01"), + (2, "ignored", "stable", "2026-09-29"), + (3, None, "next", "2026-09-30"), + (5, "new", "stable", "2026-09-30"), + ], + ), + ] + for execution_time, rows in loads: + source = ( + pd.DataFrame(rows, columns=list(source_columns_to_types)) + if use_dataframe + else adapter._values_to_sql( + rows, + target_columns_to_types=source_columns_to_types, + batch_start=0, + batch_end=len(rows), + ) + ) + if by_time: + adapter.scd_type_2_by_time( + "history", + source, + execution_time=execution_time, + updated_at_col=exp.column("updated_at"), + updated_at_as_valid_from=True, + **scd_kwargs, + ) + else: + adapter.scd_type_2_by_column( + "history", + source, + execution_time=execution_time, + check_columns=[exp.column("tracked_value")], + execution_time_as_valid_from=True, + **scd_kwargs, + ) + + assert adapter.fetchall( + "SELECT id, value, CAST(valid_from AS DATE), CAST(valid_to AS DATE) " + "FROM history ORDER BY id NULLS LAST, valid_from" + ) == [ + (1, None, date(2026, 9, 29), date(2026, 9, 30)), + (1, "A", date(2026, 9, 30), date(2026, 10, 1)), + (1, "B", date(2026, 10, 1), None), + (2, None, date(2026, 9, 29), None), + (3, "C", date(2026, 9, 29), date(2026, 9, 30)), + (3, None, date(2026, 9, 30), None), + (4, "deleted", date(2026, 9, 29), date(2026, 9, 30)), + (5, "new", date(2026, 9, 30), None), + (None, "null-key", date(2026, 9, 29), date(2026, 9, 30)), + ] + + +@pytest.mark.parametrize("by_time", [True, False]) +@pytest.mark.parametrize("use_dataframe", [True, False]) +@pytest.mark.parametrize("column_name", ["t__exists", "T__EXISTS"]) +@pytest.mark.parametrize("data_type,value", [("TEXT", "keep"), ("BOOLEAN", False)]) +def test_scd_type_2_marker_collision( + adapter: EngineAdapter, + by_time: bool, + use_dataframe: bool, + column_name: str, + data_type: str, + value: t.Any, +): + columns_to_types = { + name: exp.DataType.build(column_type) + for name, column_type in { + "id": "INT", + column_name: data_type, + "t__exists_2": data_type, + "_exists_3": data_type, + "updated_at": "TIMESTAMP", + "valid_from": "TIMESTAMP", + "valid_to": "TIMESTAMP", + }.items() + } + source_columns_to_types = { + name: column_type + for name, column_type in columns_to_types.items() + if name not in {"valid_from", "valid_to"} + } + adapter.create_table("history", columns_to_types) + for execution_time, rows in [ + ("2026-09-29", [(1, value, value, value, "2026-09-29")]), + ( + "2026-09-30", + [ + (1, None, value, value, "2026-09-30"), + (2, value, value, value, "2026-09-30"), + ], + ), + ]: + source = ( + pd.DataFrame(rows, columns=list(source_columns_to_types)) + if use_dataframe + else adapter._values_to_sql( + rows, + target_columns_to_types=source_columns_to_types, + batch_start=0, + batch_end=len(rows), + ) + ) + kwargs: t.Dict[str, t.Any] = { + "unique_key": [exp.column("id")], + "valid_from_col": exp.column("valid_from"), + "valid_to_col": exp.column("valid_to"), + "target_columns_to_types": columns_to_types, + } + if by_time: + adapter.scd_type_2_by_time( + "history", + source, + execution_time=execution_time, + updated_at_col=exp.column("updated_at"), + updated_at_as_valid_from=True, + **kwargs, + ) + else: + adapter.scd_type_2_by_column( + "history", + source, + execution_time=execution_time, + check_columns=[exp.column(column_name)], + execution_time_as_valid_from=True, + **kwargs, + ) + assert adapter.fetchall( + f'SELECT id, "{column_name}", t__exists_2, _exists_3 FROM history ORDER BY id, valid_from' + ) == [(1, value, value, value), (1, None, value, value), (2, value, value, value)] + + +@pytest.mark.parametrize("by_time", [True, False]) +@pytest.mark.parametrize("use_dataframe", [True, False]) +def test_scd_type_2_null_valid_from(adapter: EngineAdapter, by_time: bool, use_dataframe: bool): + columns_to_types = { + name: exp.DataType.build(data_type) + for name, data_type in { + "id": "INT", + "value": "TEXT", + "updated_at": "TIMESTAMP", + "valid_from": "TIMESTAMP", + "valid_to": "TIMESTAMP", + }.items() + } + source_columns_to_types = { + name: data_type + for name, data_type in columns_to_types.items() + if name not in {"valid_from", "valid_to"} + } + adapter.create_table("history", columns_to_types) + scd_kwargs: t.Dict[str, t.Any] = { + "unique_key": [exp.column("id")], + "valid_from_col": exp.column("valid_from"), + "valid_to_col": exp.column("valid_to"), + "target_columns_to_types": columns_to_types, + } + for is_initial in [True, False]: + rows: t.List[t.Tuple[t.Any, ...]] = ( + [(1, "keep", None)] if is_initial else [(2, "new", "2026-09-30")] + ) + source = ( + pd.DataFrame(rows, columns=list(source_columns_to_types)) + if use_dataframe + else adapter._values_to_sql( + rows, + target_columns_to_types=source_columns_to_types, + batch_start=0, + batch_end=1, + ) + ) + if by_time: + adapter.scd_type_2_by_time( + "history", + source, + execution_time="2026-09-29" if is_initial else "2026-09-30", + updated_at_col=exp.column("updated_at"), + updated_at_as_valid_from=True, + **scd_kwargs, + ) + else: + adapter.scd_type_2_by_column( + "history", + source, + execution_time=exp.column("updated_at") if is_initial else "2026-09-30", + check_columns=[exp.column("value")], + execution_time_as_valid_from=True, + **scd_kwargs, + ) + assert adapter.fetchall("SELECT id, value FROM history ORDER BY id") == ( + [(1, "keep")] if is_initial else [(1, "keep"), (2, "new")] + ) + + def test_replace_query_pandas(adapter: EngineAdapter, duck_conn): df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}) adapter.replace_query( diff --git a/tests/core/engine_adapter/test_spark.py b/tests/core/engine_adapter/test_spark.py index d7c3127f05..47dd6cafc9 100644 --- a/tests/core/engine_adapter/test_spark.py +++ b/tests/core/engine_adapter/test_spark.py @@ -676,6 +676,7 @@ def check_table_exists(table_name: exp.Table) -> bool: ), `joined` AS ( SELECT `source`.`_exists` AS `_exists`, + `latest`.`_exists` AS `t__exists`, `latest`.`id` AS `t_id`, `latest`.`name` AS `t_name`, `latest`.`price` AS `t_price`, @@ -692,6 +693,7 @@ def check_table_exists(table_name: exp.Table) -> bool: UNION ALL SELECT `source`.`_exists` AS `_exists`, + `latest`.`_exists` AS `t__exists`, `latest`.`id` AS `t_id`, `latest`.`name` AS `t_name`, `latest`.`price` AS `t_price`, @@ -709,10 +711,26 @@ def check_table_exists(table_name: exp.Table) -> bool: `latest`.`_exists` IS NULL ), `updated_rows` AS ( SELECT - COALESCE(`joined`.`t_id`, `joined`.`id`) AS `id`, - COALESCE(`joined`.`t_name`, `joined`.`name`) AS `name`, - COALESCE(`joined`.`t_price`, `joined`.`price`) AS `price`, - COALESCE(`joined`.`t_test_updated_at`, `joined`.`test_updated_at`) AS `test_updated_at`, + CASE + WHEN NOT `joined`.`t__exists` IS NULL + THEN `joined`.`t_id` + ELSE `joined`.`id` + END AS `id`, + CASE + WHEN NOT `joined`.`t__exists` IS NULL + THEN `joined`.`t_name` + ELSE `joined`.`name` + END AS `name`, + CASE + WHEN NOT `joined`.`t__exists` IS NULL + THEN `joined`.`t_price` + ELSE `joined`.`price` + END AS `price`, + CASE + WHEN NOT `joined`.`t__exists` IS NULL + THEN `joined`.`t_test_updated_at` + ELSE `joined`.`test_updated_at` + END AS `test_updated_at`, CASE WHEN `t_test_valid_from` IS NULL AND NOT `latest_deleted`.`_exists` IS NULL THEN CASE