Skip to content

Commit af3a80b

Browse files
brucearctorxinge.ji
andcommitted
fix(doris): fix wait-for-db FE config, normalize Identifier merge keys, and move MV tests to test_integration_doris
Co-authored-by: xinge.ji <xinge.ji@outlook.com> Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com>
1 parent 51c61e2 commit af3a80b

6 files changed

Lines changed: 179 additions & 57 deletions

File tree

‎.github/scripts/wait-for-db.sh‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -110,7 +110,7 @@ doris_ready() {
110110
if [ "$HOSTNAME" = "localhost" ]; then
111111
HOSTNAME="127.0.0.1"
112112
fi
113-
INIT_SQL="ADMIN SET FRONTEND CONFIG ('default_replication_num' = '1'); ADMIN SET FRONTEND CONFIG ('force_olap_table_replication_num' = '1'); ADMIN SET FRONTEND CONFIG ('enable_unicode_name_support' = 'true');"
113+
INIT_SQL="ADMIN SET FRONTEND CONFIG ('force_olap_table_replication_num' = '1');"
114114
if docker ps --format '{{.Names}}' 2>/dev/null | grep -qx 'doris-fe-01'; then
115115
docker exec -i doris-fe-01 mysql -h127.0.0.1 -P9030 -uroot -e "SHOW BACKENDS\G" 2>/dev/null | grep -q "Alive: true" && \
116116
docker exec -i doris-fe-01 mysql -h127.0.0.1 -P9030 -uroot -e "$INIT_SQL" 2>/dev/null
@@ -128,9 +128,7 @@ with conn.cursor() as cur:
128128
alive_idx = cols.index('Alive')
129129
if not rows or not any(str(r[alive_idx]).lower() == 'true' for r in rows):
130130
sys.exit(1)
131-
for cfg in ('default_replication_num', 'force_olap_table_replication_num'):
132-
cur.execute(f\"ADMIN SET FRONTEND CONFIG ('{cfg}' = '1')\")
133-
cur.execute(\"ADMIN SET FRONTEND CONFIG ('enable_unicode_name_support' = 'true')\")
131+
cur.execute(\"ADMIN SET FRONTEND CONFIG ('force_olap_table_replication_num' = '1')\")
134132
" 2>/dev/null
135133
fi
136134
}

‎sqlmesh/core/config/connection.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2851,6 +2851,9 @@ def init(cursor: t.Any) -> None:
28512851

28522852
return init
28532853

2854+
def get_catalog(self) -> t.Optional[str]:
2855+
return self.catalog
2856+
28542857

28552858
CONNECTION_CONFIG_TO_TYPE = {
28562859
# Map all subclasses of ConnectionConfig to the value of their `type_` field.

‎sqlmesh/core/engine_adapter/doris.py‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -891,10 +891,19 @@ def _extract_subquery_from_where(
891891
else:
892892
key_exprs = [lhs.copy()]
893893

894-
if not key_exprs or not all(isinstance(k, exp.Column) for k in key_exprs):
894+
normalized_keys: t.List[exp.Expr] = []
895+
for k in key_exprs:
896+
if isinstance(k, exp.Column):
897+
normalized_keys.append(k)
898+
elif isinstance(k, exp.Identifier):
899+
normalized_keys.append(exp.Column(this=k.copy()))
900+
else:
901+
return None
902+
903+
if not normalized_keys:
895904
return None
896905

897-
return key_exprs, subquery, extra_conditions
906+
return normalized_keys, subquery, extra_conditions
898907

899908
@staticmethod
900909
def _unwrap_concat_ws(expr: exp.Expr) -> t.Optional[exp.Expr]:

‎tests/core/engine_adapter/integration/test_integration.py‎

Lines changed: 17 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -440,6 +440,12 @@ def test_materialized_view(ctx_query_and_df: TestContext):
440440
"which this generic test does not provide; StarRocks MVs are covered by "
441441
"test_integration_starrocks.py"
442442
)
443+
if ctx.engine_adapter.dialect == "doris":
444+
pytest.skip(
445+
"Doris materialized views require engine-specific properties (SYNC source_table or MTMV DISTRIBUTED BY), "
446+
"which this generic test does not provide; Doris MVs are covered by "
447+
"test_integration_doris.py"
448+
)
443449
input_data = pd.DataFrame(
444450
[
445451
{"id": 1, "ds": "2022-01-01"},
@@ -450,22 +456,8 @@ def test_materialized_view(ctx_query_and_df: TestContext):
450456
source_table = ctx.table("source_table")
451457
ctx.engine_adapter.ctas(source_table, ctx.input_data(input_data), ctx.columns_to_types)
452458
view = ctx.table("test_view")
453-
# For Doris synchronized materialized views, we need to use only the table name without schema and add the source table name to the table properties
454-
if ctx.engine_adapter.dialect == "doris":
455-
view = exp.to_table(view.name)
456-
# For Doris, select only the first column name from columns_to_types.
457-
# Otherwise, Doris will raise an error if the materialized view has the same schema as the base table.
458-
first_column = next(iter(ctx.columns_to_types))
459-
view_query = exp.select(first_column).from_(source_table)
460-
ctx.engine_adapter.create_view(
461-
view,
462-
view_query,
463-
materialized=True,
464-
view_properties={"materialized_type": "SYNC", "source_table": source_table},
465-
)
466-
else:
467-
view_query = exp.select(*ctx.columns_to_types).from_(source_table)
468-
ctx.engine_adapter.create_view(view, view_query, materialized=True)
459+
view_query = exp.select(*ctx.columns_to_types).from_(source_table)
460+
ctx.engine_adapter.create_view(view, view_query, materialized=True)
469461
results = ctx.get_metadata_results()
470462
# Redshift considers the underlying dataset supporting materialized views as a table therefore we get 2
471463
# tables in the result
@@ -474,31 +466,13 @@ def test_materialized_view(ctx_query_and_df: TestContext):
474466
else:
475467
assert len(results.tables) == 1
476468
assert len(results.views) == 0
477-
# Doris information_schema.tables does not show materialized views and cannot query the synchronized materialized view
478-
if ctx.engine_adapter.dialect != "doris":
479-
assert len(results.materialized_views) == 1
480-
assert results.materialized_views[0] == view.name
481-
ctx.compare_with_current(view, input_data)
469+
assert len(results.materialized_views) == 1
470+
assert results.materialized_views[0] == view.name
471+
ctx.compare_with_current(view, input_data)
482472
# Make sure that dropping a materialized view also works
483-
if ctx.engine_adapter.dialect == "doris":
484-
# Wait for the materialized view to be created by retrying drop until it succeeds
485-
def drop_view_success():
486-
try:
487-
ctx.engine_adapter.drop_view(
488-
view,
489-
ignore_if_not_exists=False,
490-
materialized=True,
491-
view_properties={"materialized_type": "SYNC", "source_table": source_table},
492-
)
493-
return True
494-
except Exception:
495-
return False
496-
497-
wait_until(drop_view_success, attempts=5, wait=1)
498-
else:
499-
ctx.engine_adapter.drop_view(view, materialized=True)
500-
results = ctx.get_metadata_results()
501-
assert len(results.materialized_views) == 0
473+
ctx.engine_adapter.drop_view(view, materialized=True)
474+
results = ctx.get_metadata_results()
475+
assert len(results.materialized_views) == 0
502476

503477

504478
def test_drop_schema(ctx: TestContext):
@@ -1031,12 +1005,8 @@ def test_merge(ctx_query_and_df: TestContext):
10311005
# Athena only supports MERGE on Iceberg tables
10321006
# And it cant fall back to a logical merge on Hive tables because it cant delete records
10331007
table_format = "iceberg" if ctx.dialect == "athena" else None
1034-
# Doris needs to have a UNIQUE KEY to use MERGE
1035-
table_properties = {"unique_key": "id"} if ctx.dialect == "doris" else None
10361008

1037-
ctx.engine_adapter.create_table(
1038-
table, ctx.columns_to_types, table_format=table_format, table_properties=table_properties
1039-
)
1009+
ctx.engine_adapter.create_table(table, ctx.columns_to_types, table_format=table_format)
10401010
input_data = pd.DataFrame(
10411011
[
10421012
{"id": 1, "ds": "2022-01-01"},
@@ -1101,15 +1071,11 @@ def test_merge_source_columns(ctx_query_and_df: TestContext):
11011071
# Athena only supports MERGE on Iceberg tables
11021072
# And it cant fall back to a logical merge on Hive tables because it cant delete records
11031073
table_format = "iceberg" if ctx.dialect == "athena" else None
1104-
# Doris needs to have a UNIQUE KEY to use MERGE
1105-
table_properties = {"unique_key": "id"} if ctx.dialect == "doris" else None
11061074

11071075
columns_to_types = ctx.columns_to_types.copy()
11081076
columns_to_types["ignored_column"] = exp.DataType.build("int")
11091077

1110-
ctx.engine_adapter.create_table(
1111-
table, columns_to_types, table_format=table_format, table_properties=table_properties
1112-
)
1078+
ctx.engine_adapter.create_table(table, columns_to_types, table_format=table_format)
11131079
input_data = pd.DataFrame(
11141080
[
11151081
{"id": 1, "ds": "2022-01-01", "ignored_source": "ignored_value"},
@@ -4058,7 +4024,7 @@ def test_unicode_characters(ctx: TestContext, tmp_path: Path):
40584024
# I also think Spark may not support unicode in general but that would need to be verified.
40594025
if not ctx.engine_adapter.QUOTE_IDENTIFIERS_IN_VIEWS:
40604026
pytest.skip("Skipping as these engines have issues with unicode characters in model names")
4061-
# Doris default setting `enable_unicode_name_support=false` so it is incompatible with unicode characters in model names
4027+
# Doris static fe.conf setting `enable_unicode_name_support` defaults to false (and is not runtime-mutable)
40624028
if ctx.dialect == "doris":
40634029
pytest.skip(
40644030
"Skipping as Doris default setting has issues with unicode characters in model names"

‎tests/core/engine_adapter/integration/test_integration_doris.py‎

Lines changed: 122 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,3 +98,125 @@ def test_doris_iceberg_local_catalog(ctx: TestContext):
9898
except Exception:
9999
pass
100100
adapter.set_current_catalog("internal")
101+
102+
103+
def test_materialized_view_sync(ctx: TestContext):
104+
import pandas as pd # noqa: TID253
105+
from tests.core.engine_adapter.integration import wait_until
106+
107+
input_data = pd.DataFrame(
108+
[
109+
{"id": 1, "ds": "2022-01-01"},
110+
{"id": 2, "ds": "2022-01-02"},
111+
{"id": 3, "ds": "2022-01-03"},
112+
]
113+
)
114+
source_table = ctx.table("source_table")
115+
ctx.engine_adapter.ctas(source_table, ctx.input_data(input_data), ctx.columns_to_types)
116+
view = exp.to_table("test_sync_mv")
117+
first_column = next(iter(ctx.columns_to_types))
118+
view_query = exp.select(first_column).from_(source_table)
119+
ctx.engine_adapter.create_view(
120+
view,
121+
view_query,
122+
materialized=True,
123+
view_properties={
124+
"materialized_type": exp.Literal.string("SYNC"),
125+
"source_table": source_table,
126+
},
127+
)
128+
129+
def drop_view_success() -> bool:
130+
try:
131+
ctx.engine_adapter.drop_view(
132+
view,
133+
ignore_if_not_exists=False,
134+
materialized=True,
135+
view_properties={
136+
"materialized_type": exp.Literal.string("SYNC"),
137+
"source_table": source_table,
138+
},
139+
)
140+
return True
141+
except Exception:
142+
return False
143+
144+
wait_until(drop_view_success, attempts=5, wait=1)
145+
146+
147+
def test_materialized_view_async_mtmv(ctx: TestContext):
148+
import pandas as pd # noqa: TID253
149+
from tests.core.engine_adapter.integration import wait_until
150+
151+
input_data = pd.DataFrame(
152+
[
153+
{"id": 1, "ds": "2022-01-01"},
154+
{"id": 2, "ds": "2022-01-02"},
155+
{"id": 3, "ds": "2022-01-03"},
156+
]
157+
)
158+
source_table = ctx.table("mtmv_source")
159+
ctx.engine_adapter.ctas(source_table, ctx.input_data(input_data), ctx.columns_to_types)
160+
view = ctx.table("test_async_mtmv")
161+
view_query = exp.select(*ctx.columns_to_types).from_(source_table)
162+
ctx.engine_adapter.create_view(
163+
view,
164+
view_query,
165+
materialized=True,
166+
view_properties={
167+
"build": exp.Literal.string("IMMEDIATE"),
168+
"refresh": exp.Literal.string("COMPLETE"),
169+
"refresh_trigger": exp.Literal.string("ON MANUAL"),
170+
"distributed_by": exp.Literal.string("RANDOM BUCKETS AUTO"),
171+
},
172+
)
173+
wait_until(lambda: len(ctx.get_current_data(view)) == len(input_data), attempts=10, wait=1)
174+
ctx.compare_with_current(view, input_data)
175+
ctx.engine_adapter.drop_view(view, materialized=True)
176+
177+
178+
def test_merge_unique_key(ctx: TestContext):
179+
import pandas as pd # noqa: TID253
180+
181+
table = ctx.table("test_merge_uk")
182+
ctx.engine_adapter.create_table(
183+
table,
184+
ctx.columns_to_types,
185+
table_properties={"unique_key": exp.Literal.string("id")},
186+
)
187+
input_data = pd.DataFrame(
188+
[
189+
{"id": 1, "ds": "2022-01-01"},
190+
{"id": 2, "ds": "2022-01-02"},
191+
{"id": 3, "ds": "2022-01-03"},
192+
]
193+
)
194+
ctx.engine_adapter.merge(
195+
table,
196+
ctx.input_data(input_data),
197+
target_columns_to_types=None,
198+
unique_key=[exp.to_identifier("id")],
199+
)
200+
merge_data = pd.DataFrame(
201+
[
202+
{"id": 2, "ds": "2022-01-10"},
203+
{"id": 4, "ds": "2022-01-04"},
204+
]
205+
)
206+
ctx.engine_adapter.merge(
207+
table,
208+
ctx.input_data(merge_data),
209+
target_columns_to_types=None,
210+
unique_key=[exp.to_identifier("id")],
211+
)
212+
ctx.compare_with_current(
213+
table,
214+
pd.DataFrame(
215+
[
216+
{"id": 1, "ds": "2022-01-01"},
217+
{"id": 2, "ds": "2022-01-10"},
218+
{"id": 3, "ds": "2022-01-03"},
219+
{"id": 4, "ds": "2022-01-04"},
220+
]
221+
),
222+
)

‎tests/core/engine_adapter/test_doris.py‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1277,3 +1277,27 @@ def test_adjust_physical_properties_for_incremental(
12771277
assert to_sql_calls(adapter) == [
12781278
"CREATE TABLE IF NOT EXISTS `db`.`incremental_model` (`id` INT, `name` VARCHAR(50)) UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10"
12791279
]
1280+
1281+
1282+
def test_merge_with_identifier_unique_key(
1283+
make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter],
1284+
mocker: MockerFixture,
1285+
):
1286+
adapter = make_mocked_engine_adapter(DorisEngineAdapter)
1287+
mocker.patch("sqlmesh.core.engine_adapter.doris.random_id", return_value="abcd1234")
1288+
adapter.merge(
1289+
target_table="db.target",
1290+
source_table=parse_one("SELECT id, ds FROM db.src"),
1291+
target_columns_to_types={
1292+
"id": exp.DataType.build("INT", dialect="doris"),
1293+
"ds": exp.DataType.build("VARCHAR(32)", dialect="doris"),
1294+
},
1295+
unique_key=[exp.to_identifier("id")],
1296+
)
1297+
sql_calls = to_sql_calls(adapter)
1298+
assert any(
1299+
call.startswith(
1300+
"DELETE FROM `db`.`target` AS `_t1` USING (SELECT `id` FROM `db`.`__temp_target_abcd1234`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id`"
1301+
)
1302+
for call in sql_calls
1303+
)

0 commit comments

Comments
 (0)