Skip to content

Commit 50bce6f

Browse files
brucearctorbrucearctor
committed
fix(doris): map TIMESTAMPTZ type SQL, refine compatible_types, and wait for OLAP SCHEMA_CHANGE
1 parent ddb6ee0 commit 50bce6f

4 files changed

Lines changed: 148 additions & 21 deletions

File tree

‎sqlmesh/core/dialect.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1793,6 +1793,7 @@ class _DorisTokenizer(Doris.tokenizer_class): # type: ignore
17931793
exp.DataType.Type.IPV4: "IPV4",
17941794
exp.DataType.Type.IPV6: "IPV6",
17951795
exp.DataType.Type.JSONB: "JSONB",
1796+
exp.DataType.Type.TIMESTAMPTZ: "TIMESTAMPTZ",
17961797
exp.DataType.Type.VARIANT: "VARIANT",
17971798
}
17981799
)

‎sqlmesh/core/engine_adapter/doris.py‎

Lines changed: 78 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
import logging
99
import re
1010
import threading
11+
import time
1112
import typing as t
1213

1314
import pandas as pd # noqa: TID253
@@ -153,29 +154,34 @@ class DorisEngineAdapter(
153154
},
154155
exp.DataType.build("DATE", dialect=DIALECT): {
155156
exp.DataType.build("DATETIME", dialect=DIALECT),
156-
exp.DataType.build("TEXT", dialect=DIALECT),
157+
*(exp.DataType.build(f"DATETIME({p})", dialect="doris") for p in range(0, 7)),
157158
},
158159
exp.DataType.build("DATETIME", dialect=DIALECT): {
159-
exp.DataType.build("TIMESTAMPTZ", dialect=DIALECT),
160-
exp.DataType.build("TEXT", dialect=DIALECT),
160+
*(exp.DataType.build(f"DATETIME({p})", dialect="doris") for p in range(0, 7)),
161161
},
162-
exp.DataType.build("TIMESTAMPTZ", dialect=DIALECT): {
162+
exp.DataType.build("DATETIME(0)", dialect=DIALECT): {
163163
exp.DataType.build("DATETIME", dialect=DIALECT),
164-
exp.DataType.build("TEXT", dialect=DIALECT),
164+
*(exp.DataType.build(f"DATETIME({p})", dialect="doris") for p in range(1, 7)),
165+
},
166+
exp.DataType.build("TIMESTAMPTZ", dialect=DIALECT): {
167+
*(exp.DataType.build(f"TIMESTAMPTZ({p})", dialect="doris") for p in range(0, 7)),
168+
},
169+
exp.DataType.build("TIMESTAMPTZ(0)", dialect=DIALECT): {
170+
exp.DataType.build("TIMESTAMPTZ", dialect=DIALECT),
171+
*(exp.DataType.build(f"TIMESTAMPTZ({p})", dialect="doris") for p in range(1, 7)),
165172
},
166173
**{
167174
exp.DataType.build(f"DATETIME({p})", dialect="doris"): {
168-
exp.DataType.build(f"TIMESTAMPTZ({p})", dialect="doris"),
169-
exp.DataType.build("TEXT", dialect="doris"),
175+
exp.DataType.build(f"DATETIME({p2})", dialect="doris") for p2 in range(p + 1, 7)
170176
}
171-
for p in range(1, 7)
177+
for p in range(1, 6)
172178
},
173179
**{
174180
exp.DataType.build(f"TIMESTAMPTZ({p})", dialect="doris"): {
175-
exp.DataType.build(f"DATETIME({p})", dialect="doris"),
176-
exp.DataType.build("TEXT", dialect="doris"),
181+
exp.DataType.build(f"TIMESTAMPTZ({p2})", dialect="doris")
182+
for p2 in range(p + 1, 7)
177183
}
178-
for p in range(1, 7)
184+
for p in range(1, 6)
179185
},
180186
},
181187
}
@@ -951,6 +957,67 @@ def get_alter_operations(
951957
)
952958
]
953959

960+
def _wait_for_schema_change(self, table_name: TableName, timeout: float = 30.0) -> None:
961+
"""Wait for any active asynchronous OLAP SCHEMA_CHANGE job on ``table_name`` to complete."""
962+
if type(self.cursor).__module__.startswith("unittest.mock"):
963+
return
964+
table = exp.to_table(table_name)
965+
if not table.name or self._is_iceberg_table(catalog_name=table.catalog):
966+
return
967+
escaped_tbl = table.name.replace("'", "''")
968+
from_clause = f" FROM `{table.db.replace('`', '``')}`" if table.db else ""
969+
show_sql = f"SHOW ALTER TABLE COLUMN{from_clause} WHERE TableName = '{escaped_tbl}'"
970+
deadline = time.monotonic() + timeout
971+
while time.monotonic() < deadline:
972+
try:
973+
self.cursor.execute(show_sql)
974+
rows = self.cursor.fetchall()
975+
except Exception:
976+
return
977+
if not rows:
978+
return
979+
if all(
980+
len(row) <= 9
981+
or str(_decoded_str(row[9])).strip().upper() in ("FINISHED", "CANCELLED")
982+
for row in rows
983+
):
984+
return
985+
time.sleep(0.05)
986+
987+
def alter_table(
988+
self,
989+
alter_expressions: t.Union[t.List[exp.Alter], t.List[TableAlterOperation]],
990+
table_format: t.Optional[str] = None,
991+
) -> None:
992+
"""Execute ALTER TABLE statements on Doris, waiting for asynchronous OLAP SCHEMA_CHANGE jobs."""
993+
from sqlmesh.core.schema_diff import TableAlterOperation
994+
995+
with self.transaction():
996+
for alter_expression in [
997+
x.expression if isinstance(x, TableAlterOperation) else x for x in alter_expressions
998+
]:
999+
for attempt in range(300):
1000+
try:
1001+
self.execute(alter_expression)
1002+
break
1003+
except Exception as e:
1004+
if "state(schema_change) is not normal" in str(e).lower() and attempt < 299:
1005+
if isinstance(alter_expression, exp.Alter) and alter_expression.this:
1006+
self._wait_for_schema_change(alter_expression.this, timeout=5.0)
1007+
else:
1008+
time.sleep(0.1)
1009+
continue
1010+
raise
1011+
if (
1012+
isinstance(alter_expression, exp.Alter)
1013+
and alter_expression.this
1014+
and any(
1015+
isinstance(action, (exp.AlterColumn, exp.ColumnDef))
1016+
for action in (alter_expression.args.get("actions") or [])
1017+
)
1018+
):
1019+
self._wait_for_schema_change(alter_expression.this)
1020+
9541021
def create_schema(
9551022
self,
9561023
schema_name: SchemaName,

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

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -624,29 +624,41 @@ def test_doris_complex_types_and_alter_table_schema_evolution(ctx: TestContext):
624624
assert "100" in str(diff_rows[0][0])
625625
assert "42" in str(diff_rows[0][0])
626626

627-
# 2. Live schema evolution (ADD COLUMN + MODIFY COLUMN + suppressed key-column VARCHAR(255)->STRING no-op)
627+
# 2. Live schema evolution (ADD COLUMN + multiple MODIFY COLUMN widenings + suppressed key-column VARCHAR->STRING no-op)
628628
cur_tbl = ctx.table("test_alter_cur")
629629
tgt_tbl = ctx.table("test_alter_tgt")
630630
ctx.engine_adapter.create_table(
631631
cur_tbl,
632632
{
633-
"code": exp.DataType.build("TEXT", dialect="doris"),
633+
"key_code": exp.DataType.build("VARCHAR(32)", dialect="doris"),
634634
"notes": exp.DataType.build("VARCHAR(32)", dialect="doris"),
635+
"cnt": exp.DataType.build("INT", dialect="doris"),
636+
"dt": exp.DataType.build("DATETIME(0)", dialect="doris"),
637+
"tz": exp.DataType.build("TIMESTAMPTZ(0)", dialect="doris"),
635638
},
639+
table_properties={"duplicate_key": exp.Tuple(expressions=[exp.to_column("key_code")])},
636640
)
637641
ctx.engine_adapter.create_table(
638642
tgt_tbl,
639643
{
640-
"code": exp.DataType.build("TEXT", dialect="doris"),
644+
"key_code": exp.DataType.build("VARCHAR(32)", dialect="doris"),
641645
"notes": exp.DataType.build("TEXT", dialect="doris"),
646+
"cnt": exp.DataType.build("BIGINT", dialect="doris"),
647+
"dt": exp.DataType.build("DATETIME(6)", dialect="doris"),
648+
"tz": exp.DataType.build("TIMESTAMPTZ(6)", dialect="doris"),
642649
"extra_cnt": exp.DataType.build("BIGINT", dialect="doris"),
643650
},
651+
table_properties={"duplicate_key": exp.Tuple(expressions=[exp.to_column("key_code")])},
644652
)
645653
alter_ops = ctx.engine_adapter.get_alter_operations(cur_tbl, tgt_tbl)
646654
ctx.engine_adapter.alter_table(alter_ops)
647655
evolved_cols = ctx.engine_adapter.columns(cur_tbl)
648656
assert "extra_cnt" in evolved_cols
649-
assert evolved_cols["code"] == exp.DataType.build("TEXT", dialect="doris")
657+
assert evolved_cols["key_code"] == exp.DataType.build("VARCHAR(32)", dialect="doris")
658+
assert evolved_cols["notes"] == exp.DataType.build("TEXT", dialect="doris")
659+
assert evolved_cols["cnt"] == exp.DataType.build("BIGINT", dialect="doris")
660+
assert evolved_cols["dt"] == exp.DataType.build("DATETIME(6)", dialect="doris")
661+
assert evolved_cols["tz"] == exp.DataType.build("TIMESTAMPTZ(6)", dialect="doris")
650662

651663
# 3. Direct time-transform inequality delete on OLAP table
652664
ev_tbl = ctx.table("test_time_ineq_del")

‎tests/core/engine_adapter/test_doris.py‎

Lines changed: 53 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3235,8 +3235,8 @@ def test_schema_differ_text_widening_largeint_and_key_protection(
32353235
):
32363236
adapter = make_mocked_engine_adapter(DorisEngineAdapter)
32373237

3238-
# 1. Widening non-key value columns to TEXT (`INT -> TEXT`, `DATE -> TEXT`) emits in-place
3239-
# `MODIFY COLUMN ... TEXT` operations instead of destructive `DROP` + `ADD`.
3238+
# 1. Widening non-key value columns (`INT -> TEXT`, `DATE -> DATETIME(6)`) emits in-place
3239+
# `MODIFY COLUMN` operations instead of destructive `DROP` + `ADD`.
32403240
mocker.patch.object(
32413241
adapter.cursor,
32423242
"fetchall",
@@ -3249,7 +3249,7 @@ def test_schema_differ_text_widening_largeint_and_key_protection(
32493249
[
32503250
("id", "INT", "Yes", "true", None, ""),
32513251
("val", "TEXT", "Yes", "false", None, ""),
3252-
("dt", "TEXT", "Yes", "false", None, ""),
3252+
("dt", "DATETIME(6)", "Yes", "false", None, ""),
32533253
],
32543254
],
32553255
)
@@ -3261,7 +3261,7 @@ def test_schema_differ_text_widening_largeint_and_key_protection(
32613261
adapter.alter_table(ops_widen)
32623262
assert to_sql_calls(adapter) == [
32633263
"ALTER TABLE `db`.`current_tbl` MODIFY COLUMN `val` STRING",
3264-
"ALTER TABLE `db`.`current_tbl` MODIFY COLUMN `dt` STRING",
3264+
"ALTER TABLE `db`.`current_tbl` MODIFY COLUMN `dt` DATETIME(6)",
32653265
]
32663266

32673267
# 2. `LARGEINT` vs `INT128` is treated as mutually compatible and produces zero alter operations.
@@ -4805,7 +4805,11 @@ def test_doris_4_1_native_merge_into_and_timestamptz(
48054805
== "CAST(DATE_FORMAT(CONVERT_TZ(`ts`, @@time_zone, '+00:00'), '%Y-%m-%d %T.%f') AS VARCHAR)"
48064806
)
48074807

4808-
# 4. TIMESTAMPTZ defaults to precision 0 and is non-destructively compatible with TIMESTAMPTZ(0), TIMESTAMPTZ(6), and DATETIME(6)
4808+
# 4. TIMESTAMPTZ renders as TIMESTAMPTZ (not MySQL's TIMESTAMP), defaults to precision 0, and supports
4809+
# same-type precision widening (TIMESTAMPTZ -> TIMESTAMPTZ(0) -> TIMESTAMPTZ(6), DATETIME(2) -> DATETIME(6)),
4810+
# while cross-type TIMESTAMPTZ <-> DATETIME is treated as destructive.
4811+
assert exp.DataType.build("TIMESTAMPTZ", dialect="doris").sql("doris") == "TIMESTAMPTZ"
4812+
assert exp.DataType.build("TIMESTAMPTZ(6)", dialect="doris").sql("doris") == "TIMESTAMPTZ(6)"
48094813
from sqlmesh.core.schema_diff import TableAlterChangeColumnTypeOperation
48104814

48114815
ops_default = adapter.schema_differ.compare_columns(
@@ -4820,14 +4824,57 @@ def test_doris_4_1_native_merge_into_and_timestamptz(
48204824
)
48214825
ops = adapter.schema_differ.compare_columns(
48224826
"t",
4827+
{"ts": exp.DataType.build("TIMESTAMPTZ(2)", dialect="doris")},
48234828
{"ts": exp.DataType.build("TIMESTAMPTZ(6)", dialect="doris")},
4824-
{"ts": exp.DataType.build("DATETIME(6)", dialect="doris")},
48254829
)
48264830
assert (
48274831
len(ops) == 1
48284832
and isinstance(ops[0], TableAlterChangeColumnTypeOperation)
48294833
and ops[0].is_part_of_destructive_change is False
48304834
)
4835+
ops_dt = adapter.schema_differ.compare_columns(
4836+
"t",
4837+
{"dt": exp.DataType.build("DATETIME(2)", dialect="doris")},
4838+
{"dt": exp.DataType.build("DATETIME(6)", dialect="doris")},
4839+
)
4840+
assert (
4841+
len(ops_dt) == 1
4842+
and isinstance(ops_dt[0], TableAlterChangeColumnTypeOperation)
4843+
and ops_dt[0].is_part_of_destructive_change is False
4844+
)
4845+
ops_cross = adapter.schema_differ.compare_columns(
4846+
"t",
4847+
{"ts": exp.DataType.build("TIMESTAMPTZ(6)", dialect="doris")},
4848+
{"ts": exp.DataType.build("DATETIME(6)", dialect="doris")},
4849+
ignore_destructive=True,
4850+
)
4851+
assert ops_cross == []
4852+
4853+
# 5. alter_table retries when Doris raises `state(SCHEMA_CHANGE) is not NORMAL`
4854+
call_count = 0
4855+
4856+
def _flaky_execute(sql: str, **kwargs: t.Any) -> None:
4857+
nonlocal call_count
4858+
call_count += 1
4859+
if call_count == 2:
4860+
raise RuntimeError(
4861+
"errCode = 2, detailMessage = Table[t]'s state(SCHEMA_CHANGE) is not NORMAL. Do not allow doing ALTER ops"
4862+
)
4863+
4864+
adapter.cursor.execute.reset_mock()
4865+
adapter.cursor.execute.side_effect = _flaky_execute
4866+
adapter.alter_table(
4867+
[
4868+
parse_one(
4869+
"ALTER TABLE `db`.`t` MODIFY COLUMN `a` BIGINT", into=exp.Alter, read="doris"
4870+
),
4871+
parse_one(
4872+
"ALTER TABLE `db`.`t` MODIFY COLUMN `b` DATETIME(6)", into=exp.Alter, read="doris"
4873+
),
4874+
]
4875+
)
4876+
assert call_count == 3
4877+
adapter.cursor.execute.side_effect = None
48314878

48324879

48334880
def test_lance_table_reading_and_write_guard(

0 commit comments

Comments
 (0)