From 8a5ce1c659fa86c3575c51f0f62706bdb404d5fe Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sun, 18 Oct 2026 15:50:28 -0700 Subject: [PATCH 01/27] Feat: Add Apache Doris engine adapter with OLAP and Iceberg multi-catalog support Closes #5875 Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- .github/scripts/install-prerequisites.sh | 4 + .github/scripts/wait-for-db.sh | 28 + .github/workflows/pr.yaml | 1 + .readthedocs.yaml | 2 +- Makefile | 3 + docs/guides/configuration.md | 2 + docs/guides/connections.md | 1 + docs/installation.md | 1 + docs/integrations/engines/doris.md | 489 +++++ docs/integrations/overview.md | 1 + docs/reference/configuration.md | 1 + mkdocs.yml | 1 + pyproject.toml | 2 + sqlmesh/core/config/__init__.py | 1 + sqlmesh/core/config/connection.py | 160 ++ sqlmesh/core/engine_adapter/__init__.py | 2 + sqlmesh/core/engine_adapter/doris.py | 1948 +++++++++++++++++ .../engine_adapter/integration/__init__.py | 8 +- .../engine_adapter/integration/config.yaml | 10 + .../integration/docker/compose.doris.yaml | 43 + .../integration/test_integration.py | 87 +- .../integration/test_integration_doris.py | 100 + tests/core/engine_adapter/test_doris.py | 1279 +++++++++++ tests/core/test_connection_config.py | 114 + 24 files changed, 4269 insertions(+), 19 deletions(-) create mode 100644 docs/integrations/engines/doris.md create mode 100644 sqlmesh/core/engine_adapter/doris.py create mode 100644 tests/core/engine_adapter/integration/docker/compose.doris.yaml create mode 100644 tests/core/engine_adapter/integration/test_integration_doris.py create mode 100644 tests/core/engine_adapter/test_doris.py diff --git a/.github/scripts/install-prerequisites.sh b/.github/scripts/install-prerequisites.sh index 6997633a31..a955f0c47c 100755 --- a/.github/scripts/install-prerequisites.sh +++ b/.github/scripts/install-prerequisites.sh @@ -28,6 +28,10 @@ elif [ "$ENGINE" == "fabric" ]; then rm packages-microsoft-prod.deb ENGINE_DEPENDENCIES="msodbcsql18" +elif [ "$ENGINE" == "doris" ]; then + ENGINE_DEPENDENCIES="mysql-client" + sudo sysctl -w vm.max_map_count=2000000 + sudo swapoff -a || true fi ALL_DEPENDENCIES="$COMMON_DEPENDENCIES $ENGINE_DEPENDENCIES" diff --git a/.github/scripts/wait-for-db.sh b/.github/scripts/wait-for-db.sh index 4a076f31f8..6077324dbd 100755 --- a/.github/scripts/wait-for-db.sh +++ b/.github/scripts/wait-for-db.sh @@ -104,6 +104,34 @@ db2_ready() { done } +doris_ready() { + probe_port 9030 || return 1 + HOSTNAME=${DOCKER_HOSTNAME:-127.0.0.1} + if [ "$HOSTNAME" = "localhost" ]; then + HOSTNAME="127.0.0.1" + fi + 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');" + if command -v mysql > /dev/null 2>&1; then + mysql -h "$HOSTNAME" -P 9030 -u root -e "SHOW BACKENDS\G" 2>/dev/null | grep -q "Alive: true" && \ + mysql -h "$HOSTNAME" -P 9030 -u root -e "$INIT_SQL" 2>/dev/null + else + python3 -c " +import pymysql, sys +conn = pymysql.connect(host='$HOSTNAME', port=9030, user='root', autocommit=True) +with conn.cursor() as cur: + cur.execute('SHOW BACKENDS') + rows = cur.fetchall() + cols = [d[0] for d in cur.description] + alive_idx = cols.index('Alive') + if not rows or not any(str(r[alive_idx]).lower() == 'true' for r in rows): + sys.exit(1) + for cfg in ('default_replication_num', 'force_olap_table_replication_num'): + cur.execute(f\"ADMIN SET FRONTEND CONFIG ('{cfg}' = '1')\") + cur.execute(\"ADMIN SET FRONTEND CONFIG ('enable_unicode_name_support' = 'true')\") +" 2>/dev/null + fi +} + echo "Waiting for $ENGINE to be ready..." READINESS_FUNC="${ENGINE}_ready" diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 2c18fe1e04..75a15f2b88 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -269,6 +269,7 @@ jobs: clickhouse, risingwave, starrocks, + doris, ] env: PYTEST_XDIST_AUTO_NUM_WORKERS: 2 diff --git a/.readthedocs.yaml b/.readthedocs.yaml index ee4794538f..09262c1577 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -6,7 +6,7 @@ build: python: "3.10" jobs: pre_build: - - pip install -e ".[athena,azuresql,bigframes,bigquery,clickhouse,databricks,dbt,dlt,gcppostgres,github,llm,mssql,mysql,mwaa,postgres,redshift,slack,snowflake,starrocks,trino,web,risingwave]" + - pip install -e ".[athena,azuresql,bigframes,bigquery,clickhouse,databricks,dbt,dlt,doris,gcppostgres,github,llm,mssql,mysql,mwaa,postgres,redshift,slack,snowflake,starrocks,trino,web,risingwave]" - make api-docs mkdocs: diff --git a/Makefile b/Makefile index 0d8571b33f..b1ce245486 100644 --- a/Makefile +++ b/Makefile @@ -226,6 +226,9 @@ starrocks-test: engine-starrocks-up db2-test: engine-db2-up pytest -n auto -m "db2" --reruns 3 --junitxml=test-results/junit-db2.xml +doris-test: engine-doris-up + pytest -n auto -m "doris" --reruns 3 --junitxml=test-results/junit-doris.xml + ################# # Cloud Engines # ################# diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index 38f6814c8b..656478d8d6 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -933,6 +933,7 @@ These pages describe the connection configuration options for each execution eng * [BigQuery](../integrations/engines/bigquery.md) * [ClickHouse](../integrations/engines/clickhouse.md) * [Databricks](../integrations/engines/databricks.md) +* [Doris](../integrations/engines/doris.md) * [DuckDB](../integrations/engines/duckdb.md) * [Fabric](../integrations/engines/fabric.md) * [MotherDuck](../integrations/engines/motherduck.md) @@ -978,6 +979,7 @@ Unsupported state engines, even for development: * [Spark](../integrations/engines/spark.md) * [StarRocks](../integrations/engines/starrocks.md) * [Trino](../integrations/engines/trino.md) +* [Doris](../integrations/engines/doris.md) This example gateway configuration uses Snowflake for the data warehouse connection and Postgres for the state backend connection: diff --git a/docs/guides/connections.md b/docs/guides/connections.md index 1345905bf7..bf6d4811d6 100644 --- a/docs/guides/connections.md +++ b/docs/guides/connections.md @@ -85,6 +85,7 @@ default_gateway: local_db * [ClickHouse](../integrations/engines/clickhouse.md) * [Databricks](../integrations/engines/databricks.md) * [Db2](../integrations/engines/db2.md) +* [Doris](../integrations/engines/doris.md) * [DuckDB](../integrations/engines/duckdb.md) * [Fabric](../integrations/engines/fabric.md) * [MotherDuck](../integrations/engines/motherduck.md) diff --git a/docs/installation.md b/docs/installation.md index 6bcac928e5..221a40fcce 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -54,6 +54,7 @@ Other extras are required to use specific SQL engines, like Bigquery or Postgres | Bigquery | `pip install "sqlmesh[bigquery]"` | | ClickHouse | `pip install "sqlmesh[clickhouse]"` | | Databricks | `pip install "sqlmesh[databricks]"` | + | Doris | `pip install "sqlmesh[doris]"` | | Fabric | `pip install "sqlmesh[fabric]"` | | GCP Postgres | `pip install "sqlmesh[gcppostgres]"` | | MS SQL Server | `pip install "sqlmesh[mssql]"` | diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md new file mode 100644 index 0000000000..86def97b47 --- /dev/null +++ b/docs/integrations/engines/doris.md @@ -0,0 +1,489 @@ +# Apache Doris + +This page describes SQLMesh support for the [Apache Doris](https://doris.apache.org/) engine (versions 2.1, 3.x, 4.x, and 5.0), including native Doris OLAP tables and Apache Iceberg multi-catalog lakehouse tables. + +!!! note + Doris may not be used for the SQLMesh [state connection](../../reference/configuration.md#connections). Configure a transactional database such as Postgres, MySQL, or DuckDB (for local development) as your `state_connection`. + +## Local/Built-in Scheduler + +**Engine Adapter Type**: `doris` + +### Installation + +```bash +pip install "sqlmesh[doris]" +``` + +### Connection options + +| Option | Description | Type | Required | +| --- | --- | :---: | :---: | +| `type` | Engine type name - must be `doris` | string | Y | +| `host` | The hostname of the Doris FE (Frontend) server | string | Y | +| `user` | The username to use for authentication with Doris FE | string | Y | +| `password` | The password to use for authentication with Doris FE (Default: `None`) | string | N | +| `port` | The query port number of the Doris FE server (Default: `9030`) | int | N | +| `database` | The optional default database (schema) name | string | N | +| `catalog` | The optional default catalog name (e.g., `internal` or an external `iceberg` catalog) | string | N | +| `charset` | The character set used for the connection (Default: `utf8mb4`) | string | N | +| `collation` | The collation used for the connection | string | N | +| `connect_timeout` | Connection timeout in seconds (Default: `10`) | int | N | +| `read_timeout` | Read timeout in seconds | int | N | +| `write_timeout` | Write timeout in seconds | int | N | +| `session_variables` | Session variables to set on each new connection via `SET = ` (e.g., `workload_group`, `query_timeout`) | dict | N | +| `ssl` | Explicit SSL configuration dictionary passed to `PyMySQL` | dict | N | +| `ssl_disabled` | Whether SSL is disabled | bool | N | +| `ssl_ca` | Path to the file that contains a PEM-formatted CA certificate | string | N | +| `ssl_cert` | Path to the file that contains a PEM-formatted client certificate | string | N | +| `ssl_key` | Path to the file that contains a PEM-formatted private key for the client certificate | string | N | +| `ssl_verify_cert` | Set to `true` to verify the server's certificate against the CA certificate | bool | N | +| `ssl_verify_identity` | Set to `true` to verify the server's identity in addition to the certificate | bool | N | + +### Connection configuration example + +```yaml +gateways: + doris: + connection: + type: doris + host: fe.doris.cluster + port: 9030 + user: doris_user + password: {{ env_var('DORIS_PASSWORD') }} + database: analytics + catalog: internal + connect_timeout: 30 + session_variables: + workload_group: etl_group + query_timeout: 3600 + state_connection: + type: postgres + host: state-db.internal + port: 5432 + user: sqlmesh + password: {{ env_var('POSTGRES_PASSWORD') }} + database: sqlmesh_state +``` + +--- + +## OLAP Table Models + +Apache Doris supports three native OLAP table data models: **DUPLICATE KEY**, **UNIQUE KEY**, and **AGGREGATE KEY**. + +!!! tip "Automatic Key Column Prefix Reordering" + Doris requires key columns (`UNIQUE KEY`, `DUPLICATE KEY`, or `AGGREGATE KEY`) to appear as the **first columns** in a `CREATE TABLE` definition in the exact order of the key. SQLMesh's `DorisEngineAdapter` automatically reorders your schema definition during `CREATE TABLE` so key columns appear first, regardless of the column order in your `SELECT` query. + +### DUPLICATE KEY model (default) + +By default, OLAP tables use the `DUPLICATE KEY` model. If no key is specified, Doris uses the first column as the duplicate key and hash distribution column. + +```sql +MODEL ( + name analytics.user_events, + kind FULL, + clustered_by (user_id), + physical_properties ( + duplicate_key = ('user_id', 'event_time'), + buckets = 10 + ) +); + +SELECT + user_id::BIGINT AS user_id, + event_time::DATETIME AS event_time, + event_type::VARCHAR(64) AS event_type +FROM raw.events; +``` + +### UNIQUE KEY model (Merge-on-Write & custom sort keys) + +For `INCREMENTAL_BY_UNIQUE_KEY` models, SQLMesh automatically configures the model's `unique_key` as the Doris table's `UNIQUE KEY` and `DISTRIBUTED BY HASH(...)` key—no extra `physical_properties` required: + +```sql +MODEL ( + name analytics.orders_upsert, + kind INCREMENTAL_BY_UNIQUE_KEY ( + unique_key order_id + ), + physical_properties ( + order_by = ('updated_at', 'order_id'), + enable_unique_key_merge_on_write = 'true' + ) +); + +SELECT + order_id::BIGINT AS order_id, + customer_id::BIGINT AS customer_id, + status::VARCHAR(32) AS status, + updated_at::DATETIME AS updated_at +FROM raw.orders +WHERE updated_at BETWEEN @start_ds AND @end_ds; +``` + +You can also specify `unique_key` explicitly in `physical_properties` for `FULL` or `INCREMENTAL_BY_TIME_RANGE` models, and optionally decouple the sort key from the unique key using `order_by` (creating a `SEQUENCE MAP COL(col1, col2)` sort key on Doris 2.1+ Merge-on-Write tables). + +### AGGREGATE KEY model + +For pre-aggregated rollup tables, specify the key columns using `aggregate_key` in `physical_properties`: + +```sql +MODEL ( + name analytics.daily_site_metrics, + kind FULL, + clustered_by (site_id), + physical_properties ( + aggregate_key = ('site_id', 'event_date'), + buckets = 16 + ) +); + +SELECT + site_id::INT AS site_id, + event_date::DATE AS event_date, + SUM(pv)::BIGINT AS total_pv, + MAX(dwell_seconds)::INT AS max_dwell +FROM analytics.user_events +GROUP BY site_id, event_date; +``` + +--- + +## Distribution & Bucketing + +Doris requires a `DISTRIBUTED BY` clause on every OLAP table. SQLMesh supports three ways to configure distribution: + +1. **SQLMesh `clustered_by` + optional `buckets`** (recommended for cross-engine portability): + ```sql + MODEL ( + name analytics.events, + kind FULL, + clustered_by (user_id), + physical_properties ( + buckets = 16 + ) + ); + ``` +2. **Shorthand string `distributed_by`**: + ```sql + physical_properties ( + distributed_by = 'HASH(user_id) BUCKETS 16' + -- or 'RANDOM BUCKETS AUTO' + ) + ``` +3. **Structured dictionary `distributed_by`**: + ```sql + physical_properties ( + distributed_by = ( + kind = 'HASH', + expressions = 'user_id', + buckets = 'AUTO' + ) + ) + ``` + +If none of `clustered_by` or `distributed_by` is provided: +- When `unique_key` is present, SQLMesh defaults to `DISTRIBUTED BY HASH() BUCKETS `. +- When `buckets` is specified with `duplicate_key` or `aggregate_key`, SQLMesh emits `DISTRIBUTED BY HASH() BUCKETS ` (or `DISTRIBUTED BY RANDOM BUCKETS ` when no key is specified). + +--- + +## Partitioning + +Doris OLAP tables support **Multi-Range Batch Partitioning**, **Explicit Partition Definitions**, **Dynamic Partitioning**, and **Auto Partitioning** (Doris 2.1+). + +### Multi-range & explicit partitions + +```sql +MODEL ( + name analytics.partitioned_events, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column event_date + ), + partitioned_by RANGE(event_date), + physical_properties ( + partitions = 'FROM ("2023-01-01") TO ("2027-01-01") INTERVAL 1 MONTH' + ) +); +``` + +### Auto partitioning (`auto_partition = true`) + +In Doris 2.1+, `AUTO PARTITION` automatically creates partitions on demand as new data arrives: + +```sql +MODEL ( + name analytics.auto_partitioned_events, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column event_time + ), + partitioned_by RANGE(DATE_TRUNC(event_time, 'day')), + physical_properties ( + auto_partition = 'true' + ) +); +``` + +--- + +## Apache Iceberg & Multi-Catalog Lakehouse Integration + +Apache Doris 2.1–5.0 acts as a federated lakehouse engine that can both query and write to external **Apache Iceberg** catalogs (REST/Polaris, AWS Glue, Hive Metastore, Hadoop/S3, DLF). + +### Pattern 1: Reading Iceberg Bronze tables into Doris OLAP Silver/Gold models + +Because `DorisEngineAdapter` supports multi-catalog 3-part naming (`catalog.db.table`), your SQLMesh models can query external Iceberg tables directly while materializing sub-second serving tables in Doris's `internal` OLAP catalog: + +```sql +MODEL ( + name silver.user_sessions, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column session_date + ) +); + +SELECT + user_id, + session_id, + CAST(event_ts AS DATE) AS session_date +FROM iceberg_catalog.bronze.clickstream +WHERE event_ts BETWEEN @start_ts AND @end_ts; +``` + +### Pattern 2: Materializing Apache Iceberg tables via Doris + +SQLMesh can also create, partition, and overwrite Apache Iceberg tables directly through Doris (`ENGINE=iceberg` and `INSERT OVERWRITE TABLE ... VALUES (*)`). + +!!! important "Virtual Environments & Catalog Bootstrap on External Iceberg Catalogs" + Doris only supports `CREATE VIEW` inside the `internal` catalog; external Iceberg catalogs do not support Doris views, and Doris does not support `environment_suffix_target: catalog`. + When materializing models directly inside an external Iceberg catalog: + + - Set `virtual_environment_mode: dev_only` (and keep `environment_suffix_target: schema` or `table`) in `config.yaml` so production tables are materialized directly in the Iceberg catalog without virtual views. + - Configure `environment_catalog_mapping` in `config.yaml` to route non-`prod` virtual environment views to the `internal` catalog (otherwise `DorisEngineAdapter` will log a warning and skip view creation in external catalogs during `dev` promotion): + ```yaml + virtual_environment_mode: dev_only + environment_suffix_target: schema + environment_catalog_mapping: + "^(?!prod$).*": internal + before_all: + - > + CREATE CATALOG IF NOT EXISTS iceberg_catalog PROPERTIES ( + 'type' = 'iceberg', + 'iceberg.catalog.type' = 'rest', + 'uri' = 'http://rest-catalog:8181', + 'warehouse' = 's3://lakehouse/warehouse' + ) + ``` + +```sql +MODEL ( + name iceberg_catalog.silver.curated_events, + kind FULL, + table_format iceberg, + partitioned_by (day(event_ts), bucket(16, user_id)), + physical_properties ( + 'write-format' = 'parquet', + 'compression-codec' = 'zstd' + ) +); + +SELECT + user_id::BIGINT AS user_id, + event_ts::DATETIME AS event_ts, + payload::STRING AS payload +FROM iceberg_catalog.bronze.raw_events; +``` + +--- + +## DDL-Only / Schema Contract Management for Streaming Tables (`doris_ddl_only`) + +In real-time architectures where **Apache Flink (`flink-doris-connector`)**, **Kafka Routine Load**, **Stream Load**, or **Spark** streams data continuously into a Doris OLAP or Iceberg table, you may want SQLMesh to manage **only the table DDL (`CREATE TABLE`) and schema evolution (`ALTER TABLE`)** without executing `INSERT` or `REPLACE` statements during `sqlmesh run`. + +You can achieve this by adding a lightweight `CustomMaterialization` plugin (`materializations/doris_ddl_only.py`) to your SQLMesh project: + +```python +from __future__ import annotations + +import typing as t +from sqlmesh.core.model import Model +from sqlmesh.core.snapshot.evaluator import CustomMaterialization +from sqlmesh.utils.errors import SQLMeshError + +if t.TYPE_CHECKING: + from sqlmesh.core.engine_adapter._typing import QueryOrDF + + +class DorisDDLOnlyMaterialization(CustomMaterialization): + """Manages CREATE TABLE and ALTER TABLE schema migrations without inserting data.""" + + NAME = "doris_ddl_only" + + def create( + self, + table_name: str, + model: Model, + is_table_deployable: bool, + render_kwargs: t.Dict[str, t.Any], + skip_grants: bool = False, + **kwargs: t.Any, + ) -> None: + if not model.annotated: + raise SQLMeshError( + f"Model '{model.name}' uses 'doris_ddl_only' and must explicitly declare `columns (...)` in MODEL (...)." + ) + super().create( + table_name=table_name, + model=model, + is_table_deployable=is_table_deployable, + render_kwargs=render_kwargs, + skip_grants=skip_grants, + **kwargs, + ) + + def insert( + self, + table_name: str, + query_or_df: QueryOrDF, + model: Model, + is_first_insert: bool, + render_kwargs: t.Dict[str, t.Any], + **kwargs: t.Any, + ) -> None: + # No-op: external Flink / Routine Load / Spark pipelines write data to this table. + return + + def append( + self, + table_name: str, + query_or_df: QueryOrDF, + model: Model, + render_kwargs: t.Dict[str, t.Any], + **kwargs: t.Any, + ) -> None: + # No-op: external Flink / Routine Load / Spark pipelines write data to this table. + return +``` + +Then define your streaming target table in `models/` with `virtual_environment_mode: dev_only`: + +```sql +MODEL ( + name realtime.flink_order_events, + kind CUSTOM ( + materialization 'doris_ddl_only', + forward_only true, + disable_restatement true + ), + columns ( + order_id BIGINT, + user_id BIGINT, + amount DECIMAL(18, 2), + updated_at DATETIME + ), + physical_properties ( + unique_key = ('order_id'), + order_by = ('updated_at', 'order_id'), + buckets = 16, + enable_unique_key_merge_on_write = 'true', + enable_light_schema_change = 'true' + ) +); + +SELECT + CAST(NULL AS BIGINT) AS order_id, + CAST(NULL AS BIGINT) AS user_id, + CAST(NULL AS DECIMAL(18, 2)) AS amount, + CAST(NULL AS DATETIME) AS updated_at; +``` + +!!! tip "Zero-Downtime Schema Evolution with Flink / Routine Load" + Setting `forward_only true` and `disable_restatement true` on a `doris_ddl_only` model ensures that schema additions (`ALTER TABLE ... ADD COLUMN`) are applied in-place on the existing table (`enable_light_schema_change = 'true'`) rather than creating a new empty physical table, and prevents accidental restatement wipes. + +--- + +## Compute-Storage Decoupled Mode (Doris 3.0 / 4.x / 5.0) + +For Doris 3.0+ Compute-Storage Decoupled deployments on object storage (S3, GCS, Azure Blob, OSS) or HDFS: + +- Pass `storage_vault_name` in `physical_properties` (or `model_defaults.physical_properties` in `config.yaml`) to bind tables to a specific storage vault. +- Use `session_variables` in your `doris` connection config to bind SQLMesh workloads to a specific `workload_group` or cloud cluster. + +```sql +MODEL ( + name analytics.cloud_events, + kind FULL, + clustered_by (event_id), + physical_properties ( + storage_vault_name = 'prod_s3_vault', + file_cache_ttl_seconds = '86400' + ) +); +``` + +--- + +## Materialized Views & Indexing + +### Asynchronous Materialized Views (MTMV) + +```sql +MODEL ( + name analytics.user_summary_mv, + kind VIEW ( + materialized true + ), + partitioned_by ds, + physical_properties ( + build = 'IMMEDIATE', + refresh = 'AUTO', + refresh_trigger = 'ON SCHEDULE EVERY 1 HOUR', + distributed_by = 'HASH(user_id) BUCKETS 10' + ) +); + +SELECT + user_id, + CAST(event_time AS DATE) AS ds, + COUNT(*) AS event_count +FROM analytics.user_events +GROUP BY user_id, CAST(event_time AS DATE); +``` + +### Inverted Indexes (Doris 2.1–5.0) + +You can attach Doris inverted indexes (`USING INVERTED`) for full-text search and fast point filtering using post-statements gated on `@runtime_stage = 'creating'`: + +```sql +MODEL ( + name analytics.searchable_logs, + kind FULL +); + +SELECT + log_id::BIGINT AS log_id, + message::STRING AS message +FROM raw.logs; + +@IF( + @runtime_stage = 'creating', + CREATE INDEX idx_message ON @this_model (message) USING INVERTED PROPERTIES("parser" = "english") COMMENT 'Full-text inverted index' +); +``` + +--- + +## Migrating from `dbt-doris` + +If you are migrating an existing `dbt-doris` project to native SQLMesh models, map your `dbt-doris` config keys as follows: + +| `dbt-doris` Config | SQLMesh Model Equivalent | +| --- | --- | +| `materialized='table'` | `kind FULL` | +| `materialized='incremental'` (with `unique_key`) | `kind INCREMENTAL_BY_UNIQUE_KEY (unique_key ...)` | +| `materialized='incremental'` (partition overwrite) | `kind INCREMENTAL_BY_TIME_RANGE (time_column ...)` or `kind INCREMENTAL_BY_PARTITION` | +| `duplicate_key=['c1', 'c2']` | `physical_properties (duplicate_key = ('c1', 'c2'))` | +| `unique_key=['c1']` | `kind INCREMENTAL_BY_UNIQUE_KEY (unique_key c1)` or `physical_properties (unique_key = ('c1'))` | +| `distributed_by=['c1']`, `buckets=10` | `clustered_by (c1), physical_properties (buckets = 10)` | +| `properties={...}` | Key-value entries inside `physical_properties (...)` | diff --git a/docs/integrations/overview.md b/docs/integrations/overview.md index 1c9d56b7e2..6e0398eb5e 100644 --- a/docs/integrations/overview.md +++ b/docs/integrations/overview.md @@ -17,6 +17,7 @@ SQLMesh supports the following execution engines for running SQLMesh projects (e * [ClickHouse](./engines/clickhouse.md) (clickhouse) * [Databricks](./engines/databricks.md) (databricks) * [Db2](./engines/db2.md) (db2) +* [Doris](./engines/doris.md) (doris) * [DuckDB](./engines/duckdb.md) (duckdb) * [Fabric](./engines/fabric.md) (fabric) * [MotherDuck](./engines/motherduck.md) (motherduck) diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index ae036a7774..a31eea0ccf 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -222,6 +222,7 @@ These pages describe the connection configuration options for each execution eng * [BigQuery](../integrations/engines/bigquery.md) * [ClickHouse](../integrations/engines/clickhouse.md) * [Databricks](../integrations/engines/databricks.md) +* [Doris](../integrations/engines/doris.md) * [DuckDB](../integrations/engines/duckdb.md) * [Fabric](../integrations/engines/fabric.md) * [MotherDuck](../integrations/engines/motherduck.md) diff --git a/mkdocs.yml b/mkdocs.yml index 49c4b9163b..334feb3019 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -83,6 +83,7 @@ nav: - integrations/engines/clickhouse.md - integrations/engines/databricks.md - integrations/engines/db2.md + - integrations/engines/doris.md - integrations/engines/duckdb.md - integrations/engines/fabric.md - integrations/engines/motherduck.md diff --git a/pyproject.toml b/pyproject.toml index 378d3ffe43..41b33339e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -116,6 +116,7 @@ db2 = [ "ibm_db", "db2-sqlglot-dialect;python_version>=\"3.10\"" ] +doris = ["pymysql"] duckdb = [] fabric = ["pyodbc>=5.0.0"] fabric-mssql-python = ["mssql-python>=1.1.0;python_version>=\"3.10\""] @@ -290,6 +291,7 @@ markers = [ "trino: test for Trino (all connectors)", "risingwave: test for Risingwave", "starrocks: test for StarRocks", + "doris: test for Doris", # Other "set_default_connection", diff --git a/sqlmesh/core/config/__init__.py b/sqlmesh/core/config/__init__.py index 50d2d9a5a2..3113c61435 100644 --- a/sqlmesh/core/config/__init__.py +++ b/sqlmesh/core/config/__init__.py @@ -12,6 +12,7 @@ BigQueryConnectionConfig as BigQueryConnectionConfig, ConnectionConfig as ConnectionConfig, DatabricksConnectionConfig as DatabricksConnectionConfig, + DorisConnectionConfig as DorisConnectionConfig, DuckDBConnectionConfig as DuckDBConnectionConfig, FabricConnectionConfig as FabricConnectionConfig, GCPPostgresConnectionConfig as GCPPostgresConnectionConfig, diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 29cb1ef467..8e55c7bdd7 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -63,6 +63,8 @@ # Db2 rejects table names starting with underscore (_versions, _snapshots, # _environments, _intervals). Use a separate state_connection (e.g., DuckDB). "db2", + # OLAP primary keys and lack of transactional DDL make Doris unsuitable for state sync + "doris", } MOTHERDUCK_TOKEN_REGEX = re.compile(r"(\?|\&)(motherduck_token=)(\S*)") PASSWORD_REGEX = re.compile(r"(password=)(\S+)") @@ -2692,6 +2694,164 @@ def connect_db2(**kwargs: t.Any) -> t.Any: return connect_db2 +class DorisConnectionConfig(ConnectionConfig): + """Configuration for the Apache Doris connection. + + Apache Doris uses the MySQL network protocol and supports both internal OLAP tables + and external Multi-Catalog lakehouse catalogs (e.g., Apache Iceberg). + + Args: + host: The hostname of the Doris FE (Frontend) node. + user: The Doris username. + password: The optional Doris password. + port: The query port number of the Doris FE node. Default is 9030. + database: The optional database name. + catalog: The optional default catalog name (e.g., 'internal' or an external Iceberg catalog). + charset: The optional character set (defaults to 'utf8mb4'). + collation: The optional collation. + connect_timeout: Timeout in seconds for establishing a connection. Default is 10. + read_timeout: Timeout in seconds for reading from the connection. + write_timeout: Timeout in seconds for writing to the connection. + ssl: Optional dictionary of SSL arguments passed to PyMySQL. + ssl_ca: Path to the file that contains a PEM-formatted CA certificate. + ssl_cert: Path to the file that contains a PEM-formatted client certificate. + ssl_key: Path to the file that contains a PEM-formatted private key for the client certificate. + ssl_verify_cert: Whether to verify the server certificate against the CA certificate. + ssl_verify_identity: Whether to verify the server hostname against the certificate. + ssl_disabled: Whether to disable SSL connection. + session_variables: Optional dictionary of Doris session variables to set on each connection. + concurrent_tasks: The maximum number of tasks that can use this connection concurrently. + register_comments: Whether or not to register model comments with the SQL engine. + pre_ping: Whether or not to pre-ping the connection before starting a new statement. + """ + + host: str + user: str + password: t.Optional[str] = None + port: t.Optional[int] = 9030 + database: t.Optional[str] = None + catalog: t.Optional[str] = None + charset: t.Optional[str] = None + collation: t.Optional[str] = None + connect_timeout: t.Optional[int] = 10 + read_timeout: t.Optional[int] = None + write_timeout: t.Optional[int] = None + ssl: t.Optional[t.Union[bool, t.Dict[str, t.Any]]] = None + ssl_ca: t.Optional[str] = None + ssl_cert: t.Optional[str] = None + ssl_key: t.Optional[str] = None + ssl_verify_cert: t.Optional[bool] = None + ssl_verify_identity: t.Optional[bool] = None + ssl_disabled: t.Optional[bool] = None + session_variables: t.Optional[t.Dict[str, t.Any]] = None + + concurrent_tasks: int = 4 + register_comments: bool = True + pre_ping: bool = False + + type_: t.Literal["doris"] = Field(alias="type", default="doris") + DIALECT: t.ClassVar[t.Literal["doris"]] = "doris" + DISPLAY_NAME: t.ClassVar[t.Literal["Apache Doris"]] = "Apache Doris" + DISPLAY_ORDER: t.ClassVar[t.Literal[20]] = 20 + + _engine_import_validator = _get_engine_import_validator("pymysql", "doris") + + @field_validator("session_variables", mode="after") + @classmethod + def _validate_session_variables( + cls, v: t.Optional[t.Dict[str, t.Any]] + ) -> t.Optional[t.Dict[str, t.Any]]: + if v: + import math + + for k, val in v.items(): + if not isinstance(k, str) or not re.fullmatch(r"[a-zA-Z_][a-zA-Z0-9_]*\Z", k): + raise ConfigError(f"Invalid Doris session variable name: '{k}'") + if not isinstance(val, (str, int, float, bool)) or ( + isinstance(val, float) and (math.isnan(val) or math.isinf(val)) + ): + raise ConfigError(f"Invalid Doris session variable value for '{k}': {val!r}") + return v + + @property + def _static_connection_kwargs(self) -> t.Dict[str, t.Any]: + kwargs: t.Dict[str, t.Any] = { + "autocommit": True, + "charset": self.charset or "utf8mb4", + "local_infile": False, + } + if not self.ssl_disabled: + ssl_config: t.Dict[str, t.Any] = dict(self.ssl) if isinstance(self.ssl, dict) else {} + for key in ("ca", "cert", "key"): + val = getattr(self, f"ssl_{key}", None) + if val is not None: + ssl_config[key] = val + if ssl_config or (self.ssl is not None and self.ssl is not False): + kwargs["ssl"] = ssl_config + for key in ( + "ssl_ca", + "ssl_cert", + "ssl_key", + "ssl_verify_cert", + "ssl_verify_identity", + ): + val = getattr(self, key, None) + if val is not None: + kwargs[key] = val + return kwargs + + @property + def _connection_kwargs_keys(self) -> t.Set[str]: + is_external_catalog = bool(self.catalog and self.catalog.lower() != "internal") + return { + k + for k in ( + "host", + "user", + "password", + "port", + "database", + "charset", + "collation", + "connect_timeout", + "read_timeout", + "write_timeout", + "ssl_disabled", + ) + if getattr(self, k) is not None and not (k == "database" and is_external_catalog) + } + + @property + def _engine_adapter(self) -> t.Type[EngineAdapter]: + return engine_adapter.DorisEngineAdapter + + @property + def _connection_factory(self) -> t.Callable: + from pymysql import connect + + return connect + + @property + def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]: + if not self.catalog and not self.session_variables: + return None + + def init(cursor: t.Any) -> None: + if self.session_variables: + self._validate_session_variables(self.session_variables) + for k, v in self.session_variables.items(): + val_sql = exp.convert(v).sql("doris", comments=False) + cursor.execute(f"SET `{k}` = {val_sql}") + if self.catalog: + catalog_sql = exp.to_identifier(self.catalog, quoted=True).sql("doris") + cursor.execute(f"SWITCH {catalog_sql}") + if self.database: + db_sql = exp.to_identifier(self.database, quoted=True).sql("doris") + cursor.execute(f"USE {db_sql}") + + return init + + CONNECTION_CONFIG_TO_TYPE = { # Map all subclasses of ConnectionConfig to the value of their `type_` field. tpe.all_field_infos()["type_"].default: tpe diff --git a/sqlmesh/core/engine_adapter/__init__.py b/sqlmesh/core/engine_adapter/__init__.py index 658b149753..bcb70e87a1 100644 --- a/sqlmesh/core/engine_adapter/__init__.py +++ b/sqlmesh/core/engine_adapter/__init__.py @@ -11,6 +11,7 @@ from sqlmesh.core.engine_adapter.bigquery import BigQueryEngineAdapter from sqlmesh.core.engine_adapter.clickhouse import ClickhouseEngineAdapter from sqlmesh.core.engine_adapter.databricks import DatabricksEngineAdapter +from sqlmesh.core.engine_adapter.doris import DorisEngineAdapter from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter from sqlmesh.core.engine_adapter.mssql import MSSQLEngineAdapter from sqlmesh.core.engine_adapter.mysql import MySQLEngineAdapter @@ -50,6 +51,7 @@ "risingwave": RisingwaveEngineAdapter, "fabric": FabricEngineAdapter, "starrocks": StarRocksEngineAdapter, + "doris": DorisEngineAdapter, } # Add DB2 to the registry only when the dialect package is present diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py new file mode 100644 index 0000000000..a3ca961520 --- /dev/null +++ b/sqlmesh/core/engine_adapter/doris.py @@ -0,0 +1,1948 @@ +from __future__ import annotations + +from collections import defaultdict +import itertools +import logging +import re +import typing as t + +from sqlglot import exp, parse, parse_one + +from sqlmesh.core.dialect import to_schema +from sqlmesh.core.engine_adapter.base import ( + EngineAdapterWithIndexSupport, + InsertOverwriteStrategy, + _decoded_str, +) +from sqlmesh.core.engine_adapter.mixins import ( + GetCurrentCatalogFromFunctionMixin, + LogicalMergeMixin, + PandasNativeFetchDFSupportMixin, + RowDiffMixin, +) +from sqlmesh.core.engine_adapter.shared import ( + CatalogSupport, + CommentCreationTable, + CommentCreationView, + DataObject, + DataObjectType, + SourceQuery, + set_catalog, +) +from sqlmesh.utils import get_source_columns_to_types, random_id +from sqlmesh.utils.errors import SQLMeshError + +if t.TYPE_CHECKING: + from sqlmesh.core._typing import SchemaName, TableName + from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF + from sqlmesh.core.node import IntervalUnit + +logger = logging.getLogger(__name__) + + +@set_catalog() +class DorisEngineAdapter( + EngineAdapterWithIndexSupport, + LogicalMergeMixin, + GetCurrentCatalogFromFunctionMixin, + PandasNativeFetchDFSupportMixin, + RowDiffMixin, +): + DIALECT = "doris" + DEFAULT_BATCH_SIZE = 5000 + SUPPORTS_TRANSACTIONS = False + COMMENT_CREATION_TABLE = CommentCreationTable.IN_SCHEMA_DEF_NO_CTAS + COMMENT_CREATION_VIEW = CommentCreationView.IN_SCHEMA_DEF_NO_COMMANDS + MAX_TABLE_COMMENT_LENGTH = 2048 + MAX_COLUMN_COMMENT_LENGTH = 255 + SUPPORTS_INDEXES = True + SUPPORTS_REPLACE_TABLE = False + MAX_IDENTIFIER_LENGTH = 64 + SUPPORTS_MATERIALIZED_VIEWS = True + SUPPORTS_MATERIALIZED_VIEW_SCHEMA = True + SUPPORTS_CREATE_DROP_CATALOG = False + INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.DELETE_INSERT + CURRENT_CATALOG_EXPRESSION = exp.func("current_catalog") + SCHEMA_DIFFER_KWARGS = { + "parameterized_type_defaults": { + exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(9, 0), (0,)], + exp.DataType.build("CHAR", dialect=DIALECT).this: [(1,), (255,)], + exp.DataType.build("VARCHAR", dialect=DIALECT).this: [(65533,)], + exp.DataType.build("DATETIME", dialect=DIALECT).this: [(0,)], + exp.DataType.build("INT", dialect=DIALECT).this: [(11,)], + exp.DataType.build("BIGINT", dialect=DIALECT).this: [(20,)], + exp.DataType.build("SMALLINT", dialect=DIALECT).this: [(6,)], + exp.DataType.build("TINYINT", dialect=DIALECT).this: [(4,), (1,)], + }, + "types_with_unlimited_length": { + exp.DataType.build("TEXT", dialect=DIALECT).this: { + exp.DataType.build("VARCHAR", dialect=DIALECT).this, + exp.DataType.build("CHAR", dialect=DIALECT).this, + }, + exp.DataType.build("VARCHAR", dialect=DIALECT).this: { + exp.DataType.build("VARCHAR", dialect=DIALECT).this, + exp.DataType.build("CHAR", dialect=DIALECT).this, + exp.DataType.build("TEXT", dialect=DIALECT).this, + }, + }, + "compatible_types": { + exp.DataType.build("BOOLEAN", dialect=DIALECT): { + exp.DataType.build("TINYINT", dialect=DIALECT), + exp.DataType.build("SMALLINT", dialect=DIALECT), + exp.DataType.build("INT", dialect=DIALECT), + exp.DataType.build("BIGINT", dialect=DIALECT), + exp.DataType.build("INT128", dialect=DIALECT), + }, + exp.DataType.build("TINYINT", dialect=DIALECT): { + exp.DataType.build("SMALLINT", dialect=DIALECT), + exp.DataType.build("INT", dialect=DIALECT), + exp.DataType.build("BIGINT", dialect=DIALECT), + exp.DataType.build("INT128", dialect=DIALECT), + exp.DataType.build("DOUBLE", dialect=DIALECT), + }, + exp.DataType.build("SMALLINT", dialect=DIALECT): { + exp.DataType.build("INT", dialect=DIALECT), + exp.DataType.build("BIGINT", dialect=DIALECT), + exp.DataType.build("INT128", dialect=DIALECT), + exp.DataType.build("DOUBLE", dialect=DIALECT), + }, + exp.DataType.build("INT", dialect=DIALECT): { + exp.DataType.build("BIGINT", dialect=DIALECT), + exp.DataType.build("INT128", dialect=DIALECT), + exp.DataType.build("DOUBLE", dialect=DIALECT), + }, + exp.DataType.build("BIGINT", dialect=DIALECT): { + exp.DataType.build("INT128", dialect=DIALECT), + }, + exp.DataType.build("FLOAT", dialect=DIALECT): { + exp.DataType.build("DOUBLE", dialect=DIALECT), + }, + exp.DataType.build("DATE", dialect=DIALECT): { + exp.DataType.build("DATETIME", dialect=DIALECT), + }, + }, + } + + INVALID_KEY_TYPES = { + "FLOAT", + "DOUBLE", + "STRING", + "TEXT", + "JSON", + "JSONB", + "VARIANT", + "OBJECT", + "ARRAY", + "MAP", + "STRUCT", + "BITMAP", + "HLL", + "QUANTILE_STATE", + "AGG_STATE", + } + + @property + def catalog_support(self) -> CatalogSupport: + return CatalogSupport.FULL_SUPPORT + + @property + def default_catalog(self) -> t.Optional[str]: + return self._default_catalog or self._extra_config.get("catalog") + + def get_current_catalog(self) -> t.Optional[str]: + try: + catalog = super().get_current_catalog() + if catalog: + return _decoded_str(catalog) + except Exception: + logger.debug("current_catalog() failed on Doris, falling back to default catalog") + return self.default_catalog or "internal" + + def set_current_catalog(self, catalog: str) -> None: + catalog_sql = self._to_sql(exp.to_identifier(catalog, quoted=True)) + self.execute(exp.Command(this="SWITCH", expression=catalog_sql)) + + def _is_iceberg_table( + self, + table_format: t.Optional[str] = None, + catalog_name: t.Optional[str] = None, + table_properties: t.Optional[t.Dict[str, t.Any]] = None, + ) -> bool: + if table_format and str(table_format).strip().lower() == "iceberg": + return True + if table_properties: + engine_val = self._extract_str_value(table_properties.get("engine")) + if engine_val and engine_val.strip().lower() == "iceberg": + return True + effective_catalog = catalog_name or self.default_catalog + if effective_catalog and str(effective_catalog).strip().lower() != "internal": + return True + return False + + @staticmethod + def _extract_str_value(val: t.Any) -> t.Optional[str]: + if val is None: + return None + if isinstance(val, exp.Literal): + return str(val.this) + if isinstance(val, exp.Identifier): + return str(val.this) + if isinstance(val, exp.Column): + return str(val.name) + if isinstance(val, exp.Var): + return str(val.this) + if isinstance(val, str): + return val + if isinstance(val, exp.Expr): + return val.name or str(val) + return str(val) + + @classmethod + def _extract_column_names(cls, key_expr: t.Any) -> t.List[str]: + if key_expr is None: + return [] + if isinstance(key_expr, (exp.Tuple, exp.Array)): + names: t.List[str] = [] + for e in key_expr.expressions: + names.extend(cls._extract_column_names(e)) + return names + if isinstance(key_expr, exp.Paren): + return cls._extract_column_names(key_expr.this) + if isinstance(key_expr, (list, tuple)): + names = [] + for e in key_expr: + names.extend(cls._extract_column_names(e)) + return names + if isinstance(key_expr, exp.Column): + return [key_expr.name] + if isinstance(key_expr, (exp.Identifier, exp.Literal, exp.Var)): + raw = str(key_expr.this) + return [c.strip().strip("`") for c in raw.split(",") if c.strip()] + if isinstance(key_expr, str): + return [c.strip().strip("`") for c in key_expr.split(",") if c.strip()] + return [str(key_expr).strip().strip("`")] + + @classmethod + def _is_valid_olap_key_type(cls, dtype: exp.DataType) -> bool: + type_name = ( + dtype.this.value.upper() + if isinstance(dtype.this, exp.DataType.Type) + else str(dtype.this).upper() + ) + if type_name in cls.INVALID_KEY_TYPES: + return False + if dtype.is_type("userdefined"): + kind = str(dtype.args.get("kind") or "").upper() + if kind in cls.INVALID_KEY_TYPES: + return False + return True + + def _fetch_native_df( + self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False + ) -> DF: + self._connection_pool.get_cursor() + return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) + + def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any) -> None: + try: + super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) + except Exception as e: + err_code = e.args[0] if getattr(e, "args", None) else None + err_name = e.__class__.__name__ + is_disconnect = (err_code in (2006, 2013, 2055) and err_name == "OperationalError") or ( + err_code == 0 and err_name == "InterfaceError" + ) + if is_disconnect and not self._connection_pool.is_transaction_active: + logger.warning( + "Doris connection lost (%s: %s), reconnecting and retrying once...", + err_name, + err_code, + ) + self._connection_pool.close_cursor() + conn = self._connection_pool.get() + conn.ping(reconnect=True) + super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) + else: + raise + + def columns( + self, table_name: TableName, include_pseudo_columns: bool = False + ) -> t.Dict[str, exp.DataType]: + """Fetches column names and types for the target table, normalizing Doris type aliases.""" + self.execute(exp.Describe(this=exp.to_table(table_name), kind="TABLE")) + describe_output = self.cursor.fetchall() + cols: t.Dict[str, exp.DataType] = {} + for column_name, column_type, *_ in itertools.takewhile( + lambda row: not str(row[0]).startswith("#"), + describe_output, + ): + col_name_str = _decoded_str(column_name) if column_name else "" + raw_type = _decoded_str(column_type) if column_type else "" + if not col_name_str.strip() or not raw_type.strip(): + continue + normalized_type = re.sub( + r"\bDECIMALV[23]\b(?!\s*`?\s*:)", "DECIMAL", raw_type, flags=re.IGNORECASE + ) + normalized_type = re.sub( + r"\bDATETIMEV2\b(?!\s*`?\s*:)", "DATETIME", normalized_type, flags=re.IGNORECASE + ) + normalized_type = re.sub( + r"\bDATEV2\b(?!\s*`?\s*:)", "DATE", normalized_type, flags=re.IGNORECASE + ) + normalized_type = re.sub( + r"\bLARGEINT(?:\(\d+\)|\b)(?!\s*`?\s*:)", + "INT128", + normalized_type, + flags=re.IGNORECASE, + ) + normalized_type = re.sub( + r"\bTINYINT\(1\)", "BOOLEAN", normalized_type, flags=re.IGNORECASE + ) + normalized_type = re.sub( + r"\b(TINYINT|SMALLINT|INT|BIGINT)\(\d+\)", + r"\1", + normalized_type, + flags=re.IGNORECASE, + ) + cols[col_name_str] = exp.DataType.build(normalized_type, dialect=self.dialect, udt=True) + return cols + + def create_schema( + self, + schema_name: SchemaName, + ignore_if_exists: bool = True, + warn_on_error: bool = True, + properties: t.Optional[t.List[exp.Expr]] = None, + ) -> None: + """Create a schema (database) in Doris.""" + properties = properties or [] + return super()._create_schema( + schema_name=schema_name, + ignore_if_exists=ignore_if_exists, + warn_on_error=warn_on_error, + properties=properties, + kind="DATABASE", + ) + + def drop_schema( + self, + schema_name: SchemaName, + ignore_if_not_exists: bool = True, + cascade: bool = False, + **drop_args: t.Dict[str, exp.Expr], + ) -> None: + """Drop schema in Doris. Note: Doris doesn't support CASCADE clause.""" + return self._drop_object( + name=schema_name, + exists=ignore_if_not_exists, + kind="DATABASE", + cascade=False, + **drop_args, + ) + + def _get_data_objects( + self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None + ) -> t.List[DataObject]: + """Returns all the data objects that exist in the given schema and optional catalog.""" + schema_exp = to_schema(schema_name) + catalog = schema_exp.catalog or None + query_catalog = catalog or self.default_catalog + + query = ( + exp.select( + exp.column("table_schema").as_("schema_name"), + exp.column("table_name").as_("name"), + exp.case() + .when( + exp.column("table_type").eq("BASE TABLE"), + exp.Literal.string("table"), + ) + .when( + exp.column("table_type").eq("SYSTEM VIEW"), + exp.Literal.string("view"), + ) + .when( + exp.column("table_type").eq("VIEW"), + exp.Literal.string("view"), + ) + .when( + exp.column("table_type").eq("MATERIALIZED VIEW"), + exp.Literal.string("materialized_view"), + ) + .else_(exp.column("table_type")) + .as_("type"), + ) + .from_(exp.table_("tables", db="information_schema", catalog=query_catalog)) + .where(exp.column("table_schema").eq(schema_exp.db)) + ) + if object_names: + lowered_names = [name.lower() for name in object_names] + query = query.where(exp.func("LOWER", exp.column("table_name")).isin(*lowered_names)) + + result: t.List[DataObject] = [] + for row_schema, table_name, table_type in self.fetchall(query): + try: + schema_str = _decoded_str(row_schema) if row_schema is not None else schema_exp.db + name_str = _decoded_str(table_name) if table_name is not None else "unknown" + obj_type = _decoded_str(table_type) if table_type is not None else "table" + + upper_type = obj_type.upper() + if upper_type in ("BASE TABLE", "EXTERNAL TABLE"): + obj_type = "table" + elif upper_type in ("VIEW", "SYSTEM VIEW"): + obj_type = "view" + elif upper_type in ("MATERIALIZED VIEW", "MATERIALIZED_VIEW"): + obj_type = "materialized_view" + + result.append( + DataObject( + catalog=catalog, + schema=schema_str, + name=name_str, + type=DataObjectType.from_str(obj_type), + ) + ) + except (ValueError, AttributeError) as e: + logger.error( + f"Error processing row: {e}, row: {(row_schema, table_name, table_type)}" + ) + continue + + return result + + def create_view( + self, + view_name: TableName, + query_or_df: QueryOrDF, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + replace: bool = True, + materialized: bool = False, + materialized_properties: t.Optional[t.Dict[str, t.Any]] = None, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, + source_columns: t.Optional[t.List[str]] = None, + **create_kwargs: t.Any, + ) -> None: + target_table = exp.to_table(view_name) + effective_catalog = target_table.catalog or self.default_catalog + if effective_catalog and effective_catalog.lower() != "internal": + if materialized: + raise SQLMeshError( + f"Doris only supports creating materialized views in the 'internal' catalog, got '{effective_catalog}' " + f"for view '{target_table.sql(dialect=self.dialect)}'." + ) + logger.warning( + "Skipping view creation for '%s' because Doris does not support creating views in external catalog '%s'. " + "Configure `environment_catalog_mapping` in your SQLMesh configuration to route virtual environment views to the 'internal' catalog.", + target_table.sql(dialect=self.dialect), + effective_catalog, + ) + return + + if replace: + self.drop_view( + view_name, + ignore_if_not_exists=True, + materialized=materialized, + view_properties=view_properties, + ) + if not materialized: + return super().create_view( + view_name, + query_or_df, + target_columns_to_types, + replace=False, + materialized=False, + materialized_properties=materialized_properties, + table_description=table_description, + column_descriptions=column_descriptions, + view_properties=view_properties, + source_columns=source_columns, + **create_kwargs, + ) + self._create_materialized_view( + view_name, + query_or_df, + target_columns_to_types=target_columns_to_types, + materialized_properties=materialized_properties, + table_description=table_description, + column_descriptions=column_descriptions, + view_properties=view_properties, + source_columns=source_columns, + **create_kwargs, + ) + + def _create_materialized_view( + self, + view_name: TableName, + query_or_df: QueryOrDF, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + materialized_properties: t.Optional[t.Dict[str, t.Any]] = None, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, + source_columns: t.Optional[t.List[str]] = None, + **create_kwargs: t.Any, + ) -> None: + import pandas as pd + + query_or_df = self._native_df_to_pandas_df(query_or_df) + + if isinstance(query_or_df, pd.DataFrame): + values: t.List[t.Tuple[t.Any, ...]] = list( + query_or_df.itertuples(index=False, name=None) + ) + target_columns_to_types, source_columns = self._columns_to_types( + query_or_df, target_columns_to_types, source_columns + ) + if not target_columns_to_types: + raise SQLMeshError("columns_to_types must be provided for dataframes") + source_columns_to_types = get_source_columns_to_types( + target_columns_to_types, source_columns + ) + query_or_df = self._values_to_sql( + values, + source_columns_to_types, + batch_start=0, + batch_end=len(values), + ) + + source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( + query_or_df, + target_columns_to_types, + batch_size=0, + target_table=view_name, + source_columns=source_columns, + ) + if len(source_queries) != 1: + raise SQLMeshError("Only one source query is supported for creating materialized views") + + with source_queries[0] as query: + target_table = exp.to_table(view_name) + + schema: t.Union[exp.Table, exp.Schema] = target_table + if target_columns_to_types: + schema = self._build_schema_exp( + target_table, + target_columns_to_types, + column_descriptions, + is_view=True, + ) + + partitioned_by = None + if materialized_properties and materialized_properties.get("partitioned_by"): + partitioned_by = materialized_properties.get("partitioned_by") + elif create_kwargs.get("partitioned_by"): + partitioned_by = create_kwargs.get("partitioned_by") + elif view_properties and view_properties.get("partitioned_by"): + raw_vp = view_properties.get("partitioned_by") + if raw_vp is not None: + partitioned_by = raw_vp if isinstance(raw_vp, list) else [raw_vp] + + clustered_by = None + if materialized_properties and materialized_properties.get("clustered_by"): + clustered_by = materialized_properties.get("clustered_by") + elif create_kwargs.get("clustered_by"): + clustered_by = create_kwargs.get("clustered_by") + + properties_exp = self._build_table_properties_exp( + catalog_name=target_table.catalog, + table_properties=view_properties, + target_columns_to_types=target_columns_to_types, + table_description=( + table_description + if self.COMMENT_CREATION_VIEW.supports_schema_def and self.comments_enabled + else None + ), + partitioned_by=partitioned_by, + clustered_by=clustered_by, + partition_interval_unit=( + materialized_properties.get("partition_interval_unit") + if materialized_properties + else None + ), + table_kind="MATERIALIZED_VIEW", + ) + + is_sync_mv = bool( + view_properties + and ( + ( + self._extract_str_value(view_properties.get("materialized_type")) or "" + ).upper() + == "SYNC" + ) + ) + if is_sync_mv and isinstance(query, exp.Select): + query = query.copy() + aliased_selects: t.List[exp.Expr] = [] + for sel_expr in query.expressions: + if isinstance(sel_expr, exp.Column): + aliased_selects.append( + exp.alias_(sel_expr.copy(), f"mv_{sel_expr.name}", quoted=True) + ) + elif ( + isinstance(sel_expr, exp.Alias) + and isinstance(sel_expr.this, exp.Column) + and sel_expr.alias == sel_expr.this.name + ): + aliased_selects.append( + exp.alias_(sel_expr.this.copy(), f"mv_{sel_expr.alias}", quoted=True) + ) + else: + aliased_selects.append(sel_expr) + query.set("expressions", aliased_selects) + + if not target_table.db: + source_db: t.Optional[str] = None + if view_properties and view_properties.get("source_table"): + raw_src = view_properties["source_table"] + if isinstance(raw_src, exp.Column): + source_db = raw_src.db or None + elif isinstance(raw_src, exp.Table): + source_db = raw_src.db or None + else: + src_str = self._extract_str_value(raw_src) or str(raw_src) + source_db = exp.to_table(src_str, dialect=self.dialect).db or None + if not source_db: + first_tbl = query.find(exp.Table) + if first_tbl and first_tbl.db: + source_db = first_tbl.db + if source_db: + db_sql = exp.to_identifier(source_db, quoted=True).sql( + dialect=self.dialect, identify=True + ) + self.execute(f"USE {db_sql}") + + self.execute( + exp.Create( + this=schema, + kind="VIEW", + replace=False, + expression=query, + properties=properties_exp, + ), + quote_identifiers=self.QUOTE_IDENTIFIERS_IN_VIEWS, + ) + self._clear_data_object_cache(view_name) + + def drop_view( + self, + view_name: TableName, + ignore_if_not_exists: bool = True, + materialized: bool = False, + **kwargs: t.Any, + ) -> None: + """Drop view in Doris.""" + target_table = exp.to_table(view_name) + effective_catalog = target_table.catalog or self.default_catalog + if effective_catalog and effective_catalog.lower() != "internal": + return + + view_properties = kwargs.pop("view_properties", None) + if materialized and view_properties: + mat_type = self._extract_str_value(view_properties.get("materialized_type")) + source_table = view_properties.get("source_table") + if mat_type and mat_type.upper() == "SYNC" and source_table: + view_table_sql = target_table.sql(dialect=self.dialect, identify=True) + if isinstance(source_table, exp.Column): + source_tbl_exp = exp.table_( + source_table.name, + db=source_table.db or None, + catalog=source_table.catalog or None, + ) + elif isinstance(source_table, exp.Table): + source_tbl_exp = source_table + else: + source_table_str = self._extract_str_value(source_table) or str(source_table) + source_tbl_exp = exp.to_table(source_table_str, dialect=self.dialect) + if not target_table.db and source_tbl_exp.db: + db_sql = exp.to_identifier(source_tbl_exp.db, quoted=True).sql( + dialect=self.dialect, identify=True + ) + self.execute(f"USE {db_sql}") + source_table_sql = source_tbl_exp.sql(dialect=self.dialect, identify=True) + exists_clause = "IF EXISTS " if ignore_if_not_exists else "" + self.execute( + f"DROP MATERIALIZED VIEW {exists_clause}{view_table_sql} ON {source_table_sql}" + ) + self._clear_data_object_cache(view_name) + return + super().drop_view(view_name, ignore_if_not_exists, materialized, **kwargs) + + def _create_table_like( + self, + target_table_name: TableName, + source_table_name: TableName, + exists: bool = True, + **kwargs: t.Any, + ) -> None: + self.execute( + exp.Create( + this=exp.to_table(target_table_name), + kind="TABLE", + exists=exists, + properties=exp.Properties( + expressions=[ + exp.LikeProperty( + this=exp.to_table(source_table_name), + ), + ], + ), + ) + ) + + def _build_create_comment_table_exp( + self, table: exp.Table, table_comment: str, table_kind: str = "TABLE" + ) -> exp.Comment | str: + table_sql = table.sql(dialect=self.dialect, identify=True) + comment_sql = exp.Literal.string(self._truncate_table_comment(table_comment.strip())).sql( + dialect=self.dialect + ) + return f"ALTER TABLE {table_sql} MODIFY COMMENT {comment_sql}" + + def _build_create_comment_column_exp( + self, table: exp.Table, column_name: str, column_comment: str, table_kind: str = "TABLE" + ) -> exp.Comment | str: + table_sql = table.sql(dialect=self.dialect, identify=True) + col_sql = exp.to_identifier(column_name, quoted=True).sql( + dialect=self.dialect, identify=True + ) + comment_sql = exp.Literal.string(self._truncate_column_comment(column_comment.strip())).sql( + dialect=self.dialect + ) + return f"ALTER TABLE {table_sql} MODIFY COLUMN {col_sql} COMMENT {comment_sql}" + + def _truncate_table(self, table_name: TableName) -> None: + table = exp.to_table(table_name) + caught_err: t.Optional[Exception] = None + if not self._is_iceberg_table(catalog_name=table.catalog): + try: + super()._truncate_table(table_name) + return + except Exception as e: + err_msg = str(e).lower() + if "truncate" not in err_msg and "not support" not in err_msg: + raise e + caught_err = e + + cols_to_types = self.columns(table_name) + if cols_to_types: + valid_col_name = next( + ( + col + for col, dtype in cols_to_types.items() + if self._is_valid_olap_key_type(dtype) + ), + next(iter(cols_to_types)), + ) + target_col = exp.column(valid_col_name, quoted=True) + self.execute(exp.delete(table, where=target_col.is_(exp.null()).not_())) + self.execute(exp.delete(table, where=target_col.is_(exp.null()))) + return + if caught_err is not None: + raise caught_err + + def _insert_overwrite_by_condition( + self, + table_name: TableName, + source_queries: t.List[SourceQuery], + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + where: t.Optional[exp.Condition] = None, + insert_overwrite_strategy_override: t.Optional[InsertOverwriteStrategy] = None, + **kwargs: t.Any, + ) -> None: + table = exp.to_table(table_name) + if ( + insert_overwrite_strategy_override is None + and where is None + and self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=table.catalog, + table_properties=kwargs.get("table_properties"), + ) + ): + insert_overwrite_strategy_override = InsertOverwriteStrategy.INSERT_OVERWRITE + return super()._insert_overwrite_by_condition( + table_name, + source_queries, + target_columns_to_types=target_columns_to_types, + where=where, + insert_overwrite_strategy_override=insert_overwrite_strategy_override, + **kwargs, + ) + + def insert_overwrite_by_partition( + self, + table_name: TableName, + query_or_df: QueryOrDF, + partitioned_by: t.List[exp.Expr], + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + source_columns: t.Optional[t.List[str]] = None, + ) -> None: + parsed_exprs, _ = self._parse_partition_expressions(partitioned_by) + key_exprs: t.List[exp.Expr] = [] + for expr in parsed_exprs or partitioned_by: + if isinstance(expr, exp.Column): + key_exprs.append(expr) + else: + cols = list(expr.find_all(exp.Column)) + if cols: + key_exprs.extend(cols) + else: + key_exprs.append(expr) + super().insert_overwrite_by_partition( + table_name, + query_or_df, + key_exprs or partitioned_by, + target_columns_to_types=target_columns_to_types, + source_columns=source_columns, + ) + + def delete_from( + self, table_name: TableName, where: t.Optional[t.Union[str, exp.Expr]] = None + ) -> None: + """Delete from a table in Doris.""" + if not where or where == exp.true(): + self._truncate_table(table_name) + return + + if isinstance(where, str): + where = parse_one(where, read=self.dialect) + else: + where = where.copy() + + where = self._rewrite_between_in_where(where) + + subquery_info = self._extract_subquery_from_where(where) + if subquery_info is not None: + self._execute_delete_with_subquery(table_name, *subquery_info) + else: + super().delete_from(table_name, where) + + @staticmethod + def _rewrite_between_in_where(where: exp.Expr) -> exp.Expr: + if isinstance(where, exp.Between): + return exp.and_( + exp.GTE(this=where.this.copy(), expression=where.args["low"].copy()), + exp.LTE(this=where.this.copy(), expression=where.args["high"].copy()), + copy=False, + ) + for between in list(where.find_all(exp.Between)): + if between.find_ancestor(exp.Subquery, exp.Select, exp.Not): + continue + replacement: exp.Expr = exp.and_( + exp.GTE(this=between.this.copy(), expression=between.args["low"].copy()), + exp.LTE(this=between.this.copy(), expression=between.args["high"].copy()), + copy=False, + ) + if not isinstance(between.parent, (exp.Where, exp.Paren)): + replacement = exp.Paren(this=replacement) + between.replace(replacement) + return where + + def _extract_subquery_from_where( + self, where: exp.Expr + ) -> t.Optional[t.Tuple[t.List[exp.Expr], exp.Subquery, t.List[exp.Expr]]]: + """Extract positive IN (SELECT ...) subquery and any sibling AND conditions.""" + conjuncts: t.List[exp.Expr] = ( + list(where.flatten()) if isinstance(where, exp.And) else [where] + ) + target_in: t.Optional[exp.In] = None + subquery: t.Optional[exp.Subquery] = None + extra_conditions: t.List[exp.Expr] = [] + + for cond in conjuncts: + unwrapped_cond = cond.unnest() + if target_in is None and isinstance(unwrapped_cond, exp.In): + sq = self._get_in_subquery(unwrapped_cond) + if sq is not None: + target_in = unwrapped_cond + subquery = sq + continue + extra_conditions.append(cond) + + if target_in is None or subquery is None: + return None + + lhs = target_in.this + select_expr = subquery.this + key_exprs: t.List[exp.Expr] = [] + + lhs_concat = self._unwrap_concat_ws(lhs) + if lhs_concat is not None: + key_exprs = [e.copy() for e in lhs_concat.expressions[1:]] + if isinstance(select_expr, exp.Select) and len(select_expr.expressions) == 1: + first_sel = select_expr.expressions[0] + raw_sel = first_sel.this if isinstance(first_sel, exp.Alias) else first_sel + sel_concat = self._unwrap_concat_ws(raw_sel) + if sel_concat is not None: + select_expr = select_expr.copy() + select_expr.set("expressions", [e.copy() for e in sel_concat.expressions[1:]]) + subquery = exp.Subquery(this=select_expr) + elif isinstance(lhs, exp.Tuple): + key_exprs = [e.copy() for e in lhs.expressions] + else: + key_exprs = [lhs.copy()] + + if not key_exprs or not all(isinstance(k, exp.Column) for k in key_exprs): + return None + + return key_exprs, subquery, extra_conditions + + @staticmethod + def _unwrap_concat_ws(expr: exp.Expr) -> t.Optional[exp.Expr]: + curr = expr.unnest() + if isinstance(curr, exp.Column) and isinstance(curr.this, exp.Expr): + curr = curr.this.unnest() + if isinstance(curr, exp.ConcatWs) or ( + isinstance(curr, exp.Anonymous) and curr.name.upper() == "CONCAT_WS" + ): + return curr + return None + + @staticmethod + def _get_in_subquery(in_expr: exp.In) -> t.Optional[exp.Subquery]: + query = in_expr.args.get("query") + if isinstance(query, exp.Subquery): + return query.copy() + if isinstance(query, exp.Select): + return exp.Subquery(this=query.copy()) + if len(in_expr.expressions) == 1: + first = in_expr.expressions[0] + if isinstance(first, exp.Subquery): + return first.copy() + if isinstance(first, exp.Select): + return exp.Subquery(this=first.copy()) + return None + + def _execute_delete_with_subquery( + self, + table_name: TableName, + key_exprs: t.List[exp.Expr], + subquery: exp.Subquery, + extra_conditions: t.List[exp.Expr], + ) -> None: + """Execute DELETE FROM ... USING (SELECT ...) in Doris, with literal batch fallback for DUPLICATE KEY/Iceberg.""" + target_table = exp.to_table(table_name) + if self._is_iceberg_table(catalog_name=target_table.catalog): + self._delete_by_literal_keys(table_name, key_exprs, subquery, extra_conditions) + return + + join_conditions: t.List[exp.Expr] = [] + for idx, col in enumerate(key_exprs): + col_name = col.name + t1_col = exp.column(col_name, table="_t1") + t2_col_name = col_name + if isinstance(subquery.this, exp.Select) and idx < len(subquery.this.expressions): + sel_item = subquery.this.expressions[idx] + if sel_item.alias_or_name: + t2_col_name = sel_item.alias_or_name + t2_col = exp.column(t2_col_name, table="_t2") + join_conditions.append(t1_col.eq(t2_col)) + + for extra in extra_conditions: + qualified_extra = extra.copy() + for col in qualified_extra.find_all(exp.Column): + if not col.find_ancestor(exp.Subquery, exp.Select) and not col.table: + col.set("table", exp.to_identifier("_t1")) + join_conditions.append(qualified_extra) + + where_expr = exp.and_(*join_conditions) + + target_sql = f"{self._to_sql(target_table)} AS `_t1`" + subquery_with_alias = exp.alias_(subquery.copy(), "_t2", table=True) + subquery_sql = self._to_sql(subquery_with_alias) + where_sql = self._to_sql(where_expr) + + delete_sql = f"DELETE FROM {target_sql} USING {subquery_sql} WHERE {where_sql}" + try: + self.execute(delete_sql) + except Exception as e: + err_msg = str(e).lower() + if any( + marker in err_msg + for marker in ( + "unique key", + "olaptable", + "olap table", + "not support", + "only support", + "unsupported", + ) + ): + logger.info( + "DELETE ... USING not supported on table '%s' (%s); falling back to batched literal DELETE.", + table_name, + e, + ) + self._delete_by_literal_keys(table_name, key_exprs, subquery, extra_conditions) + else: + raise + + def _delete_by_literal_keys( + self, + table_name: TableName, + key_exprs: t.List[exp.Expr], + subquery: t.Union[exp.Subquery, exp.Select], + extra_conditions: t.Optional[t.List[exp.Expr]] = None, + ) -> None: + """Fallback for DELETE on DUPLICATE KEY / AGGREGATE KEY / Iceberg tables that prohibits OR predicates.""" + select_query = ( + subquery.this.copy() if isinstance(subquery, exp.Subquery) else subquery.copy() + ) + if isinstance(select_query, exp.Select): + select_query = select_query.distinct() + rows = self.fetchall(select_query) + if not rows: + return + + import pandas as pd + + extra_conditions = extra_conditions or [] + cols = [exp.column(k.name, quoted=True) for k in key_exprs] + batch_size = self.DEFAULT_BATCH_SIZE + + def _sanitize(val: t.Any) -> t.Any: + if val is None: + return None + try: + if pd.isna(val): + return None + except (TypeError, ValueError): + pass + return val + + sanitized_rows = [tuple(_sanitize(v) for v in row) for row in rows] + + if len(cols) == 1: + col = cols[0] + non_null_vals = [exp.convert(r[0]) for r in sanitized_rows if r[0] is not None] + has_null = any(r[0] is None for r in sanitized_rows) + + for i in range(0, len(non_null_vals), batch_size): + batch = non_null_vals[i : i + batch_size] + predicates: t.List[exp.Expr] = [col.isin(*batch), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*predicates))) + + if has_null: + predicates = [col.is_(exp.null()), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*predicates))) + return + + # Composite key (N >= 2): dynamically select the column index to batch via IN (...) + # that minimizes the number of distinct prefix groups, breaking ties in favor of the last column. + batch_idx = min( + range(len(cols)), + key=lambda idx: (len({row[:idx] + row[idx + 1 :] for row in sanitized_rows}), -idx), + ) + grouped: t.Dict[t.Tuple[t.Any, ...], t.List[t.Any]] = defaultdict(list) + for row in sanitized_rows: + prefix_key = row[:batch_idx] + row[batch_idx + 1 :] + grouped[prefix_key].append(row[batch_idx]) + + batch_col = cols[batch_idx] + prefix_cols = cols[:batch_idx] + cols[batch_idx + 1 :] + + for prefix_vals, batch_vals in grouped.items(): + prefix_preds: t.List[exp.Expr] = [ + p_col.eq(exp.convert(p_val)) if p_val is not None else p_col.is_(exp.null()) + for p_col, p_val in zip(prefix_cols, prefix_vals) + ] + non_null_batch = [exp.convert(v) for v in batch_vals if v is not None] + has_null_batch = any(v is None for v in batch_vals) + + for i in range(0, len(non_null_batch), batch_size): + batch = non_null_batch[i : i + batch_size] + preds = [*prefix_preds, batch_col.isin(*batch), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*preds))) + + if has_null_batch: + preds = [*prefix_preds, batch_col.is_(exp.null()), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*preds))) + + def adjust_physical_properties_for_incremental( + self, + physical_properties: t.Dict[str, t.Any], + *, + requires_delete_capable_table: bool, + unique_key: t.Optional[t.List[exp.Expr]], + model_name: str, + ) -> t.Dict[str, t.Any]: + """Promote ``unique_key`` from ``INCREMENTAL_BY_UNIQUE_KEY`` models into ``physical_properties`` + when no explicit OLAP key model (``unique_key``, ``duplicate_key``, ``aggregate_key``) is set. + """ + if ( + unique_key + and "unique_key" not in physical_properties + and "duplicate_key" not in physical_properties + and "aggregate_key" not in physical_properties + ): + physical_properties["unique_key"] = ( + unique_key[0] if len(unique_key) == 1 else exp.Tuple(expressions=list(unique_key)) + ) + return physical_properties + + def _create_table_from_columns( + self, + table_name: TableName, + target_columns_to_types: t.Dict[str, exp.DataType], + primary_key: t.Optional[t.Tuple[str, ...]] = None, + exists: bool = True, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + **kwargs: t.Any, + ) -> None: + table_properties = dict(kwargs.get("table_properties") or {}) + + if primary_key and "unique_key" not in table_properties: + table_properties["unique_key"] = exp.Tuple( + expressions=[exp.to_column(col) for col in primary_key] + ) + + kwargs["table_properties"] = table_properties + + super()._create_table_from_columns( + table_name=table_name, + target_columns_to_types=target_columns_to_types, + primary_key=None, + exists=exists, + table_description=table_description, + column_descriptions=column_descriptions, + **kwargs, + ) + + def _build_create_table_exp( + self, + table_name_or_schema: t.Union[exp.Schema, TableName], + expression: t.Optional[exp.Expr], + exists: bool = True, + replace: bool = False, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + table_description: t.Optional[str] = None, + table_kind: t.Optional[str] = None, + **kwargs: t.Any, + ) -> exp.Create: + kwargs.setdefault("table_properties", None) + kwargs["is_schema_ddl"] = isinstance(table_name_or_schema, exp.Schema) + + catalog_name = None + if isinstance(table_name_or_schema, exp.Schema): + if isinstance(table_name_or_schema.this, exp.Table): + catalog_name = table_name_or_schema.this.catalog + else: + catalog_name = exp.to_table(table_name_or_schema).catalog + + is_iceberg = self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=catalog_name, + table_properties=kwargs.get("table_properties"), + ) + + if not is_iceberg and target_columns_to_types: + reordered_cols = self._reorder_olap_columns(target_columns_to_types, kwargs) + if list(reordered_cols) != list(target_columns_to_types): + target_columns_to_types = reordered_cols + if isinstance(table_name_or_schema, exp.Schema): + col_def_map = { + col_def.name: col_def + for col_def in table_name_or_schema.expressions + if isinstance(col_def, exp.ColumnDef) + } + other_exprs = [ + e + for e in table_name_or_schema.expressions + if not isinstance(e, exp.ColumnDef) + ] + if len(col_def_map) == len(target_columns_to_types) and all( + c in col_def_map for c in target_columns_to_types + ): + table_name_or_schema = table_name_or_schema.copy() + table_name_or_schema.set( + "expressions", + [col_def_map[c] for c in target_columns_to_types] + other_exprs, + ) + if expression is not None and isinstance(expression, exp.Query): + expression = self._order_projections_and_filter( + expression, target_columns_to_types + ) + + return super()._build_create_table_exp( + table_name_or_schema, + expression=expression, + exists=exists, + replace=replace, + target_columns_to_types=target_columns_to_types, + table_description=table_description, + table_kind=table_kind, + **kwargs, + ) + + def _reorder_olap_columns( + self, + target_columns_to_types: t.Dict[str, exp.DataType], + kwargs: t.Dict[str, t.Any], + ) -> t.Dict[str, exp.DataType]: + table_props = kwargs.get("table_properties") or {} + explicit_key = ( + table_props.get("unique_key") + or kwargs.get("unique_key") + or table_props.get("duplicate_key") + or table_props.get("aggregate_key") + ) + col_names = list(target_columns_to_types) + if not col_names: + return target_columns_to_types + + if explicit_key: + key_cols = [ + k for k in self._extract_column_names(explicit_key) if k in target_columns_to_types + ] + if key_cols and col_names[: len(key_cols)] != key_cols: + new_order = key_cols + [c for c in col_names if c not in key_cols] + return {c: target_columns_to_types[c] for c in new_order} + return target_columns_to_types + + first_dtype = target_columns_to_types[col_names[0]] + if not self._is_valid_olap_key_type(first_dtype): + valid_col = next( + (c for c in col_names if self._is_valid_olap_key_type(target_columns_to_types[c])), + None, + ) + if valid_col: + new_order = [valid_col] + [c for c in col_names if c != valid_col] + return {c: target_columns_to_types[c] for c in new_order} + + return target_columns_to_types + + def _parse_partition_expressions( + self, partitioned_by: t.List[exp.Expr] + ) -> t.Tuple[t.List[exp.Expr], t.Optional[str]]: + """Parse partition expressions and extract partition kind and normalized expressions.""" + parsed_partitioned_by: t.List[exp.Expr] = [] + partition_kind: t.Optional[str] = None + + for expr in partitioned_by: + if isinstance(expr, exp.Anonymous) and expr.this: + func_name = str(expr.this).upper() + if func_name in ("RANGE", "LIST"): + partition_kind = func_name + for arg in expr.expressions: + if isinstance(arg, exp.Column): + parsed_partitioned_by.append(arg) + elif isinstance(arg, exp.Func) or isinstance(arg, exp.Anonymous): + parsed_partitioned_by.append(arg) + else: + parsed_partitioned_by.append(exp.to_column(str(arg))) + continue + + if isinstance(expr, exp.Literal) and getattr(expr, "is_string", False): + text = str(expr.this) + match = re.match(r"^\s*(RANGE|LIST)\s*\((.+)\)\s*$", text, flags=re.IGNORECASE) + if match: + partition_kind = match.group(1).upper() + inner = match.group(2).strip() + if ";" in inner: + raise SQLMeshError(f"Invalid partition expression: {text}") + try: + parsed_stmts = parse(f"SELECT {inner}", read=self.dialect) + if ( + len(parsed_stmts) == 1 + and isinstance(parsed_stmts[0], exp.Select) + and set(parsed_stmts[0].args.keys()) <= {"expressions"} + ): + for sel_expr in parsed_stmts[0].expressions: + parsed_partitioned_by.append(sel_expr) + continue + except Exception: + pass + inner_cols = [c.strip().strip("`") for c in inner.split(",") if c.strip()] + for col in inner_cols: + parsed_partitioned_by.append(exp.to_column(col)) + continue + + parsed_partitioned_by.append(expr) + + return parsed_partitioned_by, partition_kind + + def _build_partitioned_by_exp( + self, + partitioned_by: t.List[exp.Expr], + *, + partition_interval_unit: t.Optional[IntervalUnit] = None, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + catalog_name: t.Optional[str] = None, + **kwargs: t.Any, + ) -> t.Optional[ + t.Union[ + exp.PartitionedByProperty, + exp.PartitionByRangeProperty, + exp.PartitionByListProperty, + exp.Property, + ] + ]: + """Build Doris partitioning expression for OLAP (manual, dynamic, AUTO PARTITION) and Iceberg tables.""" + is_iceberg = bool(kwargs.get("is_iceberg")) + partitions = kwargs.get("partitions") + auto_partition = kwargs.get("auto_partition") + dynamic_partition_enabled = kwargs.get("dynamic_partition_enabled") + + partitioned_by, partition_kind = self._parse_partition_expressions(partitioned_by) + + if is_iceberg: + return exp.PartitionByListProperty( + partition_expressions=partitioned_by, + create_expressions=[], + ) + + is_auto = False + if auto_partition is not None: + auto_str = self._extract_str_value(auto_partition) + is_auto = str(auto_str).strip().lower() in ("true", "1") + + # On Doris OLAP tables, if no explicit partitions, auto_partition, or dynamic_partition.enable + # are configured, omit PARTITION BY so Doris does not create a 0-partition table that rejects inserts. + if not partitions and not is_auto and not dynamic_partition_enabled: + return None + + def to_raw_sql(expr: t.Union[exp.Literal, exp.Var, str, t.Any]) -> exp.Var: + if isinstance(expr, exp.Literal): + raw_text = str(expr.this) + elif isinstance(expr, exp.Var): + raw_text = str(expr.this) + elif isinstance(expr, str): + raw_text = expr + else: + return expr + if ";" in raw_text or "--" in raw_text or "/*" in raw_text: + raise SQLMeshError(f"Invalid partition definition expression: {raw_text}") + return exp.Var(this=raw_text, quoted=False) + + create_expressions: t.Optional[t.List[exp.Expr]] = None + if partitions: + if isinstance(partitions, exp.Tuple): + create_expressions = [to_raw_sql(e) for e in partitions.expressions] + else: + create_expressions = [to_raw_sql(partitions)] + elif is_auto or dynamic_partition_enabled: + create_expressions = [] + + inferred_list = False + if partition_kind is None and create_expressions: + try: + texts = [getattr(e, "this", "").upper() for e in create_expressions] + inferred_list = any("VALUES IN" in text for text in texts) + except Exception: + inferred_list = False + if partition_kind: + is_list = str(partition_kind).upper() == "LIST" + else: + is_list = inferred_list + + partition_prop: t.Union[exp.PartitionByListProperty, exp.PartitionByRangeProperty] + if is_list: + partition_prop = exp.PartitionByListProperty( + partition_expressions=partitioned_by, + create_expressions=create_expressions, + ) + else: + partition_prop = exp.PartitionByRangeProperty( + partition_expressions=partitioned_by, + create_expressions=create_expressions, + ) + + if is_auto: + return exp.SetConfigProperty(this=exp.var(f"AUTO {self._to_sql(partition_prop)}")) + + return partition_prop + + def _build_table_properties_exp( + self, + catalog_name: t.Optional[str] = None, + table_format: t.Optional[str] = None, + storage_format: t.Optional[str] = None, + partitioned_by: t.Optional[t.List[exp.Expr]] = None, + partition_interval_unit: t.Optional[IntervalUnit] = None, + clustered_by: t.Optional[t.List[exp.Expr]] = None, + table_properties: t.Optional[t.Dict[str, t.Any]] = None, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + table_description: t.Optional[str] = None, + table_kind: t.Optional[str] = None, + **kwargs: t.Any, + ) -> t.Optional[exp.Properties]: + """Creates a SQLGlot table properties expression for Doris DDL.""" + properties: t.List[exp.Expr] = [] + is_materialized_view = table_kind == "MATERIALIZED_VIEW" + is_schema_ddl = bool(kwargs.get("is_schema_ddl")) + + if is_materialized_view: + properties.append(exp.MaterializedProperty()) + + table_properties_copy = dict(table_properties) if table_properties else {} + is_iceberg = not is_materialized_view and self._is_iceberg_table( + table_format=table_format, + catalog_name=catalog_name, + table_properties=table_properties_copy, + ) + + # Pop Iceberg engine property if present so it is not duplicated in PROPERTIES (...) + table_properties_copy.pop("engine", None) + if is_iceberg and is_schema_ddl: + properties.append(exp.EngineProperty(this=exp.var("iceberg"))) + + if is_materialized_view: + build_val = table_properties_copy.pop("build", None) + if build_val is not None: + build_upper = self._to_upper_str(build_val) + if build_upper not in ("IMMEDIATE", "DEFERRED"): + raise SQLMeshError( + f"Invalid materialized view build mode: {build_upper}. Expected IMMEDIATE or DEFERRED." + ) + properties.append(exp.BuildProperty(this=exp.Var(this=build_upper))) + + refresh_val = table_properties_copy.pop("refresh", None) + refresh_trigger_val = table_properties_copy.pop("refresh_trigger", None) + if refresh_val is not None or refresh_trigger_val is not None: + method_var: exp.Var = exp.Var(this="AUTO") + kind_str: t.Optional[str] = None + every_lit: t.Optional[exp.Expr] = None + unit_var: t.Optional[exp.Var] = None + starts_lit: t.Optional[exp.Literal] = None + + if isinstance(refresh_val, exp.Tuple): + method_var, kind_str, every_lit, unit_var, starts_lit = ( + self._parse_refresh_tuple(refresh_val) + ) + else: + raw_refresh = ( + self._extract_str_value(refresh_val) if refresh_val is not None else "AUTO" + ) + raw_refresh_str = (raw_refresh or "AUTO").strip() + if " " in raw_refresh_str: + parts = raw_refresh_str.split(None, 1) + method_upper = parts[0].upper() + if method_upper in ("COMPLETE", "AUTO", "FAST", "FORCE", "NEVER"): + method_var = exp.Var(this=method_upper) + k, e, u, s = self._parse_trigger_string(parts[1]) + kind_str, every_lit, unit_var, starts_lit = k, e, u, s + else: + k, e, u, s = self._parse_trigger_string(raw_refresh_str) + kind_str, every_lit, unit_var, starts_lit = k, e, u, s + else: + method_upper = raw_refresh_str.upper() + if method_upper in ("MANUAL", "COMMIT"): + kind_str = method_upper + elif method_upper in ("COMPLETE", "AUTO", "FAST", "FORCE", "NEVER"): + method_var = exp.Var(this=method_upper) + else: + raise SQLMeshError( + f"Invalid materialized view refresh method: {method_upper}. Expected one of: COMPLETE, AUTO, FAST, FORCE, NEVER, MANUAL, COMMIT." + ) + + if refresh_trigger_val is not None: + trigger_text = self._extract_str_value(refresh_trigger_val) or "" + k, e, u, s = self._parse_trigger_string(trigger_text) + kind_str, every_lit, unit_var, starts_lit = k, e, u, s + + properties.append( + exp.RefreshTriggerProperty( + method=method_var, + kind=kind_str, + every=every_lit, + unit=unit_var, + starts=starts_lit, + ) + ) + + # Pop key models and order_by from table_properties_copy + unique_key = table_properties_copy.pop("unique_key", None) or kwargs.get("unique_key") + duplicate_key = table_properties_copy.pop("duplicate_key", None) + aggregate_key = table_properties_copy.pop("aggregate_key", None) + order_by = table_properties_copy.pop("order_by", None) + + unique_key_cols = self._extract_column_names(unique_key) + duplicate_key_cols = self._extract_column_names(duplicate_key) + aggregate_key_cols = self._extract_column_names(aggregate_key) + order_by_cols = self._extract_column_names(order_by) + + if not is_iceberg: + if unique_key_cols: + properties.append( + exp.UniqueKeyProperty(expressions=[exp.to_column(k) for k in unique_key_cols]) + ) + elif duplicate_key_cols: + properties.append( + exp.DuplicateKeyProperty( + expressions=[exp.to_column(k) for k in duplicate_key_cols] + ) + ) + elif aggregate_key_cols: + agg_cols_sql = ", ".join(self._to_sql(exp.to_column(k)) for k in aggregate_key_cols) + properties.append( + exp.SetConfigProperty(this=exp.var(f"AGGREGATE KEY ({agg_cols_sql})")) + ) + + if order_by_cols and not is_materialized_view: + order_cols_sql = ", ".join(self._to_sql(exp.to_column(k)) for k in order_by_cols) + properties.append( + exp.SetConfigProperty(this=exp.var(f"ORDER BY ({order_cols_sql})")) + ) + + if table_description: + properties.append( + exp.SchemaCommentProperty( + this=exp.Literal.string(self._truncate_table_comment(table_description)) + ) + ) + + # Pop partitioning-related properties + partitions = table_properties_copy.pop("partitions", None) + auto_partition = table_properties_copy.pop("auto_partition", None) + if not partitioned_by and "partitioned_by" in table_properties_copy: + raw_part_by = table_properties_copy.pop("partitioned_by") + if raw_part_by is not None: + partitioned_by = raw_part_by if isinstance(raw_part_by, list) else [raw_part_by] + else: + table_properties_copy.pop("partitioned_by", None) + + add_partition = True + if partitioned_by: + parsed_part_exprs, _ = self._parse_partition_expressions(partitioned_by) + if unique_key_cols and not is_materialized_view and not is_iceberg: + key_cols_set = set(unique_key_cols) + partition_cols = set() + for expr in parsed_part_exprs: + for col_node in expr.find_all(exp.Column): + partition_cols.add(col_node.name) + not_in_key = partition_cols - key_cols_set + if not_in_key: + logger.warning( + f"[Doris] UNIQUE KEY model: partitioned_by columns {not_in_key} not in key_cols {key_cols_set}, skip PARTITION BY." + ) + add_partition = False + + if add_partition: + if is_materialized_view: + if parsed_part_exprs: + properties.append( + exp.PartitionedByProperty( + this=exp.Schema(expressions=parsed_part_exprs), + ) + ) + else: + dyn_enable = self._extract_str_value( + table_properties_copy.get("dynamic_partition.enable") + ) + dynamic_partition_enabled = bool( + dyn_enable and dyn_enable.strip().lower() in ("true", "1") + ) + partition_expr = self._build_partitioned_by_exp( + partitioned_by, + partition_interval_unit=partition_interval_unit, + target_columns_to_types=target_columns_to_types, + catalog_name=catalog_name, + partitions=partitions, + auto_partition=auto_partition, + dynamic_partition_enabled=dynamic_partition_enabled, + is_iceberg=is_iceberg, + ) + if partition_expr: + properties.append(partition_expr) + + # Pop distribution / bucketing properties so they never leak into PROPERTIES (...) + distributed_by = table_properties_copy.pop("distributed_by", None) + top_level_buckets = table_properties_copy.pop("buckets", None) + prop_clustered_by = table_properties_copy.pop("clustered_by", None) + effective_clustered_by = clustered_by or ( + [exp.to_column(c) for c in self._extract_column_names(prop_clustered_by)] + if prop_clustered_by + else None + ) + + if not is_iceberg: + dist_prop = self._build_distributed_by_exp( + distributed_by=distributed_by, + clustered_by=effective_clustered_by, + unique_key_cols=unique_key_cols if not is_materialized_view else [], + fallback_key_cols=( + (duplicate_key_cols or aggregate_key_cols) if not is_materialized_view else [] + ), + top_level_buckets=top_level_buckets, + ) + if dist_prop is not None: + properties.append(dist_prop) + + # Default replication_num = 1 for single-BE dev/CI when not explicitly configured on OLAP tables + # Note: Only emit explicit table_properties_copy entries so unit test expectations remain clean. + if table_properties_copy: + properties.extend(self._properties_to_expressions(table_properties_copy)) + + if properties: + return exp.Properties(expressions=properties) + return None + + def _to_upper_str(self, val: t.Any) -> str: + if isinstance(val, exp.Literal): + return str(val.this).upper() + if isinstance(val, exp.Identifier): + return str(val.this).upper() + if isinstance(val, str): + return val.upper() + if isinstance(val, exp.Expr): + return val.sql(dialect=self.dialect).upper() + return str(val).upper() + + def _parse_refresh_tuple( + self, + tpl: exp.Tuple, + ) -> t.Tuple[ + exp.Var, + t.Optional[str], + t.Optional[exp.Expr], + t.Optional[exp.Var], + t.Optional[exp.Literal], + ]: + method_var: exp.Var = exp.Var(this="AUTO") + kind_str: t.Optional[str] = None + every_lit: t.Optional[exp.Expr] = None + unit_var: t.Optional[exp.Var] = None + starts_lit: t.Optional[exp.Literal] = None + for e in tpl.expressions: + if ( + isinstance(e, exp.EQ) + and isinstance(e.this, exp.Column) + and isinstance(e.this.this, exp.Identifier) + ): + key = str(e.this.this.this).strip('"').lower() + val = e.expression + if key == "method": + method_upper = self._to_upper_str(val) + if method_upper not in ("COMPLETE", "AUTO", "FAST", "FORCE", "NEVER"): + raise SQLMeshError( + f"Invalid materialized view refresh method: {method_upper}. Expected one of: COMPLETE, AUTO, FAST, FORCE, NEVER." + ) + method_var = exp.Var(this=method_upper) + elif key == "kind": + kind_upper = self._to_upper_str(val) + if kind_upper not in ("MANUAL", "COMMIT", "SCHEDULE"): + raise SQLMeshError( + f"Invalid materialized view refresh kind: {kind_upper}. Expected one of: MANUAL, COMMIT, SCHEDULE." + ) + kind_str = kind_upper + elif key == "every": + if isinstance(val, exp.Literal): + val_str = str(val.this).strip() + if not val_str.isdigit(): + raise SQLMeshError( + f"Invalid materialized view refresh every interval: {val_str}. Expected a positive integer." + ) + every_lit = exp.Literal.number(int(val_str)) + elif isinstance(val, (int, float)): + every_lit = exp.Literal.number(int(val)) + elif key == "unit": + unit_upper = self._to_upper_str(val) + if not re.match(r"^[A-Z]+$", unit_upper): + raise SQLMeshError(f"Invalid materialized view refresh unit: {unit_upper}") + unit_var = exp.Var(this=unit_upper) + elif key == "starts": + if isinstance(val, exp.Literal): + starts_lit = exp.Literal.string(str(val.this)) + else: + starts_lit = exp.Literal.string(str(val)) + return method_var, kind_str, every_lit, unit_var, starts_lit + + @staticmethod + def _parse_trigger_string( + trigger_text: str, + ) -> t.Tuple[ + t.Optional[str], + t.Optional[exp.Expr], + t.Optional[exp.Var], + t.Optional[exp.Literal], + ]: + text = trigger_text.strip() + if text.upper().startswith("ON "): + text = text[3:].strip() + kind = None + every_lit: t.Optional[exp.Expr] = None + unit_var = None + starts_lit = None + + m_kind = re.match(r"^(MANUAL|COMMIT|SCHEDULE)\b", text, flags=re.IGNORECASE) + if m_kind: + kind = m_kind.group(1).upper() + m_every = re.search(r"\bEVERY\s+(\d+)\s+([A-Za-z]+)", text, flags=re.IGNORECASE) + if m_every: + every_lit = exp.Literal.number(int(m_every.group(1))) + unit_var = exp.Var(this=m_every.group(2).upper()) + m_starts = re.search(r"\bSTARTS\s+'([^']+)'", text, flags=re.IGNORECASE) + if m_starts: + starts_lit = exp.Literal.string(m_starts.group(1)) + return kind, every_lit, unit_var, starts_lit + + def _resolve_buckets(self, raw_buckets: t.Any) -> t.Optional[exp.Expr]: + if raw_buckets is None: + return None + if isinstance(raw_buckets, exp.Literal): + val_str = str(raw_buckets.this).strip() + if val_str.upper() == "AUTO": + return exp.Var(this="AUTO") + if val_str.isdigit(): + return exp.Literal.number(int(val_str)) + if isinstance(raw_buckets, int): + return exp.Literal.number(raw_buckets) + val_str = str(self._extract_str_value(raw_buckets) or "").strip() + if val_str.upper() == "AUTO": + return exp.Var(this="AUTO") + if val_str.isdigit(): + return exp.Literal.number(int(val_str)) + raise SQLMeshError( + f"Invalid buckets value for Doris DISTRIBUTED BY: {raw_buckets}. Expected a positive integer or 'AUTO'." + ) + + def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> str: + if expression.find(exp.DataType): + expression = expression.copy() + for dtype in reversed(list(expression.find_all(exp.DataType))): + if dtype.is_type("int128"): + dtype.replace(exp.DataType.build("LARGEINT", udt=True)) + elif dtype.is_type("struct"): + new_fields: t.List[exp.Expr] = [] + for field in dtype.expressions: + if isinstance(field, exp.ColumnDef) and field.args.get("kind"): + field_id = field.this.copy() + if quote and isinstance(field_id, exp.Identifier): + field_id.set("quoted", True) + field_name_sql = super()._to_sql(field_id, quote=False, **kwargs) + field_kind_sql = self._to_sql(field.args["kind"], quote=quote, **kwargs) + new_fields.append(exp.var(f"{field_name_sql}: {field_kind_sql}")) + else: + new_fields.append(field) + dtype.set("expressions", new_fields) + return super()._to_sql(expression, quote=quote, **kwargs) + + def _build_distributed_by_exp( + self, + distributed_by: t.Any = None, + clustered_by: t.Optional[t.List[exp.Expr]] = None, + unique_key_cols: t.Optional[t.List[str]] = None, + fallback_key_cols: t.Optional[t.List[str]] = None, + top_level_buckets: t.Any = None, + ) -> t.Optional[exp.DistributedByProperty]: + default_buckets_expr = self._resolve_buckets(top_level_buckets) + + if distributed_by is not None: + distributed_info: t.Dict[str, t.Any] = {} + + # Case 1: Function expression such as HASH(a, b) + if isinstance(distributed_by, (exp.Anonymous, exp.Func)) and not isinstance( + distributed_by, (exp.Tuple, exp.Array, exp.Paren) + ): + kind_upper = distributed_by.name.upper() + if kind_upper not in ("HASH", "RANDOM"): + raise SQLMeshError( + f"Invalid distributed_by kind: {kind_upper}. Expected HASH or RANDOM." + ) + cols = [ + exp.to_column(c) for c in self._extract_column_names(distributed_by.expressions) + ] + return exp.DistributedByProperty( + kind=exp.Var(this=kind_upper), + expressions=cols, + buckets=default_buckets_expr or exp.Literal.number(10), + order=None, + ) + + # Case 2: String shorthand such as "HASH(id) BUCKETS 10", "RANDOM BUCKETS AUTO", or "id" + if isinstance(distributed_by, str) or ( + isinstance(distributed_by, exp.Literal) + and getattr(distributed_by, "is_string", False) + ): + text = ( + str(distributed_by.this).strip() + if isinstance(distributed_by, exp.Literal) + else distributed_by.strip() + ) + m = re.match( + r"^(HASH|RANDOM)\s*(?:\(([^)]*)\))?\s*(?:BUCKETS\s+(AUTO|\d+))?$", + text, + flags=re.IGNORECASE, + ) + if m: + kind_upper = m.group(1).upper() + cols_str = m.group(2) + b_str = m.group(3) + parsed_cols = ( + [ + exp.to_column(c.strip().strip("`")) + for c in cols_str.split(",") + if c.strip() + ] + if cols_str + else None + ) + b_expr = self._resolve_buckets(b_str) if b_str else default_buckets_expr + return exp.DistributedByProperty( + kind=exp.Var(this=kind_upper), + expressions=parsed_cols, + buckets=b_expr, + order=None, + ) + cols = [exp.to_column(c) for c in self._extract_column_names(text)] + if cols: + return exp.DistributedByProperty( + kind=exp.Var(this="HASH"), + expressions=cols, + buckets=default_buckets_expr or exp.Literal.number(10), + order=None, + ) + + # Case 3: Tuple or Array (either key=value pairs or list of column names from dbt-doris) + if isinstance(distributed_by, (exp.Tuple, exp.Array)): + has_eq = any(isinstance(e, exp.EQ) for e in distributed_by.expressions) + if has_eq: + for expr in distributed_by.expressions: + if isinstance(expr, exp.EQ) and hasattr(expr.this, "this"): + key = str(expr.this.this).strip('"').strip("'").lower() + if isinstance(expr.expression, exp.Literal): + distributed_info[key] = expr.expression.this + elif isinstance(expr.expression, (exp.Array, exp.Tuple)): + distributed_info[key] = self._extract_column_names(expr.expression) + else: + distributed_info[key] = self._extract_str_value(expr.expression) + else: + cols = [exp.to_column(c) for c in self._extract_column_names(distributed_by)] + if cols: + return exp.DistributedByProperty( + kind=exp.Var(this="HASH"), + expressions=cols, + buckets=default_buckets_expr or exp.Literal.number(10), + order=None, + ) + elif isinstance(distributed_by, exp.Paren) and isinstance(distributed_by.this, exp.EQ): + expr = distributed_by.this + if hasattr(expr.this, "this"): + key = str(expr.this.this).strip('"').strip("'").lower() + if isinstance(expr.expression, exp.Literal): + distributed_info[key] = expr.expression.this + else: + distributed_info[key] = self._extract_str_value(expr.expression) + elif isinstance(distributed_by, dict): + distributed_info = distributed_by + + if distributed_info: + raw_kind = distributed_info.get("kind") + expressions = distributed_info.get("expressions") + buckets = distributed_info.get("buckets", top_level_buckets) + + if raw_kind: + kind_upper = str(raw_kind).strip().upper() + if kind_upper not in ("HASH", "RANDOM"): + raise SQLMeshError( + f"Invalid distributed_by kind: {kind_upper}. Expected HASH or RANDOM." + ) + buckets_expr = self._resolve_buckets(buckets) + + expressions_list = None + if expressions: + col_names = self._extract_column_names(expressions) + expressions_list = [exp.to_column(c) for c in col_names] + + return exp.DistributedByProperty( + kind=exp.Var(this=kind_upper), + expressions=expressions_list, + buckets=buckets_expr, + order=None, + ) + + # Fallback 1: Top-level clustered_by (or physical_properties clustered_by) + if clustered_by: + cols = [ + exp.to_column(c) for expr in clustered_by for c in self._extract_column_names(expr) + ] + if cols: + return exp.DistributedByProperty( + kind=exp.Var(this="HASH"), + expressions=cols, + buckets=default_buckets_expr or exp.Literal.number(10), + order=None, + ) + + # Fallback 2: Default DISTRIBUTED BY HASH () BUCKETS + if unique_key_cols: + first_col = unique_key_cols[0] + logger.info( + f"[Doris] Adding default distributed_by using unique_key column: {first_col}" + ) + return exp.DistributedByProperty( + expressions=[exp.to_column(first_col)], + kind="HASH", + buckets=default_buckets_expr or exp.Literal.number(10), + ) + + # Fallback 3: Standalone buckets with duplicate_key/aggregate_key or RANDOM + if default_buckets_expr is not None: + if fallback_key_cols: + return exp.DistributedByProperty( + kind=exp.Var(this="HASH"), + expressions=[exp.to_column(fallback_key_cols[0])], + buckets=default_buckets_expr, + order=None, + ) + return exp.DistributedByProperty( + kind=exp.Var(this="RANDOM"), + expressions=None, + buckets=default_buckets_expr, + order=None, + ) + + return None + + def _get_temp_table( + self, table: TableName, table_only: bool = False, quoted: bool = True + ) -> exp.Table: + """Returns the temp table expression, truncated to fit Doris's 64-char identifier limit.""" + table_copy = t.cast(exp.Table, exp.to_table(table).copy()) + table_copy.set( + "this", + exp.to_identifier( + f"__temp_{table_copy.name[:48]}_{random_id(short=True)}", quoted=quoted + ), + ) + + if table_only: + table_copy.set("db", None) + table_copy.set("catalog", None) + + return table_copy + + def _properties_to_expressions(self, properties: t.Dict[str, t.Any]) -> t.List[exp.Expr]: + """Convert a dictionary of properties to a list of Doris PROPERTIES ('k'='v') expressions.""" + expressions: t.List[exp.Expr] = [] + skip_keys = { + "build", + "refresh", + "refresh_trigger", + "distributed_by", + "clustered_by", + "buckets", + "partitioned_by", + "partitions", + "auto_partition", + "unique_key", + "duplicate_key", + "aggregate_key", + "order_by", + "engine", + "properties", + "materialized_type", + "source_table", + } + for key, value in properties.items(): + if key in skip_keys or value is None: + continue + raw_str = self._extract_str_value(value) + str_literal = exp.Literal.string(raw_str if raw_str is not None else str(value)) + expressions.append(exp.Property(this=exp.Literal.string(str(key)), value=str_literal)) + return expressions diff --git a/tests/core/engine_adapter/integration/__init__.py b/tests/core/engine_adapter/integration/__init__.py index fe1951add6..7ce05f451e 100644 --- a/tests/core/engine_adapter/integration/__init__.py +++ b/tests/core/engine_adapter/integration/__init__.py @@ -80,6 +80,7 @@ def pytest_marks(self) -> t.List[MarkDecorator]: IntegrationTestEngine("clickhouse", catalog_types=["standalone", "cluster"]), IntegrationTestEngine("risingwave"), IntegrationTestEngine("starrocks"), + IntegrationTestEngine("doris"), # Cloud engines that need paid accounts / special credentials IntegrationTestEngine("clickhouse_cloud", cloud=True), IntegrationTestEngine("redshift", cloud=True), @@ -276,6 +277,7 @@ def timestamp_columns(self) -> t.List[str]: if v.sql().lower().startswith("timestamp") or (v.sql().lower() == "datetime" and self.dialect == "bigquery") or (v.sql().lower() == "datetime" and self.dialect == "starrocks") + or (v.sql().lower() == "datetime" and self.dialect == "doris") ] @property @@ -462,7 +464,7 @@ def get_table_comment( AND pgc.relkind = '{"v" if table_kind == "VIEW" else "r"}' ; """ - elif self.dialect in ["mysql", "snowflake", "starrocks"]: + elif self.dialect in ["mysql", "snowflake", "starrocks", "doris"]: # Snowflake treats all identifiers as uppercase unless they are lowercase and quoted. # They are lowercase and quoted in sushi but not in the inline tests. if self.dialect == "snowflake" and snowflake_capitalize_ids: @@ -473,6 +475,7 @@ def get_table_comment( "mysql": "table_comment", "snowflake": "comment", "starrocks": "table_comment", + "doris": "table_comment", } query = f""" @@ -594,7 +597,7 @@ def get_column_comments( AND pgc.relkind = '{"v" if table_kind == "VIEW" else "r"}' ; """ - elif self.dialect in ["mysql", "snowflake", "trino", "starrocks"]: + elif self.dialect in ["mysql", "snowflake", "trino", "starrocks", "doris"]: # Snowflake treats all identifiers as uppercase unless they are lowercase and quoted. # They are lowercase and quoted in sushi but not in the inline tests. if self.dialect == "snowflake" and snowflake_capitalize_ids: @@ -606,6 +609,7 @@ def get_column_comments( "snowflake": "comment", "trino": "comment", "starrocks": "column_comment", + "doris": "column_comment", } query = f""" diff --git a/tests/core/engine_adapter/integration/config.yaml b/tests/core/engine_adapter/integration/config.yaml index 24ba2b2a6d..931f5e0343 100644 --- a/tests/core/engine_adapter/integration/config.yaml +++ b/tests/core/engine_adapter/integration/config.yaml @@ -128,6 +128,16 @@ gateways: check_import: false state_connection: type: duckdb + inttest_doris: + connection: + type: doris + host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }} + port: 9030 + user: root + password: "" + check_import: false + state_connection: + type: duckdb # Cloud databases diff --git a/tests/core/engine_adapter/integration/docker/compose.doris.yaml b/tests/core/engine_adapter/integration/docker/compose.doris.yaml new file mode 100644 index 0000000000..bb1b77d44a --- /dev/null +++ b/tests/core/engine_adapter/integration/docker/compose.doris.yaml @@ -0,0 +1,43 @@ +services: + fe-01: + image: apache/doris:fe-4.0.8 + container_name: doris-fe-01 + hostname: fe-01 + environment: + - FE_SERVERS=fe1:172.30.80.2:9010 + - FE_ID=1 + ports: + - "8030:8030" + - "9030:9030" + volumes: + - iceberg-warehouse:/tmp/doris_iceberg_warehouse + networks: + doris_net: + ipv4_address: 172.30.80.2 + + be-01: + image: apache/doris:be-4.0.8 + container_name: doris-be-01 + hostname: be-01 + depends_on: + - fe-01 + environment: + - FE_SERVERS=fe1:172.30.80.2:9010 + - BE_ADDR=172.30.80.3:9050 + ports: + - "8040:8040" + - "9050:9050" + volumes: + - iceberg-warehouse:/tmp/doris_iceberg_warehouse + networks: + doris_net: + ipv4_address: 172.30.80.3 + +volumes: + iceberg-warehouse: + +networks: + doris_net: + ipam: + config: + - subnet: 172.30.80.0/24 diff --git a/tests/core/engine_adapter/integration/test_integration.py b/tests/core/engine_adapter/integration/test_integration.py index 6a2905d64d..a2cdfb00b3 100644 --- a/tests/core/engine_adapter/integration/test_integration.py +++ b/tests/core/engine_adapter/integration/test_integration.py @@ -101,6 +101,7 @@ def test_catalog_operations(ctx: TestContext): if ( ctx.engine_adapter.catalog_support.is_unsupported or ctx.engine_adapter.catalog_support.is_single_catalog_only + or ctx.dialect == "doris" ): pytest.skip( f"Engine adapter {ctx.engine_adapter.dialect} doesn't support catalog operations" @@ -154,7 +155,7 @@ def create_objects_and_validate(schema_name: str): assert len(results.materialized_views) == 0 assert len(results.non_temp_tables) == 2 - if ctx.engine_adapter.catalog_support.is_unsupported: + if ctx.engine_adapter.catalog_support.is_unsupported or ctx.dialect == "doris": pytest.skip( f"Engine adapter {ctx.engine_adapter.dialect} doesn't support catalog operations" ) @@ -449,8 +450,22 @@ def test_materialized_view(ctx_query_and_df: TestContext): source_table = ctx.table("source_table") ctx.engine_adapter.ctas(source_table, ctx.input_data(input_data), ctx.columns_to_types) view = ctx.table("test_view") - view_query = exp.select(*ctx.columns_to_types).from_(source_table) - ctx.engine_adapter.create_view(view, view_query, materialized=True) + # 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 + if ctx.engine_adapter.dialect == "doris": + view = exp.to_table(view.name) + # For Doris, select only the first column name from columns_to_types. + # Otherwise, Doris will raise an error if the materialized view has the same schema as the base table. + first_column = next(iter(ctx.columns_to_types)) + view_query = exp.select(first_column).from_(source_table) + ctx.engine_adapter.create_view( + view, + view_query, + materialized=True, + view_properties={"materialized_type": "SYNC", "source_table": source_table}, + ) + else: + view_query = exp.select(*ctx.columns_to_types).from_(source_table) + ctx.engine_adapter.create_view(view, view_query, materialized=True) results = ctx.get_metadata_results() # Redshift considers the underlying dataset supporting materialized views as a table therefore we get 2 # tables in the result @@ -459,13 +474,31 @@ def test_materialized_view(ctx_query_and_df: TestContext): else: assert len(results.tables) == 1 assert len(results.views) == 0 - assert len(results.materialized_views) == 1 - assert results.materialized_views[0] == view.name - ctx.compare_with_current(view, input_data) + # Doris information_schema.tables does not show materialized views and cannot query the synchronized materialized view + if ctx.engine_adapter.dialect != "doris": + assert len(results.materialized_views) == 1 + assert results.materialized_views[0] == view.name + ctx.compare_with_current(view, input_data) # Make sure that dropping a materialized view also works - ctx.engine_adapter.drop_view(view, materialized=True) - results = ctx.get_metadata_results() - assert len(results.materialized_views) == 0 + if ctx.engine_adapter.dialect == "doris": + # Wait for the materialized view to be created by retrying drop until it succeeds + def drop_view_success(): + try: + ctx.engine_adapter.drop_view( + view, + ignore_if_not_exists=False, + materialized=True, + view_properties={"materialized_type": "SYNC", "source_table": source_table}, + ) + return True + except Exception: + return False + + wait_until(drop_view_success, attempts=5, wait=1) + else: + ctx.engine_adapter.drop_view(view, materialized=True) + results = ctx.get_metadata_results() + assert len(results.materialized_views) == 0 def test_drop_schema(ctx: TestContext): @@ -998,8 +1031,12 @@ def test_merge(ctx_query_and_df: TestContext): # Athena only supports MERGE on Iceberg tables # And it cant fall back to a logical merge on Hive tables because it cant delete records table_format = "iceberg" if ctx.dialect == "athena" else None + # Doris needs to have a UNIQUE KEY to use MERGE + table_properties = {"unique_key": "id"} if ctx.dialect == "doris" else None - ctx.engine_adapter.create_table(table, ctx.columns_to_types, table_format=table_format) + ctx.engine_adapter.create_table( + table, ctx.columns_to_types, table_format=table_format, table_properties=table_properties + ) input_data = pd.DataFrame( [ {"id": 1, "ds": "2022-01-01"}, @@ -1064,11 +1101,15 @@ def test_merge_source_columns(ctx_query_and_df: TestContext): # Athena only supports MERGE on Iceberg tables # And it cant fall back to a logical merge on Hive tables because it cant delete records table_format = "iceberg" if ctx.dialect == "athena" else None + # Doris needs to have a UNIQUE KEY to use MERGE + table_properties = {"unique_key": "id"} if ctx.dialect == "doris" else None columns_to_types = ctx.columns_to_types.copy() columns_to_types["ignored_column"] = exp.DataType.build("int") - ctx.engine_adapter.create_table(table, columns_to_types, table_format=table_format) + ctx.engine_adapter.create_table( + table, columns_to_types, table_format=table_format, table_properties=table_properties + ) input_data = pd.DataFrame( [ {"id": 1, "ds": "2022-01-01", "ignored_source": "ignored_value"}, @@ -2684,6 +2725,7 @@ def test_dialects(ctx: TestContext): "spark": pd.Timestamp("2020-01-01 00:00:00"), "databricks": pd.Timestamp("2020-01-01 00:00:00"), "starrocks": pd.Timestamp("2020-01-01 00:00:00"), + "doris": pd.Timestamp("2020-01-01 00:00:00"), }, ), ( @@ -3246,8 +3288,8 @@ def test_value_normalization( if column_type == exp.DataType.Type.TIMESTAMPTZ: if ctx.dialect == "trino" and ctx.engine_adapter.current_catalog_type == "hive": pytest.skip("Trino on Hive doesn't support TIMESTAMP WITH TIME ZONE fields") - if ctx.dialect == "fabric": - pytest.skip("Fabric doesn't support TIMESTAMP WITH TIME ZONE fields") + if ctx.dialect in ("fabric", "doris"): + pytest.skip(f"{ctx.dialect} doesn't support TIMESTAMP WITH TIME ZONE fields") if not isinstance(ctx.engine_adapter, RowDiffMixin): pytest.skip( @@ -3258,8 +3300,8 @@ def test_value_normalization( # resolve dialect-specific types if column_type in (DATA_TYPE.DATETIME, DATA_TYPE.TIMESTAMP, DATA_TYPE.TIMESTAMPTZ): - if ctx.dialect in ("mysql", "trino"): - # MySQL needs DATETIME(6) instead of DATETIME or subseconds will be truncated. + if ctx.dialect in ("mysql", "trino", "doris"): + # MySQL and Doris need DATETIME(6) instead of DATETIME or subseconds will be truncated. # It also needs TIMESTAMP(6) as the column type for CREATE TABLE or the truncation will occur full_column_type = exp.DataType.build( column_type, @@ -3283,7 +3325,11 @@ def test_value_normalization( exp.to_table(ctx.table("test_value_normalization")), dialect=ctx.dialect ) columns_to_types_normalized = { - normalize_identifiers(k, dialect=ctx.dialect).sql(dialect=ctx.dialect): v + ( + normalize_identifiers(k, dialect=ctx.dialect).name + if ctx.dialect == "doris" + else normalize_identifiers(k, dialect=ctx.dialect).sql(dialect=ctx.dialect) + ): v for k, v in columns_to_types.items() } @@ -3941,6 +3987,10 @@ def test_materialized_view_evaluation(ctx: TestContext): "which this generic test does not provide; StarRocks MVs are covered by " "test_integration_starrocks.py" ) + elif dialect == "doris": + pytest.skip( + "Skipping doris as synchronous materialized views do not support specifying schema" + ) model_name = ctx.table("test_tbl") mview_name = ctx.table("test_mview") @@ -4008,6 +4058,11 @@ def test_unicode_characters(ctx: TestContext, tmp_path: Path): # I also think Spark may not support unicode in general but that would need to be verified. if not ctx.engine_adapter.QUOTE_IDENTIFIERS_IN_VIEWS: pytest.skip("Skipping as these engines have issues with unicode characters in model names") + # Doris default setting `enable_unicode_name_support=false` so it is incompatible with unicode characters in model names + if ctx.dialect == "doris": + pytest.skip( + "Skipping as Doris default setting has issues with unicode characters in model names" + ) model_name = "客户数据" table = ctx.table(model_name).sql(dialect=ctx.dialect) diff --git a/tests/core/engine_adapter/integration/test_integration_doris.py b/tests/core/engine_adapter/integration/test_integration_doris.py new file mode 100644 index 0000000000..3e29bb53c5 --- /dev/null +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import typing as t + +import pytest +from pytest import FixtureRequest +from sqlglot import exp + +from sqlmesh.core.engine_adapter import DorisEngineAdapter +from tests.core.engine_adapter.integration import ( + ENGINES_BY_NAME, + IntegrationTestEngine, + TestContext, + generate_pytest_params, +) + + +@pytest.fixture(params=list(generate_pytest_params(ENGINES_BY_NAME["doris"]))) +def ctx( + request: FixtureRequest, + create_test_context: t.Callable[[IntegrationTestEngine, str, str], t.Iterable[TestContext]], +) -> t.Iterable[TestContext]: + yield from create_test_context(*request.param) + + +@pytest.fixture +def engine_adapter(ctx: TestContext) -> DorisEngineAdapter: + assert isinstance(ctx.engine_adapter, DorisEngineAdapter) + return ctx.engine_adapter + + +def test_engine_adapter(ctx: TestContext): + assert isinstance(ctx.engine_adapter, DorisEngineAdapter) + assert ctx.engine_adapter.fetchone("select 1") == (1,) + + +def test_doris_iceberg_local_catalog(ctx: TestContext): + adapter = ctx.engine_adapter + assert isinstance(adapter, DorisEngineAdapter) + + catalog_name = f"local_iceberg_{ctx.test_id}" + schema_name = f"iceberg_db_{ctx.test_id}" + table_fqn = f"{catalog_name}.{schema_name}.events" + + try: + adapter.execute( + f""" + CREATE CATALOG IF NOT EXISTS `{catalog_name}` PROPERTIES ( + 'type' = 'iceberg', + 'iceberg.catalog.type' = 'hadoop', + 'warehouse' = 'file:///tmp/doris_iceberg_warehouse/{catalog_name}' + ) + """, + quote_identifiers=False, + ) + adapter.create_schema(f"{catalog_name}.{schema_name}") + adapter.create_table( + table_fqn, + { + "id": exp.DataType.build("int"), + "ds": exp.DataType.build("date"), + }, + partitioned_by=[exp.to_column("ds")], + ) + + cols = adapter.columns(table_fqn) + assert "id" in cols + assert "ds" in cols + + data_objects = adapter.get_data_objects(f"{catalog_name}.{schema_name}") + assert any(obj.name == "events" for obj in data_objects) + + schema_target_fqn = f"{catalog_name}.{schema_name}.events_target" + adapter.create_table( + schema_target_fqn, + { + "id": exp.DataType.build("int"), + "ds": exp.DataType.build("date"), + "region": exp.DataType.build("varchar(32)"), + }, + partitioned_by=[exp.to_column("ds")], + ) + alter_ops = adapter.get_alter_operations(table_fqn, schema_target_fqn) + assert len(alter_ops) == 1 + adapter.alter_table(alter_ops) + + updated_cols = adapter.columns(table_fqn) + assert "region" in updated_cols + finally: + try: + adapter.drop_schema( + f"{catalog_name}.{schema_name}", ignore_if_not_exists=True, cascade=True + ) + except Exception: + pass + try: + adapter.execute(f"DROP CATALOG IF EXISTS `{catalog_name}`", quote_identifiers=False) + except Exception: + pass + adapter.set_current_catalog("internal") diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py new file mode 100644 index 0000000000..098f52406a --- /dev/null +++ b/tests/core/engine_adapter/test_doris.py @@ -0,0 +1,1279 @@ +import typing as t + +import pytest +from pytest_mock.plugin import MockerFixture +from sqlglot import expressions as exp +from sqlglot import parse_one + +from sqlmesh.core.engine_adapter.doris import DorisEngineAdapter +from sqlmesh.core.engine_adapter.shared import CatalogSupport, DataObject, DataObjectType +from tests.core.engine_adapter import to_sql_calls + +pytestmark = [pytest.mark.doris, pytest.mark.engine] + + +def test_create_view(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_view("test_view", parse_one("SELECT a FROM tbl")) + adapter.create_view("test_view", parse_one("SELECT a FROM tbl"), replace=False) + # check view properties + adapter.create_view( + "test_view", + parse_one("SELECT a FROM tbl"), + replace=False, + view_properties={"a": exp.convert(1)}, + ) + + assert to_sql_calls(adapter) == [ + "DROP VIEW IF EXISTS `test_view`", + "CREATE VIEW `test_view` AS SELECT `a` FROM `tbl`", + "CREATE VIEW `test_view` AS SELECT `a` FROM `tbl`", + "CREATE VIEW `test_view` AS SELECT `a` FROM `tbl`", + ] + + +def test_create_view_with_comment(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_view( + "test_view", + parse_one("SELECT a FROM tbl"), + replace=False, + target_columns_to_types={"a": exp.DataType.build("INT")}, + table_description="test_description", + column_descriptions={"a": "test_column_description"}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE VIEW `test_view` (`a` COMMENT 'test_column_description') COMMENT 'test_description' AS SELECT `a` FROM `tbl`", + ] + + +def test_create_materialized_view(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_view( + "test_view", + parse_one("SELECT a FROM tbl"), + materialized=True, + target_columns_to_types={"a": exp.DataType.build("INT")}, + ) + adapter.create_view( + "test_view", + parse_one("SELECT a FROM tbl"), + replace=False, + materialized=True, + target_columns_to_types={"a": exp.DataType.build("INT")}, + ) + + assert to_sql_calls(adapter) == [ + "DROP MATERIALIZED VIEW IF EXISTS `test_view`", + "CREATE MATERIALIZED VIEW `test_view` (`a`) AS SELECT `a` FROM `tbl`", + "CREATE MATERIALIZED VIEW `test_view` (`a`) AS SELECT `a` FROM `tbl`", + ] + + adapter.cursor.reset_mock() + adapter.create_view( + "test_view", + parse_one("SELECT a, b FROM tbl"), + replace=False, + materialized=True, + target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}, + column_descriptions={"a": "test_column_description", "b": "test_column_description"}, + ) + adapter.create_view( + "test_view", parse_one("SELECT a, b FROM tbl"), replace=False, materialized=True + ) + + assert to_sql_calls(adapter) == [ + "CREATE MATERIALIZED VIEW `test_view` (`a` COMMENT 'test_column_description', `b` COMMENT 'test_column_description') AS SELECT `a`, `b` FROM `tbl`", + "CREATE MATERIALIZED VIEW `test_view` AS SELECT `a`, `b` FROM `tbl`", + ] + + +def test_create_table(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT")}, + column_descriptions={"a": "test_column_description"}, + table_properties={"unique_key": exp.Column(this=exp.Identifier(this="a", quoted=True))}, + ) + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT")}, + column_descriptions={"a": "test_column_description"}, + table_properties={"duplicate_key": exp.Column(this=exp.Identifier(this="a", quoted=True))}, + ) + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT")}, + table_description="test_description", + column_descriptions={"a": "test_column_description"}, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT COMMENT 'test_column_description') UNIQUE KEY (`a`) DISTRIBUTED BY HASH (`a`) BUCKETS 10", + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT COMMENT 'test_column_description') DUPLICATE KEY (`a`)", + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT COMMENT 'test_column_description') COMMENT 'test_description'", + ] + + adapter.cursor.reset_mock() + + +def test_create_table_like(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.create_table_like("target_table", "source_table") + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `target_table` LIKE `source_table`", + ] + + +def test_create_schema(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_schema("test_schema") + + assert to_sql_calls(adapter) == [ + "CREATE DATABASE IF NOT EXISTS `test_schema`", + ] + + +def test_drop_schema(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.drop_schema("test_schema") + adapter.drop_schema("test_schema", ignore_if_not_exists=False) + + assert to_sql_calls(adapter) == [ + "DROP DATABASE IF EXISTS `test_schema`", + "DROP DATABASE `test_schema`", + ] + + +def test_rename_table(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.rename_table("old_table", "new_table") + adapter.cursor.execute.assert_called_once_with("ALTER TABLE `old_table` RENAME `new_table`") + + +def test_replace_by_key( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + temp_table = parse_one("temp_table") + mocker.patch.object(adapter, "_get_temp_table", return_value=temp_table) + mocker.patch.object(adapter, "ctas") + mocker.patch.object(adapter, "delete_from") + mocker.patch.object(adapter, "drop_table") + + columns_to_types = {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")} + source_query = parse_one("SELECT a, b FROM src") + target_table = "target_table" + + adapter._replace_by_key(target_table, source_query, columns_to_types, [exp.column("a")], True) + adapter._replace_by_key(target_table, source_query, columns_to_types, [exp.column("a")], False) + adapter._replace_by_key( + target_table, source_query, columns_to_types, [exp.column("a"), exp.column("b")], True + ) + adapter._replace_by_key( + target_table, source_query, columns_to_types, [exp.column("a"), exp.column("b")], False + ) + + assert to_sql_calls(adapter, identify=True) == [ + "INSERT INTO `target_table` (`a`, `b`) SELECT `a`, `b` FROM (SELECT `a` AS `a`, `b` AS `b`, ROW_NUMBER() OVER (PARTITION BY `a` ORDER BY `a`) AS _row_number FROM `temp_table`) AS _t WHERE _row_number = 1", + "INSERT INTO `target_table` (`a`, `b`) SELECT `a`, `b` FROM `temp_table`", + "INSERT INTO `target_table` (`a`, `b`) SELECT `a`, `b` FROM (SELECT `a` AS `a`, `b` AS `b`, ROW_NUMBER() OVER (PARTITION BY `a`, `b` ORDER BY `a`, `b`) AS _row_number FROM `temp_table`) AS _t WHERE _row_number = 1", + "INSERT INTO `target_table` (`a`, `b`) SELECT `a`, `b` FROM `temp_table`", + ] + + +def test_create_index(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.create_index("test_table", "test_index", ("cola",)) + adapter.cursor.execute.assert_called_once_with( + "CREATE INDEX IF NOT EXISTS `test_index` ON `test_table`(`cola`)" + ) + + +def test_create_table_with_distributed_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + distributed_by = { + "expressions": ["a", "b"], + "kind": "HASH", + "buckets": 8, + } + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}, + table_properties={"distributed_by": distributed_by}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT, `b` INT) DISTRIBUTED BY HASH (`a`, `b`) BUCKETS 8", + ] + + adapter.cursor.execute.reset_mock() + + distributed_by = { + "expressions": None, + "kind": "RANDOM", + "buckets": None, + } + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}, + table_properties={"distributed_by": distributed_by}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT, `b` INT) DISTRIBUTED BY RANDOM", + ] + + adapter.cursor.execute.reset_mock() + + distributed_by = { + "expressions": ["a"], + "kind": "HASH", + "buckets": "AUTO", + } + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}, + table_properties={"distributed_by": distributed_by}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT, `b` INT) DISTRIBUTED BY HASH (`a`) BUCKETS AUTO", + ] + + +def test_create_table_with_properties( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT")}, + table_properties={ + "refresh_interval": "86400", + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT) PROPERTIES ('refresh_interval'='86400')", + ] + + +def test_create_table_with_partitioned_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("DATE")}, + partitioned_by=[exp.Literal.string("RANGE(b)")], + table_properties={ + "partitions": exp.Literal.string( + "FROM ('2000-11-14') TO ('2021-11-14') INTERVAL 2 YEAR" + ) + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT, `b` DATE) PARTITION BY RANGE (`b`) (FROM ('2000-11-14') TO ('2021-11-14') INTERVAL 2 YEAR)", + ] + + +def test_create_table_with_range_partitioned_by_with_partitions( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "waiter_id": exp.DataType.build("INT"), + "customer_id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATETIME"), + }, + partitioned_by=[exp.Literal.string("RANGE(ds)")], + table_properties={ + "partitions": exp.Tuple( + expressions=[ + exp.Literal.string('PARTITION `p2023` VALUES [("2023-01-01"), ("2024-01-01"))'), + exp.Literal.string('PARTITION `p2024` VALUES [("2024-01-01"), ("2025-01-01"))'), + exp.Literal.string('PARTITION `p2025` VALUES [("2025-01-01"), ("2026-01-01"))'), + exp.Literal.string("PARTITION `other` VALUES LESS THAN MAXVALUE"), + ] + ), + "distributed_by": exp.Tuple( + expressions=[ + exp.EQ( + this=exp.Column(this=exp.Identifier(this="kind", quoted=True)), + expression=exp.Literal.string("HASH"), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="expressions", quoted=True)), + expression=exp.Column(this=exp.Identifier(this="id", quoted=True)), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="buckets", quoted=True)), + expression=exp.Literal.number(10), + ), + ] + ), + "replication_allocation": exp.Literal.string("tag.location.default: 3"), + "in_memory": exp.Literal.string("false"), + "storage_format": exp.Literal.string("V2"), + "disable_auto_compaction": exp.Literal.string("false"), + }, + ) + + expected_sql = ( + "CREATE TABLE IF NOT EXISTS `test_table` " + "(`id` INT, `waiter_id` INT, `customer_id` INT, `ds` DATETIME) " + "PARTITION BY RANGE (`ds`) " + '(PARTITION `p2023` VALUES [("2023-01-01"), ("2024-01-01")), ' + 'PARTITION `p2024` VALUES [("2024-01-01"), ("2025-01-01")), ' + 'PARTITION `p2025` VALUES [("2025-01-01"), ("2026-01-01")), ' + "PARTITION `other` VALUES LESS THAN MAXVALUE) " + "DISTRIBUTED BY HASH (`id`) BUCKETS 10 " + "PROPERTIES (" + "'replication_allocation'='tag.location.default: 3', " + "'in_memory'='false', " + "'storage_format'='V2', " + "'disable_auto_compaction'='false')" + ) + + assert to_sql_calls(adapter) == [expected_sql] + + +def test_create_table_with_list_partitioned_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "status": exp.DataType.build("VARCHAR(10)"), + }, + partitioned_by=[exp.Literal.string("LIST(status)")], + table_properties={ + "partitions": exp.Tuple( + expressions=[ + exp.Literal.string('PARTITION `active` VALUES IN ("active", "pending")'), + exp.Literal.string('PARTITION `inactive` VALUES IN ("inactive", "disabled")'), + ] + ), + }, + ) + + expected_sql = ( + "CREATE TABLE IF NOT EXISTS `test_table` " + "(`id` INT, `status` VARCHAR(10)) " + "PARTITION BY LIST (`status`) " + '(PARTITION `active` VALUES IN ("active", "pending"), ' + 'PARTITION `inactive` VALUES IN ("inactive", "disabled"))' + ) + + assert to_sql_calls(adapter) == [expected_sql] + + +def test_create_table_with_range_partitioned_by_anonymous_function( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + """Test that RANGE(ds) function call syntax generates correct SQL without duplicate RANGE.""" + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATETIME"), + }, + # This simulates how partitioned_by RANGE(ds) gets parsed from model definition + partitioned_by=[exp.Anonymous(this="RANGE", expressions=[exp.to_column("ds")])], + table_properties={ + "partitions": exp.Literal.string( + 'FROM ("2000-11-14") TO ("2099-11-14") INTERVAL 1 MONTH' + ) + }, + ) + + expected_sql = ( + "CREATE TABLE IF NOT EXISTS `test_table` " + "(`id` INT, `ds` DATETIME) " + "PARTITION BY RANGE (`ds`) " + '(FROM ("2000-11-14") TO ("2099-11-14") INTERVAL 1 MONTH)' + ) + + assert to_sql_calls(adapter) == [expected_sql] + + +def test_create_materialized_view_with_duplicate_key( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_view( + "test_mv", + parse_one("SELECT id, status, COUNT(*) as cnt FROM orders GROUP BY id, status"), + materialized=True, + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "status": exp.DataType.build("VARCHAR(10)"), + "cnt": exp.DataType.build("BIGINT"), + }, + view_properties={ + "duplicate_key": exp.Tuple( + expressions=[ + exp.to_column("id"), + exp.to_column("status"), + ] + ), + }, + ) + + expected_sqls = [ + "DROP MATERIALIZED VIEW IF EXISTS `test_mv`", + ( + "CREATE MATERIALIZED VIEW `test_mv` " + "(`id`, `status`, `cnt`) " + "DUPLICATE KEY (`id`, `status`) " + "AS SELECT `id`, `status`, COUNT(*) AS `cnt` FROM `orders` GROUP BY `id`, `status`" + ), + ] + + assert to_sql_calls(adapter) == expected_sqls + + +def test_create_full_materialized_view( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + view_properties: t.Dict[str, exp.Expr] = { + "build": exp.Literal.string("IMMEDIATE"), + "refresh": exp.Tuple( + expressions=[ + exp.EQ( + this=exp.Column(this=exp.Identifier(this="method", quoted=True)), + expression=exp.Literal.string("AUTO"), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="kind", quoted=True)), + expression=exp.Literal.string("SCHEDULE"), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="every", quoted=True)), + expression=exp.Literal.number(1), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="unit", quoted=True)), + expression=exp.Literal.string("DAY"), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="starts", quoted=True)), + expression=exp.Literal.string("2024-12-01 20:30:00"), + ), + ], + ), + "distributed_by": exp.Tuple( + expressions=[ + exp.EQ( + this=exp.Column(this=exp.Identifier(this="kind", quoted=True)), + expression=exp.Literal.string("HASH"), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="expressions", quoted=True)), + expression=exp.Column(this=exp.Identifier(this="orderkey", quoted=True)), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="buckets", quoted=True)), + expression=exp.Literal.number(2), + ), + ] + ), + "unique_key": exp.to_column("orderkey"), + "replication_num": exp.Literal.string("1"), + } + materialized_properties = { + "partitioned_by": [ + parse_one("DATE_TRUNC(o_orderdate, 'MONTH')", dialect="doris"), + ], + "clustered_by": [], + "partition_interval_unit": None, + } + columns_to_types = { + "orderdate": exp.DataType.build("DATE"), + "orderkey": exp.DataType.build("INT"), + "partkey": exp.DataType.build("INT"), + } + column_descriptions = { + "orderdate": "order date", + "orderkey": "order key", + "partkey": "part key", + } + query = parse_one( + """ + SELECT + o_orderdate, + l_orderkey, + l_partkey + FROM + orders + LEFT JOIN lineitem ON l_orderkey = o_orderkey + LEFT JOIN partsupp ON ps_partkey = l_partkey + and l_suppkey = ps_suppkey + """ + ) + adapter.create_view( + "complete_mv", + query, + replace=False, + materialized=True, + target_columns_to_types=columns_to_types, + column_descriptions=column_descriptions, + table_description="test_description", + view_properties=view_properties, + materialized_properties=materialized_properties, + ) + expected_sqls = [ + "CREATE MATERIALIZED VIEW `complete_mv` (`orderdate` COMMENT 'order date', `orderkey` COMMENT 'order key', `partkey` COMMENT 'part key') " + "BUILD IMMEDIATE REFRESH AUTO ON SCHEDULE EVERY 1 DAY STARTS '2024-12-01 20:30:00' KEY (`orderkey`) COMMENT 'test_description' " + "PARTITION BY (DATE_TRUNC(`o_orderdate`, 'MONTH')) " + "DISTRIBUTED BY HASH (`orderkey`) BUCKETS 2 PROPERTIES ('replication_num'='1') " + "AS SELECT `o_orderdate`, `l_orderkey`, `l_partkey` FROM `orders` LEFT JOIN `lineitem` ON `l_orderkey` = `o_orderkey` LEFT JOIN `partsupp` ON `ps_partkey` = `l_partkey` AND `l_suppkey` = `ps_suppkey`", + ] + sql_calls = to_sql_calls(adapter) + + # Remove extra spaces for comparison + def norm(s): + return " ".join(s.split()) + + for expected_sql in expected_sqls: + assert any(norm(expected_sql) == norm(sql) for sql in sql_calls), ( + f"Expected SQL not found.\nExpected: {expected_sql}\nGot: {sql_calls}" + ) + + +def test_create_table_with_single_string_distributed_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + """Test creating table with distributed_by where expressions is a single string.""" + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + distributed_by = { + "expressions": "recordid", # Single string instead of array + "kind": "HASH", + "buckets": 10, + } + adapter.create_table( + "test_table", + target_columns_to_types={ + "recordid": exp.DataType.build("INT"), + "name": exp.DataType.build("VARCHAR"), + }, + table_properties={"distributed_by": distributed_by}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`recordid` INT, `name` VARCHAR) DISTRIBUTED BY HASH (`recordid`) BUCKETS 10", + ] + + +def test_delete_from_with_subquery(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Test DELETE FROM with IN subquery + adapter.delete_from( + "test_schema_xtll416o.test_table", + "id IN (SELECT id FROM test_schema_xtll416o.temp_test_table_ik29031e)", + ) + + # Test DELETE FROM with NOT IN subquery (must NOT rewrite to USING ... WHERE <>) + adapter.delete_from( + "test_schema_xtll416o.test_table", + "id NOT IN (SELECT id FROM test_schema_xtll416o.temp_test_table_ik29031e)", + ) + + # Test simple DELETE FROM (should use base implementation) + adapter.delete_from("test_schema_xtll416o.test_table", "id = 1") + + assert to_sql_calls(adapter) == [ + "DELETE FROM `test_schema_xtll416o`.`test_table` AS `_t1` USING (SELECT `id` FROM `test_schema_xtll416o`.`temp_test_table_ik29031e`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id`", + "DELETE FROM `test_schema_xtll416o`.`test_table` WHERE NOT `id` IN (SELECT `id` FROM `test_schema_xtll416o`.`temp_test_table_ik29031e`)", + "DELETE FROM `test_schema_xtll416o`.`test_table` WHERE `id` = 1", + ] + + +def test_delete_from_with_complex_subquery( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Test DELETE FROM with complex subquery (multiple columns) + adapter.delete_from( + "test_schema_xtll416o.test_table", + "(id, name) IN (SELECT id, name FROM test_schema_xtll416o.temp_test_table_ik29031e WHERE active = 1)", + ) + + # Test DELETE FROM with subquery that has WHERE clause + adapter.delete_from( + "test_schema_xtll416o.test_table", + "id IN (SELECT id FROM test_schema_xtll416o.temp_test_table_ik29031e WHERE created_date > '2024-01-01')", + ) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `test_schema_xtll416o`.`test_table` AS `_t1` USING (SELECT `id`, `name` FROM `test_schema_xtll416o`.`temp_test_table_ik29031e` WHERE `active` = 1) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`name` = `_t2`.`name`", + "DELETE FROM `test_schema_xtll416o`.`test_table` AS `_t1` USING (SELECT `id` FROM `test_schema_xtll416o`.`temp_test_table_ik29031e` WHERE `created_date` > '2024-01-01') AS `_t2` WHERE `_t1`.`id` = `_t2`.`id`", + ] + + +def test_delete_from_fallback_to_base( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Test DELETE FROM with simple conditions (should use base implementation) + adapter.delete_from("test_table", "id = 1") + adapter.delete_from("test_table", "name = 'test' AND status = 'active'") + adapter.delete_from("test_table", "id IN (1, 2, 3)") # Simple IN with values, not subquery + + assert to_sql_calls(adapter) == [ + "DELETE FROM `test_table` WHERE `id` = 1", + "DELETE FROM `test_table` WHERE `name` = 'test' AND `status` = 'active'", + "DELETE FROM `test_table` WHERE `id` IN (1, 2, 3)", + ] + + +def test_delete_from_full_table(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Test DELETE FROM with WHERE TRUE (should use TRUNCATE) + adapter.delete_from("test_table", exp.true()) + + assert to_sql_calls(adapter) == [ + "TRUNCATE TABLE `test_table`", + ] + + +def test_create_table_olap_key_prefix_reordering_and_order_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "events", + target_columns_to_types={ + "payload": exp.DataType.build("VARCHAR(255)"), + "id": exp.DataType.build("BIGINT"), + "ds": exp.DataType.build("DATE"), + "updated_at": exp.DataType.build("DATETIME"), + }, + unique_key=[exp.to_column("id"), exp.to_column("ds")], + table_properties={ + "order_by": exp.Tuple(expressions=[exp.to_column("updated_at"), exp.to_column("id")]), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `events` (`id` BIGINT, `ds` DATE, `payload` VARCHAR(255), `updated_at` DATETIME) UNIQUE KEY (`id`, `ds`) ORDER BY (`updated_at`, `id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10", + ] + + +def test_create_table_aggregate_key_and_clustered_by_buckets( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "metrics_agg", + target_columns_to_types={ + "total": exp.DataType.build("BIGINT"), + "tenant_id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATE"), + }, + clustered_by=[exp.to_column("tenant_id")], + table_properties={ + "aggregate_key": exp.Tuple( + expressions=[exp.to_column("tenant_id"), exp.to_column("ds")] + ), + "buckets": exp.Literal.number(16), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `metrics_agg` (`tenant_id` INT, `ds` DATE, `total` BIGINT) AGGREGATE KEY (`tenant_id`, `ds`) DISTRIBUTED BY HASH (`tenant_id`) BUCKETS 16", + ] + + +def test_create_table_auto_partition_and_unannotated_partitioned_by_guard( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + # Without partitions or auto_partition, OLAP partitioned_by is safely omitted + adapter.create_table( + "olap_no_partitions", + target_columns_to_types={"id": exp.DataType.build("INT"), "ds": exp.DataType.build("DATE")}, + partitioned_by=[exp.to_column("ds")], + ) + # With auto_partition = true, emits AUTO PARTITION BY RANGE (...) () + adapter.create_table( + "olap_auto_partition", + target_columns_to_types={"id": exp.DataType.build("INT"), "ds": exp.DataType.build("DATE")}, + partitioned_by=[parse_one("RANGE(DATE_TRUNC(ds, 'DAY'))", dialect="doris")], + table_properties={"auto_partition": exp.Literal.string("true")}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `olap_no_partitions` (`id` INT, `ds` DATE)", + "CREATE TABLE IF NOT EXISTS `olap_auto_partition` (`id` INT, `ds` DATE) AUTO PARTITION BY RANGE (DATE_TRUNC(`ds`, 'DAY')) ()", + ] + + +def test_iceberg_multi_catalog_ddl_and_dml( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + assert adapter.catalog_support == CatalogSupport.FULL_SUPPORT + assert adapter.default_catalog is None + + # 1. Explicit table_format='iceberg' on 2-part table name emits ENGINE=iceberg and PARTITION BY LIST (...) () + adapter.create_table( + "silver.iceberg_events", + target_columns_to_types={ + "id": exp.DataType.build("BIGINT"), + "ds": exp.DataType.build("DATE"), + "region": exp.DataType.build("STRING"), + }, + table_format="iceberg", + partitioned_by=[parse_one("DAY(ds)", dialect="doris"), exp.to_column("region")], + unique_key=[exp.to_column("id")], + table_properties={"write-format": exp.Literal.string("parquet")}, + ) + + # 2. 3-part external Iceberg catalog CTAS omits ENGINE=iceberg (invalid in Doris CTAS grammar) + adapter.ctas( + "iceberg_cat.silver.events_ctas", + parse_one("SELECT id, ds FROM bronze.events"), + target_columns_to_types={ + "id": exp.DataType.build("BIGINT"), + "ds": exp.DataType.build("DATE"), + }, + ) + + # 3. replace_query on existing Iceberg table uses atomic INSERT OVERWRITE TABLE + mocker.patch.object( + adapter, + "get_data_object", + return_value=DataObject( + catalog="iceberg_cat", schema="silver", name="events_ctas", type=DataObjectType.TABLE + ), + ) + adapter.replace_query( + "iceberg_cat.silver.events_ctas", + parse_one("SELECT id, ds FROM bronze.events"), + target_columns_to_types={ + "id": exp.DataType.build("BIGINT"), + "ds": exp.DataType.build("DATE"), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `silver`.`iceberg_events` (`id` BIGINT, `ds` DATE, `region` STRING) ENGINE=iceberg PARTITION BY LIST (DAY(`ds`), `region`) () PROPERTIES ('write-format'='parquet')", + "CREATE TABLE IF NOT EXISTS `iceberg_cat`.`silver`.`events_ctas` AS SELECT CAST(`id` AS BIGINT) AS `id`, CAST(`ds` AS DATE) AS `ds` FROM (SELECT `id`, `ds` FROM `bronze`.`events`) AS `_subquery`", + "INSERT OVERWRITE TABLE `iceberg_cat`.`silver`.`events_ctas` (`id`, `ds`) SELECT `id`, `ds` FROM `bronze`.`events`", + ] + + +def test_columns_normalization_largeint_and_struct( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.cursor.fetchall.return_value = [ + ("id", "LARGEINT", "No", "true", None, ""), + ("price", "DECIMALV3(9, 0)", "Yes", "false", None, ""), + ("is_active", "TINYINT(1)", "Yes", "false", None, ""), + ("count_val", "INT(11)", "Yes", "false", None, ""), + ("meta", "STRUCT", "Yes", "false", None, ""), + ] + + cols = adapter.columns("test_db.test_table") + assert cols == { + "id": exp.DataType.build("INT128"), + "price": exp.DataType.build("DECIMAL(9, 0)", dialect="doris"), + "is_active": exp.DataType.build("BOOLEAN", dialect="doris"), + "count_val": exp.DataType.build("INT", dialect="doris"), + "meta": exp.DataType.build("STRUCT", dialect="doris"), + } + + adapter.cursor.reset_mock() + adapter.create_table( + "test_db.struct_tbl", + target_columns_to_types={"id": cols["id"], "meta": cols["meta"]}, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_db`.`struct_tbl` (`id` LARGEINT, `meta` STRUCT<`f1`: INT, `f2`: VARCHAR(20)>)", + ] + + +def test_get_data_objects_with_catalog_and_mv( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter, patch_get_data_objects=False) + mocker.patch.object( + adapter, + "fetchall", + return_value=[ + ("silver", "tbl1", "table"), + ("silver", "mv1", "materialized_view"), + ], + ) + + objs = adapter._get_data_objects("iceberg_cat.silver") + assert len(objs) == 2 + assert objs[0].type == DataObjectType.TABLE + assert objs[0].catalog == "iceberg_cat" + assert objs[0].schema_name == "silver" + assert objs[0].name == "tbl1" + assert objs[1].type == DataObjectType.MATERIALIZED_VIEW + assert objs[1].name == "mv1" + + +def test_comments_escaping(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + table_comment_sql = adapter._build_create_comment_table_exp( + exp.to_table("db.tbl"), " O'Reilly's \"table\" \\ test " + ) + col_comment_sql = adapter._build_create_comment_column_exp( + exp.to_table("db.tbl"), "col'1", ' Col\'s "comment" ' + ) + assert ( + table_comment_sql + == "ALTER TABLE `db`.`tbl` MODIFY COMMENT 'O''Reilly''s \"table\" \\\\ test'" + ) + assert ( + col_comment_sql + == "ALTER TABLE `db`.`tbl` MODIFY COLUMN `col'1` COMMENT 'Col''s \"comment\"'" + ) + + +def test_delete_from_between_concat_ws_and_literal_fallback( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. BETWEEN rewrite in delete_from + adapter.delete_from("test_table", "ds BETWEEN '2024-01-01' AND '2024-01-31'") + + # 2. _replace_by_key / LogicalMergeMixin CONCAT_WS unwrap into multi-column USING + key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", exp.column("id"), exp.column("ds")) + merge_where = key_exp.isin(query=exp.select(key_exp).from_("src_table")) + adapter.delete_from("target_table", where=merge_where) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `test_table` WHERE `ds` >= '2024-01-01' AND `ds` <= '2024-01-31'", + "DELETE FROM `target_table` AS `_t1` USING (SELECT `id`, `ds` FROM `src_table`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`ds` = `_t2`.`ds`", + ] + + # 3. Fallback to zero-OR batched _delete_by_literal_keys when USING fails on DUPLICATE KEY / Iceberg + adapter.cursor.reset_mock() + mocker.patch.object( + adapter, + "fetchall", + return_value=[("tenant_a", 1), ("tenant_a", 2), ("tenant_a", None)], + ) + adapter._delete_by_literal_keys( + "dup_table", + [exp.column("tenant"), exp.column("id")], + parse_one("SELECT tenant, id FROM src_table", into=exp.Select), + ) + assert to_sql_calls(adapter) == [ + "DELETE FROM `dup_table` WHERE `tenant` = 'tenant_a' AND `id` IN (1, 2)", + "DELETE FROM `dup_table` WHERE `tenant` = 'tenant_a' AND `id` IS NULL", + ] + + +def test_execute_operational_error_reconnect( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + class OperationalError(Exception): + pass + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + conn_mock = mocker.MagicMock() + mocker.patch.object(adapter._connection_pool, "get", return_value=conn_mock) + + call_count = 0 + + def side_effect(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise OperationalError(2013, "Lost connection to MySQL server during query") + + adapter.cursor.execute.side_effect = side_effect + adapter._execute("SELECT 1") + + assert call_count == 2 + conn_mock.ping.assert_called_once_with(reconnect=True) + + +def test_external_catalog_view_skip_and_materialized_error( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + from sqlmesh.utils.errors import SQLMeshError + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Regular view creation and deletion on external catalog is a safe no-op (warns instead of failing dev promotion) + adapter.create_view("iceberg_cat.silver__dev.events", parse_one("SELECT 1 AS id")) + adapter.drop_view("iceberg_cat.silver__dev.events") + assert to_sql_calls(adapter) == [] + + # Materialized view on external catalog still raises SQLMeshError + with pytest.raises( + SQLMeshError, match="Doris only supports creating materialized views in the 'internal'" + ): + adapter.create_view( + "iceberg_cat.silver.events_mv", + parse_one("SELECT 1 AS id"), + materialized=True, + ) + + +def test_duplicate_key_with_buckets_and_random_buckets_fallback( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "dup_bucketed", + target_columns_to_types={"id": exp.DataType.build("INT"), "val": exp.DataType.build("INT")}, + table_properties={ + "duplicate_key": exp.Tuple(expressions=[exp.to_column("id")]), + "buckets": exp.Literal.number(16), + }, + ) + adapter.create_table( + "random_bucketed", + target_columns_to_types={"id": exp.DataType.build("INT"), "val": exp.DataType.build("INT")}, + table_properties={ + "buckets": exp.Literal.number(8), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `dup_bucketed` (`id` INT, `val` INT) DUPLICATE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 16", + "CREATE TABLE IF NOT EXISTS `random_bucketed` (`id` INT, `val` INT) DISTRIBUTED BY RANDOM BUCKETS 8", + ] + + +def test_truncate_table_fallback_on_unsupported_truncate( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch.object( + adapter, + "columns", + return_value={"id": exp.DataType.build("BIGINT"), "ds": exp.DataType.build("DATE")}, + ) + + call_count = 0 + + def side_effect(sql, *args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1 and "TRUNCATE" in str(sql): + raise Exception("errCode = 2, detailMessage = Not support truncate table on iceberg") + + adapter.cursor.execute.side_effect = side_effect + adapter._truncate_table("silver.iceberg_two_part") + + assert to_sql_calls(adapter) == [ + "TRUNCATE TABLE `silver`.`iceberg_two_part`", + "DELETE FROM `silver`.`iceberg_two_part` WHERE NOT `id` IS NULL", + "DELETE FROM `silver`.`iceberg_two_part` WHERE `id` IS NULL", + ] + + +def test_row_diff_mixin_concat_columns( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + concat_expr = adapter.concat_columns( + { + "id": exp.DataType.build("INT"), + "flag": exp.DataType.build("BOOLEAN"), + "ts": exp.DataType.build("DATETIME"), + } + ) + assert ( + concat_expr.sql("doris") + == "CONCAT(COALESCE(CAST(CAST(id AS BIGINT) AS VARCHAR), ''), ',', COALESCE(CAST(CAST(flag AS INT) AS VARCHAR), ''), ',', COALESCE(CAST(DATE_FORMAT(ts, '%Y-%m-%d %T.%f') AS VARCHAR), ''))" + ) + + +def test_temp_table_length_and_catalog_switch_and_wrapped_concat_ws( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. _get_temp_table never exceeds Doris's 64-character identifier limit + long_table = "db." + ("a" * 60) + temp_tbl = adapter._get_temp_table(long_table) + assert len(temp_tbl.name) <= 64 + + # 2. set_current_catalog preserves backtick quoting on hyphenated catalog names + adapter.set_current_catalog("iceberg-cat") + + # 3. delete_from unwraps parenthesized CONCAT_WS(...) + key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", exp.column("id"), exp.column("ds")) + wrapped_where = exp.Paren(this=key_exp).isin(query=exp.select(key_exp).from_("src_table")) + adapter.delete_from("target_table", where=wrapped_where) + + assert to_sql_calls(adapter) == [ + "SWITCH `iceberg-cat`", + "DELETE FROM `target_table` AS `_t1` USING (SELECT `id`, `ds` FROM `src_table`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`ds` = `_t2`.`ds`", + ] + + +def test_insert_overwrite_by_partition_unwraps_wrappers_and_transforms( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch.object( + adapter, + "_get_temp_table", + return_value=exp.table_("__temp_tbl", db="db"), + ) + + adapter.insert_overwrite_by_partition( + "db.events", + parse_one("SELECT id, ds FROM db.src", read="doris"), + partitioned_by=[ + parse_one("RANGE(DATE_TRUNC(ds, 'day'))", read="doris"), + ], + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATETIME"), + }, + ) + + sql_calls = to_sql_calls(adapter) + assert ( + "DELETE FROM `db`.`events` WHERE DATE_TRUNC(`ds`, 'day') IN (SELECT DISTINCT DATE_TRUNC(`ds`, 'day') FROM `db`.`__temp_tbl`)" + not in sql_calls + ) + assert any( + "SELECT DISTINCT `ds` FROM `db`.`__temp_tbl`" in call + or "USING (SELECT DISTINCT `ds` FROM `db`.`__temp_tbl`)" in call + for call in sql_calls + ) + + +def test_sync_materialized_view_drop_and_replace( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.create_view( + "db.sync_mv", + parse_one( + "SELECT region, SUM(amount) AS total FROM db.orders GROUP BY region", read="doris" + ), + replace=True, + materialized=True, + view_properties={ + "materialized_type": exp.Literal.string("sync"), + "source_table": exp.Literal.string("db.orders"), + }, + ) + + assert to_sql_calls(adapter) == [ + "DROP MATERIALIZED VIEW IF EXISTS `db`.`sync_mv` ON `db`.`orders`", + "CREATE MATERIALIZED VIEW `db`.`sync_mv` AS SELECT `region` AS `mv_region`, SUM(`amount`) AS `total` FROM `db`.`orders` GROUP BY `region`", + ] + + adapter.cursor.execute.reset_mock() + + adapter.create_view( + "sync_mv", + parse_one("SELECT id FROM db.orders", read="doris"), + replace=True, + materialized=True, + view_properties={ + "materialized_type": exp.Literal.string("SYNC"), + "source_table": exp.to_table("db.orders"), + }, + ) + + assert to_sql_calls(adapter) == [ + "USE `db`", + "DROP MATERIALIZED VIEW IF EXISTS `sync_mv` ON `db`.`orders`", + "USE `db`", + "CREATE MATERIALIZED VIEW `sync_mv` AS SELECT `id` AS `mv_id` FROM `db`.`orders`", + ] + + +def test_delete_with_subquery_fallback_and_cardinality_aware_composite_batching( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # When DELETE ... USING fails on DUPLICATE/AGGREGATE KEY OLAP tables, fallback to literal batching + # and pick the highest-cardinality column (id: 3 values) to batch in IN (...) while grouping by low-cardinality ds (1 value). + mocker.patch.object( + adapter, + "fetchall", + return_value=[(1, "2024-01-01"), (2, "2024-01-01"), (3, "2024-01-01")], + ) + + first_delete = True + + def execute_side_effect(sql, *args, **kwargs): + nonlocal first_delete + if first_delete and "USING" in str(sql): + first_delete = False + raise Exception("errCode = 2, detailMessage = delete using only support unique key") + + adapter.cursor.execute.side_effect = execute_side_effect + + key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", exp.column("id"), exp.column("ds")) + where = exp.Paren(this=key_exp).isin(query=exp.select(key_exp).from_("src_table")) + adapter.delete_from("target_table", where=where) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `target_table` AS `_t1` USING (SELECT `id`, `ds` FROM `src_table`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`ds` = `_t2`.`ds`", + "DELETE FROM `target_table` WHERE `ds` = '2024-01-01' AND `id` IN (1, 2, 3)", + ] + + +def test_rewrite_between_in_where_preserves_not_between_and_wraps_or( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.delete_from( + "target_table", + where=parse_one( + "(ds BETWEEN '2024-01-01' AND '2024-01-07' OR status = 'expired') AND id NOT BETWEEN 10 AND 20", + read="doris", + ), + ) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `target_table` WHERE ((`ds` >= '2024-01-01' AND `ds` <= '2024-01-07') OR `status` = 'expired') AND NOT `id` BETWEEN 10 AND 20", + ] + + +@pytest.mark.parametrize( + "key_cols,rows", + [ + (["id"], [(1,), (None,), (2,), (3,)]), + (["id", "ds"], [(1, "2024-01-01"), (2, "2024-01-01"), (None, "2024-01-01"), (1, None)]), + ( + ["tenant", "region", "id"], + [ + ("t1", "us", 10), + ("t1", "us", 20), + ("t1", "us", None), + ("t1", None, 30), + (None, None, None), + ], + ), + ], +) +def test_delete_by_literal_keys_doris_delete_handler_invariants( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], + mocker: MockerFixture, + key_cols: t.List[str], + rows: t.List[t.Tuple[t.Any, ...]], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch.object(adapter, "fetchall", return_value=rows) + + adapter._delete_by_literal_keys( + exp.to_table("db.target_tbl"), + [exp.column(c) for c in key_cols], + parse_one(f"SELECT {', '.join(key_cols)} FROM db.src_tbl", read="doris", into=exp.Select), + ) + + sql_calls = to_sql_calls(adapter) + assert len(sql_calls) >= 1 + for sql in sql_calls: + parsed = parse_one(sql, read="doris") + assert isinstance(parsed, exp.Delete) + # Doris DeleteHandler forbids OR and BETWEEN in DELETE WHERE predicates + assert parsed.find(exp.Or) is None + assert parsed.find(exp.Between) is None + + +@pytest.mark.parametrize( + "non_key_type", + ["FLOAT", "DOUBLE", "ARRAY", "MAP", "STRUCT", "JSON", "VARIANT"], +) +def test_olap_key_type_and_prefix_invariants( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], + non_key_type: str, +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "db.complex_tbl", + target_columns_to_types={ + "bad_first_col": exp.DataType.build(non_key_type, dialect="doris"), + "valid_key": exp.DataType.build("BIGINT", dialect="doris"), + }, + ) + + sql_calls = to_sql_calls(adapter) + assert len(sql_calls) == 1 + # valid_key must be moved to column index 0 ahead of non-keyable bad_first_col + assert sql_calls[0].startswith( + "CREATE TABLE IF NOT EXISTS `db`.`complex_tbl` (`valid_key` BIGINT, `bad_first_col`" + ) + + +def test_adjust_physical_properties_for_incremental( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Single unique_key promoted when no explicit key is present + props = adapter.adjust_physical_properties_for_incremental( + {}, + requires_delete_capable_table=True, + unique_key=[exp.to_column("id")], + model_name="db.incremental_model", + ) + assert props == {"unique_key": exp.to_column("id")} + + # Composite unique_key promoted as exp.Tuple + props_composite = adapter.adjust_physical_properties_for_incremental( + {}, + requires_delete_capable_table=True, + unique_key=[exp.to_column("id"), exp.to_column("ds")], + model_name="db.incremental_model", + ) + assert props_composite == { + "unique_key": exp.Tuple(expressions=[exp.to_column("id"), exp.to_column("ds")]) + } + + # Explicit duplicate_key in physical_properties is preserved without overwriting + props_explicit = adapter.adjust_physical_properties_for_incremental( + {"duplicate_key": exp.to_column("id")}, + requires_delete_capable_table=True, + unique_key=[exp.to_column("id")], + model_name="db.incremental_model", + ) + assert props_explicit == {"duplicate_key": exp.to_column("id")} + + # Verify end-to-end CREATE TABLE with promoted unique_key reorders key columns first and emits UNIQUE KEY + adapter.create_table( + "db.incremental_model", + target_columns_to_types={ + "name": exp.DataType.build("VARCHAR(50)", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + }, + table_properties=props, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`incremental_model` (`id` INT, `name` VARCHAR(50)) UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10" + ] diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py index c506d401d9..800fb5fb19 100644 --- a/tests/core/test_connection_config.py +++ b/tests/core/test_connection_config.py @@ -15,6 +15,7 @@ ClickhouseConnectionConfig, ConnectionConfig, DatabricksConnectionConfig, + DorisConnectionConfig, DuckDBAttachOptions, DuckDBConnectionConfig, FabricConnectionConfig, @@ -2479,3 +2480,116 @@ def test_starrocks(make_config): assert advanced_config.charset == "utf8mb4" assert advanced_config.ssl_disabled is True assert advanced_config.concurrent_tasks == 10 + + +def test_doris(make_config): + """Test DorisConnectionConfig basic functionality and multi-catalog/session options""" + # Basic configuration + config = make_config( + type="doris", + host="localhost", + user="root", + password="password", + port=9030, + database="demo", + check_import=False, + ) + assert isinstance(config, DorisConnectionConfig) + assert config.type_ == "doris" + assert config.host == "localhost" + assert config.user == "root" + assert config.password == "password" + assert config.port == 9030 + assert config.database == "demo" + assert config.catalog is None + assert config.get_catalog() is None + assert config.DIALECT == "doris" + assert config.DISPLAY_NAME == "Apache Doris" + assert config.DISPLAY_ORDER == 20 + assert config.is_recommended_for_state_sync is False + assert config.is_forbidden_for_state_sync is True + assert config._cursor_init is None + + # Test with minimal configuration (using default port and optional password) + minimal_config = make_config( + type="doris", + host="fe.doris.cluster", + user="doris_user", + check_import=False, + ) + assert isinstance(minimal_config, DorisConnectionConfig) + assert minimal_config.port == 9030 # Default Doris FE port + assert minimal_config.host == "fe.doris.cluster" + assert minimal_config.user == "doris_user" + assert minimal_config.password is None + + # Test with Iceberg catalog, session variables, and TLS options + advanced_config = make_config( + type="doris", + host="doris-fe", + user="admin", + password="admin123", + port=9030, + database="analytics", + catalog="iceberg_cat", + charset="utf8mb4", + ssl_ca="/path/to/ca.pem", + ssl_verify_cert=True, + session_variables={"query_timeout": 300, "workload_group": "etl_group"}, + concurrent_tasks=10, + check_import=False, + ) + assert isinstance(advanced_config, DorisConnectionConfig) + assert advanced_config.charset == "utf8mb4" + assert advanced_config.catalog == "iceberg_cat" + assert advanced_config.get_catalog() == "iceberg_cat" + assert advanced_config.concurrent_tasks == 10 + assert "database" not in advanced_config._static_connection_kwargs + assert advanced_config._static_connection_kwargs["ssl"] == {"ca": "/path/to/ca.pem"} + assert advanced_config._static_connection_kwargs["ssl_verify_cert"] is True + + mock_cursor = MagicMock() + assert advanced_config._cursor_init is not None + advanced_config._cursor_init(mock_cursor) + assert [call.args[0] for call in mock_cursor.execute.call_args_list] == [ + "SET `query_timeout` = 300", + "SET `workload_group` = 'etl_group'", + "SWITCH `iceberg_cat`", + "USE `analytics`", + ] + + ssl_bool_config = make_config( + type="doris", + host="doris-fe", + user="admin", + ssl=True, + check_import=False, + ) + assert ssl_bool_config._static_connection_kwargs["ssl"] == {} + + ssl_false_config = make_config( + type="doris", + host="doris-fe", + user="admin", + ssl=False, + check_import=False, + ) + assert "ssl" not in ssl_false_config._static_connection_kwargs + + with pytest.raises(ConfigError, match="Invalid Doris session variable name"): + make_config( + type="doris", + host="doris-fe", + user="admin", + session_variables={"bad key; DROP": 1}, + check_import=False, + ) + + with pytest.raises(ConfigError, match="Invalid Doris session variable value"): + make_config( + type="doris", + host="doris-fe", + user="admin", + session_variables={"timeout": float("nan")}, + check_import=False, + ) From 8eeba87c6b867306e5a2ad1fc122777590258e4f Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sun, 18 Oct 2026 17:54:26 -0700 Subject: [PATCH 02/27] refactor(doris): scope sync MV USE db, unify _to_sql, and document Doris constraints Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 51 +++++++++++++++------------- 1 file changed, 27 insertions(+), 24 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index a3ca961520..32429ce3c7 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -6,6 +6,7 @@ import re import typing as t +import pandas as pd # noqa: TID253 from sqlglot import exp, parse, parse_one from sqlmesh.core.dialect import to_schema @@ -430,12 +431,12 @@ def create_view( if materialized: raise SQLMeshError( f"Doris only supports creating materialized views in the 'internal' catalog, got '{effective_catalog}' " - f"for view '{target_table.sql(dialect=self.dialect)}'." + f"for view '{self._to_sql(target_table)}'." ) logger.warning( "Skipping view creation for '%s' because Doris does not support creating views in external catalog '%s'. " "Configure `environment_catalog_mapping` in your SQLMesh configuration to route virtual environment views to the 'internal' catalog.", - target_table.sql(dialect=self.dialect), + self._to_sql(target_table), effective_catalog, ) return @@ -485,10 +486,11 @@ def _create_materialized_view( source_columns: t.Optional[t.List[str]] = None, **create_kwargs: t.Any, ) -> None: - import pandas as pd - query_or_df = self._native_df_to_pandas_df(query_or_df) + # Inline DataFrames as a SELECT ... FROM VALUES query rather than delegating DataFrame + # handling to _get_source_queries_and_columns_to_types, which would otherwise materialize + # a temporary table that gets dropped as soon as the SourceQuery context manager exits. if isinstance(query_or_df, pd.DataFrame): values: t.List[t.Tuple[t.Any, ...]] = list( query_or_df.itertuples(index=False, name=None) @@ -594,7 +596,10 @@ def _create_materialized_view( aliased_selects.append(sel_expr) query.set("expressions", aliased_selects) - if not target_table.db: + # Doris synchronous materialized views are single-table rollup indexes whose DDL syntax + # (`CREATE MATERIALIZED VIEW AS SELECT ... FROM .`) forbids a + # `.` qualifier on `` and requires an active database on the session. + if is_sync_mv and not target_table.db: source_db: t.Optional[str] = None if view_properties and view_properties.get("source_table"): raw_src = view_properties["source_table"] @@ -610,9 +615,7 @@ def _create_materialized_view( if first_tbl and first_tbl.db: source_db = first_tbl.db if source_db: - db_sql = exp.to_identifier(source_db, quoted=True).sql( - dialect=self.dialect, identify=True - ) + db_sql = self._to_sql(exp.to_identifier(source_db, quoted=True)) self.execute(f"USE {db_sql}") self.execute( @@ -645,7 +648,7 @@ def drop_view( mat_type = self._extract_str_value(view_properties.get("materialized_type")) source_table = view_properties.get("source_table") if mat_type and mat_type.upper() == "SYNC" and source_table: - view_table_sql = target_table.sql(dialect=self.dialect, identify=True) + view_table_sql = self._to_sql(target_table) if isinstance(source_table, exp.Column): source_tbl_exp = exp.table_( source_table.name, @@ -657,12 +660,12 @@ def drop_view( else: source_table_str = self._extract_str_value(source_table) or str(source_table) source_tbl_exp = exp.to_table(source_table_str, dialect=self.dialect) + # Unqualified synchronous materialized views require the source table's database + # to be active on the session before executing `DROP MATERIALIZED VIEW ON `. if not target_table.db and source_tbl_exp.db: - db_sql = exp.to_identifier(source_tbl_exp.db, quoted=True).sql( - dialect=self.dialect, identify=True - ) + db_sql = self._to_sql(exp.to_identifier(source_tbl_exp.db, quoted=True)) self.execute(f"USE {db_sql}") - source_table_sql = source_tbl_exp.sql(dialect=self.dialect, identify=True) + source_table_sql = self._to_sql(source_tbl_exp) exists_clause = "IF EXISTS " if ignore_if_not_exists else "" self.execute( f"DROP MATERIALIZED VIEW {exists_clause}{view_table_sql} ON {source_table_sql}" @@ -696,21 +699,19 @@ def _create_table_like( def _build_create_comment_table_exp( self, table: exp.Table, table_comment: str, table_kind: str = "TABLE" ) -> exp.Comment | str: - table_sql = table.sql(dialect=self.dialect, identify=True) - comment_sql = exp.Literal.string(self._truncate_table_comment(table_comment.strip())).sql( - dialect=self.dialect + table_sql = self._to_sql(table) + comment_sql = self._to_sql( + exp.Literal.string(self._truncate_table_comment(table_comment.strip())) ) return f"ALTER TABLE {table_sql} MODIFY COMMENT {comment_sql}" def _build_create_comment_column_exp( self, table: exp.Table, column_name: str, column_comment: str, table_kind: str = "TABLE" ) -> exp.Comment | str: - table_sql = table.sql(dialect=self.dialect, identify=True) - col_sql = exp.to_identifier(column_name, quoted=True).sql( - dialect=self.dialect, identify=True - ) - comment_sql = exp.Literal.string(self._truncate_column_comment(column_comment.strip())).sql( - dialect=self.dialect + table_sql = self._to_sql(table) + col_sql = self._to_sql(exp.to_identifier(column_name, quoted=True)) + comment_sql = self._to_sql( + exp.Literal.string(self._truncate_column_comment(column_comment.strip())) ) return f"ALTER TABLE {table_sql} MODIFY COLUMN {col_sql} COMMENT {comment_sql}" @@ -727,6 +728,10 @@ def _truncate_table(self, table_name: TableName) -> None: raise e caught_err = e + # When TRUNCATE TABLE is unsupported (e.g., on Iceberg tables), clear all rows via DELETE. + # Doris's DeleteHandler rejects DELETE without WHERE, constant WHERE TRUE, OR predicates, + # and floating-point/complex column predicates, so we pick a key-compatible column and + # execute two complementary `IS NOT NULL` and `IS NULL` deletes. cols_to_types = self.columns(table_name) if cols_to_types: valid_col_name = next( @@ -998,8 +1003,6 @@ def _delete_by_literal_keys( if not rows: return - import pandas as pd - extra_conditions = extra_conditions or [] cols = [exp.column(k.name, quoted=True) for k in key_exprs] batch_size = self.DEFAULT_BATCH_SIZE From 51c61e2f73a999ec5f5a4d3f1d35cb85f9c1de24 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sun, 18 Oct 2026 18:12:38 -0700 Subject: [PATCH 03/27] refactor(doris): extract static helpers and prefer docker exec in wait-for-db.sh Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- .github/scripts/wait-for-db.sh | 5 ++- sqlmesh/core/engine_adapter/doris.py | 54 ++++++++++++++-------------- 2 files changed, 32 insertions(+), 27 deletions(-) diff --git a/.github/scripts/wait-for-db.sh b/.github/scripts/wait-for-db.sh index 6077324dbd..cda12bd801 100755 --- a/.github/scripts/wait-for-db.sh +++ b/.github/scripts/wait-for-db.sh @@ -111,7 +111,10 @@ doris_ready() { HOSTNAME="127.0.0.1" fi 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');" - if command -v mysql > /dev/null 2>&1; then + if docker ps --format '{{.Names}}' 2>/dev/null | grep -qx 'doris-fe-01'; then + docker exec -i doris-fe-01 mysql -h127.0.0.1 -P9030 -uroot -e "SHOW BACKENDS\G" 2>/dev/null | grep -q "Alive: true" && \ + docker exec -i doris-fe-01 mysql -h127.0.0.1 -P9030 -uroot -e "$INIT_SQL" 2>/dev/null + elif command -v mysql > /dev/null 2>&1; then mysql -h "$HOSTNAME" -P 9030 -u root -e "SHOW BACKENDS\G" 2>/dev/null | grep -q "Alive: true" && \ mysql -h "$HOSTNAME" -P 9030 -u root -e "$INIT_SQL" 2>/dev/null else diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 32429ce3c7..97b30ccafc 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -986,6 +986,17 @@ def _execute_delete_with_subquery( else: raise + @staticmethod + def _sanitize_key_value(val: t.Any) -> t.Any: + if val is None: + return None + try: + if pd.isna(val): + return None + except (TypeError, ValueError): + pass + return val + def _delete_by_literal_keys( self, table_name: TableName, @@ -1007,17 +1018,7 @@ def _delete_by_literal_keys( cols = [exp.column(k.name, quoted=True) for k in key_exprs] batch_size = self.DEFAULT_BATCH_SIZE - def _sanitize(val: t.Any) -> t.Any: - if val is None: - return None - try: - if pd.isna(val): - return None - except (TypeError, ValueError): - pass - return val - - sanitized_rows = [tuple(_sanitize(v) for v in row) for row in rows] + sanitized_rows = [tuple(self._sanitize_key_value(v) for v in row) for row in rows] if len(cols) == 1: col = cols[0] @@ -1269,6 +1270,20 @@ def _parse_partition_expressions( return parsed_partitioned_by, partition_kind + @staticmethod + def _to_raw_partition_sql(expr: t.Union[exp.Literal, exp.Var, str, t.Any]) -> exp.Var: + if isinstance(expr, exp.Literal): + raw_text = str(expr.this) + elif isinstance(expr, exp.Var): + raw_text = str(expr.this) + elif isinstance(expr, str): + raw_text = expr + else: + return expr + if ";" in raw_text or "--" in raw_text or "/*" in raw_text: + raise SQLMeshError(f"Invalid partition definition expression: {raw_text}") + return exp.Var(this=raw_text, quoted=False) + def _build_partitioned_by_exp( self, partitioned_by: t.List[exp.Expr], @@ -1309,25 +1324,12 @@ def _build_partitioned_by_exp( if not partitions and not is_auto and not dynamic_partition_enabled: return None - def to_raw_sql(expr: t.Union[exp.Literal, exp.Var, str, t.Any]) -> exp.Var: - if isinstance(expr, exp.Literal): - raw_text = str(expr.this) - elif isinstance(expr, exp.Var): - raw_text = str(expr.this) - elif isinstance(expr, str): - raw_text = expr - else: - return expr - if ";" in raw_text or "--" in raw_text or "/*" in raw_text: - raise SQLMeshError(f"Invalid partition definition expression: {raw_text}") - return exp.Var(this=raw_text, quoted=False) - create_expressions: t.Optional[t.List[exp.Expr]] = None if partitions: if isinstance(partitions, exp.Tuple): - create_expressions = [to_raw_sql(e) for e in partitions.expressions] + create_expressions = [self._to_raw_partition_sql(e) for e in partitions.expressions] else: - create_expressions = [to_raw_sql(partitions)] + create_expressions = [self._to_raw_partition_sql(partitions)] elif is_auto or dynamic_partition_enabled: create_expressions = [] From af3a80bc30748762ad785cc47461bd434cd37839 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sun, 18 Oct 2026 18:45:49 -0700 Subject: [PATCH 04/27] 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 Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- .github/scripts/wait-for-db.sh | 6 +- sqlmesh/core/config/connection.py | 3 + sqlmesh/core/engine_adapter/doris.py | 13 +- .../integration/test_integration.py | 68 +++------- .../integration/test_integration_doris.py | 122 ++++++++++++++++++ tests/core/engine_adapter/test_doris.py | 24 ++++ 6 files changed, 179 insertions(+), 57 deletions(-) diff --git a/.github/scripts/wait-for-db.sh b/.github/scripts/wait-for-db.sh index cda12bd801..33105ffd5b 100755 --- a/.github/scripts/wait-for-db.sh +++ b/.github/scripts/wait-for-db.sh @@ -110,7 +110,7 @@ doris_ready() { if [ "$HOSTNAME" = "localhost" ]; then HOSTNAME="127.0.0.1" fi - 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');" + INIT_SQL="ADMIN SET FRONTEND CONFIG ('force_olap_table_replication_num' = '1');" if docker ps --format '{{.Names}}' 2>/dev/null | grep -qx 'doris-fe-01'; then docker exec -i doris-fe-01 mysql -h127.0.0.1 -P9030 -uroot -e "SHOW BACKENDS\G" 2>/dev/null | grep -q "Alive: true" && \ 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: alive_idx = cols.index('Alive') if not rows or not any(str(r[alive_idx]).lower() == 'true' for r in rows): sys.exit(1) - for cfg in ('default_replication_num', 'force_olap_table_replication_num'): - cur.execute(f\"ADMIN SET FRONTEND CONFIG ('{cfg}' = '1')\") - cur.execute(\"ADMIN SET FRONTEND CONFIG ('enable_unicode_name_support' = 'true')\") + cur.execute(\"ADMIN SET FRONTEND CONFIG ('force_olap_table_replication_num' = '1')\") " 2>/dev/null fi } diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 8e55c7bdd7..1334810a1c 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -2851,6 +2851,9 @@ def init(cursor: t.Any) -> None: return init + def get_catalog(self) -> t.Optional[str]: + return self.catalog + CONNECTION_CONFIG_TO_TYPE = { # Map all subclasses of ConnectionConfig to the value of their `type_` field. diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 97b30ccafc..64c59e12c5 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -891,10 +891,19 @@ def _extract_subquery_from_where( else: key_exprs = [lhs.copy()] - if not key_exprs or not all(isinstance(k, exp.Column) for k in key_exprs): + normalized_keys: t.List[exp.Expr] = [] + for k in key_exprs: + if isinstance(k, exp.Column): + normalized_keys.append(k) + elif isinstance(k, exp.Identifier): + normalized_keys.append(exp.Column(this=k.copy())) + else: + return None + + if not normalized_keys: return None - return key_exprs, subquery, extra_conditions + return normalized_keys, subquery, extra_conditions @staticmethod def _unwrap_concat_ws(expr: exp.Expr) -> t.Optional[exp.Expr]: diff --git a/tests/core/engine_adapter/integration/test_integration.py b/tests/core/engine_adapter/integration/test_integration.py index a2cdfb00b3..948e7c8237 100644 --- a/tests/core/engine_adapter/integration/test_integration.py +++ b/tests/core/engine_adapter/integration/test_integration.py @@ -440,6 +440,12 @@ def test_materialized_view(ctx_query_and_df: TestContext): "which this generic test does not provide; StarRocks MVs are covered by " "test_integration_starrocks.py" ) + if ctx.engine_adapter.dialect == "doris": + pytest.skip( + "Doris materialized views require engine-specific properties (SYNC source_table or MTMV DISTRIBUTED BY), " + "which this generic test does not provide; Doris MVs are covered by " + "test_integration_doris.py" + ) input_data = pd.DataFrame( [ {"id": 1, "ds": "2022-01-01"}, @@ -450,22 +456,8 @@ def test_materialized_view(ctx_query_and_df: TestContext): source_table = ctx.table("source_table") ctx.engine_adapter.ctas(source_table, ctx.input_data(input_data), ctx.columns_to_types) view = ctx.table("test_view") - # 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 - if ctx.engine_adapter.dialect == "doris": - view = exp.to_table(view.name) - # For Doris, select only the first column name from columns_to_types. - # Otherwise, Doris will raise an error if the materialized view has the same schema as the base table. - first_column = next(iter(ctx.columns_to_types)) - view_query = exp.select(first_column).from_(source_table) - ctx.engine_adapter.create_view( - view, - view_query, - materialized=True, - view_properties={"materialized_type": "SYNC", "source_table": source_table}, - ) - else: - view_query = exp.select(*ctx.columns_to_types).from_(source_table) - ctx.engine_adapter.create_view(view, view_query, materialized=True) + view_query = exp.select(*ctx.columns_to_types).from_(source_table) + ctx.engine_adapter.create_view(view, view_query, materialized=True) results = ctx.get_metadata_results() # Redshift considers the underlying dataset supporting materialized views as a table therefore we get 2 # tables in the result @@ -474,31 +466,13 @@ def test_materialized_view(ctx_query_and_df: TestContext): else: assert len(results.tables) == 1 assert len(results.views) == 0 - # Doris information_schema.tables does not show materialized views and cannot query the synchronized materialized view - if ctx.engine_adapter.dialect != "doris": - assert len(results.materialized_views) == 1 - assert results.materialized_views[0] == view.name - ctx.compare_with_current(view, input_data) + assert len(results.materialized_views) == 1 + assert results.materialized_views[0] == view.name + ctx.compare_with_current(view, input_data) # Make sure that dropping a materialized view also works - if ctx.engine_adapter.dialect == "doris": - # Wait for the materialized view to be created by retrying drop until it succeeds - def drop_view_success(): - try: - ctx.engine_adapter.drop_view( - view, - ignore_if_not_exists=False, - materialized=True, - view_properties={"materialized_type": "SYNC", "source_table": source_table}, - ) - return True - except Exception: - return False - - wait_until(drop_view_success, attempts=5, wait=1) - else: - ctx.engine_adapter.drop_view(view, materialized=True) - results = ctx.get_metadata_results() - assert len(results.materialized_views) == 0 + ctx.engine_adapter.drop_view(view, materialized=True) + results = ctx.get_metadata_results() + assert len(results.materialized_views) == 0 def test_drop_schema(ctx: TestContext): @@ -1031,12 +1005,8 @@ def test_merge(ctx_query_and_df: TestContext): # Athena only supports MERGE on Iceberg tables # And it cant fall back to a logical merge on Hive tables because it cant delete records table_format = "iceberg" if ctx.dialect == "athena" else None - # Doris needs to have a UNIQUE KEY to use MERGE - table_properties = {"unique_key": "id"} if ctx.dialect == "doris" else None - ctx.engine_adapter.create_table( - table, ctx.columns_to_types, table_format=table_format, table_properties=table_properties - ) + ctx.engine_adapter.create_table(table, ctx.columns_to_types, table_format=table_format) input_data = pd.DataFrame( [ {"id": 1, "ds": "2022-01-01"}, @@ -1101,15 +1071,11 @@ def test_merge_source_columns(ctx_query_and_df: TestContext): # Athena only supports MERGE on Iceberg tables # And it cant fall back to a logical merge on Hive tables because it cant delete records table_format = "iceberg" if ctx.dialect == "athena" else None - # Doris needs to have a UNIQUE KEY to use MERGE - table_properties = {"unique_key": "id"} if ctx.dialect == "doris" else None columns_to_types = ctx.columns_to_types.copy() columns_to_types["ignored_column"] = exp.DataType.build("int") - ctx.engine_adapter.create_table( - table, columns_to_types, table_format=table_format, table_properties=table_properties - ) + ctx.engine_adapter.create_table(table, columns_to_types, table_format=table_format) input_data = pd.DataFrame( [ {"id": 1, "ds": "2022-01-01", "ignored_source": "ignored_value"}, @@ -4058,7 +4024,7 @@ def test_unicode_characters(ctx: TestContext, tmp_path: Path): # I also think Spark may not support unicode in general but that would need to be verified. if not ctx.engine_adapter.QUOTE_IDENTIFIERS_IN_VIEWS: pytest.skip("Skipping as these engines have issues with unicode characters in model names") - # Doris default setting `enable_unicode_name_support=false` so it is incompatible with unicode characters in model names + # Doris static fe.conf setting `enable_unicode_name_support` defaults to false (and is not runtime-mutable) if ctx.dialect == "doris": pytest.skip( "Skipping as Doris default setting has issues with unicode characters in model names" diff --git a/tests/core/engine_adapter/integration/test_integration_doris.py b/tests/core/engine_adapter/integration/test_integration_doris.py index 3e29bb53c5..9f580b5d58 100644 --- a/tests/core/engine_adapter/integration/test_integration_doris.py +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -98,3 +98,125 @@ def test_doris_iceberg_local_catalog(ctx: TestContext): except Exception: pass adapter.set_current_catalog("internal") + + +def test_materialized_view_sync(ctx: TestContext): + import pandas as pd # noqa: TID253 + from tests.core.engine_adapter.integration import wait_until + + input_data = pd.DataFrame( + [ + {"id": 1, "ds": "2022-01-01"}, + {"id": 2, "ds": "2022-01-02"}, + {"id": 3, "ds": "2022-01-03"}, + ] + ) + source_table = ctx.table("source_table") + ctx.engine_adapter.ctas(source_table, ctx.input_data(input_data), ctx.columns_to_types) + view = exp.to_table("test_sync_mv") + first_column = next(iter(ctx.columns_to_types)) + view_query = exp.select(first_column).from_(source_table) + ctx.engine_adapter.create_view( + view, + view_query, + materialized=True, + view_properties={ + "materialized_type": exp.Literal.string("SYNC"), + "source_table": source_table, + }, + ) + + def drop_view_success() -> bool: + try: + ctx.engine_adapter.drop_view( + view, + ignore_if_not_exists=False, + materialized=True, + view_properties={ + "materialized_type": exp.Literal.string("SYNC"), + "source_table": source_table, + }, + ) + return True + except Exception: + return False + + wait_until(drop_view_success, attempts=5, wait=1) + + +def test_materialized_view_async_mtmv(ctx: TestContext): + import pandas as pd # noqa: TID253 + from tests.core.engine_adapter.integration import wait_until + + input_data = pd.DataFrame( + [ + {"id": 1, "ds": "2022-01-01"}, + {"id": 2, "ds": "2022-01-02"}, + {"id": 3, "ds": "2022-01-03"}, + ] + ) + source_table = ctx.table("mtmv_source") + ctx.engine_adapter.ctas(source_table, ctx.input_data(input_data), ctx.columns_to_types) + view = ctx.table("test_async_mtmv") + view_query = exp.select(*ctx.columns_to_types).from_(source_table) + ctx.engine_adapter.create_view( + view, + view_query, + materialized=True, + view_properties={ + "build": exp.Literal.string("IMMEDIATE"), + "refresh": exp.Literal.string("COMPLETE"), + "refresh_trigger": exp.Literal.string("ON MANUAL"), + "distributed_by": exp.Literal.string("RANDOM BUCKETS AUTO"), + }, + ) + wait_until(lambda: len(ctx.get_current_data(view)) == len(input_data), attempts=10, wait=1) + ctx.compare_with_current(view, input_data) + ctx.engine_adapter.drop_view(view, materialized=True) + + +def test_merge_unique_key(ctx: TestContext): + import pandas as pd # noqa: TID253 + + table = ctx.table("test_merge_uk") + ctx.engine_adapter.create_table( + table, + ctx.columns_to_types, + table_properties={"unique_key": exp.Literal.string("id")}, + ) + input_data = pd.DataFrame( + [ + {"id": 1, "ds": "2022-01-01"}, + {"id": 2, "ds": "2022-01-02"}, + {"id": 3, "ds": "2022-01-03"}, + ] + ) + ctx.engine_adapter.merge( + table, + ctx.input_data(input_data), + target_columns_to_types=None, + unique_key=[exp.to_identifier("id")], + ) + merge_data = pd.DataFrame( + [ + {"id": 2, "ds": "2022-01-10"}, + {"id": 4, "ds": "2022-01-04"}, + ] + ) + ctx.engine_adapter.merge( + table, + ctx.input_data(merge_data), + target_columns_to_types=None, + unique_key=[exp.to_identifier("id")], + ) + ctx.compare_with_current( + table, + pd.DataFrame( + [ + {"id": 1, "ds": "2022-01-01"}, + {"id": 2, "ds": "2022-01-10"}, + {"id": 3, "ds": "2022-01-03"}, + {"id": 4, "ds": "2022-01-04"}, + ] + ), + ) diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 098f52406a..83d696a7ec 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -1277,3 +1277,27 @@ def test_adjust_physical_properties_for_incremental( assert to_sql_calls(adapter) == [ "CREATE TABLE IF NOT EXISTS `db`.`incremental_model` (`id` INT, `name` VARCHAR(50)) UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10" ] + + +def test_merge_with_identifier_unique_key( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], + mocker: MockerFixture, +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch("sqlmesh.core.engine_adapter.doris.random_id", return_value="abcd1234") + adapter.merge( + target_table="db.target", + source_table=parse_one("SELECT id, ds FROM db.src"), + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "ds": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + unique_key=[exp.to_identifier("id")], + ) + sql_calls = to_sql_calls(adapter) + assert any( + call.startswith( + "DELETE FROM `db`.`target` AS `_t1` USING (SELECT `id` FROM `db`.`__temp_target_abcd1234`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id`" + ) + for call in sql_calls + ) From d03b3d97ffde9331ff5e8a6530a17b76b3b6b4f9 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:24:11 -0700 Subject: [PATCH 05/27] chore(ci): optimize doris docker startup and memory footprint in ci Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- .github/scripts/install-prerequisites.sh | 2 - .github/scripts/wait-for-db.sh | 63 +++++++++++-------- .../integration/docker/compose.doris.yaml | 40 ++++++++---- .../integration/test_integration_doris.py | 4 +- 4 files changed, 66 insertions(+), 43 deletions(-) diff --git a/.github/scripts/install-prerequisites.sh b/.github/scripts/install-prerequisites.sh index a955f0c47c..09998836d0 100755 --- a/.github/scripts/install-prerequisites.sh +++ b/.github/scripts/install-prerequisites.sh @@ -29,9 +29,7 @@ elif [ "$ENGINE" == "fabric" ]; then ENGINE_DEPENDENCIES="msodbcsql18" elif [ "$ENGINE" == "doris" ]; then - ENGINE_DEPENDENCIES="mysql-client" sudo sysctl -w vm.max_map_count=2000000 - sudo swapoff -a || true fi ALL_DEPENDENCIES="$COMMON_DEPENDENCIES $ENGINE_DEPENDENCIES" diff --git a/.github/scripts/wait-for-db.sh b/.github/scripts/wait-for-db.sh index 33105ffd5b..e4e6605b05 100755 --- a/.github/scripts/wait-for-db.sh +++ b/.github/scripts/wait-for-db.sh @@ -105,32 +105,43 @@ db2_ready() { } doris_ready() { - probe_port 9030 || return 1 - HOSTNAME=${DOCKER_HOSTNAME:-127.0.0.1} - if [ "$HOSTNAME" = "localhost" ]; then - HOSTNAME="127.0.0.1" - fi - INIT_SQL="ADMIN SET FRONTEND CONFIG ('force_olap_table_replication_num' = '1');" - if docker ps --format '{{.Names}}' 2>/dev/null | grep -qx 'doris-fe-01'; then - docker exec -i doris-fe-01 mysql -h127.0.0.1 -P9030 -uroot -e "SHOW BACKENDS\G" 2>/dev/null | grep -q "Alive: true" && \ - docker exec -i doris-fe-01 mysql -h127.0.0.1 -P9030 -uroot -e "$INIT_SQL" 2>/dev/null - elif command -v mysql > /dev/null 2>&1; then - mysql -h "$HOSTNAME" -P 9030 -u root -e "SHOW BACKENDS\G" 2>/dev/null | grep -q "Alive: true" && \ - mysql -h "$HOSTNAME" -P 9030 -u root -e "$INIT_SQL" 2>/dev/null - else - python3 -c " -import pymysql, sys -conn = pymysql.connect(host='$HOSTNAME', port=9030, user='root', autocommit=True) -with conn.cursor() as cur: - cur.execute('SHOW BACKENDS') - rows = cur.fetchall() - cols = [d[0] for d in cur.description] - alive_idx = cols.index('Alive') - if not rows or not any(str(r[alive_idx]).lower() == 'true' for r in rows): - sys.exit(1) - cur.execute(\"ADMIN SET FRONTEND CONFIG ('force_olap_table_replication_num' = '1')\") -" 2>/dev/null - fi + probe_port 9030 + + echo "Waiting for Doris FE to accept SQL connections..." + while true; do + if docker exec -i doris-fe mysql -h127.0.0.1 -P9030 -uroot -e "SHOW FRONTENDS\G" 2>/dev/null | grep -q "Alive: true"; then + echo "Doris FE is alive" + break + fi + sleep 2 + done + + docker exec -i doris-fe mysql -h127.0.0.1 -P9030 -uroot -e "ALTER SYSTEM ADD BACKEND '172.30.80.3:9050';" 2>/dev/null || true + + echo "Checking for 1 alive Doris backends..." + while true; do + echo "Checking Doris backends..." + ALIVE_BACKENDS=$(docker exec -i doris-fe mysql -h127.0.0.1 -P9030 -uroot -e "show backends \G" 2>/dev/null | grep -c "^ *Alive: true *$") + + # fallback value if failed to get number + if ! [[ "$ALIVE_BACKENDS" =~ ^[0-9]+$ ]]; then + echo "WARN: Unable to parse number of alive backends, got: '$ALIVE_BACKENDS'" + ALIVE_BACKENDS=0 + fi + + echo "Found $ALIVE_BACKENDS alive backends" + + if [ "$ALIVE_BACKENDS" -ge 1 ]; then + echo "Doris has 1 or more alive backends" + break + fi + + echo "Waiting for more backends to become alive..." + sleep 3 + done + + # set default replication num to 1 (there is only one be in the docker compose file) + docker exec -i doris-fe mysql -h127.0.0.1 -P9030 -uroot -e "ADMIN SET FRONTEND CONFIG ('force_olap_table_replication_num' = '1');" } echo "Waiting for $ENGINE to be ready..." diff --git a/tests/core/engine_adapter/integration/docker/compose.doris.yaml b/tests/core/engine_adapter/integration/docker/compose.doris.yaml index bb1b77d44a..50376667c5 100644 --- a/tests/core/engine_adapter/integration/docker/compose.doris.yaml +++ b/tests/core/engine_adapter/integration/docker/compose.doris.yaml @@ -1,11 +1,16 @@ services: - fe-01: + doris-fe: image: apache/doris:fe-4.0.8 - container_name: doris-fe-01 - hostname: fe-01 - environment: - - FE_SERVERS=fe1:172.30.80.2:9010 - - FE_ID=1 + container_name: doris-fe + hostname: doris-fe + entrypoint: + - /bin/bash + - -c + command: + - | + sed -i 's/-Xmx8192m -Xms8192m/-Xmx2048m -Xms512m/g' /opt/apache-doris/fe/conf/fe.conf + echo "priority_networks = 172.30.80.0/24" >> /opt/apache-doris/fe/conf/fe.conf + exec bash /opt/apache-doris/fe/bin/start_fe.sh --console ports: - "8030:8030" - "9030:9030" @@ -15,15 +20,24 @@ services: doris_net: ipv4_address: 172.30.80.2 - be-01: + doris-be: image: apache/doris:be-4.0.8 - container_name: doris-be-01 - hostname: be-01 + container_name: doris-be + hostname: doris-be depends_on: - - fe-01 - environment: - - FE_SERVERS=fe1:172.30.80.2:9010 - - BE_ADDR=172.30.80.3:9050 + - doris-fe + entrypoint: + - /bin/bash + - -c + command: + - | + export SKIP_CHECK_ULIMIT=true + sed -i 's/-Xmx2048m/-Xmx1024m/g' /opt/apache-doris/be/conf/be.conf + echo "priority_networks = 172.30.80.0/24" >> /opt/apache-doris/be/conf/be.conf + echo "mem_limit = 2G" >> /opt/apache-doris/be/conf/be.conf + echo "storage_flood_stage_usage_percent = 99" >> /opt/apache-doris/be/conf/be.conf + echo "storage_flood_stage_left_capacity_bytes = 104857600" >> /opt/apache-doris/be/conf/be.conf + exec bash /opt/apache-doris/be/bin/start_be.sh --console ports: - "8040:8040" - "9050:9050" diff --git a/tests/core/engine_adapter/integration/test_integration_doris.py b/tests/core/engine_adapter/integration/test_integration_doris.py index 9f580b5d58..3092d70c11 100644 --- a/tests/core/engine_adapter/integration/test_integration_doris.py +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -141,7 +141,7 @@ def drop_view_success() -> bool: except Exception: return False - wait_until(drop_view_success, attempts=5, wait=1) + wait_until(drop_view_success, attempts=30, wait=1) def test_materialized_view_async_mtmv(ctx: TestContext): @@ -170,7 +170,7 @@ def test_materialized_view_async_mtmv(ctx: TestContext): "distributed_by": exp.Literal.string("RANDOM BUCKETS AUTO"), }, ) - wait_until(lambda: len(ctx.get_current_data(view)) == len(input_data), attempts=10, wait=1) + wait_until(lambda: len(ctx.get_current_data(view)) == len(input_data), attempts=30, wait=1) ctx.compare_with_current(view, input_data) ctx.engine_adapter.drop_view(view, materialized=True) From 3c58131bed47093c831aa1dc9364cc5f17af541a Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Wed, 21 Oct 2026 09:01:51 -0700 Subject: [PATCH 06/27] fix(doris): address review feedback on create_view, reconnect retry, and docs Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/integrations/engines/doris.md | 5 +- sqlmesh/core/engine_adapter/doris.py | 71 +++++++++++++++++-------- tests/core/engine_adapter/test_doris.py | 45 +++++++++++++--- 3 files changed, 92 insertions(+), 29 deletions(-) diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index 86def97b47..26b4a8275d 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -1,5 +1,8 @@ # Apache Doris +!!! info + The Doris engine adapter is a community contribution. Due to this, only limited community support is available. + This page describes SQLMesh support for the [Apache Doris](https://doris.apache.org/) engine (versions 2.1, 3.x, 4.x, and 5.0), including native Doris OLAP tables and Apache Iceberg multi-catalog lakehouse tables. !!! note @@ -260,7 +263,7 @@ SQLMesh can also create, partition, and overwrite Apache Iceberg tables directly When materializing models directly inside an external Iceberg catalog: - Set `virtual_environment_mode: dev_only` (and keep `environment_suffix_target: schema` or `table`) in `config.yaml` so production tables are materialized directly in the Iceberg catalog without virtual views. - - Configure `environment_catalog_mapping` in `config.yaml` to route non-`prod` virtual environment views to the `internal` catalog (otherwise `DorisEngineAdapter` will log a warning and skip view creation in external catalogs during `dev` promotion): + - Configure `environment_catalog_mapping` in `config.yaml` to route non-`prod` virtual environment views to the `internal` catalog (otherwise `DorisEngineAdapter` will raise an error when attempting to create views in an external catalog): ```yaml virtual_environment_mode: dev_only environment_suffix_target: schema diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 64c59e12c5..c83e5d726f 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -244,7 +244,26 @@ def _fetch_native_df( self._connection_pool.get_cursor() return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) + _SAFE_RETRY_SQL_PREFIX_RE = re.compile( + r"\A(?:\s|/\*.*?\*/|--[^\r\n]*(?:[\r\n]|\Z))*(?:SELECT|SHOW|DESCRIBE|DESC|EXPLAIN|USE|SWITCH|SET)\b", + re.IGNORECASE | re.DOTALL, + ) + + @classmethod + def _is_safe_to_retry_sql(cls, sql: str) -> bool: + """Returns True only for read-only or idempotent session statements that cannot duplicate writes.""" + return bool(cls._SAFE_RETRY_SQL_PREFIX_RE.match(sql)) + + def ping(self) -> None: + self._connection_pool.get().ping(reconnect=False) + def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any) -> None: + if self._pre_ping: + try: + self.ping() + except Exception: + logger.info("Connection to the database was lost. Reconnecting...") + self._connection_pool.close() try: super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) except Exception as e: @@ -253,9 +272,13 @@ def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any is_disconnect = (err_code in (2006, 2013, 2055) and err_name == "OperationalError") or ( err_code == 0 and err_name == "InterfaceError" ) - if is_disconnect and not self._connection_pool.is_transaction_active: + if ( + is_disconnect + and not self._connection_pool.is_transaction_active + and self._is_safe_to_retry_sql(sql) + ): logger.warning( - "Doris connection lost (%s: %s), reconnecting and retrying once...", + "Doris connection lost (%s: %s) during read/session statement, reconnecting and retrying once...", err_name, err_code, ) @@ -425,35 +448,29 @@ def create_view( source_columns: t.Optional[t.List[str]] = None, **create_kwargs: t.Any, ) -> None: + """ + Doris view creation behavior: + - Regular VIEW: supports atomic `CREATE OR REPLACE VIEW` (since Doris 2.0) via `super().create_view`. + - MATERIALIZED VIEW: does NOT support `CREATE OR REPLACE MATERIALIZED VIEW`, so `replace=True` + drops any existing object/materialized view first before creating the new one. + """ target_table = exp.to_table(view_name) effective_catalog = target_table.catalog or self.default_catalog if effective_catalog and effective_catalog.lower() != "internal": - if materialized: - raise SQLMeshError( - f"Doris only supports creating materialized views in the 'internal' catalog, got '{effective_catalog}' " - f"for view '{self._to_sql(target_table)}'." - ) - logger.warning( - "Skipping view creation for '%s' because Doris does not support creating views in external catalog '%s'. " - "Configure `environment_catalog_mapping` in your SQLMesh configuration to route virtual environment views to the 'internal' catalog.", - self._to_sql(target_table), - effective_catalog, + view_kind = "materialized views" if materialized else "views" + raise SQLMeshError( + f"Doris only supports creating {view_kind} in the 'internal' catalog, got '{effective_catalog}' " + f"for view '{self._to_sql(target_table)}'. Configure `virtual_environment_mode: dev_only` and " + f"`environment_catalog_mapping` in your SQLMesh configuration to route virtual environment views " + f"to the 'internal' catalog." ) - return - if replace: - self.drop_view( - view_name, - ignore_if_not_exists=True, - materialized=materialized, - view_properties=view_properties, - ) if not materialized: return super().create_view( view_name, query_or_df, target_columns_to_types, - replace=False, + replace=replace, materialized=False, materialized_properties=materialized_properties, table_description=table_description, @@ -462,6 +479,18 @@ def create_view( source_columns=source_columns, **create_kwargs, ) + + # Doris does not support `CREATE OR REPLACE MATERIALIZED VIEW`, so drop first when `replace=True`. + if replace: + self.drop_data_object_on_type_mismatch( + self.get_data_object(view_name), DataObjectType.MATERIALIZED_VIEW + ) + self.drop_view( + view_name, + ignore_if_not_exists=True, + materialized=True, + view_properties=view_properties, + ) self._create_materialized_view( view_name, query_or_df, diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 83d696a7ec..aeb04d3e6d 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -25,8 +25,7 @@ def test_create_view(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdap ) assert to_sql_calls(adapter) == [ - "DROP VIEW IF EXISTS `test_view`", - "CREATE VIEW `test_view` AS SELECT `a` FROM `tbl`", + "CREATE OR REPLACE VIEW `test_view` AS SELECT `a` FROM `tbl`", "CREATE VIEW `test_view` AS SELECT `a` FROM `tbl`", "CREATE VIEW `test_view` AS SELECT `a` FROM `tbl`", ] @@ -34,6 +33,14 @@ def test_create_view(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdap def test_create_view_with_comment(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_view( + "test_view", + parse_one("SELECT a FROM tbl"), + replace=True, + target_columns_to_types={"a": exp.DataType.build("INT")}, + table_description="test_description", + column_descriptions={"a": "test_column_description"}, + ) adapter.create_view( "test_view", parse_one("SELECT a FROM tbl"), @@ -44,6 +51,7 @@ def test_create_view_with_comment(make_mocked_engine_adapter: t.Callable[..., Do ) assert to_sql_calls(adapter) == [ + "CREATE OR REPLACE VIEW `test_view` (`a` COMMENT 'test_column_description') COMMENT 'test_description' AS SELECT `a` FROM `tbl`", "CREATE VIEW `test_view` (`a` COMMENT 'test_column_description') COMMENT 'test_description' AS SELECT `a` FROM `tbl`", ] @@ -916,22 +924,45 @@ def side_effect(*args, **kwargs): assert call_count == 2 conn_mock.ping.assert_called_once_with(reconnect=True) + # Write statements (e.g. INSERT) must NOT be retried after OperationalError(2013) to avoid duplicating rows + conn_mock.reset_mock() + write_calls = 0 + + def write_side_effect(*args, **kwargs): + nonlocal write_calls + write_calls += 1 + raise OperationalError(2013, "Lost connection to MySQL server during query") + + adapter.cursor.execute.side_effect = write_side_effect + with pytest.raises(OperationalError): + adapter._execute("INSERT INTO `tbl` VALUES (1)") + + assert write_calls == 1 + conn_mock.ping.assert_not_called() -def test_external_catalog_view_skip_and_materialized_error( + +def test_external_catalog_view_error_and_drop_noop( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], ): from sqlmesh.utils.errors import SQLMeshError adapter = make_mocked_engine_adapter(DorisEngineAdapter) - # Regular view creation and deletion on external catalog is a safe no-op (warns instead of failing dev promotion) - adapter.create_view("iceberg_cat.silver__dev.events", parse_one("SELECT 1 AS id")) + # Regular view creation on an external catalog raises SQLMeshError pointing to virtual_environment_mode and environment_catalog_mapping + with pytest.raises( + SQLMeshError, + match=r"Doris only supports creating views in the 'internal' catalog, got 'iceberg_cat'.*virtual_environment_mode: dev_only.*environment_catalog_mapping", + ): + adapter.create_view("iceberg_cat.silver__dev.events", parse_one("SELECT 1 AS id")) + + # Drop view on external catalog remains a safe no-op during cleanup adapter.drop_view("iceberg_cat.silver__dev.events") assert to_sql_calls(adapter) == [] - # Materialized view on external catalog still raises SQLMeshError + # Materialized view on external catalog also raises SQLMeshError with pytest.raises( - SQLMeshError, match="Doris only supports creating materialized views in the 'internal'" + SQLMeshError, + match="Doris only supports creating materialized views in the 'internal' catalog", ): adapter.create_view( "iceberg_cat.silver.events_mv", From a3b2d77fcbada0b68874469addf8590aa4833a2f Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Wed, 21 Oct 2026 12:10:19 -0700 Subject: [PATCH 07/27] fix(doris): harden reconnect retry against timeouts and validate MV before drop Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/integrations/engines/doris.md | 1 + sqlmesh/core/engine_adapter/doris.py | 41 ++++++++++++------- .../integration/test_integration.py | 11 ++++- tests/core/engine_adapter/test_doris.py | 36 +++++++++++++++- 4 files changed, 71 insertions(+), 18 deletions(-) diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index 26b4a8275d..b5ab0271dd 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -34,6 +34,7 @@ pip install "sqlmesh[doris]" | `connect_timeout` | Connection timeout in seconds (Default: `10`) | int | N | | `read_timeout` | Read timeout in seconds | int | N | | `write_timeout` | Write timeout in seconds | int | N | +| `pre_ping` | Whether to ping the connection before executing a statement and reconnect if lost (Default: `False`) | bool | N | | `session_variables` | Session variables to set on each new connection via `SET = ` (e.g., `workload_group`, `query_timeout`) | dict | N | | `ssl` | Explicit SSL configuration dictionary passed to `PyMySQL` | dict | N | | `ssl_disabled` | Whether SSL is disabled | bool | N | diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index c83e5d726f..22c1dd6835 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -245,13 +245,13 @@ def _fetch_native_df( return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) _SAFE_RETRY_SQL_PREFIX_RE = re.compile( - r"\A(?:\s|/\*.*?\*/|--[^\r\n]*(?:[\r\n]|\Z))*(?:SELECT|SHOW|DESCRIBE|DESC|EXPLAIN|USE|SWITCH|SET)\b", + r"\A(?:\s|/\*.*?\*/|--[^\r\n]*(?:[\r\n]|\Z))*(?:SELECT|WITH|SHOW|DESCRIBE|DESC|EXPLAIN|USE|SWITCH)\b", re.IGNORECASE | re.DOTALL, ) @classmethod def _is_safe_to_retry_sql(cls, sql: str) -> bool: - """Returns True only for read-only or idempotent session statements that cannot duplicate writes.""" + """Returns True only for read-only or idempotent catalog/database switch statements that cannot duplicate writes.""" return bool(cls._SAFE_RETRY_SQL_PREFIX_RE.match(sql)) def ping(self) -> None: @@ -272,8 +272,13 @@ def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any is_disconnect = (err_code in (2006, 2013, 2055) and err_name == "OperationalError") or ( err_code == 0 and err_name == "InterfaceError" ) + is_timeout = ( + isinstance(getattr(e, "__context__", None), TimeoutError) + or "timed out" in str(e).lower() + ) if ( is_disconnect + and not is_timeout and not self._connection_pool.is_transaction_active and self._is_safe_to_retry_sql(sql) ): @@ -287,6 +292,8 @@ def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any conn.ping(reconnect=True) super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) else: + if is_disconnect and not self._connection_pool.is_transaction_active: + self._connection_pool.close() raise def columns( @@ -450,9 +457,9 @@ def create_view( ) -> None: """ Doris view creation behavior: - - Regular VIEW: supports atomic `CREATE OR REPLACE VIEW` (since Doris 2.0) via `super().create_view`. + - Regular VIEW: supports atomic `CREATE OR REPLACE VIEW` via `super().create_view`. - MATERIALIZED VIEW: does NOT support `CREATE OR REPLACE MATERIALIZED VIEW`, so `replace=True` - drops any existing object/materialized view first before creating the new one. + drops any existing object/materialized view first (after validating the new query and properties). """ target_table = exp.to_table(view_name) effective_catalog = target_table.catalog or self.default_catalog @@ -480,21 +487,11 @@ def create_view( **create_kwargs, ) - # Doris does not support `CREATE OR REPLACE MATERIALIZED VIEW`, so drop first when `replace=True`. - if replace: - self.drop_data_object_on_type_mismatch( - self.get_data_object(view_name), DataObjectType.MATERIALIZED_VIEW - ) - self.drop_view( - view_name, - ignore_if_not_exists=True, - materialized=True, - view_properties=view_properties, - ) self._create_materialized_view( view_name, query_or_df, target_columns_to_types=target_columns_to_types, + replace=replace, materialized_properties=materialized_properties, table_description=table_description, column_descriptions=column_descriptions, @@ -508,6 +505,7 @@ def _create_materialized_view( view_name: TableName, query_or_df: QueryOrDF, target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + replace: bool = True, materialized_properties: t.Optional[t.Dict[str, t.Any]] = None, table_description: t.Optional[str] = None, column_descriptions: t.Optional[t.Dict[str, str]] = None, @@ -625,6 +623,19 @@ def _create_materialized_view( aliased_selects.append(sel_expr) query.set("expressions", aliased_selects) + # Doris does not support `CREATE OR REPLACE MATERIALIZED VIEW`, so drop first when `replace=True` + # after query and table property validation has succeeded. + if replace: + self.drop_data_object_on_type_mismatch( + self.get_data_object(view_name), DataObjectType.MATERIALIZED_VIEW + ) + self.drop_view( + view_name, + ignore_if_not_exists=True, + materialized=True, + view_properties=view_properties, + ) + # Doris synchronous materialized views are single-table rollup indexes whose DDL syntax # (`CREATE MATERIALIZED VIEW AS SELECT ... FROM .
`) forbids a # `.` qualifier on `` and requires an active database on the session. diff --git a/tests/core/engine_adapter/integration/test_integration.py b/tests/core/engine_adapter/integration/test_integration.py index 948e7c8237..5d33dea1e2 100644 --- a/tests/core/engine_adapter/integration/test_integration.py +++ b/tests/core/engine_adapter/integration/test_integration.py @@ -101,11 +101,14 @@ def test_catalog_operations(ctx: TestContext): if ( ctx.engine_adapter.catalog_support.is_unsupported or ctx.engine_adapter.catalog_support.is_single_catalog_only - or ctx.dialect == "doris" ): pytest.skip( f"Engine adapter {ctx.engine_adapter.dialect} doesn't support catalog operations" ) + if ctx.dialect == "doris": + pytest.skip( + "Doris external catalogs require connector properties (`CREATE CATALOG ... PROPERTIES (...)`); covered in test_integration_doris.py" + ) # use a unique name so that integration tests on cloud databases can run in parallel catalog_name = "testing" if not ctx.is_remote else ctx.add_test_suffix("testing") @@ -155,10 +158,14 @@ def create_objects_and_validate(schema_name: str): assert len(results.materialized_views) == 0 assert len(results.non_temp_tables) == 2 - if ctx.engine_adapter.catalog_support.is_unsupported or ctx.dialect == "doris": + if ctx.engine_adapter.catalog_support.is_unsupported: pytest.skip( f"Engine adapter {ctx.engine_adapter.dialect} doesn't support catalog operations" ) + if ctx.dialect == "doris": + pytest.skip( + "Doris external catalogs require connector properties and do not support views; covered in test_integration_doris.py" + ) if ctx.dialect == "spark": pytest.skip( "Currently local spark is configured to have iceberg be the testing catalog and drop cascade doesn't work on iceberg. Skipping until we have time to fix." diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index aeb04d3e6d..582744f250 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -909,6 +909,7 @@ class OperationalError(Exception): adapter = make_mocked_engine_adapter(DorisEngineAdapter) conn_mock = mocker.MagicMock() mocker.patch.object(adapter._connection_pool, "get", return_value=conn_mock) + close_spy = mocker.spy(adapter._connection_pool, "close") call_count = 0 @@ -919,13 +920,34 @@ def side_effect(*args, **kwargs): raise OperationalError(2013, "Lost connection to MySQL server during query") adapter.cursor.execute.side_effect = side_effect - adapter._execute("SELECT 1") + adapter._execute("WITH `cte` AS (SELECT 1) SELECT * FROM `cte`") assert call_count == 2 conn_mock.ping.assert_called_once_with(reconnect=True) + # Read timeout (2013 with TimeoutError / "timed out") must NOT be retried even for SELECT + conn_mock.reset_mock() + close_spy.reset_mock() + timeout_calls = 0 + + def timeout_side_effect(*args, **kwargs): + nonlocal timeout_calls + timeout_calls += 1 + raise OperationalError( + 2013, "Lost connection to MySQL server during query (timed out)" + ) from TimeoutError("timed out") + + adapter.cursor.execute.side_effect = timeout_side_effect + with pytest.raises(OperationalError): + adapter._execute("SELECT 1") + + assert timeout_calls == 1 + conn_mock.ping.assert_not_called() + close_spy.assert_called_once() + # Write statements (e.g. INSERT) must NOT be retried after OperationalError(2013) to avoid duplicating rows conn_mock.reset_mock() + close_spy.reset_mock() write_calls = 0 def write_side_effect(*args, **kwargs): @@ -939,6 +961,7 @@ def write_side_effect(*args, **kwargs): assert write_calls == 1 conn_mock.ping.assert_not_called() + close_spy.assert_called_once() def test_external_catalog_view_error_and_drop_noop( @@ -970,6 +993,17 @@ def test_external_catalog_view_error_and_drop_noop( materialized=True, ) + # Materialized view with invalid property raises SQLMeshError BEFORE dropping existing MV + with pytest.raises(SQLMeshError, match="Invalid materialized view build mode"): + adapter.create_view( + "silver.events_mv", + parse_one("SELECT 1 AS id"), + replace=True, + materialized=True, + view_properties={"build": exp.Literal.string("INVALID_MODE")}, + ) + assert to_sql_calls(adapter) == [] + def test_duplicate_key_with_buckets_and_random_buckets_fallback( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], From d54cd012bdc5fbd4ae1a79e1e2a7642c1cb787c4 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Wed, 21 Oct 2026 14:43:29 -0700 Subject: [PATCH 08/27] fix(doris): harden AST parsing, MTMV metadata, key validation, and docs Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/concepts/models/overview.md | 1 + docs/guides/configuration.md | 2 +- docs/integrations/engines/doris.md | 48 +++- docs/quickstart/cli.md | 1 + sqlmesh/core/engine_adapter/doris.py | 243 ++++++++++++++---- .../engine_adapter/integration/__init__.py | 16 ++ .../integration/test_integration.py | 8 +- .../integration/test_integration_doris.py | 4 + tests/core/engine_adapter/test_doris.py | 114 ++++++++ 9 files changed, 375 insertions(+), 62 deletions(-) diff --git a/docs/concepts/models/overview.md b/docs/concepts/models/overview.md index 1fd16a96fc..ace9900b62 100644 --- a/docs/concepts/models/overview.md +++ b/docs/concepts/models/overview.md @@ -181,6 +181,7 @@ This table lists each engine's support for `TABLE` and `VIEW` object comments: | BigQuery | Y | Y | | ClickHouse | Y | Y | | Databricks | Y | Y | +| Doris | Y | Y | | DuckDB <=0.9 | N | N | | DuckDB >=0.10 | Y | Y | | MySQL | Y | Y | diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index 656478d8d6..a34fd189f1 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -976,10 +976,10 @@ Other state engines with fast and reliable database transactions (less tested th Unsupported state engines, even for development: * [ClickHouse](../integrations/engines/clickhouse.md) +* [Doris](../integrations/engines/doris.md) * [Spark](../integrations/engines/spark.md) * [StarRocks](../integrations/engines/starrocks.md) * [Trino](../integrations/engines/trino.md) -* [Doris](../integrations/engines/doris.md) This example gateway configuration uses Snowflake for the data warehouse connection and Postgres for the state backend connection: diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index b5ab0271dd..3cd891788b 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -36,13 +36,15 @@ pip install "sqlmesh[doris]" | `write_timeout` | Write timeout in seconds | int | N | | `pre_ping` | Whether to ping the connection before executing a statement and reconnect if lost (Default: `False`) | bool | N | | `session_variables` | Session variables to set on each new connection via `SET = ` (e.g., `workload_group`, `query_timeout`) | dict | N | -| `ssl` | Explicit SSL configuration dictionary passed to `PyMySQL` | dict | N | +| `ssl` | Explicit SSL configuration (`True`/`False` or dictionary) passed to `PyMySQL` | bool \| dict | N | | `ssl_disabled` | Whether SSL is disabled | bool | N | | `ssl_ca` | Path to the file that contains a PEM-formatted CA certificate | string | N | | `ssl_cert` | Path to the file that contains a PEM-formatted client certificate | string | N | | `ssl_key` | Path to the file that contains a PEM-formatted private key for the client certificate | string | N | | `ssl_verify_cert` | Set to `true` to verify the server's certificate against the CA certificate | bool | N | | `ssl_verify_identity` | Set to `true` to verify the server's identity in addition to the certificate | bool | N | +| `concurrent_tasks` | The maximum number of tasks that can use this connection concurrently (Default: `4`) | int | N | +| `register_comments` | Whether or not to register model comments with the SQL engine (Default: `True`) | bool | N | ### Connection configuration example @@ -68,6 +70,11 @@ gateways: user: sqlmesh password: {{ env_var('POSTGRES_PASSWORD') }} database: sqlmesh_state + +default_gateway: doris + +model_defaults: + dialect: doris ``` --- @@ -126,15 +133,15 @@ FROM raw.orders WHERE updated_at BETWEEN @start_ds AND @end_ds; ``` -You can also specify `unique_key` explicitly in `physical_properties` for `FULL` or `INCREMENTAL_BY_TIME_RANGE` models, and optionally decouple the sort key from the unique key using `order_by` (creating a `SEQUENCE MAP COL(col1, col2)` sort key on Doris 2.1+ Merge-on-Write tables). +You can also specify `unique_key` explicitly in `physical_properties` for `FULL` or `INCREMENTAL_BY_TIME_RANGE` models, and optionally decouple the sort key from the unique key using `order_by` (emitting `ORDER BY (col1, col2)` on Doris 2.1+ Merge-on-Write tables). ### AGGREGATE KEY model -For pre-aggregated rollup tables, specify the key columns using `aggregate_key` in `physical_properties`: +For pre-aggregated rollup tables, specify the key columns using `aggregate_key` in `physical_properties`. Note that Doris requires non-key value columns in an `AGGREGATE KEY` table to specify an aggregation function (`SUM`, `MAX`, `MIN`, `REPLACE`, etc.) in the column definition unless all columns are part of `aggregate_key` (or the table is created via a custom DDL / pre-statement): ```sql MODEL ( - name analytics.daily_site_metrics, + name analytics.daily_site_keys, kind FULL, clustered_by (site_id), physical_properties ( @@ -145,9 +152,7 @@ MODEL ( SELECT site_id::INT AS site_id, - event_date::DATE AS event_date, - SUM(pv)::BIGINT AS total_pv, - MAX(dwell_seconds)::INT AS max_dwell + event_date::DATE AS event_date FROM analytics.user_events GROUP BY site_id, event_date; ``` @@ -197,6 +202,10 @@ If none of `clustered_by` or `distributed_by` is provided: Doris OLAP tables support **Multi-Range Batch Partitioning**, **Explicit Partition Definitions**, **Dynamic Partitioning**, and **Auto Partitioning** (Doris 2.1+). +!!! note "OLAP Partitioning Requirements" + - On Doris OLAP tables, `PARTITION BY` requires at least one of `partitions`, `auto_partition = 'true'`, or `'dynamic_partition.enable' = 'true'` in `physical_properties`. If none is provided, SQLMesh omits the `PARTITION BY` clause so Doris does not create a zero-partition table that rejects inserts. + - On `UNIQUE KEY` tables, Doris requires all `partitioned_by` columns to be a subset of `unique_key`. + ### Multi-range & explicit partitions ```sql @@ -257,7 +266,7 @@ WHERE event_ts BETWEEN @start_ts AND @end_ts; ### Pattern 2: Materializing Apache Iceberg tables via Doris -SQLMesh can also create, partition, and overwrite Apache Iceberg tables directly through Doris (`ENGINE=iceberg` and `INSERT OVERWRITE TABLE ... VALUES (*)`). +SQLMesh can also create, partition, and overwrite Apache Iceberg tables directly through Doris (`ENGINE=iceberg` and `INSERT OVERWRITE TABLE`). !!! important "Virtual Environments & Catalog Bootstrap on External Iceberg Catalogs" Doris only supports `CREATE VIEW` inside the `internal` catalog; external Iceberg catalogs do not support Doris views, and Doris does not support `environment_suffix_target: catalog`. @@ -455,6 +464,29 @@ FROM analytics.user_events GROUP BY user_id, CAST(event_time AS DATE); ``` +### Synchronous Materialized Views (`materialized_type = 'SYNC'`) + +Doris also supports single-table synchronous materialized views (rollup indexes maintained atomically on base table writes) by setting `materialized_type = 'SYNC'` and `source_table` in `physical_properties`: + +```sql +MODEL ( + name analytics.user_events_sync_mv, + kind VIEW ( + materialized true + ), + physical_properties ( + materialized_type = 'SYNC', + source_table = 'analytics.user_events' + ) +); + +SELECT + user_id, + COUNT(event_type) AS event_count +FROM analytics.user_events +GROUP BY user_id; +``` + ### Inverted Indexes (Doris 2.1–5.0) You can attach Doris inverted indexes (`USING INVERTED`) for full-text search and fast point filtering using post-statements gated on `@runtime_stage = 'creating'`: diff --git a/docs/quickstart/cli.md b/docs/quickstart/cli.md index 6aa4fc34ed..e7f55b7467 100644 --- a/docs/quickstart/cli.md +++ b/docs/quickstart/cli.md @@ -118,6 +118,7 @@ Choose your SQL engine: [16] RisingWave [17] Fabric [18] StarRocks + [20] Apache Doris Enter a number: 1 ``` diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 22c1dd6835..8866960e6f 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -241,6 +241,12 @@ def _is_valid_olap_key_type(cls, dtype: exp.DataType) -> bool: def _fetch_native_df( self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False ) -> DF: + if self._pre_ping: + try: + self.ping() + except Exception: + logger.info("Connection to the database was lost. Reconnecting...") + self._connection_pool.close() self._connection_pool.get_cursor() return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) @@ -288,9 +294,13 @@ def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any err_code, ) self._connection_pool.close_cursor() - conn = self._connection_pool.get() - conn.ping(reconnect=True) - super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) + try: + conn = self._connection_pool.get() + conn.ping(reconnect=True) + super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) + except Exception: + self._connection_pool.close() + raise else: if is_disconnect and not self._connection_pool.is_transaction_active: self._connection_pool.close() @@ -439,6 +449,29 @@ def _get_data_objects( ) continue + # In Doris 2.1+, asynchronous materialized views (MTMVs) are stored as OLAP base tables + # and report table_type = 'BASE TABLE' in information_schema.tables. Reclassify them via + # the mv_infos() table-valued function so drop_data_object_on_type_mismatch and metadata + # inspection recognize them as MATERIALIZED_VIEW instead of issuing DROP TABLE. + if (not query_catalog or query_catalog.lower() == "internal") and any( + obj.type.is_table for obj in result + ): + try: + db_literal = self._to_sql(exp.Literal.string(schema_exp.db), quote=False) + mv_rows = self.fetchall( + f"SELECT Name FROM mv_infos('database' = {db_literal})", + quote_identifiers=False, + ) + mv_names = { + _decoded_str(row[0]).lower() for row in mv_rows if len(row) == 1 and row[0] + } + if mv_names: + for obj in result: + if obj.name.lower() in mv_names: + obj.type = DataObjectType.MATERIALIZED_VIEW + except Exception: + logger.debug("mv_infos lookup failed on Doris, skipping MTMV reclassification") + return result def create_view( @@ -550,8 +583,18 @@ def _create_materialized_view( with source_queries[0] as query: target_table = exp.to_table(view_name) + is_sync_mv = bool( + view_properties + and ( + ( + self._extract_str_value(view_properties.get("materialized_type")) or "" + ).upper() + == "SYNC" + ) + ) + schema: t.Union[exp.Table, exp.Schema] = target_table - if target_columns_to_types: + if target_columns_to_types and not is_sync_mv: schema = self._build_schema_exp( target_table, target_columns_to_types, @@ -594,15 +637,6 @@ def _create_materialized_view( table_kind="MATERIALIZED_VIEW", ) - is_sync_mv = bool( - view_properties - and ( - ( - self._extract_str_value(view_properties.get("materialized_type")) or "" - ).upper() - == "SYNC" - ) - ) if is_sync_mv and isinstance(query, exp.Select): query = query.copy() aliased_selects: t.List[exp.Expr] = [] @@ -849,7 +883,7 @@ def delete_from( self, table_name: TableName, where: t.Optional[t.Union[str, exp.Expr]] = None ) -> None: """Delete from a table in Doris.""" - if not where or where == exp.true(): + if not where: self._truncate_table(table_name) return @@ -858,6 +892,13 @@ def delete_from( else: where = where.copy() + where = self._simplify_boolean_where(where) + if where == exp.true(): + self._truncate_table(table_name) + return + if where == exp.false(): + return + where = self._rewrite_between_in_where(where) subquery_info = self._extract_subquery_from_where(where) @@ -866,6 +907,41 @@ def delete_from( else: super().delete_from(table_name, where) + @classmethod + def _simplify_boolean_where(cls, expr: exp.Expr) -> exp.Expr: + if isinstance(expr, (exp.Subquery, exp.Select)): + return expr + for k, v in list(expr.args.items()): + if isinstance(v, exp.Expr): + expr.set(k, cls._simplify_boolean_where(v)) + elif isinstance(v, list): + expr.set( + k, + [ + cls._simplify_boolean_where(item) if isinstance(item, exp.Expr) else item + for item in v + ], + ) + if isinstance(expr, exp.Paren) and isinstance(expr.this, exp.Boolean): + return expr.this + if isinstance(expr, exp.Not) and isinstance(expr.this, exp.Boolean): + return exp.false() if expr.this == exp.true() else exp.true() + if isinstance(expr, exp.And): + if expr.this == exp.false() or expr.expression == exp.false(): + return exp.false() + if expr.this == exp.true(): + return expr.expression + if expr.expression == exp.true(): + return expr.this + if isinstance(expr, exp.Or): + if expr.this == exp.true() or expr.expression == exp.true(): + return exp.true() + if expr.this == exp.false(): + return expr.expression + if expr.expression == exp.false(): + return expr.this + return expr + @staticmethod def _rewrite_between_in_where(where: exp.Expr) -> exp.Expr: if isinstance(where, exp.Between): @@ -891,6 +967,7 @@ def _extract_subquery_from_where( self, where: exp.Expr ) -> t.Optional[t.Tuple[t.List[exp.Expr], exp.Subquery, t.List[exp.Expr]]]: """Extract positive IN (SELECT ...) subquery and any sibling AND conditions.""" + where = where.unnest() conjuncts: t.List[exp.Expr] = ( list(where.flatten()) if isinstance(where, exp.And) else [where] ) @@ -1149,7 +1226,12 @@ def _create_table_from_columns( ) -> None: table_properties = dict(kwargs.get("table_properties") or {}) - if primary_key and "unique_key" not in table_properties: + if ( + primary_key + and "unique_key" not in table_properties + and "duplicate_key" not in table_properties + and "aggregate_key" not in table_properties + ): table_properties["unique_key"] = exp.Tuple( expressions=[exp.to_column(col) for col in primary_key] ) @@ -1217,9 +1299,26 @@ def _build_create_table_exp( [col_def_map[c] for c in target_columns_to_types] + other_exprs, ) if expression is not None and isinstance(expression, exp.Query): - expression = self._order_projections_and_filter( - expression, target_columns_to_types - ) + reordered_in_place = False + if isinstance(expression, exp.Select): + proj_by_name = { + proj.alias_or_name: proj + for proj in expression.expressions + if proj.alias_or_name + } + if len(proj_by_name) == len(expression.expressions) == len( + target_columns_to_types + ) and all(c in proj_by_name for c in target_columns_to_types): + expression = expression.copy() + expression.set( + "expressions", + [proj_by_name[c] for c in target_columns_to_types], + ) + reordered_in_place = True + if not reordered_in_place: + expression = self._order_projections_and_filter( + expression, target_columns_to_types + ) return super()._build_create_table_exp( table_name_or_schema, @@ -1238,20 +1337,38 @@ def _reorder_olap_columns( kwargs: t.Dict[str, t.Any], ) -> t.Dict[str, exp.DataType]: table_props = kwargs.get("table_properties") or {} - explicit_key = ( - table_props.get("unique_key") - or kwargs.get("unique_key") - or table_props.get("duplicate_key") - or table_props.get("aggregate_key") - ) + unique_key = table_props.get("unique_key") or kwargs.get("unique_key") + duplicate_key = table_props.get("duplicate_key") + aggregate_key = table_props.get("aggregate_key") + + specified_key_types = [ + name + for name, val in ( + ("unique_key", unique_key), + ("duplicate_key", duplicate_key), + ("aggregate_key", aggregate_key), + ) + if self._extract_column_names(val) + ] + if len(specified_key_types) > 1: + raise SQLMeshError( + f"Doris tables support at most one key model (unique_key, duplicate_key, or aggregate_key), " + f"got multiple: {', '.join(specified_key_types)}." + ) + + explicit_key = unique_key or duplicate_key or aggregate_key col_names = list(target_columns_to_types) if not col_names: return target_columns_to_types if explicit_key: - key_cols = [ - k for k in self._extract_column_names(explicit_key) if k in target_columns_to_types - ] + raw_key_cols = self._extract_column_names(explicit_key) + missing_cols = [k for k in raw_key_cols if k not in target_columns_to_types] + if missing_cols: + raise SQLMeshError( + f"Doris key columns {missing_cols} not found in table schema columns {col_names}." + ) + key_cols = [k for k in raw_key_cols if k in target_columns_to_types] if key_cols and col_names[: len(key_cols)] != key_cols: new_order = key_cols + [c for c in col_names if c not in key_cols] return {c: target_columns_to_types[c] for c in new_order} @@ -1277,18 +1394,18 @@ def _parse_partition_expressions( partition_kind: t.Optional[str] = None for expr in partitioned_by: - if isinstance(expr, exp.Anonymous) and expr.this: - func_name = str(expr.this).upper() - if func_name in ("RANGE", "LIST"): - partition_kind = func_name - for arg in expr.expressions: - if isinstance(arg, exp.Column): - parsed_partitioned_by.append(arg) - elif isinstance(arg, exp.Func) or isinstance(arg, exp.Anonymous): - parsed_partitioned_by.append(arg) - else: - parsed_partitioned_by.append(exp.to_column(str(arg))) - continue + if isinstance(expr, exp.List) or ( + isinstance(expr, exp.Anonymous) + and expr.this + and str(expr.this).upper() in ("RANGE", "LIST") + ): + partition_kind = "LIST" if isinstance(expr, exp.List) else str(expr.this).upper() + for arg in expr.expressions: + if isinstance(arg, (exp.Column, exp.Func, exp.Anonymous)): + parsed_partitioned_by.append(arg) + else: + parsed_partitioned_by.append(exp.to_column(str(arg))) + continue if isinstance(expr, exp.Literal) and getattr(expr, "is_string", False): text = str(expr.this) @@ -1521,6 +1638,28 @@ def _build_table_properties_exp( order_by_cols = self._extract_column_names(order_by) if not is_iceberg: + specified_key_types = [ + name + for name, cols in ( + ("unique_key", unique_key_cols), + ("duplicate_key", duplicate_key_cols), + ("aggregate_key", aggregate_key_cols), + ) + if cols + ] + if len(specified_key_types) > 1: + raise SQLMeshError( + f"Doris tables support at most one key model (unique_key, duplicate_key, or aggregate_key), " + f"got multiple: {', '.join(specified_key_types)}." + ) + active_key_cols = unique_key_cols or duplicate_key_cols or aggregate_key_cols + if active_key_cols and target_columns_to_types: + missing_cols = [k for k in active_key_cols if k not in target_columns_to_types] + if missing_cols: + raise SQLMeshError( + f"Doris key columns {missing_cols} not found in table schema columns {list(target_columns_to_types)}." + ) + if unique_key_cols: properties.append( exp.UniqueKeyProperty(expressions=[exp.to_column(k) for k in unique_key_cols]) @@ -1789,11 +1928,15 @@ def _build_distributed_by_exp( if distributed_by is not None: distributed_info: t.Dict[str, t.Any] = {} - # Case 1: Function expression such as HASH(a, b) + # Case 1: Function expression such as HASH(a, b) or RANDOM() if isinstance(distributed_by, (exp.Anonymous, exp.Func)) and not isinstance( distributed_by, (exp.Tuple, exp.Array, exp.Paren) ): - kind_upper = distributed_by.name.upper() + kind_upper = ( + "RANDOM" + if isinstance(distributed_by, exp.Rand) + else distributed_by.name.upper() + ) if kind_upper not in ("HASH", "RANDOM"): raise SQLMeshError( f"Invalid distributed_by kind: {kind_upper}. Expected HASH or RANDOM." @@ -1803,20 +1946,24 @@ def _build_distributed_by_exp( ] return exp.DistributedByProperty( kind=exp.Var(this=kind_upper), - expressions=cols, - buckets=default_buckets_expr or exp.Literal.number(10), + expressions=cols if cols else None, + buckets=default_buckets_expr or (exp.Literal.number(10) if cols else None), order=None, ) - # Case 2: String shorthand such as "HASH(id) BUCKETS 10", "RANDOM BUCKETS AUTO", or "id" - if isinstance(distributed_by, str) or ( + # Case 2: String shorthand or bare column/identifier such as "HASH(id) BUCKETS 10", "RANDOM BUCKETS AUTO", or id + if isinstance(distributed_by, (str, exp.Column, exp.Identifier, exp.Var)) or ( isinstance(distributed_by, exp.Literal) and getattr(distributed_by, "is_string", False) ): text = ( - str(distributed_by.this).strip() - if isinstance(distributed_by, exp.Literal) - else distributed_by.strip() + distributed_by.strip() + if isinstance(distributed_by, str) + else ( + str(distributed_by.this).strip() + if isinstance(distributed_by, exp.Literal) + else distributed_by.name.strip() + ) ) m = re.match( r"^(HASH|RANDOM)\s*(?:\(([^)]*)\))?\s*(?:BUCKETS\s+(AUTO|\d+))?$", diff --git a/tests/core/engine_adapter/integration/__init__.py b/tests/core/engine_adapter/integration/__init__.py index 7ce05f451e..3aa2639a92 100644 --- a/tests/core/engine_adapter/integration/__init__.py +++ b/tests/core/engine_adapter/integration/__init__.py @@ -758,6 +758,17 @@ def create_catalog(self, catalog_name: str): self.engine_adapter.execute(f'CREATE DATABASE IF NOT EXISTS "{catalog_name}"') except Exception: pass + elif self.dialect == "doris": + self.engine_adapter.execute( + f""" + CREATE CATALOG IF NOT EXISTS `{catalog_name}` PROPERTIES ( + 'type' = 'iceberg', + 'iceberg.catalog.type' = 'hadoop', + 'warehouse' = 'file:///tmp/doris_iceberg_warehouse/{catalog_name}' + ) + """, + quote_identifiers=False, + ) self._catalogs.append(catalog_name) @@ -769,6 +780,11 @@ def drop_catalog(self, catalog_name: str): elif self.dialect == "fabric": # Use the engine adapter's built-in catalog dropping functionality self.engine_adapter.drop_catalog(catalog_name) + elif self.dialect == "doris": + self.engine_adapter.set_current_catalog("internal") + self.engine_adapter.execute( + f"DROP CATALOG IF EXISTS `{catalog_name}`", quote_identifiers=False + ) else: self.engine_adapter.execute(f'DROP DATABASE IF EXISTS "{catalog_name}"') diff --git a/tests/core/engine_adapter/integration/test_integration.py b/tests/core/engine_adapter/integration/test_integration.py index 5d33dea1e2..dbad929aa6 100644 --- a/tests/core/engine_adapter/integration/test_integration.py +++ b/tests/core/engine_adapter/integration/test_integration.py @@ -105,10 +105,6 @@ def test_catalog_operations(ctx: TestContext): pytest.skip( f"Engine adapter {ctx.engine_adapter.dialect} doesn't support catalog operations" ) - if ctx.dialect == "doris": - pytest.skip( - "Doris external catalogs require connector properties (`CREATE CATALOG ... PROPERTIES (...)`); covered in test_integration_doris.py" - ) # use a unique name so that integration tests on cloud databases can run in parallel catalog_name = "testing" if not ctx.is_remote else ctx.add_test_suffix("testing") @@ -3299,6 +3295,7 @@ def test_value_normalization( ) columns_to_types_normalized = { ( + # "value" is in DorisGenerator.RESERVED_KEYWORDS, so .sql(dialect="doris") returns "`value`" normalize_identifiers(k, dialect=ctx.dialect).name if ctx.dialect == "doris" else normalize_identifiers(k, dialect=ctx.dialect).sql(dialect=ctx.dialect) @@ -3962,7 +3959,8 @@ def test_materialized_view_evaluation(ctx: TestContext): ) elif dialect == "doris": pytest.skip( - "Skipping doris as synchronous materialized views do not support specifying schema" + "Doris materialized views require DISTRIBUTED BY or physical_properties configuration; " + "Doris MVs are covered by test_integration_doris.py" ) model_name = ctx.table("test_tbl") diff --git a/tests/core/engine_adapter/integration/test_integration_doris.py b/tests/core/engine_adapter/integration/test_integration_doris.py index 3092d70c11..8c1ffdc47b 100644 --- a/tests/core/engine_adapter/integration/test_integration_doris.py +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -172,7 +172,11 @@ def test_materialized_view_async_mtmv(ctx: TestContext): ) wait_until(lambda: len(ctx.get_current_data(view)) == len(input_data), attempts=30, wait=1) ctx.compare_with_current(view, input_data) + results = ctx.get_metadata_results() + assert exp.to_table(view).name in results.materialized_views + assert exp.to_table(view).name not in results.tables ctx.engine_adapter.drop_view(view, materialized=True) + assert len(ctx.get_metadata_results().materialized_views) == 0 def test_merge_unique_key(ctx: TestContext): diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 582744f250..b3ad1720ad 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -7,6 +7,7 @@ from sqlmesh.core.engine_adapter.doris import DorisEngineAdapter from sqlmesh.core.engine_adapter.shared import CatalogSupport, DataObject, DataObjectType +from sqlmesh.utils.errors import SQLMeshError from tests.core.engine_adapter import to_sql_calls pytestmark = [pytest.mark.doris, pytest.mark.engine] @@ -1366,3 +1367,116 @@ def test_merge_with_identifier_unique_key( ) for call in sql_calls ) + + +def test_list_partition_ast_and_random_distributed_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.create_table( + "db.list_tbl", + target_columns_to_types={ + "region": exp.DataType.build("VARCHAR(16)", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + }, + partitioned_by=[parse_one("LIST(region)", read="doris")], + table_properties={ + "partitions": exp.Literal.string("PARTITION p_na VALUES IN ('US', 'CA')"), + "distributed_by": parse_one("RANDOM()", read="doris"), + "buckets": exp.Literal.number(8), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`list_tbl` (`region` VARCHAR(16), `id` INT) PARTITION BY LIST (`region`) (PARTITION p_na VALUES IN ('US', 'CA')) DISTRIBUTED BY RANDOM BUCKETS 8" + ] + + +def test_key_model_mutual_exclusion_and_missing_column_validation( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + with pytest.raises(SQLMeshError, match="at most one key model"): + adapter.create_table( + "db.bad_keys", + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + table_properties={ + "unique_key": exp.to_column("id"), + "duplicate_key": exp.to_column("id"), + }, + ) + + with pytest.raises(SQLMeshError, match="not found in table schema columns"): + adapter.create_table( + "db.missing_key", + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + table_properties={ + "unique_key": exp.to_column("nonexistent_col"), + }, + ) + + +def test_ctas_reorders_projections_in_place_without_double_subquery( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.ctas( + "db.ctas_reordered", + parse_one("SELECT name, id FROM db.src", read="doris"), + target_columns_to_types={ + "name": exp.DataType.build("VARCHAR(50)", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + }, + table_properties={"unique_key": exp.to_column("id")}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`ctas_reordered` UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10 AS SELECT CAST(`id` AS INT) AS `id`, CAST(`name` AS VARCHAR(50)) AS `name` FROM (SELECT `name`, `id` FROM `db`.`src`) AS `_subquery`" + ] + + +def test_get_data_objects_reclassifies_mtmv_via_mv_infos( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter, patch_get_data_objects=False) + mocker.patch.object( + adapter, + "fetchall", + side_effect=[ + [("silver", "regular_tbl", "BASE TABLE"), ("silver", "async_mv", "BASE TABLE")], + [("async_mv",)], + ], + ) + + objects = adapter._get_data_objects("silver") + assert [(o.name, o.type) for o in objects] == [ + ("regular_tbl", DataObjectType.TABLE), + ("async_mv", DataObjectType.MATERIALIZED_VIEW), + ] + + +def test_delete_from_boolean_literal_simplification( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # FALSE is a no-op + adapter.delete_from("db.tbl", where=exp.false()) + assert to_sql_calls(adapter) == [] + + # Compound AND TRUE is simplified away so Doris DeleteHandler does not reject boolean literal + adapter.delete_from("db.tbl", where="id = 1 AND TRUE") + assert to_sql_calls(adapter) == ["DELETE FROM `db`.`tbl` WHERE `id` = 1"] + + adapter.cursor.execute.reset_mock() + adapter.delete_from("db.tbl", where="TRUE") + assert to_sql_calls(adapter) == ["TRUNCATE TABLE `db`.`tbl`"] From 1b8fffb16b36189ee71210e37b3b5e82e7239294 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:11:46 -0700 Subject: [PATCH 09/27] fix(doris): support case-insensitive key matching and validate order_by Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- .github/scripts/wait-for-db.sh | 2 +- sqlmesh/core/engine_adapter/doris.py | 50 ++++++++++++++++++------- tests/core/engine_adapter/test_doris.py | 28 ++++++++++++++ 3 files changed, 65 insertions(+), 15 deletions(-) diff --git a/.github/scripts/wait-for-db.sh b/.github/scripts/wait-for-db.sh index e4e6605b05..489cdd4583 100755 --- a/.github/scripts/wait-for-db.sh +++ b/.github/scripts/wait-for-db.sh @@ -121,7 +121,7 @@ doris_ready() { echo "Checking for 1 alive Doris backends..." while true; do echo "Checking Doris backends..." - ALIVE_BACKENDS=$(docker exec -i doris-fe mysql -h127.0.0.1 -P9030 -uroot -e "show backends \G" 2>/dev/null | grep -c "^ *Alive: true *$") + ALIVE_BACKENDS=$(docker exec -i doris-fe mysql -h127.0.0.1 -P9030 -uroot -e "show backends \G" 2>/dev/null | grep -c "^ *Alive: true *$" || true) # fallback value if failed to get number if ! [[ "$ALIVE_BACKENDS" =~ ^[0-9]+$ ]]; then diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 8866960e6f..1a70a88bc1 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -1281,7 +1281,7 @@ def _build_create_table_exp( target_columns_to_types = reordered_cols if isinstance(table_name_or_schema, exp.Schema): col_def_map = { - col_def.name: col_def + col_def.name.lower(): col_def for col_def in table_name_or_schema.expressions if isinstance(col_def, exp.ColumnDef) } @@ -1291,28 +1291,28 @@ def _build_create_table_exp( if not isinstance(e, exp.ColumnDef) ] if len(col_def_map) == len(target_columns_to_types) and all( - c in col_def_map for c in target_columns_to_types + c.lower() in col_def_map for c in target_columns_to_types ): table_name_or_schema = table_name_or_schema.copy() table_name_or_schema.set( "expressions", - [col_def_map[c] for c in target_columns_to_types] + other_exprs, + [col_def_map[c.lower()] for c in target_columns_to_types] + other_exprs, ) if expression is not None and isinstance(expression, exp.Query): reordered_in_place = False if isinstance(expression, exp.Select): proj_by_name = { - proj.alias_or_name: proj + proj.alias_or_name.lower(): proj for proj in expression.expressions if proj.alias_or_name } if len(proj_by_name) == len(expression.expressions) == len( target_columns_to_types - ) and all(c in proj_by_name for c in target_columns_to_types): + ) and all(c.lower() in proj_by_name for c in target_columns_to_types): expression = expression.copy() expression.set( "expressions", - [proj_by_name[c] for c in target_columns_to_types], + [proj_by_name[c.lower()] for c in target_columns_to_types], ) reordered_in_place = True if not reordered_in_place: @@ -1362,15 +1362,17 @@ def _reorder_olap_columns( return target_columns_to_types if explicit_key: + col_by_lower = {c.lower(): c for c in col_names} raw_key_cols = self._extract_column_names(explicit_key) - missing_cols = [k for k in raw_key_cols if k not in target_columns_to_types] + missing_cols = [k for k in raw_key_cols if k.lower() not in col_by_lower] if missing_cols: raise SQLMeshError( f"Doris key columns {missing_cols} not found in table schema columns {col_names}." ) - key_cols = [k for k in raw_key_cols if k in target_columns_to_types] + key_cols = [col_by_lower[k.lower()] for k in raw_key_cols if k.lower() in col_by_lower] if key_cols and col_names[: len(key_cols)] != key_cols: - new_order = key_cols + [c for c in col_names if c not in key_cols] + key_cols_set = set(key_cols) + new_order = key_cols + [c for c in col_names if c not in key_cols_set] return {c: target_columns_to_types[c] for c in new_order} return target_columns_to_types @@ -1652,13 +1654,22 @@ def _build_table_properties_exp( f"Doris tables support at most one key model (unique_key, duplicate_key, or aggregate_key), " f"got multiple: {', '.join(specified_key_types)}." ) + col_by_lower = ( + {c.lower(): c for c in target_columns_to_types} if target_columns_to_types else {} + ) active_key_cols = unique_key_cols or duplicate_key_cols or aggregate_key_cols - if active_key_cols and target_columns_to_types: - missing_cols = [k for k in active_key_cols if k not in target_columns_to_types] + if active_key_cols and col_by_lower and not is_materialized_view: + missing_cols = [k for k in active_key_cols if k.lower() not in col_by_lower] if missing_cols: raise SQLMeshError( - f"Doris key columns {missing_cols} not found in table schema columns {list(target_columns_to_types)}." + f"Doris key columns {missing_cols} not found in table schema columns {list(target_columns_to_types or ())}." ) + if unique_key_cols: + unique_key_cols = [col_by_lower[k.lower()] for k in unique_key_cols] + elif duplicate_key_cols: + duplicate_key_cols = [col_by_lower[k.lower()] for k in duplicate_key_cols] + elif aggregate_key_cols: + aggregate_key_cols = [col_by_lower[k.lower()] for k in aggregate_key_cols] if unique_key_cols: properties.append( @@ -1677,6 +1688,17 @@ def _build_table_properties_exp( ) if order_by_cols and not is_materialized_view: + if not unique_key_cols: + raise SQLMeshError( + "Doris table property 'order_by' is only supported on UNIQUE KEY tables." + ) + if col_by_lower: + missing_order_cols = [k for k in order_by_cols if k.lower() not in col_by_lower] + if missing_order_cols: + raise SQLMeshError( + f"Doris order_by columns {missing_order_cols} not found in table schema columns {list(target_columns_to_types or ())}." + ) + order_by_cols = [col_by_lower[k.lower()] for k in order_by_cols] order_cols_sql = ", ".join(self._to_sql(exp.to_column(k)) for k in order_by_cols) properties.append( exp.SetConfigProperty(this=exp.var(f"ORDER BY ({order_cols_sql})")) @@ -1703,11 +1725,11 @@ def _build_table_properties_exp( if partitioned_by: parsed_part_exprs, _ = self._parse_partition_expressions(partitioned_by) if unique_key_cols and not is_materialized_view and not is_iceberg: - key_cols_set = set(unique_key_cols) + key_cols_set = {k.lower() for k in unique_key_cols} partition_cols = set() for expr in parsed_part_exprs: for col_node in expr.find_all(exp.Column): - partition_cols.add(col_node.name) + partition_cols.add(col_node.name.lower()) not_in_key = partition_cols - key_cols_set if not_in_key: logger.warning( diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index b3ad1720ad..e1046b7560 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -1423,6 +1423,34 @@ def test_key_model_mutual_exclusion_and_missing_column_validation( }, ) + with pytest.raises(SQLMeshError, match="only supported on UNIQUE KEY tables"): + adapter.create_table( + "db.order_by_without_uk", + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + table_properties={ + "duplicate_key": exp.to_column("id"), + "order_by": exp.to_column("name"), + }, + ) + + # Case-insensitive key column matching resolves to the canonical schema column name + adapter.create_table( + "db.case_insensitive_uk", + target_columns_to_types={ + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + }, + table_properties={ + "unique_key": exp.to_column("ID"), + }, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`case_insensitive_uk` (`id` INT, `name` VARCHAR(32)) UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10" + ] + def test_ctas_reorders_projections_in_place_without_double_subquery( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], From c874f4471be0b17295f77a29026f79e2be05b7d8 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:30:23 -0700 Subject: [PATCH 10/27] fix(doris): reorder unannotated CTAS key projections and validate distributed_by Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 43 ++++++++++++++++++++++++- tests/core/engine_adapter/test_doris.py | 42 ++++++++++++++++++++++-- 2 files changed, 82 insertions(+), 3 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 1a70a88bc1..3df812e717 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -1208,9 +1208,12 @@ def adjust_physical_properties_for_incremental( and "unique_key" not in physical_properties and "duplicate_key" not in physical_properties and "aggregate_key" not in physical_properties + and all(isinstance(k, (exp.Column, exp.Identifier)) for k in unique_key) ): + cleaned_keys = [exp.column(k.name) for k in unique_key] + physical_properties = dict(physical_properties) physical_properties["unique_key"] = ( - unique_key[0] if len(unique_key) == 1 else exp.Tuple(expressions=list(unique_key)) + cleaned_keys[0] if len(cleaned_keys) == 1 else exp.Tuple(expressions=cleaned_keys) ) return physical_properties @@ -1319,6 +1322,24 @@ def _build_create_table_exp( expression = self._order_projections_and_filter( expression, target_columns_to_types ) + elif not is_iceberg and isinstance(expression, exp.Select): + proj_by_name = { + proj.alias_or_name.lower(): proj + for proj in expression.expressions + if proj.alias_or_name and not isinstance(proj, exp.Star) + } + if proj_by_name and len(proj_by_name) == len(expression.expressions): + select_cols_to_types = { + proj.alias_or_name: exp.DataType.build("UNKNOWN") + for proj in expression.expressions + } + reordered_select_cols = self._reorder_olap_columns(select_cols_to_types, kwargs) + if list(reordered_select_cols) != list(select_cols_to_types): + expression = expression.copy() + expression.set( + "expressions", + [proj_by_name[c.lower()] for c in reordered_select_cols], + ) return super()._build_create_table_exp( table_name_or_schema, @@ -1786,6 +1807,26 @@ def _build_table_properties_exp( top_level_buckets=top_level_buckets, ) if dist_prop is not None: + if not is_materialized_view: + dist_kind = str(dist_prop.args.get("kind") or "").upper() + if unique_key_cols and dist_kind == "RANDOM": + raise SQLMeshError( + "Doris UNIQUE KEY tables do not support DISTRIBUTED BY RANDOM. " + "Use DISTRIBUTED BY HASH (...) instead." + ) + if col_by_lower and dist_prop.expressions: + dist_col_names = [c.name for c in dist_prop.expressions] + missing_dist_cols = [ + c for c in dist_col_names if c.lower() not in col_by_lower + ] + if missing_dist_cols: + raise SQLMeshError( + f"Doris distributed_by columns {missing_dist_cols} not found in table schema columns {list(target_columns_to_types or ())}." + ) + dist_prop.set( + "expressions", + [exp.to_column(col_by_lower[c.lower()]) for c in dist_col_names], + ) properties.append(dist_prop) # Default replication_num = 1 for single-BE dev/CI when not explicitly configured on OLAP tables diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index e1046b7560..14c22249fe 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -1445,12 +1445,42 @@ def test_key_model_mutual_exclusion_and_missing_column_validation( }, table_properties={ "unique_key": exp.to_column("ID"), + "distributed_by": "HASH(ID) BUCKETS 8", }, ) assert to_sql_calls(adapter) == [ - "CREATE TABLE IF NOT EXISTS `db`.`case_insensitive_uk` (`id` INT, `name` VARCHAR(32)) UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10" + "CREATE TABLE IF NOT EXISTS `db`.`case_insensitive_uk` (`id` INT, `name` VARCHAR(32)) UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 8" ] + with pytest.raises( + SQLMeshError, match="Doris UNIQUE KEY tables do not support DISTRIBUTED BY RANDOM" + ): + adapter.create_table( + "db.uk_random_dist", + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + table_properties={ + "unique_key": exp.to_column("id"), + "distributed_by": "RANDOM BUCKETS 8", + }, + ) + + with pytest.raises( + SQLMeshError, match="Doris distributed_by columns .* not found in table schema columns" + ): + adapter.create_table( + "db.missing_dist_col", + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + table_properties={ + "distributed_by": "HASH(missing_col) BUCKETS 8", + }, + ) + def test_ctas_reorders_projections_in_place_without_double_subquery( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], @@ -1467,8 +1497,16 @@ def test_ctas_reorders_projections_in_place_without_double_subquery( table_properties={"unique_key": exp.to_column("id")}, ) + # Unannotated CTAS (target_columns_to_types=None) also reorders Select projections for key prefix + adapter.ctas( + "db.ctas_unannotated_reordered", + parse_one("SELECT name, id FROM db.src", read="doris"), + table_properties={"unique_key": exp.to_column("id")}, + ) + assert to_sql_calls(adapter) == [ - "CREATE TABLE IF NOT EXISTS `db`.`ctas_reordered` UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10 AS SELECT CAST(`id` AS INT) AS `id`, CAST(`name` AS VARCHAR(50)) AS `name` FROM (SELECT `name`, `id` FROM `db`.`src`) AS `_subquery`" + "CREATE TABLE IF NOT EXISTS `db`.`ctas_reordered` UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10 AS SELECT CAST(`id` AS INT) AS `id`, CAST(`name` AS VARCHAR(50)) AS `name` FROM (SELECT `name`, `id` FROM `db`.`src`) AS `_subquery`", + "CREATE TABLE IF NOT EXISTS `db`.`ctas_unannotated_reordered` UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10 AS SELECT `id`, `name` FROM `db`.`src`", ] From f45ee2bdd0ecf567d8d6b6988d0401e661784bcc Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Wed, 21 Oct 2026 16:43:15 -0700 Subject: [PATCH 11/27] fix(doris): disable MV recreation on evaluation and resolve table refs in physical properties Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 4 ++++ tests/core/engine_adapter/test_doris.py | 8 ++++++-- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 3df812e717..92a4133055 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -61,9 +61,13 @@ class DorisEngineAdapter( MAX_IDENTIFIER_LENGTH = 64 SUPPORTS_MATERIALIZED_VIEWS = True SUPPORTS_MATERIALIZED_VIEW_SCHEMA = True + RECREATE_MATERIALIZED_VIEW_ON_EVALUATION = False SUPPORTS_CREATE_DROP_CATALOG = False INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.DELETE_INSERT CURRENT_CATALOG_EXPRESSION = exp.func("current_catalog") + RESOLVE_TABLE_REFS_IN_PHYSICAL_PROPERTIES: t.FrozenSet[str] = frozenset( + {"excluded_trigger_tables", "source_table"} + ) SCHEMA_DIFFER_KWARGS = { "parameterized_type_defaults": { exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(9, 0), (0,)], diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 14c22249fe..274195d939 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -1301,12 +1301,16 @@ def test_adjust_physical_properties_for_incremental( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], ): adapter = make_mocked_engine_adapter(DorisEngineAdapter) + assert adapter.RECREATE_MATERIALIZED_VIEW_ON_EVALUATION is False + assert adapter.RESOLVE_TABLE_REFS_IN_PHYSICAL_PROPERTIES == frozenset( + {"excluded_trigger_tables", "source_table"} + ) - # Single unique_key promoted when no explicit key is present + # Single unique_key promoted when no explicit key is present (table qualifiers stripped) props = adapter.adjust_physical_properties_for_incremental( {}, requires_delete_capable_table=True, - unique_key=[exp.to_column("id")], + unique_key=[exp.to_column("src.id")], model_name="db.incremental_model", ) assert props == {"unique_key": exp.to_column("id")} From c5bdff7c8cab577da1d415b2a92037550f73c7dd Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Wed, 21 Oct 2026 17:00:34 -0700 Subject: [PATCH 12/27] fix(doris): fix MODEL physical_properties parsing, UDT key check, and sync MV source_table Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 92 +++++++------- tests/core/engine_adapter/test_doris.py | 155 +++++++++++++++++++++++- 2 files changed, 202 insertions(+), 45 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 92a4133055..24c9e39866 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -188,12 +188,16 @@ def _is_iceberg_table( def _extract_str_value(val: t.Any) -> t.Optional[str]: if val is None: return None + if isinstance(val, exp.Boolean): + return "true" if val.this else "false" + if isinstance(val, bool): + return "true" if val else "false" if isinstance(val, exp.Literal): return str(val.this) if isinstance(val, exp.Identifier): return str(val.this) - if isinstance(val, exp.Column): - return str(val.name) + if isinstance(val, (exp.Column, exp.Table)): + return ".".join(part.name for part in val.parts) if isinstance(val, exp.Var): return str(val.this) if isinstance(val, str): @@ -236,8 +240,14 @@ def _is_valid_olap_key_type(cls, dtype: exp.DataType) -> bool: ) if type_name in cls.INVALID_KEY_TYPES: return False - if dtype.is_type("userdefined"): - kind = str(dtype.args.get("kind") or "").upper() + if dtype.this == exp.DataType.Type.USERDEFINED: + kind = ( + str(dtype.args.get("kind") or dtype.sql(dialect=cls.DIALECT)) + .upper() + .split("(", 1)[0] + .strip() + .strip("`") + ) if kind in cls.INVALID_KEY_TYPES: return False return True @@ -680,14 +690,7 @@ def _create_materialized_view( if is_sync_mv and not target_table.db: source_db: t.Optional[str] = None if view_properties and view_properties.get("source_table"): - raw_src = view_properties["source_table"] - if isinstance(raw_src, exp.Column): - source_db = raw_src.db or None - elif isinstance(raw_src, exp.Table): - source_db = raw_src.db or None - else: - src_str = self._extract_str_value(raw_src) or str(raw_src) - source_db = exp.to_table(src_str, dialect=self.dialect).db or None + source_db = self._to_source_table(view_properties["source_table"]).db or None if not source_db: first_tbl = query.find(exp.Table) if first_tbl and first_tbl.db: @@ -708,6 +711,22 @@ def _create_materialized_view( ) self._clear_data_object_cache(view_name) + def _to_source_table(self, raw_src: t.Any) -> exp.Table: + if isinstance(raw_src, exp.Table): + tbl = raw_src.copy() + elif isinstance(raw_src, exp.Column): + tbl = exp.table_( + raw_src.name, + db=raw_src.table or None, + catalog=raw_src.db or None, + ) + else: + src_str = self._extract_str_value(raw_src) or str(raw_src) + tbl = exp.to_table(src_str, dialect=self.dialect) + if tbl.catalog and tbl.catalog.lower() == "internal": + tbl.set("catalog", None) + return tbl + def drop_view( self, view_name: TableName, @@ -727,17 +746,7 @@ def drop_view( source_table = view_properties.get("source_table") if mat_type and mat_type.upper() == "SYNC" and source_table: view_table_sql = self._to_sql(target_table) - if isinstance(source_table, exp.Column): - source_tbl_exp = exp.table_( - source_table.name, - db=source_table.db or None, - catalog=source_table.catalog or None, - ) - elif isinstance(source_table, exp.Table): - source_tbl_exp = source_table - else: - source_table_str = self._extract_str_value(source_table) or str(source_table) - source_tbl_exp = exp.to_table(source_table_str, dialect=self.dialect) + source_tbl_exp = self._to_source_table(source_table) # Unqualified synchronous materialized views require the source table's database # to be active on the session before executing `DROP MATERIALIZED VIEW ON `. if not target_table.db and source_tbl_exp.db: @@ -1080,8 +1089,12 @@ def _execute_delete_with_subquery( for extra in extra_conditions: qualified_extra = extra.copy() for col in qualified_extra.find_all(exp.Column): - if not col.find_ancestor(exp.Subquery, exp.Select) and not col.table: + if not col.find_ancestor(exp.Subquery, exp.Select) and ( + not col.table or col.table.lower() == target_table.name.lower() + ): col.set("table", exp.to_identifier("_t1")) + col.set("db", None) + col.set("catalog", None) join_conditions.append(qualified_extra) where_expr = exp.and_(*join_conditions) @@ -1843,15 +1856,10 @@ def _build_table_properties_exp( return None def _to_upper_str(self, val: t.Any) -> str: - if isinstance(val, exp.Literal): - return str(val.this).upper() - if isinstance(val, exp.Identifier): - return str(val.this).upper() - if isinstance(val, str): - return val.upper() - if isinstance(val, exp.Expr): - return val.sql(dialect=self.dialect).upper() - return str(val).upper() + extracted = self._extract_str_value(val) + if extracted is not None: + return extracted.strip().strip("`").upper() + return str(val).strip().strip("`").upper() def _parse_refresh_tuple( self, @@ -1869,12 +1877,10 @@ def _parse_refresh_tuple( unit_var: t.Optional[exp.Var] = None starts_lit: t.Optional[exp.Literal] = None for e in tpl.expressions: - if ( - isinstance(e, exp.EQ) - and isinstance(e.this, exp.Column) - and isinstance(e.this.this, exp.Identifier) - ): - key = str(e.this.this.this).strip('"').lower() + if isinstance(e, exp.EQ): + key = ( + (self._extract_str_value(e.this) or "").strip('"').strip("'").strip("`").lower() + ) val = e.expression if key == "method": method_upper = self._to_upper_str(val) @@ -1906,10 +1912,10 @@ def _parse_refresh_tuple( raise SQLMeshError(f"Invalid materialized view refresh unit: {unit_upper}") unit_var = exp.Var(this=unit_upper) elif key == "starts": - if isinstance(val, exp.Literal): - starts_lit = exp.Literal.string(str(val.this)) - else: - starts_lit = exp.Literal.string(str(val)) + starts_str = self._extract_str_value(val) + starts_lit = exp.Literal.string( + starts_str if starts_str is not None else str(val) + ) return method_var, kind_str, every_lit, unit_var, starts_lit @staticmethod diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 274195d939..6069d2d12d 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -1178,6 +1178,27 @@ def test_sync_materialized_view_drop_and_replace( "CREATE MATERIALIZED VIEW `sync_mv` AS SELECT `id` AS `mv_id` FROM `db`.`orders`", ] + adapter.cursor.execute.reset_mock() + + # source_table parsed as exp.Column (from MODEL physical_properties) and 3-part internal.db.orders + adapter.create_view( + "sync_mv", + parse_one("SELECT id FROM orders", read="doris"), + replace=True, + materialized=True, + view_properties={ + "materialized_type": exp.Literal.string("SYNC"), + "source_table": parse_one("internal.db.orders", read="doris"), + }, + ) + + assert to_sql_calls(adapter) == [ + "USE `db`", + "DROP MATERIALIZED VIEW IF EXISTS `sync_mv` ON `db`.`orders`", + "USE `db`", + "CREATE MATERIALIZED VIEW `sync_mv` AS SELECT `id` AS `mv_id` FROM `orders`", + ] + def test_delete_with_subquery_fallback_and_cardinality_aware_composite_batching( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture @@ -1274,7 +1295,19 @@ def test_delete_by_literal_keys_doris_delete_handler_invariants( @pytest.mark.parametrize( "non_key_type", - ["FLOAT", "DOUBLE", "ARRAY", "MAP", "STRUCT", "JSON", "VARIANT"], + [ + "FLOAT", + "DOUBLE", + "ARRAY", + "MAP", + "STRUCT", + "JSON", + "VARIANT", + "BITMAP", + "HLL", + "QUANTILE_STATE", + "AGG_STATE(sum(INT))", + ], ) def test_olap_key_type_and_prefix_invariants( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], @@ -1284,7 +1317,7 @@ def test_olap_key_type_and_prefix_invariants( adapter.create_table( "db.complex_tbl", target_columns_to_types={ - "bad_first_col": exp.DataType.build(non_key_type, dialect="doris"), + "bad_first_col": exp.DataType.build(non_key_type, dialect="doris", udt=True), "valid_key": exp.DataType.build("BIGINT", dialect="doris"), }, ) @@ -1550,3 +1583,121 @@ def test_delete_from_boolean_literal_simplification( adapter.cursor.execute.reset_mock() adapter.delete_from("db.tbl", where="TRUE") assert to_sql_calls(adapter) == ["TRUNCATE TABLE `db`.`tbl`"] + + # Target-qualified extra condition in DELETE ... USING is rewritten to alias _t1 + adapter.cursor.execute.reset_mock() + adapter.delete_from( + "db.target", + where="target.ds = '2024-01-01' AND id IN (SELECT id FROM db.src)", + ) + assert to_sql_calls(adapter) == [ + "DELETE FROM `db`.`target` AS `_t1` USING (SELECT `id` FROM `db`.`src`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`ds` = '2024-01-01'" + ] + + +def test_resolve_table_refs_in_physical_properties_for_mv_and_sync_mv( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + from sqlmesh.core.dialect import parse + from sqlmesh.core.model import SqlModel, load_sql_based_model + from sqlmesh.core.snapshot import Snapshot, SnapshotChangeCategory + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + base_model = t.cast( + SqlModel, + load_sql_based_model( + parse( + """ + MODEL ( + name doris_db.orders, + kind FULL, + dialect doris, + columns (id INT, amount INT) + ); + SELECT 1 AS id, 10 AS amount; + """, + default_dialect="doris", + ) + ), + ) + base_snapshot = Snapshot.from_node(base_model, nodes={}, ttl="in 1 week") + base_snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + physical_tbl = exp.to_table(base_snapshot.table_name()) + expected_physical = f"{physical_tbl.db}.{physical_tbl.name}" + snapshots = {base_snapshot.name: base_snapshot} + + # 1. Async MTMV with excluded_trigger_tables resolves model ref to physical snapshot table + async_mv_model = t.cast( + SqlModel, + load_sql_based_model( + parse( + """ + MODEL ( + name doris_db.orders_mtmv, + kind VIEW (materialized true), + dialect doris, + columns (id INT, amount INT), + physical_properties ( + build = IMMEDIATE, + refresh = AUTO, + excluded_trigger_tables = doris_db.orders + ) + ); + SELECT id, amount FROM doris_db.orders; + """, + default_dialect="doris", + ) + ), + ) + rendered_async_props = async_mv_model.render_physical_properties( + snapshots=snapshots, engine_adapter=adapter + ) + adapter.create_view( + async_mv_model.name, + async_mv_model.render_query() or parse_one("SELECT id, amount FROM doris_db.orders"), + replace=False, + materialized=True, + target_columns_to_types=async_mv_model.columns_to_types, + view_properties=rendered_async_props, + ) + async_ddl = to_sql_calls(adapter)[-1] + assert f"'excluded_trigger_tables'='{expected_physical}'" in async_ddl + + # 2. Sync MV with source_table resolves model ref to physical snapshot table for DROP ... ON + adapter.cursor.execute.reset_mock() + sync_mv_model = t.cast( + SqlModel, + load_sql_based_model( + parse( + """ + MODEL ( + name doris_db.orders_sync_mv, + kind VIEW (materialized true), + dialect doris, + physical_properties ( + materialized_type = SYNC, + source_table = doris_db.orders + ) + ); + SELECT id, SUM(amount) AS total FROM doris_db.orders GROUP BY id; + """, + default_dialect="doris", + ) + ), + ) + rendered_sync_props = sync_mv_model.render_physical_properties( + snapshots=snapshots, engine_adapter=adapter + ) + adapter.create_view( + sync_mv_model.name, + sync_mv_model.render_query(snapshots=snapshots) + or parse_one("SELECT id, SUM(amount) AS total FROM doris_db.orders GROUP BY id"), + replace=True, + materialized=True, + view_properties=rendered_sync_props, + ) + sync_calls = to_sql_calls(adapter) + assert sync_calls[0] == ( + f"DROP MATERIALIZED VIEW IF EXISTS `doris_db`.`orders_sync_mv` ON `{physical_tbl.db}`.`{physical_tbl.name}`" + ) From eed2be151c1aa2d075a7ef4432cf15786a8748ef Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Wed, 21 Oct 2026 18:43:09 -0700 Subject: [PATCH 13/27] fix(doris): validate refresh_trigger and document MV physical table resolution Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/integrations/engines/doris.md | 7 +++- sqlmesh/core/engine_adapter/doris.py | 51 +++++++++++++++++++------ tests/core/engine_adapter/test_doris.py | 11 ++++++ 3 files changed, 56 insertions(+), 13 deletions(-) diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index 3cd891788b..abf8fd5b12 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -452,7 +452,8 @@ MODEL ( build = 'IMMEDIATE', refresh = 'AUTO', refresh_trigger = 'ON SCHEDULE EVERY 1 HOUR', - distributed_by = 'HASH(user_id) BUCKETS 10' + distributed_by = 'HASH(user_id) BUCKETS 10', + excluded_trigger_tables = 'analytics.user_events' ) ); @@ -464,6 +465,10 @@ FROM analytics.user_events GROUP BY user_id, CAST(event_time AS DATE); ``` +!!! note "Automatic Physical Table Resolution & Non-Destructive Evaluation" + - Doris materialized views are only created on initial deployment or when a rebuild is required (`RECREATE_MATERIALIZED_VIEW_ON_EVALUATION = False`), preserving MTMV refresh state across routine `sqlmesh run` invocations. + - When `excluded_trigger_tables` (on async MTMVs) or `source_table` (on synchronous MVs) references a SQLMesh-managed model, SQLMesh automatically resolves the model name to its underlying physical snapshot table name. + ### Synchronous Materialized Views (`materialized_type = 'SYNC'`) Doris also supports single-table synchronous materialized views (rollup indexes maintained atomically on base table writes) by setting `materialized_type = 'SYNC'` and `source_table` in `physical_properties`: diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 24c9e39866..5d1dc01244 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -1621,10 +1621,19 @@ def _build_table_properties_exp( unit_var: t.Optional[exp.Var] = None starts_lit: t.Optional[exp.Literal] = None + if isinstance(refresh_val, exp.Paren) and isinstance(refresh_val.this, exp.EQ): + refresh_val = exp.Tuple(expressions=[refresh_val.this]) + elif isinstance(refresh_val, exp.EQ): + refresh_val = exp.Tuple(expressions=[refresh_val]) + if isinstance(refresh_val, exp.Tuple): method_var, kind_str, every_lit, unit_var, starts_lit = ( self._parse_refresh_tuple(refresh_val) ) + if refresh_trigger_val is not None: + trigger_text = self._extract_str_value(refresh_trigger_val) or "" + k, e, u, s = self._parse_trigger_string(trigger_text) + kind_str, every_lit, unit_var, starts_lit = k, e, u, s else: raw_refresh = ( self._extract_str_value(refresh_val) if refresh_val is not None else "AUTO" @@ -1936,8 +1945,11 @@ def _parse_trigger_string( starts_lit = None m_kind = re.match(r"^(MANUAL|COMMIT|SCHEDULE)\b", text, flags=re.IGNORECASE) - if m_kind: - kind = m_kind.group(1).upper() + if not m_kind: + raise SQLMeshError( + f"Invalid materialized view refresh trigger: {trigger_text}. Expected MANUAL, COMMIT, or SCHEDULE." + ) + kind = m_kind.group(1).upper() m_every = re.search(r"\bEVERY\s+(\d+)\s+([A-Za-z]+)", text, flags=re.IGNORECASE) if m_every: every_lit = exp.Literal.number(int(m_every.group(1))) @@ -2077,8 +2089,14 @@ def _build_distributed_by_exp( has_eq = any(isinstance(e, exp.EQ) for e in distributed_by.expressions) if has_eq: for expr in distributed_by.expressions: - if isinstance(expr, exp.EQ) and hasattr(expr.this, "this"): - key = str(expr.this.this).strip('"').strip("'").lower() + if isinstance(expr, exp.EQ): + key = ( + (self._extract_str_value(expr.this) or "") + .strip('"') + .strip("'") + .strip("`") + .lower() + ) if isinstance(expr.expression, exp.Literal): distributed_info[key] = expr.expression.this elif isinstance(expr.expression, (exp.Array, exp.Tuple)): @@ -2094,14 +2112,23 @@ def _build_distributed_by_exp( buckets=default_buckets_expr or exp.Literal.number(10), order=None, ) - elif isinstance(distributed_by, exp.Paren) and isinstance(distributed_by.this, exp.EQ): - expr = distributed_by.this - if hasattr(expr.this, "this"): - key = str(expr.this.this).strip('"').strip("'").lower() - if isinstance(expr.expression, exp.Literal): - distributed_info[key] = expr.expression.this - else: - distributed_info[key] = self._extract_str_value(expr.expression) + elif isinstance(distributed_by, exp.EQ) or ( + isinstance(distributed_by, exp.Paren) and isinstance(distributed_by.this, exp.EQ) + ): + expr = ( + distributed_by.this if isinstance(distributed_by, exp.Paren) else distributed_by + ) + key = ( + (self._extract_str_value(expr.this) or "") + .strip('"') + .strip("'") + .strip("`") + .lower() + ) + if isinstance(expr.expression, exp.Literal): + distributed_info[key] = expr.expression.this + else: + distributed_info[key] = self._extract_str_value(expr.expression) elif isinstance(distributed_by, dict): distributed_info = distributed_by diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 6069d2d12d..bd8e96863b 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -1005,6 +1005,17 @@ def test_external_catalog_view_error_and_drop_noop( ) assert to_sql_calls(adapter) == [] + # Invalid refresh_trigger also raises SQLMeshError instead of silently falling back to AUTO + with pytest.raises(SQLMeshError, match="Invalid materialized view refresh trigger"): + adapter.create_view( + "silver.events_mv", + parse_one("SELECT 1 AS id"), + replace=True, + materialized=True, + view_properties={"refresh_trigger": exp.Literal.string("ON DEMAND")}, + ) + assert to_sql_calls(adapter) == [] + def test_duplicate_key_with_buckets_and_random_buckets_fallback( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], From 3d575a5504c50e2df74bd4dd54abc61c678c0f16 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Wed, 21 Oct 2026 19:57:43 -0700 Subject: [PATCH 14/27] fix(doris): harden key/partition/delete invariants and expand unit + integration tests Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 185 ++++++++++++++++-- .../integration/test_integration_doris.py | 139 +++++++++++++ tests/core/engine_adapter/test_doris.py | 68 ++++++- 3 files changed, 371 insertions(+), 21 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 5d1dc01244..917742aa93 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -918,6 +918,14 @@ def delete_from( if subquery_info is not None: self._execute_delete_with_subquery(table_name, *subquery_info) else: + target_table = exp.to_table(table_name) + for col in where.find_all(exp.Column): + if not col.find_ancestor(exp.Subquery, exp.Select) and ( + col.table and col.table.lower() == target_table.name.lower() + ): + col.set("table", None) + col.set("db", None) + col.set("catalog", None) super().delete_from(table_name, where) @classmethod @@ -1157,7 +1165,19 @@ def _delete_by_literal_keys( if not rows: return - extra_conditions = extra_conditions or [] + target_table = exp.to_table(table_name) + cleaned_extra_conditions: t.List[exp.Expr] = [] + for extra in extra_conditions or []: + cleaned_extra = extra.copy() + for col in cleaned_extra.find_all(exp.Column): + if not col.find_ancestor(exp.Subquery, exp.Select) and ( + col.table and col.table.lower() == target_table.name.lower() + ): + col.set("table", None) + col.set("db", None) + col.set("catalog", None) + cleaned_extra_conditions.append(cleaned_extra) + extra_conditions = cleaned_extra_conditions cols = [exp.column(k.name, quoted=True) for k in key_exprs] batch_size = self.DEFAULT_BATCH_SIZE @@ -1358,6 +1378,51 @@ def _build_create_table_exp( [proj_by_name[c.lower()] for c in reordered_select_cols], ) + # Doris AUTO PARTITION BY RANGE requires the partition column(s) to be defined as NOT NULL. + if not is_iceberg and isinstance(table_name_or_schema, exp.Schema): + table_props = kwargs.get("table_properties") or {} + auto_partition = table_props.get("auto_partition") + is_auto = bool( + auto_partition is not None + and str(self._extract_str_value(auto_partition) or "").strip().lower() + in ("true", "1") + ) + raw_part_by = kwargs.get("partitioned_by") or table_props.get("partitioned_by") + if is_auto and raw_part_by: + part_list = raw_part_by if isinstance(raw_part_by, list) else [raw_part_by] + parsed_part_exprs, part_kind = self._parse_partition_expressions(part_list) + if (part_kind or "RANGE").upper() != "LIST": + auto_range_cols = { + col.name.lower() + for expr in parsed_part_exprs + for col in expr.find_all(exp.Column) + } + if auto_range_cols: + table_name_or_schema = table_name_or_schema.copy() + updated_exprs: t.List[exp.Expr] = [] + for expr in table_name_or_schema.expressions: + if ( + isinstance(expr, exp.ColumnDef) + and expr.name.lower() in auto_range_cols + ): + existing_constraints = list(expr.args.get("constraints") or []) + has_not_null = any( + isinstance(c, exp.NotNullColumnConstraint) + or isinstance( + getattr(c, "kind", None), exp.NotNullColumnConstraint + ) + for c in existing_constraints + ) + if not has_not_null: + expr = expr.copy() + expr.set( + "constraints", + [exp.ColumnConstraint(kind=exp.NotNullColumnConstraint())] + + existing_constraints, + ) + updated_exprs.append(expr) + table_name_or_schema.set("expressions", updated_exprs) + return super()._build_create_table_exp( table_name_or_schema, expression=expression, @@ -1421,8 +1486,44 @@ def _reorder_olap_columns( None, ) if valid_col: - new_order = [valid_col] + [c for c in col_names if c != valid_col] - return {c: target_columns_to_types[c] for c in new_order} + col_names = [valid_col] + [c for c in col_names if c != valid_col] + + # Ensure any OLAP partition columns with valid key types precede non-keyable columns + raw_part_by = kwargs.get("partitioned_by") or table_props.get("partitioned_by") + if raw_part_by: + part_list = raw_part_by if isinstance(raw_part_by, list) else [raw_part_by] + parsed_part_exprs, _ = self._parse_partition_expressions(part_list) + col_by_lower = {c.lower(): c for c in col_names} + part_cols = [ + col_by_lower[col.name.lower()] + for expr in parsed_part_exprs + for col in expr.find_all(exp.Column) + if col.name.lower() in col_by_lower + and self._is_valid_olap_key_type( + target_columns_to_types[col_by_lower[col.name.lower()]] + ) + ] + first_non_key_idx = next( + ( + i + for i, c in enumerate(col_names) + if not self._is_valid_olap_key_type(target_columns_to_types[c]) + ), + None, + ) + if ( + part_cols + and first_non_key_idx is not None + and any(col_names.index(p) > first_non_key_idx for p in part_cols) + ): + moved_parts = [p for p in part_cols if col_names.index(p) > first_non_key_idx] + moved_set = set(moved_parts) + before_non_key = [c for c in col_names[:first_non_key_idx] if c not in moved_set] + after_non_key = [c for c in col_names[first_non_key_idx:] if c not in moved_set] + col_names = before_non_key + moved_parts + after_non_key + + if col_names != list(target_columns_to_types): + return {c: target_columns_to_types[c] for c in col_names} return target_columns_to_types @@ -1513,6 +1614,17 @@ def _build_partitioned_by_exp( dynamic_partition_enabled = kwargs.get("dynamic_partition_enabled") partitioned_by, partition_kind = self._parse_partition_expressions(partitioned_by) + if target_columns_to_types: + col_by_lower = {c.lower(): c for c in target_columns_to_types} + normalized_exprs: t.List[exp.Expr] = [] + for expr in partitioned_by: + expr_copy = expr.copy() + for col_node in expr_copy.find_all(exp.Column): + canonical = col_by_lower.get(col_node.name.lower()) + if canonical and canonical != col_node.name: + col_node.set("this", exp.to_identifier(canonical)) + normalized_exprs.append(expr_copy) + partitioned_by = normalized_exprs if is_iceberg: return exp.PartitionByListProperty( @@ -1685,6 +1797,9 @@ def _build_table_properties_exp( duplicate_key_cols = self._extract_column_names(duplicate_key) aggregate_key_cols = self._extract_column_names(aggregate_key) order_by_cols = self._extract_column_names(order_by) + col_by_lower = ( + {c.lower(): c for c in target_columns_to_types} if target_columns_to_types else {} + ) if not is_iceberg: specified_key_types = [ @@ -1701,9 +1816,6 @@ def _build_table_properties_exp( f"Doris tables support at most one key model (unique_key, duplicate_key, or aggregate_key), " f"got multiple: {', '.join(specified_key_types)}." ) - col_by_lower = ( - {c.lower(): c for c in target_columns_to_types} if target_columns_to_types else {} - ) active_key_cols = unique_key_cols or duplicate_key_cols or aggregate_key_cols if active_key_cols and col_by_lower and not is_materialized_view: missing_cols = [k for k in active_key_cols if k.lower() not in col_by_lower] @@ -1771,16 +1883,18 @@ def _build_table_properties_exp( add_partition = True if partitioned_by: parsed_part_exprs, _ = self._parse_partition_expressions(partitioned_by) - if unique_key_cols and not is_materialized_view and not is_iceberg: - key_cols_set = {k.lower() for k in unique_key_cols} + strict_part_key_cols = unique_key_cols or aggregate_key_cols + if strict_part_key_cols and not is_materialized_view and not is_iceberg: + key_cols_set = {k.lower() for k in strict_part_key_cols} partition_cols = set() for expr in parsed_part_exprs: for col_node in expr.find_all(exp.Column): partition_cols.add(col_node.name.lower()) not_in_key = partition_cols - key_cols_set if not_in_key: + key_label = "UNIQUE KEY" if unique_key_cols else "AGGREGATE KEY" logger.warning( - f"[Doris] UNIQUE KEY model: partitioned_by columns {not_in_key} not in key_cols {key_cols_set}, skip PARTITION BY." + f"[Doris] {key_label} model: partitioned_by columns {not_in_key} not in key_cols {key_cols_set}, skip PARTITION BY." ) add_partition = False @@ -1810,6 +1924,19 @@ def _build_table_properties_exp( is_iceberg=is_iceberg, ) if partition_expr: + if col_by_lower: + part_col_nodes = [ + col_node + for expr in parsed_part_exprs + for col_node in expr.find_all(exp.Column) + ] + missing_part_cols = [ + c.name for c in part_col_nodes if c.name.lower() not in col_by_lower + ] + if missing_part_cols: + raise SQLMeshError( + f"Doris partitioned_by columns {missing_part_cols} not found in table schema columns {list(target_columns_to_types or ())}." + ) properties.append(partition_expr) # Pop distribution / bucketing properties so they never leak into PROPERTIES (...) @@ -1823,21 +1950,23 @@ def _build_table_properties_exp( ) if not is_iceberg: + strict_key_cols = ( + (unique_key_cols or aggregate_key_cols) if not is_materialized_view else [] + ) dist_prop = self._build_distributed_by_exp( distributed_by=distributed_by, clustered_by=effective_clustered_by, - unique_key_cols=unique_key_cols if not is_materialized_view else [], - fallback_key_cols=( - (duplicate_key_cols or aggregate_key_cols) if not is_materialized_view else [] - ), + unique_key_cols=strict_key_cols, + fallback_key_cols=duplicate_key_cols if not is_materialized_view else [], top_level_buckets=top_level_buckets, ) if dist_prop is not None: if not is_materialized_view: dist_kind = str(dist_prop.args.get("kind") or "").upper() - if unique_key_cols and dist_kind == "RANDOM": + if strict_key_cols and dist_kind == "RANDOM": + key_kind_label = "UNIQUE KEY" if unique_key_cols else "AGGREGATE KEY" raise SQLMeshError( - "Doris UNIQUE KEY tables do not support DISTRIBUTED BY RANDOM. " + f"Doris {key_kind_label} tables do not support DISTRIBUTED BY RANDOM. " "Use DISTRIBUTED BY HASH (...) instead." ) if col_by_lower and dist_prop.expressions: @@ -1853,6 +1982,19 @@ def _build_table_properties_exp( "expressions", [exp.to_column(col_by_lower[c.lower()]) for c in dist_col_names], ) + if strict_key_cols and dist_prop.expressions: + strict_key_set = {k.lower() for k in strict_key_cols} + non_key_dist_cols = [ + c.name + for c in dist_prop.expressions + if c.name.lower() not in strict_key_set + ] + if non_key_dist_cols: + key_kind_label = "UNIQUE KEY" if unique_key_cols else "AGGREGATE KEY" + raise SQLMeshError( + f"Doris {key_kind_label} tables require distributed_by columns to be key columns, " + f"got {non_key_dist_cols} which are not in {strict_key_cols}." + ) properties.append(dist_prop) # Default replication_num = 1 for single-BE dev/CI when not explicitly configured on OLAP tables @@ -2159,9 +2301,16 @@ def _build_distributed_by_exp( # Fallback 1: Top-level clustered_by (or physical_properties clustered_by) if clustered_by: - cols = [ - exp.to_column(c) for expr in clustered_by for c in self._extract_column_names(expr) - ] + col_names = [c for expr in clustered_by for c in self._extract_column_names(expr)] + if unique_key_cols: + strict_key_set = {k.lower() for k in unique_key_cols} + not_in_key = [c for c in col_names if c.lower() not in strict_key_set] + if not_in_key: + logger.warning( + f"[Doris] clustered_by columns {not_in_key} not in key columns {unique_key_cols}, falling back to key column distribution." + ) + col_names = [] + cols = [exp.to_column(c) for c in col_names] if cols: return exp.DistributedByProperty( kind=exp.Var(this="HASH"), diff --git a/tests/core/engine_adapter/integration/test_integration_doris.py b/tests/core/engine_adapter/integration/test_integration_doris.py index 8c1ffdc47b..cae796ec0a 100644 --- a/tests/core/engine_adapter/integration/test_integration_doris.py +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -224,3 +224,142 @@ def test_merge_unique_key(ctx: TestContext): ] ), ) + + +def test_auto_partition_and_insert_overwrite_by_partition(ctx: TestContext): + import pandas as pd # noqa: TID253 + from sqlglot import parse_one + + table = ctx.table("test_auto_partition") + columns_to_types = { + "payload": exp.DataType.build("STRING", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + } + partitioned_by = [parse_one("RANGE(DATE_TRUNC(ds, 'day'))", read="doris")] + ctx.engine_adapter.create_table( + table, + columns_to_types, + partitioned_by=partitioned_by, + table_properties={"auto_partition": exp.Literal.string("true")}, + ) + initial_data = pd.DataFrame( + [ + {"payload": "a", "id": 1, "ds": "2022-01-01"}, + {"payload": "b", "id": 2, "ds": "2022-01-02"}, + ] + ) + ctx.engine_adapter.insert_append( + table, + ctx.input_data(initial_data, columns_to_types), + target_columns_to_types=columns_to_types, + ) + overwrite_data = pd.DataFrame( + [ + {"payload": "b_updated", "id": 20, "ds": "2022-01-02"}, + {"payload": "c", "id": 3, "ds": "2022-01-03"}, + ] + ) + ctx.engine_adapter.insert_overwrite_by_partition( + table, + ctx.input_data(overwrite_data, columns_to_types), + partitioned_by=partitioned_by, + target_columns_to_types=columns_to_types, + ) + rows = ctx.engine_adapter.fetchall( + exp.select("id", exp.cast(exp.column("ds"), "VARCHAR").as_("ds"), "payload") + .from_(table) + .order_by("id") + ) + assert list(rows) == [ + (1, "2022-01-01", "a"), + (3, "2022-01-03", "c"), + (20, "2022-01-02", "b_updated"), + ] + + +def test_unique_key_order_by_and_key_reordering(ctx: TestContext): + import pandas as pd # noqa: TID253 + + table = ctx.table("test_uk_order_by") + columns_to_types = { + "payload": exp.DataType.build("VARCHAR(64)", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + "seq": exp.DataType.build("INT", dialect="doris"), + } + ctx.engine_adapter.create_table( + table, + columns_to_types, + unique_key=[exp.to_column("id"), exp.to_column("ds")], + table_properties={ + "order_by": exp.Tuple(expressions=[exp.to_column("seq"), exp.to_column("id")]), + }, + ) + df = pd.DataFrame( + [ + {"payload": "first", "id": 1, "ds": "2022-01-01", "seq": 10}, + {"payload": "second", "id": 2, "ds": "2022-01-02", "seq": 5}, + ] + ) + ctx.engine_adapter.insert_append( + table, + ctx.input_data(df, columns_to_types), + target_columns_to_types=columns_to_types, + ) + rows = ctx.engine_adapter.fetchall( + exp.select("id", exp.cast(exp.column("ds"), "VARCHAR").as_("ds"), "seq", "payload") + .from_(table) + .order_by("id") + ) + assert list(rows) == [ + (1, "2022-01-01", 10, "first"), + (2, "2022-01-02", 5, "second"), + ] + + +def test_duplicate_key_delete_with_subquery_and_between(ctx: TestContext): + import pandas as pd # noqa: TID253 + + target_table = ctx.table("test_dup_delete_target") + src_table = ctx.table("test_dup_delete_src") + ctx.engine_adapter.create_table( + target_table, + ctx.columns_to_types, + table_properties={ + "duplicate_key": exp.Tuple(expressions=[exp.to_column("id"), exp.to_column("ds")]), + }, + ) + input_data = pd.DataFrame( + [ + {"id": 1, "ds": "2022-01-01"}, + {"id": 2, "ds": "2022-01-02"}, + {"id": 3, "ds": "2022-01-03"}, + {"id": 4, "ds": "2022-01-04"}, + ] + ) + ctx.engine_adapter.insert_append( + target_table, + ctx.input_data(input_data), + target_columns_to_types=ctx.columns_to_types, + ) + # 1. BETWEEN rewrite in DELETE WHERE + ctx.engine_adapter.delete_from(target_table, "ds BETWEEN '2022-01-01' AND '2022-01-01'") + + # 2. Subquery DELETE on DUPLICATE KEY table (triggers literal batch fallback since USING requires UNIQUE KEY) + src_data = pd.DataFrame([{"id": 2, "ds": "2022-01-02"}]) + ctx.engine_adapter.ctas(src_table, ctx.input_data(src_data), ctx.columns_to_types) + key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", exp.column("id"), exp.column("ds")) + ctx.engine_adapter.delete_from( + target_table, + where=key_exp.isin(query=exp.select(key_exp).from_(src_table)), + ) + ctx.compare_with_current( + target_table, + pd.DataFrame( + [ + {"id": 3, "ds": "2022-01-03"}, + {"id": 4, "ds": "2022-01-04"}, + ] + ), + ) diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index bd8e96863b..78f84a25fc 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -724,17 +724,21 @@ def test_create_table_auto_partition_and_unannotated_partitioned_by_guard( target_columns_to_types={"id": exp.DataType.build("INT"), "ds": exp.DataType.build("DATE")}, partitioned_by=[exp.to_column("ds")], ) - # With auto_partition = true, emits AUTO PARTITION BY RANGE (...) () + # With auto_partition = true, emits AUTO PARTITION BY RANGE (...) () and marks range partition column NOT NULL adapter.create_table( "olap_auto_partition", - target_columns_to_types={"id": exp.DataType.build("INT"), "ds": exp.DataType.build("DATE")}, + target_columns_to_types={ + "payload": exp.DataType.build("STRING", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + }, partitioned_by=[parse_one("RANGE(DATE_TRUNC(ds, 'DAY'))", dialect="doris")], table_properties={"auto_partition": exp.Literal.string("true")}, ) assert to_sql_calls(adapter) == [ "CREATE TABLE IF NOT EXISTS `olap_no_partitions` (`id` INT, `ds` DATE)", - "CREATE TABLE IF NOT EXISTS `olap_auto_partition` (`id` INT, `ds` DATE) AUTO PARTITION BY RANGE (DATE_TRUNC(`ds`, 'DAY')) ()", + "CREATE TABLE IF NOT EXISTS `olap_auto_partition` (`id` INT, `ds` DATE NOT NULL, `payload` STRING) AUTO PARTITION BY RANGE (DATE_TRUNC(`ds`, 'DAY')) ()", ] @@ -1529,6 +1533,54 @@ def test_key_model_mutual_exclusion_and_missing_column_validation( }, ) + # Explicit distributed_by on UNIQUE KEY / AGGREGATE KEY must be a subset of key columns + with pytest.raises( + SQLMeshError, + match="Doris UNIQUE KEY tables require distributed_by columns to be key columns", + ): + adapter.create_table( + "db.uk_non_key_dist", + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + table_properties={ + "unique_key": exp.to_column("id"), + "distributed_by": "HASH(name) BUCKETS 8", + }, + ) + + # Cross-engine clustered_by on non-key column for UNIQUE KEY table falls back to key column + adapter.cursor.execute.reset_mock() + adapter.create_table( + "db.uk_clustered_by_fallback", + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + clustered_by=[exp.to_column("name")], + table_properties={ + "unique_key": exp.to_column("id"), + }, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`uk_clustered_by_fallback` (`id` INT, `name` VARCHAR(32)) UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10" + ] + + # Missing partitioned_by column raises SQLMeshError when partitioning is active + with pytest.raises( + SQLMeshError, match="Doris partitioned_by columns .* not found in table schema columns" + ): + adapter.create_table( + "db.missing_part_col", + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + }, + partitioned_by=[parse_one("RANGE(missing_ds)", read="doris")], + table_properties={"auto_partition": exp.Literal.string("true")}, + ) + def test_ctas_reorders_projections_in_place_without_double_subquery( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], @@ -1605,6 +1657,16 @@ def test_delete_from_boolean_literal_simplification( "DELETE FROM `db`.`target` AS `_t1` USING (SELECT `id` FROM `db`.`src`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`ds` = '2024-01-01'" ] + # Target-qualified condition in simple DELETE FROM is stripped to unqualified column for Doris DeleteHandler + adapter.cursor.execute.reset_mock() + adapter.delete_from( + "db.target", + where="target.ds = '2024-01-01' AND target.id = 5", + ) + assert to_sql_calls(adapter) == [ + "DELETE FROM `db`.`target` WHERE `ds` = '2024-01-01' AND `id` = 5" + ] + def test_resolve_table_refs_in_physical_properties_for_mv_and_sync_mv( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], From 6fe3f5f7287afa442c95f1e178910584c414f735 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Thu, 22 Oct 2026 08:45:30 -0700 Subject: [PATCH 15/27] fix(doris): harden DataFrame/seed key coercion, Iceberg bucket/truncate, CTAS auto-partition, and docs Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/integrations/engines/doris.md | 4 +- docs/quickstart/cli.md | 1 + sqlmesh/core/engine_adapter/doris.py | 266 +++++++++++++++++++----- tests/core/engine_adapter/test_doris.py | 231 ++++++++++++++++++++ 4 files changed, 450 insertions(+), 52 deletions(-) diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index abf8fd5b12..8617491142 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -204,7 +204,7 @@ Doris OLAP tables support **Multi-Range Batch Partitioning**, **Explicit Partiti !!! note "OLAP Partitioning Requirements" - On Doris OLAP tables, `PARTITION BY` requires at least one of `partitions`, `auto_partition = 'true'`, or `'dynamic_partition.enable' = 'true'` in `physical_properties`. If none is provided, SQLMesh omits the `PARTITION BY` clause so Doris does not create a zero-partition table that rejects inserts. - - On `UNIQUE KEY` tables, Doris requires all `partitioned_by` columns to be a subset of `unique_key`. + - On `UNIQUE KEY` and `AGGREGATE KEY` tables, Doris requires all `partitioned_by` columns to be a subset of the key columns. ### Multi-range & explicit partitions @@ -509,7 +509,7 @@ FROM raw.logs; @IF( @runtime_stage = 'creating', - CREATE INDEX idx_message ON @this_model (message) USING INVERTED PROPERTIES("parser" = "english") COMMENT 'Full-text inverted index' + ALTER TABLE @this_model ADD INDEX idx_message (message) USING INVERTED COMMENT 'Full-text inverted index' ); ``` diff --git a/docs/quickstart/cli.md b/docs/quickstart/cli.md index e7f55b7467..4b300c9018 100644 --- a/docs/quickstart/cli.md +++ b/docs/quickstart/cli.md @@ -118,6 +118,7 @@ Choose your SQL engine: [16] RisingWave [17] Fabric [18] StarRocks + [19] Db2 [20] Apache Doris Enter a number: 1 diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 917742aa93..5540f62459 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -50,7 +50,7 @@ class DorisEngineAdapter( RowDiffMixin, ): DIALECT = "doris" - DEFAULT_BATCH_SIZE = 5000 + DEFAULT_BATCH_SIZE = 500 SUPPORTS_TRANSACTIONS = False COMMENT_CREATION_TABLE = CommentCreationTable.IN_SCHEMA_DEF_NO_CTAS COMMENT_CREATION_VIEW = CommentCreationView.IN_SCHEMA_DEF_NO_COMMANDS @@ -133,6 +133,15 @@ class DorisEngineAdapter( "DOUBLE", "STRING", "TEXT", + "TINYTEXT", + "MEDIUMTEXT", + "LONGTEXT", + "BLOB", + "TINYBLOB", + "MEDIUMBLOB", + "LONGBLOB", + "BINARY", + "VARBINARY", "JSON", "JSONB", "VARIANT", @@ -399,7 +408,7 @@ def _get_data_objects( self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None ) -> t.List[DataObject]: """Returns all the data objects that exist in the given schema and optional catalog.""" - schema_exp = to_schema(schema_name) + schema_exp = to_schema(schema_name) if schema_name else exp.Table() catalog = schema_exp.catalog or None query_catalog = catalog or self.default_catalog @@ -428,7 +437,11 @@ def _get_data_objects( .as_("type"), ) .from_(exp.table_("tables", db="information_schema", catalog=query_catalog)) - .where(exp.column("table_schema").eq(schema_exp.db)) + .where( + exp.column("table_schema").eq( + schema_exp.db if schema_exp.db else exp.func("DATABASE") + ) + ) ) if object_names: lowered_names = [name.lower() for name in object_names] @@ -471,9 +484,13 @@ def _get_data_objects( obj.type.is_table for obj in result ): try: - db_literal = self._to_sql(exp.Literal.string(schema_exp.db), quote=False) + db_arg = ( + self._to_sql(exp.Literal.string(schema_exp.db), quote=False) + if schema_exp.db + else "DATABASE()" + ) mv_rows = self.fetchall( - f"SELECT Name FROM mv_infos('database' = {db_literal})", + f"SELECT Name FROM mv_infos('database' = {db_arg})", quote_identifiers=False, ) mv_names = { @@ -677,11 +694,20 @@ def _create_materialized_view( self.drop_data_object_on_type_mismatch( self.get_data_object(view_name), DataObjectType.MATERIALIZED_VIEW ) + drop_view_props = dict(view_properties) if view_properties else None + if ( + is_sync_mv + and drop_view_props is not None + and not drop_view_props.get("source_table") + ): + first_tbl = query.find(exp.Table) + if first_tbl: + drop_view_props["source_table"] = first_tbl self.drop_view( view_name, ignore_if_not_exists=True, materialized=True, - view_properties=view_properties, + view_properties=drop_view_props, ) # Doris synchronous materialized views are single-table rollup indexes whose DDL syntax @@ -874,16 +900,28 @@ def insert_overwrite_by_partition( source_columns: t.Optional[t.List[str]] = None, ) -> None: parsed_exprs, _ = self._parse_partition_expressions(partitioned_by) - key_exprs: t.List[exp.Expr] = [] + raw_key_exprs: t.List[exp.Expr] = [] for expr in parsed_exprs or partitioned_by: if isinstance(expr, exp.Column): - key_exprs.append(expr) + raw_key_exprs.append(expr) else: cols = list(expr.find_all(exp.Column)) if cols: - key_exprs.extend(cols) + raw_key_exprs.extend(cols) else: - key_exprs.append(expr) + raw_key_exprs.append(expr) + key_exprs: t.List[exp.Expr] = [] + seen_cols: t.Set[str] = set() + for expr in raw_key_exprs: + if isinstance(expr, exp.Column): + col_lower = expr.name.lower() + if col_lower not in seen_cols: + seen_cols.add(col_lower) + key_exprs.append( + exp.column(expr.name, quoted=bool(getattr(expr.this, "quoted", False))) + ) + else: + key_exprs.append(expr) super().insert_overwrite_by_partition( table_name, query_or_df, @@ -945,6 +983,13 @@ def _simplify_boolean_where(cls, expr: exp.Expr) -> exp.Expr: ) if isinstance(expr, exp.Paren) and isinstance(expr.this, exp.Boolean): return expr.this + if ( + isinstance(expr, exp.EQ) + and isinstance(expr.this, exp.Literal) + and isinstance(expr.expression, exp.Literal) + and expr.this.is_number == expr.expression.is_number + ): + return exp.true() if expr.this == expr.expression else exp.false() if isinstance(expr, exp.Not) and isinstance(expr.this, exp.Boolean): return exp.false() if expr.this == exp.true() else exp.true() if isinstance(expr, exp.And): @@ -1083,14 +1128,26 @@ def _execute_delete_with_subquery( return join_conditions: t.List[exp.Expr] = [] + first_select = subquery.this + while isinstance(first_select, exp.SetOperation): + first_select = first_select.this for idx, col in enumerate(key_exprs): col_name = col.name t1_col = exp.column(col_name, table="_t1") t2_col_name = col_name - if isinstance(subquery.this, exp.Select) and idx < len(subquery.this.expressions): - sel_item = subquery.this.expressions[idx] + if isinstance(first_select, exp.Select) and idx < len(first_select.expressions): + sel_item = first_select.expressions[idx] if sel_item.alias_or_name: t2_col_name = sel_item.alias_or_name + elif isinstance(subquery.this, exp.Select): + alias_name = f"_col_{idx}" + subquery = subquery.copy() + sub_select = t.cast(exp.Select, subquery.this) + new_exprs = list(sub_select.expressions) + new_exprs[idx] = exp.alias_(sel_item.copy(), alias_name, quoted=True) + sub_select.set("expressions", new_exprs) + first_select = sub_select + t2_col_name = alias_name t2_col = exp.column(t2_col_name, table="_t2") join_conditions.append(t1_col.eq(t2_col)) @@ -1161,7 +1218,7 @@ def _delete_by_literal_keys( ) if isinstance(select_query, exp.Select): select_query = select_query.distinct() - rows = self.fetchall(select_query) + rows = self.fetchall(select_query, quote_identifiers=True) if not rows: return @@ -1316,8 +1373,11 @@ def _build_create_table_exp( ) if not is_iceberg and target_columns_to_types: + original_cols = dict(target_columns_to_types) reordered_cols = self._reorder_olap_columns(target_columns_to_types, kwargs) - if list(reordered_cols) != list(target_columns_to_types): + order_changed = list(reordered_cols) != list(original_cols) + types_changed = any(reordered_cols[c] != original_cols[c] for c in reordered_cols) + if order_changed or types_changed: target_columns_to_types = reordered_cols if isinstance(table_name_or_schema, exp.Schema): col_def_map = { @@ -1333,12 +1393,19 @@ def _build_create_table_exp( if len(col_def_map) == len(target_columns_to_types) and all( c.lower() in col_def_map for c in target_columns_to_types ): + new_col_defs: t.List[exp.Expr] = [] + for c, new_dtype in target_columns_to_types.items(): + col_def = col_def_map[c.lower()] + if col_def.args.get("kind") != new_dtype: + col_def = col_def.copy() + col_def.set("kind", new_dtype) + new_col_defs.append(col_def) table_name_or_schema = table_name_or_schema.copy() table_name_or_schema.set( "expressions", - [col_def_map[c.lower()] for c in target_columns_to_types] + other_exprs, + new_col_defs + other_exprs, ) - if expression is not None and isinstance(expression, exp.Query): + if order_changed and expression is not None and isinstance(expression, exp.Query): reordered_in_place = False if isinstance(expression, exp.Select): proj_by_name = { @@ -1359,27 +1426,11 @@ def _build_create_table_exp( expression = self._order_projections_and_filter( expression, target_columns_to_types ) - elif not is_iceberg and isinstance(expression, exp.Select): - proj_by_name = { - proj.alias_or_name.lower(): proj - for proj in expression.expressions - if proj.alias_or_name and not isinstance(proj, exp.Star) - } - if proj_by_name and len(proj_by_name) == len(expression.expressions): - select_cols_to_types = { - proj.alias_or_name: exp.DataType.build("UNKNOWN") - for proj in expression.expressions - } - reordered_select_cols = self._reorder_olap_columns(select_cols_to_types, kwargs) - if list(reordered_select_cols) != list(select_cols_to_types): - expression = expression.copy() - expression.set( - "expressions", - [proj_by_name[c.lower()] for c in reordered_select_cols], - ) + elif not is_iceberg and isinstance(expression, (exp.Select, exp.SetOperation)): + expression = self._reorder_unannotated_ctas_query(expression, kwargs) # Doris AUTO PARTITION BY RANGE requires the partition column(s) to be defined as NOT NULL. - if not is_iceberg and isinstance(table_name_or_schema, exp.Schema): + if not is_iceberg: table_props = kwargs.get("table_properties") or {} auto_partition = table_props.get("auto_partition") is_auto = bool( @@ -1397,7 +1448,15 @@ def _build_create_table_exp( for expr in parsed_part_exprs for col in expr.find_all(exp.Column) } - if auto_range_cols: + strict_key = ( + table_props.get("unique_key") + or kwargs.get("unique_key") + or table_props.get("aggregate_key") + ) + strict_key_set = {k.lower() for k in self._extract_column_names(strict_key)} + if strict_key_set and not auto_range_cols.issubset(strict_key_set): + auto_range_cols = set() + if auto_range_cols and isinstance(table_name_or_schema, exp.Schema): table_name_or_schema = table_name_or_schema.copy() updated_exprs: t.List[exp.Expr] = [] for expr in table_name_or_schema.expressions: @@ -1422,6 +1481,34 @@ def _build_create_table_exp( ) updated_exprs.append(expr) table_name_or_schema.set("expressions", updated_exprs) + elif auto_range_cols and isinstance(expression, exp.Select): + expression = expression.copy() + new_select_exprs: t.List[exp.Expr] = [] + for proj in expression.expressions: + if proj.alias_or_name.lower() in auto_range_cols: + if isinstance(proj, exp.Alias): + inner = proj.this + if not ( + isinstance(inner, exp.Anonymous) + and inner.name.upper() == "NON_NULLABLE" + ): + proj = proj.copy() + proj.set( + "this", + exp.Anonymous( + this="NON_NULLABLE", expressions=[inner.copy()] + ), + ) + elif isinstance(proj, exp.Column): + proj = exp.alias_( + exp.Anonymous( + this="NON_NULLABLE", expressions=[proj.copy()] + ), + proj.name, + quoted=bool(getattr(proj.this, "quoted", True)), + ) + new_select_exprs.append(proj) + expression.set("expressions", new_select_exprs) return super()._build_create_table_exp( table_name_or_schema, @@ -1434,6 +1521,45 @@ def _build_create_table_exp( **kwargs, ) + def _reorder_unannotated_ctas_query( + self, expression: exp.Expr, kwargs: t.Dict[str, t.Any] + ) -> exp.Expr: + if isinstance(expression, exp.SetOperation): + expression = expression.copy() + expression.set("this", self._reorder_unannotated_ctas_query(expression.this, kwargs)) + expression.set( + "expression", self._reorder_unannotated_ctas_query(expression.expression, kwargs) + ) + return expression + if isinstance(expression, exp.Select): + proj_by_name = { + proj.alias_or_name.lower(): proj + for proj in expression.expressions + if proj.alias_or_name and not isinstance(proj, exp.Star) + } + if proj_by_name and len(proj_by_name) == len(expression.expressions): + select_cols_to_types = { + proj.alias_or_name: exp.DataType.build("UNKNOWN") + for proj in expression.expressions + } + reordered_select_cols = self._reorder_olap_columns(select_cols_to_types, kwargs) + if list(reordered_select_cols) != list(select_cols_to_types): + expression = expression.copy() + expression.set( + "expressions", + [proj_by_name[c.lower()] for c in reordered_select_cols], + ) + return expression + + @classmethod + def _is_coercible_string_key_type(cls, dtype: exp.DataType) -> bool: + type_name = ( + dtype.this.value.upper() + if isinstance(dtype.this, exp.DataType.Type) + else str(dtype.this).upper() + ) + return type_name in ("TEXT", "STRING", "TINYTEXT", "MEDIUMTEXT", "LONGTEXT") + def _reorder_olap_columns( self, target_columns_to_types: t.Dict[str, exp.DataType], @@ -1473,20 +1599,44 @@ def _reorder_olap_columns( f"Doris key columns {missing_cols} not found in table schema columns {col_names}." ) key_cols = [col_by_lower[k.lower()] for k in raw_key_cols if k.lower() in col_by_lower] - if key_cols and col_names[: len(key_cols)] != key_cols: - key_cols_set = set(key_cols) - new_order = key_cols + [c for c in col_names if c not in key_cols_set] - return {c: target_columns_to_types[c] for c in new_order} - return target_columns_to_types + updated_types = dict(target_columns_to_types) + for k in key_cols: + k_dtype = updated_types[k] + if not self._is_valid_olap_key_type(k_dtype): + if self._is_coercible_string_key_type(k_dtype): + updated_types[k] = exp.DataType.build("VARCHAR(255)", dialect=self.dialect) + else: + raise SQLMeshError( + f"Doris key column '{k}' has type '{k_dtype.sql(dialect=self.dialect)}' which is not supported as an OLAP key column." + ) + key_cols_set = set(key_cols) + new_order = ( + key_cols + [c for c in col_names if c not in key_cols_set] + if key_cols and col_names[: len(key_cols)] != key_cols + else col_names + ) + return {c: updated_types[c] for c in new_order} - first_dtype = target_columns_to_types[col_names[0]] + updated_types = dict(target_columns_to_types) + first_dtype = updated_types[col_names[0]] if not self._is_valid_olap_key_type(first_dtype): valid_col = next( - (c for c in col_names if self._is_valid_olap_key_type(target_columns_to_types[c])), + (c for c in col_names if self._is_valid_olap_key_type(updated_types[c])), None, ) if valid_col: col_names = [valid_col] + [c for c in col_names if c != valid_col] + else: + coercible_col = next( + (c for c in col_names if self._is_coercible_string_key_type(updated_types[c])), + None, + ) + if coercible_col: + updated_types[coercible_col] = exp.DataType.build( + "VARCHAR(255)", dialect=self.dialect + ) + if col_names[0] != coercible_col: + col_names = [coercible_col] + [c for c in col_names if c != coercible_col] # Ensure any OLAP partition columns with valid key types precede non-keyable columns raw_part_by = kwargs.get("partitioned_by") or table_props.get("partitioned_by") @@ -1499,15 +1649,13 @@ def _reorder_olap_columns( for expr in parsed_part_exprs for col in expr.find_all(exp.Column) if col.name.lower() in col_by_lower - and self._is_valid_olap_key_type( - target_columns_to_types[col_by_lower[col.name.lower()]] - ) + and self._is_valid_olap_key_type(updated_types[col_by_lower[col.name.lower()]]) ] first_non_key_idx = next( ( i for i, c in enumerate(col_names) - if not self._is_valid_olap_key_type(target_columns_to_types[c]) + if not self._is_valid_olap_key_type(updated_types[c]) ), None, ) @@ -1522,8 +1670,8 @@ def _reorder_olap_columns( after_non_key = [c for c in col_names[first_non_key_idx:] if c not in moved_set] col_names = before_non_key + moved_parts + after_non_key - if col_names != list(target_columns_to_types): - return {c: target_columns_to_types[c] for c in col_names} + if col_names != list(target_columns_to_types) or updated_types != target_columns_to_types: + return {c: updated_types[c] for c in col_names} return target_columns_to_types @@ -1627,8 +1775,26 @@ def _build_partitioned_by_exp( partitioned_by = normalized_exprs if is_iceberg: + iceberg_part_exprs = [ + expr.copy().transform( + lambda node: ( + exp.Anonymous( + this="BUCKET", + expressions=[node.expression.copy(), node.this.copy()], + ) + if isinstance(node, exp.PartitionedByBucket) + else exp.Anonymous( + this="TRUNCATE", + expressions=[node.expression.copy(), node.this.copy()], + ) + if isinstance(node, exp.PartitionByTruncate) + else node + ) + ) + for expr in partitioned_by + ] return exp.PartitionByListProperty( - partition_expressions=partitioned_by, + partition_expressions=iceberg_part_exprs, create_expressions=[], ) diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 78f84a25fc..a60298422d 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -1774,3 +1774,234 @@ def test_resolve_table_refs_in_physical_properties_for_mv_and_sync_mv( assert sync_calls[0] == ( f"DROP MATERIALIZED VIEW IF EXISTS `doris_db`.`orders_sync_mv` ON `{physical_tbl.db}`.`{physical_tbl.name}`" ) + + +def test_dataframe_seed_all_string_coercion_and_batch_size( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + import pandas as pd + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + assert adapter.DEFAULT_BATCH_SIZE == 500 + + # 1. All-TEXT schema (e.g. from columns_to_types_from_df on a string CSV seed) coerces first column to VARCHAR(255) + adapter.create_table( + "db.country_seed", + target_columns_to_types={ + "country_code": exp.DataType.build("text"), + "country_name": exp.DataType.build("text"), + }, + ) + + # 2. Schema with DOUBLE + TEXT moves TEXT column to position 0 and coerces it to VARCHAR(255) + adapter.create_table( + "db.metric_seed", + target_columns_to_types={ + "score": exp.DataType.build("double"), + "metric_name": exp.DataType.build("text"), + }, + ) + + # 3. Explicit unique_key on a TEXT column coerces the key column to VARCHAR(255) + adapter.create_table( + "db.explicit_text_uk", + target_columns_to_types={ + "payload": exp.DataType.build("text"), + "code": exp.DataType.build("text"), + }, + table_properties={"unique_key": exp.to_column("code")}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`country_seed` (`country_code` VARCHAR(255), `country_name` STRING)", + "CREATE TABLE IF NOT EXISTS `db`.`metric_seed` (`metric_name` VARCHAR(255), `score` DOUBLE)", + "CREATE TABLE IF NOT EXISTS `db`.`explicit_text_uk` (`code` VARCHAR(255), `payload` STRING) UNIQUE KEY (`code`) DISTRIBUTED BY HASH (`code`) BUCKETS 10", + ] + + # 4. Explicit unique_key on a non-coercible invalid key type (DOUBLE) raises SQLMeshError + with pytest.raises(SQLMeshError, match="not supported as an OLAP key column"): + adapter.create_table( + "db.bad_double_uk", + target_columns_to_types={ + "score": exp.DataType.build("DOUBLE", dialect="doris"), + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + table_properties={"unique_key": exp.to_column("score")}, + ) + + # 5. DataFrame CTAS with pretty_sql=True and 500 rows succeeds without RecursionError + pretty_adapter = make_mocked_engine_adapter(DorisEngineAdapter, pretty_sql=True) + df = pd.DataFrame({"code": [f"c_{i}" for i in range(500)], "val": list(range(500))}) + pretty_adapter.ctas("db.df_seed", df) + df_sql_calls = to_sql_calls(pretty_adapter) + assert len(df_sql_calls) == 1 + assert "CREATE TABLE IF NOT EXISTS `db`.`df_seed`" in df_sql_calls[0] + + +def test_iceberg_bucket_and_truncate_partition_transforms( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + from sqlmesh.core.dialect import parse + from sqlmesh.core.model import SqlModel, load_sql_based_model + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + iceberg_model = t.cast( + SqlModel, + load_sql_based_model( + parse( + """ + MODEL ( + name iceberg_catalog.silver.curated_events, + kind FULL, + dialect doris, + table_format iceberg, + partitioned_by (day(event_ts), bucket(16, user_id), truncate(10, region)) + ); + SELECT + user_id::BIGINT AS user_id, + event_ts::DATETIME AS event_ts, + region::STRING AS region + FROM iceberg_catalog.bronze.raw_events; + """, + default_dialect="doris", + ) + ), + ) + + adapter.create_table( + iceberg_model.name, + target_columns_to_types=iceberg_model.columns_to_types or {}, + table_format=iceberg_model.table_format, + partitioned_by=iceberg_model.partitioned_by, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `iceberg_catalog`.`silver`.`curated_events` (`user_id` BIGINT, `event_ts` DATETIME, `region` STRING) ENGINE=iceberg PARTITION BY LIST (DAY(`event_ts`), BUCKET(16, `user_id`), TRUNCATE(10, `region`)) ()" + ] + + +def test_ctas_auto_partition_non_nullable_and_union_key_reordering( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. CTAS with AUTO PARTITION BY RANGE wraps partition column projection in NON_NULLABLE(...) + adapter.ctas( + "db.ctas_auto_part", + parse_one("SELECT payload, id, ds FROM db.src", read="doris"), + target_columns_to_types={ + "payload": exp.DataType.build("STRING", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + }, + partitioned_by=[parse_one("RANGE(DATE_TRUNC(ds, 'DAY'))", read="doris")], + table_properties={"auto_partition": exp.Literal.string("true")}, + ) + + # 2. Unannotated CTAS with UNION ALL reorders projections in all branches for unique_key prefix + adapter.ctas( + "db.ctas_union_reordered", + parse_one( + "SELECT name, id FROM db.src1 UNION ALL SELECT name, id FROM db.src2", read="doris" + ), + table_properties={"unique_key": exp.to_column("id")}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`ctas_auto_part` AUTO PARTITION BY RANGE (DATE_TRUNC(`ds`, 'DAY')) () AS SELECT CAST(`id` AS INT) AS `id`, NON_NULLABLE(CAST(`ds` AS DATE)) AS `ds`, CAST(`payload` AS STRING) AS `payload` FROM (SELECT `payload`, `id`, `ds` FROM `db`.`src`) AS `_subquery`", + "CREATE TABLE IF NOT EXISTS `db`.`ctas_union_reordered` UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10 AS SELECT `id`, `name` FROM `db`.`src1` UNION ALL SELECT `id`, `name` FROM `db`.`src2`", + ] + + +def test_delete_tautology_and_union_subquery_and_dedup_partition_overwrite( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Literal tautology 1=1 routes to TRUNCATE TABLE; contradiction 1=0 is a no-op + adapter.delete_from("db.tbl", where="1 = 1") + adapter.delete_from("db.tbl", where="1 = 0") + + # 2. IN subquery with UNION resolves the projected column name from the first SELECT branch + adapter.delete_from( + "db.tbl", + where="id IN (SELECT src_id FROM db.a UNION SELECT src_id FROM db.b)", + ) + + # 3. IN subquery with unaliased expression aliases the projection for _t2 join + adapter.delete_from( + "db.tbl", + where="id IN (SELECT other_id + 1 FROM db.a)", + ) + + assert to_sql_calls(adapter) == [ + "TRUNCATE TABLE `db`.`tbl`", + "DELETE FROM `db`.`tbl` AS `_t1` USING (SELECT `src_id` FROM `db`.`a` UNION SELECT `src_id` FROM `db`.`b`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`src_id`", + "DELETE FROM `db`.`tbl` AS `_t1` USING (SELECT `other_id` + 1 AS `_col_0` FROM `db`.`a`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`_col_0`", + ] + + # 4. insert_overwrite_by_partition deduplicates columns and strips table qualifiers + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter, + "_get_temp_table", + return_value=exp.table_("__temp_tbl", db="db"), + ) + adapter.insert_overwrite_by_partition( + "db.events", + parse_one("SELECT id, ds FROM db.src", read="doris"), + partitioned_by=[ + parse_one("YEAR(src.ds)", read="doris"), + parse_one("MONTH(src.ds)", read="doris"), + ], + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATE"), + }, + ) + overwrite_calls = to_sql_calls(adapter) + assert any( + "DELETE FROM `db`.`events` AS `_t1` USING (SELECT DISTINCT `ds` FROM `db`.`__temp_tbl`) AS `_t2` WHERE `_t1`.`ds` = `_t2`.`ds`" + == call + for call in overwrite_calls + ) + + +def test_sync_mv_inferred_source_table_on_replace_and_unqualified_get_data_objects( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter, patch_get_data_objects=False) + + # 1. Sync MV with replace=True and no explicit source_table in view_properties infers base table from query + mocker.patch.object(adapter, "get_data_object", return_value=None) + adapter.create_view( + "sync_rollup_mv", + parse_one("SELECT id, SUM(amount) AS total FROM prod_db.orders GROUP BY id", read="doris"), + replace=True, + materialized=True, + view_properties={"materialized_type": exp.var("SYNC")}, + ) + assert to_sql_calls(adapter) == [ + "USE `prod_db`", + "DROP MATERIALIZED VIEW IF EXISTS `sync_rollup_mv` ON `prod_db`.`orders`", + "USE `prod_db`", + "CREATE MATERIALIZED VIEW `sync_rollup_mv` AS SELECT `id` AS `mv_id`, SUM(`amount`) AS `total` FROM `prod_db`.`orders` GROUP BY `id`", + ] + + # 2. _get_data_objects with empty schema uses DATABASE() instead of empty string literal + adapter.cursor.execute.reset_mock() + fetchall_spy = mocker.patch.object( + adapter, + "fetchall", + side_effect=[ + [("active_db", "tbl1", "BASE TABLE")], + [], + ], + ) + objs = adapter._get_data_objects("") + assert len(objs) == 1 + assert objs[0].schema_name == "active_db" + first_query_sql = fetchall_spy.call_args_list[0].args[0].sql("doris") + assert "table_schema = DATABASE()" in first_query_sql + second_query_sql = fetchall_spy.call_args_list[1].args[0] + assert "mv_infos('database' = DATABASE())" in second_query_sql From 33c908ebab4345b74104df087d59d62cfc7c7503 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Thu, 22 Oct 2026 10:52:05 -0700 Subject: [PATCH 16/27] fix(doris): resolve CTAS key coercion, self-ref replace, Iceberg CTAS, and DELETE edge cases Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 334 ++++++++++++++++++++---- tests/core/engine_adapter/test_doris.py | 151 +++++++++++ 2 files changed, 441 insertions(+), 44 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 5540f62459..330003ca67 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -8,6 +8,7 @@ import pandas as pd # noqa: TID253 from sqlglot import exp, parse, parse_one +from sqlglot.optimizer.qualify_columns import quote_identifiers from sqlmesh.core.dialect import to_schema from sqlmesh.core.engine_adapter.base import ( @@ -30,7 +31,7 @@ SourceQuery, set_catalog, ) -from sqlmesh.utils import get_source_columns_to_types, random_id +from sqlmesh.utils import columns_to_types_all_known, get_source_columns_to_types, random_id from sqlmesh.utils.errors import SQLMeshError if t.TYPE_CHECKING: @@ -271,7 +272,34 @@ def _fetch_native_df( logger.info("Connection to the database was lost. Reconnecting...") self._connection_pool.close() self._connection_pool.get_cursor() - return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) + try: + return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) + except Exception as e: + err_code = e.args[0] if getattr(e, "args", None) else None + err_name = e.__class__.__name__ + is_disconnect = (err_code in (2006, 2013, 2055) and err_name == "OperationalError") or ( + err_code == 0 and err_name == "InterfaceError" + ) + is_timeout = ( + isinstance(getattr(e, "__context__", None), TimeoutError) + or "timed out" in str(e).lower() + ) + if is_disconnect and not is_timeout and not self._connection_pool.is_transaction_active: + logger.warning( + "Doris connection lost (%s: %s) during fetchdf, reconnecting and retrying once...", + err_name, + err_code, + ) + self._connection_pool.close() + try: + self._connection_pool.get_cursor() + return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) + except Exception: + self._connection_pool.close() + raise + if is_disconnect and not self._connection_pool.is_transaction_active: + self._connection_pool.close() + raise _SAFE_RETRY_SQL_PREFIX_RE = re.compile( r"\A(?:\s|/\*.*?\*/|--[^\r\n]*(?:[\r\n]|\Z))*(?:SELECT|WITH|SHOW|DESCRIBE|DESC|EXPLAIN|USE|SWITCH)\b", @@ -955,16 +983,22 @@ def delete_from( subquery_info = self._extract_subquery_from_where(where) if subquery_info is not None: self._execute_delete_with_subquery(table_name, *subquery_info) - else: - target_table = exp.to_table(table_name) - for col in where.find_all(exp.Column): - if not col.find_ancestor(exp.Subquery, exp.Select) and ( - col.table and col.table.lower() == target_table.name.lower() - ): - col.set("table", None) - col.set("db", None) - col.set("catalog", None) - super().delete_from(table_name, where) + return + + tuple_in_info = self._extract_literal_tuple_in_from_where(where) + if tuple_in_info is not None: + self._delete_key_rows(table_name, *tuple_in_info, is_ast_values=True) + return + + target_table = exp.to_table(table_name) + for col in where.find_all(exp.Column): + if not col.find_ancestor(exp.Subquery, exp.Select) and ( + col.table and col.table.lower() == target_table.name.lower() + ): + col.set("table", None) + col.set("db", None) + col.set("catalog", None) + super().delete_from(table_name, where) @classmethod def _simplify_boolean_where(cls, expr: exp.Expr) -> exp.Expr: @@ -1060,15 +1094,17 @@ def _extract_subquery_from_where( lhs_concat = self._unwrap_concat_ws(lhs) if lhs_concat is not None: + if not (isinstance(select_expr, exp.Select) and len(select_expr.expressions) == 1): + return None + first_sel = select_expr.expressions[0] + raw_sel = first_sel.this if isinstance(first_sel, exp.Alias) else first_sel + sel_concat = self._unwrap_concat_ws(raw_sel) + if sel_concat is None or len(sel_concat.expressions) != len(lhs_concat.expressions): + return None key_exprs = [e.copy() for e in lhs_concat.expressions[1:]] - if isinstance(select_expr, exp.Select) and len(select_expr.expressions) == 1: - first_sel = select_expr.expressions[0] - raw_sel = first_sel.this if isinstance(first_sel, exp.Alias) else first_sel - sel_concat = self._unwrap_concat_ws(raw_sel) - if sel_concat is not None: - select_expr = select_expr.copy() - select_expr.set("expressions", [e.copy() for e in sel_concat.expressions[1:]]) - subquery = exp.Subquery(this=select_expr) + select_expr = select_expr.copy() + select_expr.set("expressions", [e.copy() for e in sel_concat.expressions[1:]]) + subquery = exp.Subquery(this=select_expr) elif isinstance(lhs, exp.Tuple): key_exprs = [e.copy() for e in lhs.expressions] else: @@ -1088,6 +1124,58 @@ def _extract_subquery_from_where( return normalized_keys, subquery, extra_conditions + def _extract_literal_tuple_in_from_where( + self, where: exp.Expr + ) -> t.Optional[t.Tuple[t.List[exp.Expr], t.List[t.Tuple[exp.Expr, ...]], t.List[exp.Expr]]]: + """Extract positive literal tuple IN predicate `(a, b) IN ((1, 2), (3, 4))` and sibling AND conditions.""" + where = where.unnest() + conjuncts: t.List[exp.Expr] = ( + list(where.flatten()) if isinstance(where, exp.And) else [where] + ) + target_in: t.Optional[exp.In] = None + extra_conditions: t.List[exp.Expr] = [] + + for cond in conjuncts: + unwrapped_cond = cond.unnest() + if ( + target_in is None + and isinstance(unwrapped_cond, exp.In) + and self._get_in_subquery(unwrapped_cond) is None + and isinstance(unwrapped_cond.this, exp.Tuple) + and len(unwrapped_cond.this.expressions) >= 2 + and unwrapped_cond.expressions + and all( + isinstance(item, exp.Tuple) + and len(item.expressions) == len(unwrapped_cond.this.expressions) + for item in unwrapped_cond.expressions + ) + ): + target_in = unwrapped_cond + continue + extra_conditions.append(cond) + + if target_in is None: + return None + + normalized_keys: t.List[exp.Expr] = [] + for k in target_in.this.expressions: + if isinstance(k, exp.Column): + normalized_keys.append(k.copy()) + elif isinstance(k, exp.Identifier): + normalized_keys.append(exp.Column(this=k.copy())) + else: + return None + + tuple_rows: t.List[t.Tuple[exp.Expr, ...]] = [ + tuple(e.copy() for e in item.expressions) + for item in target_in.expressions + if isinstance(item, exp.Tuple) + ] + if not tuple_rows: + return None + + return normalized_keys, tuple_rows, extra_conditions + @staticmethod def _unwrap_concat_ws(expr: exp.Expr) -> t.Optional[exp.Expr]: curr = expr.unnest() @@ -1131,23 +1219,27 @@ def _execute_delete_with_subquery( first_select = subquery.this while isinstance(first_select, exp.SetOperation): first_select = first_select.this + used_t2_aliases: t.Set[str] = set() for idx, col in enumerate(key_exprs): col_name = col.name t1_col = exp.column(col_name, table="_t1") t2_col_name = col_name if isinstance(first_select, exp.Select) and idx < len(first_select.expressions): sel_item = first_select.expressions[idx] - if sel_item.alias_or_name: - t2_col_name = sel_item.alias_or_name + candidate_name = sel_item.alias_or_name + if candidate_name and candidate_name.lower() not in used_t2_aliases: + t2_col_name = candidate_name elif isinstance(subquery.this, exp.Select): alias_name = f"_col_{idx}" subquery = subquery.copy() sub_select = t.cast(exp.Select, subquery.this) new_exprs = list(sub_select.expressions) - new_exprs[idx] = exp.alias_(sel_item.copy(), alias_name, quoted=True) + inner_expr = sel_item.this if isinstance(sel_item, exp.Alias) else sel_item + new_exprs[idx] = exp.alias_(inner_expr.copy(), alias_name, quoted=True) sub_select.set("expressions", new_exprs) first_select = sub_select t2_col_name = alias_name + used_t2_aliases.add(t2_col_name.lower()) t2_col = exp.column(t2_col_name, table="_t2") join_conditions.append(t1_col.eq(t2_col)) @@ -1219,6 +1311,26 @@ def _delete_by_literal_keys( if isinstance(select_query, exp.Select): select_query = select_query.distinct() rows = self.fetchall(select_query, quote_identifiers=True) + if not rows: + return + sanitized_rows = [tuple(self._sanitize_key_value(v) for v in row) for row in rows] + self._delete_key_rows( + table_name, + key_exprs, + sanitized_rows, + extra_conditions=extra_conditions, + is_ast_values=False, + ) + + def _delete_key_rows( + self, + table_name: TableName, + key_exprs: t.List[exp.Expr], + rows: t.Sequence[t.Tuple[t.Any, ...]], + extra_conditions: t.Optional[t.List[exp.Expr]] = None, + *, + is_ast_values: bool = False, + ) -> None: if not rows: return @@ -1238,12 +1350,22 @@ def _delete_by_literal_keys( cols = [exp.column(k.name, quoted=True) for k in key_exprs] batch_size = self.DEFAULT_BATCH_SIZE - sanitized_rows = [tuple(self._sanitize_key_value(v) for v in row) for row in rows] + def _is_null_val(val: t.Any) -> bool: + if val is None: + return True + if is_ast_values and isinstance(val, exp.Null): + return True + return False + + def _to_expr(val: t.Any) -> exp.Expr: + if is_ast_values and isinstance(val, exp.Expr): + return val.copy() + return exp.convert(val) if len(cols) == 1: col = cols[0] - non_null_vals = [exp.convert(r[0]) for r in sanitized_rows if r[0] is not None] - has_null = any(r[0] is None for r in sanitized_rows) + non_null_vals = [_to_expr(r[0]) for r in rows if not _is_null_val(r[0])] + has_null = any(_is_null_val(r[0]) for r in rows) for i in range(0, len(non_null_vals), batch_size): batch = non_null_vals[i : i + batch_size] @@ -1259,10 +1381,10 @@ def _delete_by_literal_keys( # that minimizes the number of distinct prefix groups, breaking ties in favor of the last column. batch_idx = min( range(len(cols)), - key=lambda idx: (len({row[:idx] + row[idx + 1 :] for row in sanitized_rows}), -idx), + key=lambda idx: (len({row[:idx] + row[idx + 1 :] for row in rows}), -idx), ) grouped: t.Dict[t.Tuple[t.Any, ...], t.List[t.Any]] = defaultdict(list) - for row in sanitized_rows: + for row in rows: prefix_key = row[:batch_idx] + row[batch_idx + 1 :] grouped[prefix_key].append(row[batch_idx]) @@ -1271,11 +1393,11 @@ def _delete_by_literal_keys( for prefix_vals, batch_vals in grouped.items(): prefix_preds: t.List[exp.Expr] = [ - p_col.eq(exp.convert(p_val)) if p_val is not None else p_col.is_(exp.null()) + p_col.eq(_to_expr(p_val)) if not _is_null_val(p_val) else p_col.is_(exp.null()) for p_col, p_val in zip(prefix_cols, prefix_vals) ] - non_null_batch = [exp.convert(v) for v in batch_vals if v is not None] - has_null_batch = any(v is None for v in batch_vals) + non_null_batch = [_to_expr(v) for v in batch_vals if not _is_null_val(v)] + has_null_batch = any(_is_null_val(v) for v in batch_vals) for i in range(0, len(non_null_batch), batch_size): batch = non_null_batch[i : i + batch_size] @@ -1345,6 +1467,92 @@ def _create_table_from_columns( **kwargs, ) + def _create_table_from_source_queries( + self, + table_name: TableName, + source_queries: t.List[SourceQuery], + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + exists: bool = True, + replace: bool = False, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + table_kind: t.Optional[str] = None, + track_rows_processed: bool = True, + **kwargs: t.Any, + ) -> None: + table = exp.to_table(table_name) + effective_catalog = table.catalog or self.default_catalog + is_iceberg = self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=table.catalog, + table_properties=kwargs.get("table_properties"), + ) + is_internal_iceberg = is_iceberg and ( + not effective_catalog or effective_catalog.lower() == "internal" + ) + + if not is_iceberg and target_columns_to_types: + target_columns_to_types = self._reorder_olap_columns(target_columns_to_types, kwargs) + + target_columns_to_types_known = bool( + target_columns_to_types and columns_to_types_all_known(target_columns_to_types) + ) + first_query = source_queries[0].query_factory() if source_queries else None + self_referencing = bool( + first_query + and any( + quote_identifiers(t_node.copy()) == quote_identifiers(table.copy()) + for t_node in first_query.find_all(exp.Table) + ) + ) + + # 1. Doris's CTAS grammar forbids `ENGINE=iceberg`, so internal/2-part Iceberg tables with known + # column types must be created via schema DDL (`_create_table_from_columns`) + `INSERT INTO`. + # 2. When `replace_query` is called with a self-referencing query on a non-existent table, + # `EngineAdapter.replace_query` creates the empty table via `_create_table_from_columns` + # first and then calls `_create_table_from_source_queries(..., replace=False)`. Executing + # `CREATE TABLE IF NOT EXISTS ... AS SELECT` would be a no-op, so we insert the source queries. + if ( + is_internal_iceberg or (self_referencing and not replace) + ) and target_columns_to_types_known: + assert target_columns_to_types is not None + if not self_referencing: + if replace: + self.drop_table(table_name, exists=True) + self._create_table_from_columns( + table_name, + target_columns_to_types, + exists=exists, + table_description=table_description, + column_descriptions=column_descriptions, + **kwargs, + ) + for source_query in source_queries: + with source_query as query: + query = self._order_projections_and_filter( + query, target_columns_to_types, coerce_types=True + ) + self._insert_append_query( + table_name, + query, + target_columns_to_types, + track_rows_processed=track_rows_processed, + ) + return + + super()._create_table_from_source_queries( + table_name, + source_queries, + target_columns_to_types=target_columns_to_types, + exists=exists, + replace=replace, + table_description=table_description, + column_descriptions=column_descriptions, + table_kind=table_kind, + track_rows_processed=track_rows_processed, + **kwargs, + ) + def _build_create_table_exp( self, table_name_or_schema: t.Union[exp.Schema, TableName], @@ -1405,7 +1613,11 @@ def _build_create_table_exp( "expressions", new_col_defs + other_exprs, ) - if order_changed and expression is not None and isinstance(expression, exp.Query): + if ( + (order_changed or types_changed) + and expression is not None + and isinstance(expression, exp.Query) + ): reordered_in_place = False if isinstance(expression, exp.Select): proj_by_name = { @@ -1417,14 +1629,31 @@ def _build_create_table_exp( target_columns_to_types ) and all(c.lower() in proj_by_name for c in target_columns_to_types): expression = expression.copy() - expression.set( - "expressions", - [proj_by_name[c.lower()] for c in target_columns_to_types], - ) + updated_projs: t.List[exp.Expr] = [] + for c, new_dtype in target_columns_to_types.items(): + proj = proj_by_name[c.lower()].copy() + if types_changed and original_cols[c] != new_dtype: + if isinstance(proj, exp.Alias) and isinstance( + proj.this, exp.Cast + ): + proj.this.set("to", new_dtype.copy()) + elif isinstance(proj, exp.Alias): + proj.set( + "this", + exp.cast(proj.this.copy(), to=new_dtype.copy()), + ) + elif isinstance(proj, exp.Column): + proj = exp.alias_( + exp.cast(proj.copy(), to=new_dtype.copy()), + proj.name, + quoted=bool(getattr(proj.this, "quoted", True)), + ) + updated_projs.append(proj) + expression.set("expressions", updated_projs) reordered_in_place = True if not reordered_in_place: expression = self._order_projections_and_filter( - expression, target_columns_to_types + expression, target_columns_to_types, coerce_types=types_changed ) elif not is_iceberg and isinstance(expression, (exp.Select, exp.SetOperation)): expression = self._reorder_unannotated_ctas_query(expression, kwargs) @@ -2292,20 +2521,37 @@ def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> expression = expression.copy() for dtype in reversed(list(expression.find_all(exp.DataType))): if dtype.is_type("int128"): - dtype.replace(exp.DataType.build("LARGEINT", udt=True)) + replacement = exp.DataType.build("LARGEINT", udt=True) + if dtype is expression: + expression = replacement + else: + dtype.replace(replacement) elif dtype.is_type("struct"): - new_fields: t.List[exp.Expr] = [] + new_field_sqls: t.List[str] = [] for field in dtype.expressions: if isinstance(field, exp.ColumnDef) and field.args.get("kind"): field_id = field.this.copy() if quote and isinstance(field_id, exp.Identifier): field_id.set("quoted", True) - field_name_sql = super()._to_sql(field_id, quote=False, **kwargs) - field_kind_sql = self._to_sql(field.args["kind"], quote=quote, **kwargs) - new_fields.append(exp.var(f"{field_name_sql}: {field_kind_sql}")) + field_name_sql = super()._to_sql( + field_id, quote=False, pretty=False, **kwargs + ) + field_kind_sql = super()._to_sql( + field.args["kind"], quote=quote, pretty=False, **kwargs + ) + new_field_sqls.append(f"{field_name_sql}: {field_kind_sql}") else: - new_fields.append(field) - dtype.set("expressions", new_fields) + new_field_sqls.append( + super()._to_sql(field, quote=quote, pretty=False, **kwargs) + ) + struct_replacement = exp.DataType( + this=exp.DataType.Type.USERDEFINED, + kind=f"STRUCT<{', '.join(new_field_sqls)}>", + ) + if dtype is expression: + expression = struct_replacement + else: + dtype.replace(struct_replacement) return super()._to_sql(expression, quote=quote, **kwargs) def _build_distributed_by_exp( diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index a60298422d..8053218e78 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -2005,3 +2005,154 @@ def test_sync_mv_inferred_source_table_on_replace_and_unqualified_get_data_objec assert "table_schema = DATABASE()" in first_query_sql second_query_sql = fetchall_spy.call_args_list[1].args[0] assert "mv_infos('database' = DATABASE())" in second_query_sql + + +def test_ctas_string_key_coercion_and_internal_iceberg_ctas( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. CTAS with explicit unique_key on TEXT column coerces CAST(code AS VARCHAR(255)) and reorders key first + adapter.ctas( + "db.seed_uk_ctas", + parse_one("SELECT payload, code FROM db.raw_seed", read="doris"), + target_columns_to_types={ + "payload": exp.DataType.build("text"), + "code": exp.DataType.build("text"), + }, + table_properties={"unique_key": exp.to_column("code")}, + ) + + # 2. CTAS with all-TEXT columns (implicit DUPLICATE KEY on first column) coerces first column CAST to VARCHAR(255) + adapter.ctas( + "db.seed_dup_ctas", + parse_one("SELECT country_code, country_name FROM db.raw_countries", read="doris"), + target_columns_to_types={ + "country_code": exp.DataType.build("text"), + "country_name": exp.DataType.build("text"), + }, + ) + + # 3. 2-part internal Iceberg CTAS (table_format="iceberg") creates table with ENGINE=iceberg and inserts rows + adapter.ctas( + "silver.iceberg_events", + parse_one("SELECT id, ds FROM bronze.events", read="doris"), + target_columns_to_types={ + "id": exp.DataType.build("BIGINT"), + "ds": exp.DataType.build("DATE"), + }, + table_format="iceberg", + partitioned_by=[parse_one("DAY(ds)", dialect="doris")], + table_properties={"write-format": exp.Literal.string("parquet")}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`seed_uk_ctas` UNIQUE KEY (`code`) DISTRIBUTED BY HASH (`code`) BUCKETS 10 AS SELECT CAST(`code` AS VARCHAR(255)) AS `code`, CAST(`payload` AS STRING) AS `payload` FROM (SELECT `payload`, `code` FROM `db`.`raw_seed`) AS `_subquery`", + "CREATE TABLE IF NOT EXISTS `db`.`seed_dup_ctas` AS SELECT CAST(`country_code` AS VARCHAR(255)) AS `country_code`, CAST(`country_name` AS STRING) AS `country_name` FROM (SELECT `country_code`, `country_name` FROM `db`.`raw_countries`) AS `_subquery`", + "CREATE TABLE IF NOT EXISTS `silver`.`iceberg_events` (`id` BIGINT, `ds` DATE) ENGINE=iceberg PARTITION BY LIST (DAY(`ds`)) () PROPERTIES ('write-format'='parquet')", + "INSERT INTO `silver`.`iceberg_events` (`id`, `ds`) SELECT CAST(`id` AS BIGINT) AS `id`, CAST(`ds` AS DATE) AS `ds` FROM (SELECT `id`, `ds` FROM `bronze`.`events`) AS `_subquery`", + ] + + +def test_self_referencing_replace_query_on_new_table( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch.object(adapter, "get_data_object", return_value=None) + + adapter.replace_query( + "db.self_ref_tbl", + parse_one("SELECT id, name FROM db.self_ref_tbl", read="doris"), + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "name": exp.DataType.build("VARCHAR(50)"), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`self_ref_tbl` (`id` INT, `name` VARCHAR(50))", + "INSERT INTO `db`.`self_ref_tbl` (`id`, `name`) SELECT CAST(`id` AS INT) AS `id`, CAST(`name` AS VARCHAR(50)) AS `name` FROM (SELECT `id`, `name` FROM `db`.`self_ref_tbl`) AS `_subquery`", + ] + + +def test_delete_literal_tuple_in_and_subquery_duplicate_aliases_and_concat_ws_guard( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Literal tuple IN (e.g. from SQLMesh state sync snapshot_id_filter) rewrites to grouped single-column IN + adapter.delete_from( + "sqlmesh._snapshots", + parse_one( + "(name, identifier) IN (('model_a', 'id_1'), ('model_a', 'id_2'), ('model_b', 'id_3'))", + read="doris", + ), + ) + + # 2. Composite subquery selecting duplicate short column names (`u.id, m.id`) aliases the second column (`_col_1`) + adapter.delete_from( + "db.pair_links", + parse_one( + "(a_id, b_id) IN (SELECT u.id, m.id FROM db.users u JOIN db.members m ON u.org_id = m.org_id)", + read="doris", + ), + ) + + # 3. CONCAT_WS on LHS with a single non-CONCAT_WS column in RHS subquery does not unwrap into mismatched arity + adapter.delete_from( + "db.target_tbl", + parse_one( + "CONCAT_WS('__', a, b) IN (SELECT combined_key FROM db.staging_keys)", + read="doris", + ), + ) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `sqlmesh`.`_snapshots` WHERE `name` = 'model_a' AND `identifier` IN ('id_1', 'id_2')", + "DELETE FROM `sqlmesh`.`_snapshots` WHERE `name` = 'model_b' AND `identifier` IN ('id_3')", + "DELETE FROM `db`.`pair_links` AS `_t1` USING (SELECT `u`.`id`, `m`.`id` AS `_col_1` FROM `db`.`users` AS `u` JOIN `db`.`members` AS `m` ON `u`.`org_id` = `m`.`org_id`) AS `_t2` WHERE `_t1`.`a_id` = `_t2`.`id` AND `_t1`.`b_id` = `_t2`.`_col_1`", + "DELETE FROM `db`.`target_tbl` WHERE CONCAT_WS('__', `a`, `b`) IN (SELECT `combined_key` FROM `db`.`staging_keys`)", + ] + + +def test_to_sql_root_int128_and_nested_struct_pretty_sql_and_fetchdf_reconnect( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + import pandas as pd + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Root exp.DataType INT128 replacement + assert adapter._to_sql(exp.DataType.build("INT128")) == "LARGEINT" + + # 2. Nested STRUCT with INT128 under _pretty_sql = True + adapter._pretty_sql = True + nested_struct = exp.DataType.build( + "STRUCT>", + dialect="doris", + ) + assert ( + adapter._to_sql(nested_struct) + == "STRUCT<`outer_id`: LARGEINT, `nested`: STRUCT<`inner_id`: LARGEINT, `label`: VARCHAR(20)>>" + ) + adapter._pretty_sql = False + + # 3. _fetch_native_df retries once on non-timeout OperationalError(2013) disconnect + class OperationalError(Exception): + pass + + close_spy = mocker.spy(adapter._connection_pool, "close") + fetch_calls = 0 + + def fake_read_sql_query(sql, con, *args, **kwargs): + nonlocal fetch_calls + fetch_calls += 1 + if fetch_calls == 1: + raise OperationalError(2013, "Lost connection to MySQL server during query") + return pd.DataFrame({"id": [1]}) + + mocker.patch("pandas.io.sql.read_sql_query", side_effect=fake_read_sql_query) + df = adapter.fetchdf("SELECT 1 AS id") + assert len(df) == 1 + assert fetch_calls == 2 + assert close_spy.call_count == 1 From b32787ec33d20050628a94c59cef059acfe18e6d Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:09:15 -0700 Subject: [PATCH 17/27] feat(doris): strip FOR UPDATE locks, guard composite create_index, and support partition_by alias Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 75 +++++++++++++++++++++---- tests/core/engine_adapter/test_doris.py | 52 +++++++++++++++++ 2 files changed, 116 insertions(+), 11 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 330003ca67..97b32a4580 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -666,8 +666,12 @@ def _create_materialized_view( partitioned_by = materialized_properties.get("partitioned_by") elif create_kwargs.get("partitioned_by"): partitioned_by = create_kwargs.get("partitioned_by") - elif view_properties and view_properties.get("partitioned_by"): - raw_vp = view_properties.get("partitioned_by") + elif view_properties and ( + view_properties.get("partitioned_by") or view_properties.get("partition_by") + ): + raw_vp = view_properties.get("partitioned_by") or view_properties.get( + "partition_by" + ) if raw_vp is not None: partitioned_by = raw_vp if isinstance(raw_vp, list) else [raw_vp] @@ -763,6 +767,17 @@ def _create_materialized_view( ), quote_identifiers=self.QUOTE_IDENTIFIERS_IN_VIEWS, ) + build_mode = ( + self._to_upper_str(view_properties["build"]) + if view_properties and view_properties.get("build") is not None + else None + ) + if ( + not is_sync_mv + and build_mode == "DEFERRED" + and bool((materialized_properties or {}).get("has_audits")) + ): + self.execute(f"REFRESH MATERIALIZED VIEW {self._to_sql(target_table)} AUTO") self._clear_data_object_cache(view_name) def _to_source_table(self, raw_src: t.Any) -> exp.Table: @@ -815,6 +830,21 @@ def drop_view( return super().drop_view(view_name, ignore_if_not_exists, materialized, **kwargs) + def create_index( + self, + table_name: TableName, + index_name: str, + columns: t.Tuple[str, ...], + exists: bool = True, + ) -> None: + if len(columns) != 1: + logger.warning( + f"[Doris] Skipping CREATE INDEX {index_name} on {table_name}: " + f"Doris only supports single-column indexes, got columns {columns}." + ) + return + super().create_index(table_name, index_name, columns, exists=exists) + def _create_table_like( self, target_table_name: TableName, @@ -1667,7 +1697,11 @@ def _build_create_table_exp( and str(self._extract_str_value(auto_partition) or "").strip().lower() in ("true", "1") ) - raw_part_by = kwargs.get("partitioned_by") or table_props.get("partitioned_by") + raw_part_by = ( + kwargs.get("partitioned_by") + or table_props.get("partitioned_by") + or table_props.get("partition_by") + ) if is_auto and raw_part_by: part_list = raw_part_by if isinstance(raw_part_by, list) else [raw_part_by] parsed_part_exprs, part_kind = self._parse_partition_expressions(part_list) @@ -1868,7 +1902,11 @@ def _reorder_olap_columns( col_names = [coercible_col] + [c for c in col_names if c != coercible_col] # Ensure any OLAP partition columns with valid key types precede non-keyable columns - raw_part_by = kwargs.get("partitioned_by") or table_props.get("partitioned_by") + raw_part_by = ( + kwargs.get("partitioned_by") + or table_props.get("partitioned_by") + or table_props.get("partition_by") + ) if raw_part_by: part_list = raw_part_by if isinstance(raw_part_by, list) else [raw_part_by] parsed_part_exprs, _ = self._parse_partition_expressions(part_list) @@ -1911,7 +1949,14 @@ def _parse_partition_expressions( parsed_partitioned_by: t.List[exp.Expr] = [] partition_kind: t.Optional[str] = None - for expr in partitioned_by: + flattened_partitioned_by: t.List[exp.Expr] = [] + for item in partitioned_by: + if isinstance(item, (exp.Tuple, exp.Array)): + flattened_partitioned_by.extend(item.expressions) + else: + flattened_partitioned_by.append(item) + + for expr in flattened_partitioned_by: if isinstance(expr, exp.List) or ( isinstance(expr, exp.Anonymous) and expr.this @@ -2268,12 +2313,14 @@ def _build_table_properties_exp( # Pop partitioning-related properties partitions = table_properties_copy.pop("partitions", None) auto_partition = table_properties_copy.pop("auto_partition", None) - if not partitioned_by and "partitioned_by" in table_properties_copy: - raw_part_by = table_properties_copy.pop("partitioned_by") - if raw_part_by is not None: - partitioned_by = raw_part_by if isinstance(raw_part_by, list) else [raw_part_by] - else: - table_properties_copy.pop("partitioned_by", None) + raw_part_by = table_properties_copy.pop("partitioned_by", None) + raw_part_by_alias = table_properties_copy.pop("partition_by", None) + if not partitioned_by: + chosen_part_by = raw_part_by if raw_part_by is not None else raw_part_by_alias + if chosen_part_by is not None: + partitioned_by = ( + chosen_part_by if isinstance(chosen_part_by, list) else [chosen_part_by] + ) add_partition = True if partitioned_by: @@ -2517,6 +2564,12 @@ def _resolve_buckets(self, raw_buckets: t.Any) -> t.Optional[exp.Expr]: ) def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> str: + if expression.find(exp.Lock): + expression = expression.copy() + for select_node in expression.find_all(exp.Select): + if select_node.args.get("locks"): + select_node.set("locks", None) + logger.debug("[Doris] Removed unsupported FOR UPDATE lock from SELECT statement") if expression.find(exp.DataType): expression = expression.copy() for dtype in reversed(list(expression.find_all(exp.DataType))): diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 8053218e78..dae48805ff 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -2156,3 +2156,55 @@ def fake_read_sql_query(sql, con, *args, **kwargs): assert len(df) == 1 assert fetch_calls == 2 assert close_spy.call_count == 1 + + +def test_strip_for_update_and_multi_col_index_and_partition_by_alias_and_deferred_mv_audit_refresh( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Strip FOR UPDATE locks from SELECT statements (e.g. EnvironmentState.finalize) + locked_select = ( + exp.select("plan_id") + .from_("sqlmesh._environments") + .where(exp.column("name").eq("prod")) + .lock() + ) + adapter.execute(locked_select) + + # 2. Multi-column composite index creation is skipped gracefully (Doris only supports single-column indexes) + adapter.create_index("sqlmesh._snapshots", "idx_snapshots", ("name", "identifier")) + + # 3. partition_by alias in table_properties works for OLAP AUTO PARTITION BY RANGE + adapter.create_table( + "db.alias_part_tbl", + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATE"), + }, + table_properties={ + "partition_by": parse_one("RANGE(DATE_TRUNC(ds, 'day'))", read="doris"), + "auto_partition": exp.Literal.string("true"), + }, + ) + + # 4. Async materialized view with build='DEFERRED' and has_audits=True triggers REFRESH MATERIALIZED VIEW ... AUTO + adapter.create_view( + "db.deferred_audited_mv", + parse_one("SELECT id, ds FROM db.alias_part_tbl", read="doris"), + replace=False, + materialized=True, + materialized_properties={"has_audits": True}, + view_properties={ + "build": exp.Literal.string("DEFERRED"), + "refresh": exp.Literal.string("COMPLETE ON MANUAL"), + "partition_by": exp.Tuple(expressions=[exp.column("ds")]), + }, + ) + + assert to_sql_calls(adapter) == [ + "SELECT `plan_id` FROM `sqlmesh`.`_environments` WHERE `name` = 'prod'", + "CREATE TABLE IF NOT EXISTS `db`.`alias_part_tbl` (`id` INT, `ds` DATE NOT NULL) AUTO PARTITION BY RANGE (DATE_TRUNC(`ds`, 'DAY')) ()", + "CREATE MATERIALIZED VIEW `db`.`deferred_audited_mv` BUILD DEFERRED REFRESH COMPLETE ON MANUAL PARTITION BY (`ds`) AS SELECT `id`, `ds` FROM `db`.`alias_part_tbl`", + "REFRESH MATERIALIZED VIEW `db`.`deferred_audited_mv` AUTO", + ] From ecdd4c4c4aaf22eb7b48eb151598aeb67d3bc20a Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Thu, 22 Oct 2026 16:44:37 -0700 Subject: [PATCH 18/27] fix(doris): support OR key deletes, bare boolean conjuncts, and OLAP key alter guard Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 217 +++++++++++++++++++++++- tests/core/engine_adapter/test_doris.py | 74 ++++++++ 2 files changed, 288 insertions(+), 3 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 97b32a4580..66dd650f15 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -38,6 +38,7 @@ from sqlmesh.core._typing import SchemaName, TableName from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF from sqlmesh.core.node import IntervalUnit + from sqlmesh.core.schema_diff import TableAlterOperation logger = logging.getLogger(__name__) @@ -58,6 +59,7 @@ class DorisEngineAdapter( MAX_TABLE_COMMENT_LENGTH = 2048 MAX_COLUMN_COMMENT_LENGTH = 255 SUPPORTS_INDEXES = True + SUPPORTS_TUPLE_IN = False SUPPORTS_REPLACE_TABLE = False MAX_IDENTIFIER_LENGTH = 64 SUPPORTS_MATERIALIZED_VIEWS = True @@ -361,17 +363,37 @@ def columns( self, table_name: TableName, include_pseudo_columns: bool = False ) -> t.Dict[str, exp.DataType]: """Fetches column names and types for the target table, normalizing Doris type aliases.""" + cols, key_cols = self._describe_columns_with_keys(table_name) + self._last_described_key_cols: t.Set[str] = key_cols + return cols + + def _describe_columns_with_keys( + self, table_name: TableName + ) -> t.Tuple[t.Dict[str, exp.DataType], t.Set[str]]: self.execute(exp.Describe(this=exp.to_table(table_name), kind="TABLE")) describe_output = self.cursor.fetchall() cols: t.Dict[str, exp.DataType] = {} - for column_name, column_type, *_ in itertools.takewhile( - lambda row: not str(row[0]).startswith("#"), + key_cols: t.Set[str] = set() + for row in itertools.takewhile( + lambda r: not str(r[0]).startswith("#"), describe_output, ): + column_name = row[0] if len(row) > 0 else None + column_type = row[1] if len(row) > 1 else None + is_key = row[3] if len(row) > 3 else None col_name_str = _decoded_str(column_name) if column_name else "" raw_type = _decoded_str(column_type) if column_type else "" if not col_name_str.strip() or not raw_type.strip(): continue + if is_key is not None and _decoded_str(is_key).strip().lower() in ( + "true", + "pri", + "uni", + "dup", + "agg", + "yes", + ): + key_cols.add(col_name_str.lower()) normalized_type = re.sub( r"\bDECIMALV[23]\b(?!\s*`?\s*:)", "DECIMAL", raw_type, flags=re.IGNORECASE ) @@ -397,7 +419,59 @@ def columns( flags=re.IGNORECASE, ) cols[col_name_str] = exp.DataType.build(normalized_type, dialect=self.dialect, udt=True) - return cols + if cols and not key_cols: + first_col = next(iter(cols)) + key_cols.add(first_col.lower()) + return cols, key_cols + + def get_alter_operations( + self, + current_table_name: TableName, + target_table_name: TableName, + *, + ignore_destructive: bool = False, + ignore_additive: bool = False, + ) -> t.List[TableAlterOperation]: + from sqlmesh.core.schema_diff import ( + TableAlterChangeColumnTypeOperation, + TableAlterOperation, + ) + + current_table = exp.to_table(current_table_name) + is_iceberg = self._is_iceberg_table(catalog_name=current_table.catalog) + if is_iceberg: + return super().get_alter_operations( + current_table_name, + target_table_name, + ignore_destructive=ignore_destructive, + ignore_additive=ignore_additive, + ) + + self._last_described_key_cols = set() + current_cols = self.columns(current_table_name) + current_key_cols = set(getattr(self, "_last_described_key_cols", set())) + if current_cols and not current_key_cols: + current_key_cols = {next(iter(current_cols)).lower()} + target_cols = self.columns(target_table_name) + operations = t.cast( + t.List[TableAlterOperation], + self.schema_differ.compare_columns( + current_table_name, + current_cols, + target_cols, + ignore_destructive=ignore_destructive, + ignore_additive=ignore_additive, + ), + ) + return [ + op + for op in operations + if not ( + isinstance(op, TableAlterChangeColumnTypeOperation) + and op.column.name.lower() in current_key_cols + and not self._is_valid_olap_key_type(op.column_type) + ) + ] def create_schema( self, @@ -1009,6 +1083,7 @@ def delete_from( return where = self._rewrite_between_in_where(where) + where = self._normalize_boolean_column_conjuncts(where) subquery_info = self._extract_subquery_from_where(where) if subquery_info is not None: @@ -1020,6 +1095,11 @@ def delete_from( self._delete_key_rows(table_name, *tuple_in_info, is_ast_values=True) return + or_key_info = self._extract_or_key_rows_from_where(where) + if or_key_info is not None: + self._delete_key_rows(table_name, *or_key_info, is_ast_values=True) + return + target_table = exp.to_table(table_name) for col in where.find_all(exp.Column): if not col.find_ancestor(exp.Subquery, exp.Select) and ( @@ -1206,6 +1286,136 @@ def _extract_literal_tuple_in_from_where( return normalized_keys, tuple_rows, extra_conditions + @classmethod + def _normalize_boolean_column_conjuncts(cls, where: exp.Expr) -> exp.Expr: + """Rewrite top-level bare boolean column conjuncts (`AND is_dev`, `AND NOT is_dev`) to + explicit binary comparisons (`is_dev = TRUE`, `is_dev = FALSE`) required by Doris DeleteHandler. + """ + unwrapped = where.unnest() + if isinstance(unwrapped, exp.Column): + return unwrapped.copy().eq(exp.true()) + if isinstance(unwrapped, exp.Not) and isinstance(unwrapped.this.unnest(), exp.Column): + return unwrapped.this.unnest().copy().eq(exp.false()) + if isinstance(unwrapped, exp.And): + conjuncts = list(unwrapped.flatten()) + changed = False + new_conjuncts: t.List[exp.Expr] = [] + for cond in conjuncts: + inner = cond.unnest() + if isinstance(inner, exp.Column): + new_conjuncts.append(inner.copy().eq(exp.true())) + changed = True + elif isinstance(inner, exp.Not) and isinstance(inner.this.unnest(), exp.Column): + new_conjuncts.append(inner.this.unnest().copy().eq(exp.false())) + changed = True + else: + new_conjuncts.append(cond) + if changed: + return exp.and_(*new_conjuncts, copy=False) + return where + + def _extract_or_key_rows_from_where( + self, where: exp.Expr + ) -> t.Optional[t.Tuple[t.List[exp.Expr], t.List[t.Tuple[exp.Expr, ...]], t.List[exp.Expr]]]: + """Extract a top-level OR of equality/IS NULL key predicates (e.g. from `snapshot_id_filter` + when `SUPPORTS_TUPLE_IN = False`, or `delete_expired_environments`) and sibling AND conditions. + """ + where = where.unnest() + conjuncts: t.List[exp.Expr] = ( + list(where.flatten()) if isinstance(where, exp.And) else [where] + ) + target_or: t.Optional[exp.Or] = None + extra_conditions: t.List[exp.Expr] = [] + + for cond in conjuncts: + unwrapped_cond = cond.unnest() + if target_or is None and isinstance(unwrapped_cond, exp.Or): + target_or = unwrapped_cond + continue + extra_conditions.append(cond) + + if target_or is None: + return None + + disjuncts = list(target_or.flatten()) + if len(disjuncts) < 2: + return None + + canonical_col_order: t.Optional[t.List[str]] = None + canonical_cols: t.List[exp.Expr] = [] + tuple_rows: t.List[t.Tuple[exp.Expr, ...]] = [] + + for branch in disjuncts: + branch_unwrapped = branch.unnest() + branch_conjuncts = ( + list(branch_unwrapped.flatten()) + if isinstance(branch_unwrapped, exp.And) + else [branch_unwrapped] + ) + col_to_val: t.Dict[str, exp.Expr] = {} + col_to_expr: t.Dict[str, exp.Expr] = {} + for item in branch_conjuncts: + item_unwrapped = item.unnest() + if isinstance(item_unwrapped, exp.EQ): + left, right = item_unwrapped.this, item_unwrapped.expression + if isinstance(left, (exp.Column, exp.Identifier)) and isinstance( + right, (exp.Literal, exp.Boolean, exp.Null) + ): + col_expr = ( + left.copy() + if isinstance(left, exp.Column) + else exp.Column(this=left.copy()) + ) + val_expr = right.copy() + elif isinstance(right, (exp.Column, exp.Identifier)) and isinstance( + left, (exp.Literal, exp.Boolean, exp.Null) + ): + col_expr = ( + right.copy() + if isinstance(right, exp.Column) + else exp.Column(this=right.copy()) + ) + val_expr = left.copy() + else: + return None + elif isinstance(item_unwrapped, exp.Is) and isinstance( + item_unwrapped.expression, exp.Null + ): + left = item_unwrapped.this + if isinstance(left, (exp.Column, exp.Identifier)): + col_expr = ( + left.copy() + if isinstance(left, exp.Column) + else exp.Column(this=left.copy()) + ) + val_expr = exp.null() + else: + return None + else: + return None + + col_key = col_expr.name.lower() + if not col_key or col_key in col_to_val: + return None + col_to_val[col_key] = val_expr + col_to_expr[col_key] = col_expr + + if not col_to_val: + return None + + if canonical_col_order is None: + canonical_col_order = list(col_to_val.keys()) + canonical_cols = [col_to_expr[k] for k in canonical_col_order] + elif set(col_to_val.keys()) != set(canonical_col_order): + return None + + tuple_rows.append(tuple(col_to_val[k] for k in canonical_col_order)) + + if not canonical_cols or not tuple_rows: + return None + + return canonical_cols, tuple_rows, extra_conditions + @staticmethod def _unwrap_concat_ws(expr: exp.Expr) -> t.Optional[exp.Expr]: curr = expr.unnest() @@ -2842,6 +3052,7 @@ def _properties_to_expressions(self, properties: t.Dict[str, t.Any]) -> t.List[e "distributed_by", "clustered_by", "buckets", + "partition_by", "partitioned_by", "partitions", "auto_partition", diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index dae48805ff..0795782c4d 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -2208,3 +2208,77 @@ def test_strip_for_update_and_multi_col_index_and_partition_by_alias_and_deferre "CREATE MATERIALIZED VIEW `db`.`deferred_audited_mv` BUILD DEFERRED REFRESH COMPLETE ON MANUAL PARTITION BY (`ds`) AS SELECT `id`, `ds` FROM `db`.`alias_part_tbl`", "REFRESH MATERIALIZED VIEW `db`.`deferred_audited_mv` AUTO", ] + + +def test_state_sync_or_key_delete_and_boolean_column_and_alter_key_coercion_guard( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + from sqlmesh.core.snapshot import SnapshotId, SnapshotNameVersion + from sqlmesh.core.state_sync.db.utils import snapshot_id_filter, snapshot_name_version_filter + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + assert adapter.SUPPORTS_TUPLE_IN is False + + # 1. snapshot_id_filter with SUPPORTS_TUPLE_IN=False produces OR of AND equalities; + # delete_from decomposes it into zero-OR grouped IN deletes. + sids = [ + SnapshotId(name="db.model_a", identifier="id_1"), + SnapshotId(name="db.model_a", identifier="id_2"), + SnapshotId(name="db.model_b", identifier="id_3"), + ] + for where in snapshot_id_filter(adapter, sids, batch_size=100): + adapter.delete_from("sqlmesh._snapshots", where=where) + + # 2. IntervalState._delete_intervals_by_dev_version passes `where.and_(exp.column("is_dev"))` + snvs = [ + SnapshotNameVersion(name="db.model_a", version="v1"), + SnapshotNameVersion(name="db.model_a", version="v2"), + ] + for where in snapshot_name_version_filter( + adapter, snvs, version_column_name="dev_version", alias=None, batch_size=100 + ): + adapter.delete_from("sqlmesh._intervals", where.and_(exp.column("is_dev"))) + + # 3. EnvironmentState.delete_expired_environments passes single-column OR `environment_name = 'dev1' OR environment_name = 'dev2'` + adapter.delete_from( + "sqlmesh._environment_statements", + where=exp.or_( + exp.column("environment_name").eq(exp.Literal.string("dev_1")), + exp.column("environment_name").eq(exp.Literal.string("dev_2")), + ), + ) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `sqlmesh`.`_snapshots` WHERE `name` = 'db.model_a' AND `identifier` IN ('id_1', 'id_2')", + "DELETE FROM `sqlmesh`.`_snapshots` WHERE `name` = 'db.model_b' AND `identifier` IN ('id_3')", + "DELETE FROM `sqlmesh`.`_intervals` WHERE `name` = 'db.model_a' AND `dev_version` IN ('v1', 'v2') AND `is_dev` = TRUE", + "DELETE FROM `sqlmesh`.`_environment_statements` WHERE `environment_name` IN ('dev_1', 'dev_2')", + ] + + # 4. get_alter_operations suppresses invalid VARCHAR(255) -> STRING type change on an OLAP key column + # while preserving ADD COLUMN and non-key column widening. + adapter.cursor.reset_mock() + mocker.patch.object( + adapter.cursor, + "fetchall", + side_effect=[ + # DESCRIBE db.current_tbl: `code` is the OLAP key column coerced to VARCHAR(255), `notes` is a non-key VARCHAR(100) + [ + ("code", "VARCHAR(255)", "Yes", "true", None, ""), + ("notes", "VARCHAR(100)", "Yes", "false", None, ""), + ], + # DESCRIBE db.target_tbl: `id` was added, `code` is STRING, `notes` widened to STRING + [ + ("id", "INT", "Yes", "true", None, ""), + ("code", "TEXT", "Yes", "false", None, ""), + ("notes", "TEXT", "Yes", "false", None, ""), + ], + ], + ) + ops = adapter.get_alter_operations("db.current_tbl", "db.target_tbl") + adapter.cursor.reset_mock() + adapter.alter_table(ops) + assert to_sql_calls(adapter) == [ + "ALTER TABLE `db`.`current_tbl` ADD COLUMN `id` INT", + "ALTER TABLE `db`.`current_tbl` MODIFY COLUMN `notes` STRING", + ] From f187846db7b695c8e672175c3e59db4f67d38ee0 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Thu, 22 Oct 2026 18:19:51 -0700 Subject: [PATCH 19/27] feat(doris): support unannotated view column comments and StarRocks property parity Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 155 ++++++++++++++++++++---- tests/core/engine_adapter/test_doris.py | 98 +++++++++++++++ 2 files changed, 231 insertions(+), 22 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 66dd650f15..c1a9aeaa99 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -200,6 +200,12 @@ def _is_iceberg_table( def _extract_str_value(val: t.Any) -> t.Optional[str]: if val is None: return None + if isinstance(val, exp.Paren): + return DorisEngineAdapter._extract_str_value(val.this) + if isinstance(val, (exp.Tuple, exp.Array)): + return ",".join( + s for e in val.expressions if (s := DorisEngineAdapter._extract_str_value(e)) + ) if isinstance(val, exp.Boolean): return "true" if val.this else "false" if isinstance(val, bool): @@ -218,6 +224,22 @@ def _extract_str_value(val: t.Any) -> t.Optional[str]: return val.name or str(val) return str(val) + @staticmethod + def _split_cleaned_column_names(raw: str) -> t.List[str]: + s = raw.strip() + while len(s) >= 2 and ( + (s[0] == "(" and s[-1] == ")") + or (s[0] == "'" and s[-1] == "'") + or (s[0] == '"' and s[-1] == '"') + ): + s = s[1:-1].strip() + result: t.List[str] = [] + for part in s.split(","): + cleaned = part.strip().strip("'\"").strip("()").strip().strip("`'\"").strip() + if cleaned: + result.append(cleaned) + return result + @classmethod def _extract_column_names(cls, key_expr: t.Any) -> t.List[str]: if key_expr is None: @@ -235,13 +257,12 @@ def _extract_column_names(cls, key_expr: t.Any) -> t.List[str]: names.extend(cls._extract_column_names(e)) return names if isinstance(key_expr, exp.Column): - return [key_expr.name] + return cls._split_cleaned_column_names(key_expr.name) if isinstance(key_expr, (exp.Identifier, exp.Literal, exp.Var)): - raw = str(key_expr.this) - return [c.strip().strip("`") for c in raw.split(",") if c.strip()] + return cls._split_cleaned_column_names(str(key_expr.this)) if isinstance(key_expr, str): - return [c.strip().strip("`") for c in key_expr.split(",") if c.strip()] - return [str(key_expr).strip().strip("`")] + return cls._split_cleaned_column_names(key_expr) + return cls._split_cleaned_column_names(str(key_expr)) @classmethod def _is_valid_olap_key_type(cls, dtype: exp.DataType) -> bool: @@ -639,6 +660,16 @@ def create_view( ) if not materialized: + if ( + not target_columns_to_types + and column_descriptions + and isinstance(query_or_df, exp.Query) + and query_or_df.named_selects + and "*" not in query_or_df.named_selects + ): + target_columns_to_types = { + col: exp.DataType.build("UNKNOWN") for col in query_or_df.named_selects + } return super().create_view( view_name, query_or_df, @@ -726,11 +757,24 @@ def _create_materialized_view( ) ) + schema_cols_to_types = target_columns_to_types + if ( + not schema_cols_to_types + and column_descriptions + and not is_sync_mv + and isinstance(query, exp.Query) + and query.named_selects + and "*" not in query.named_selects + ): + schema_cols_to_types = { + col: exp.DataType.build("UNKNOWN") for col in query.named_selects + } + schema: t.Union[exp.Table, exp.Schema] = target_table - if target_columns_to_types and not is_sync_mv: + if schema_cols_to_types and not is_sync_mv: schema = self._build_schema_exp( target_table, - target_columns_to_types, + schema_cols_to_types, column_descriptions, is_view=True, ) @@ -740,14 +784,6 @@ def _create_materialized_view( partitioned_by = materialized_properties.get("partitioned_by") elif create_kwargs.get("partitioned_by"): partitioned_by = create_kwargs.get("partitioned_by") - elif view_properties and ( - view_properties.get("partitioned_by") or view_properties.get("partition_by") - ): - raw_vp = view_properties.get("partitioned_by") or view_properties.get( - "partition_by" - ) - if raw_vp is not None: - partitioned_by = raw_vp if isinstance(raw_vp, list) else [raw_vp] clustered_by = None if materialized_properties and materialized_properties.get("clustered_by"): @@ -2039,6 +2075,28 @@ def _reorder_olap_columns( kwargs: t.Dict[str, t.Any], ) -> t.Dict[str, exp.DataType]: table_props = kwargs.get("table_properties") or {} + if table_props.get("primary_key") is not None or kwargs.get("primary_key") is not None: + raise SQLMeshError( + "Doris does not support 'primary_key' in physical_properties; use 'unique_key', 'duplicate_key', or 'aggregate_key' instead." + ) + if table_props.get("partition") is not None: + raise SQLMeshError( + "Invalid property 'partition' in physical_properties; did you mean 'partition_by', 'partitioned_by', or 'partitions'?" + ) + if ( + sum( + x is not None + for x in ( + kwargs.get("partitioned_by") or None, + table_props.get("partitioned_by"), + table_props.get("partition_by"), + ) + ) + > 1 + ): + raise SQLMeshError( + "Multiple partition definitions found: specify only one of model 'partitioned_by' or physical_properties 'partitioned_by' / 'partition_by'." + ) unique_key = table_props.get("unique_key") or kwargs.get("unique_key") duplicate_key = table_props.get("duplicate_key") aggregate_key = table_props.get("aggregate_key") @@ -2438,6 +2496,17 @@ def _build_table_properties_exp( ) # Pop key models and order_by from table_properties_copy + if ( + table_properties_copy.pop("primary_key", None) is not None + or kwargs.get("primary_key") is not None + ): + raise SQLMeshError( + "Doris does not support 'primary_key' in physical_properties; use 'unique_key', 'duplicate_key', or 'aggregate_key' instead." + ) + if table_properties_copy.get("partition") is not None: + raise SQLMeshError( + "Invalid property 'partition' in physical_properties; did you mean 'partition_by', 'partitioned_by', or 'partitions'?" + ) unique_key = table_properties_copy.pop("unique_key", None) or kwargs.get("unique_key") duplicate_key = table_properties_copy.pop("duplicate_key", None) aggregate_key = table_properties_copy.pop("aggregate_key", None) @@ -2525,6 +2594,20 @@ def _build_table_properties_exp( auto_partition = table_properties_copy.pop("auto_partition", None) raw_part_by = table_properties_copy.pop("partitioned_by", None) raw_part_by_alias = table_properties_copy.pop("partition_by", None) + if ( + sum( + x is not None + for x in ( + partitioned_by or None, + raw_part_by, + raw_part_by_alias, + ) + ) + > 1 + ): + raise SQLMeshError( + "Multiple partition definitions found: specify only one of model 'partitioned_by' or physical_properties 'partitioned_by' / 'partition_by'." + ) if not partitioned_by: chosen_part_by = raw_part_by if raw_part_by is not None else raw_part_by_alias if chosen_part_by is not None: @@ -2846,6 +2929,14 @@ def _build_distributed_by_exp( cols = [ exp.to_column(c) for c in self._extract_column_names(distributed_by.expressions) ] + if kind_upper == "RANDOM" and cols: + raise SQLMeshError( + "Doris DISTRIBUTED BY RANDOM does not accept column expressions." + ) + if kind_upper == "HASH" and not cols: + raise SQLMeshError( + "Doris DISTRIBUTED BY HASH requires at least one column expression." + ) return exp.DistributedByProperty( kind=exp.Var(this=kind_upper), expressions=cols if cols else None, @@ -2868,7 +2959,7 @@ def _build_distributed_by_exp( ) ) m = re.match( - r"^(HASH|RANDOM)\s*(?:\(([^)]*)\))?\s*(?:BUCKETS\s+(AUTO|\d+))?$", + r"^(HASH|RANDOM)(?:\s*\(([^)]*)\)|\b)\s*(?:BUCKETS\s+(\S+))?$", text, flags=re.IGNORECASE, ) @@ -2877,14 +2968,18 @@ def _build_distributed_by_exp( cols_str = m.group(2) b_str = m.group(3) parsed_cols = ( - [ - exp.to_column(c.strip().strip("`")) - for c in cols_str.split(",") - if c.strip() - ] + [exp.to_column(c) for c in self._extract_column_names(cols_str)] if cols_str else None ) + if kind_upper == "RANDOM" and parsed_cols: + raise SQLMeshError( + "Doris DISTRIBUTED BY RANDOM does not accept column expressions." + ) + if kind_upper == "HASH" and not parsed_cols: + raise SQLMeshError( + "Doris DISTRIBUTED BY HASH requires at least one column expression." + ) b_expr = self._resolve_buckets(b_str) if b_str else default_buckets_expr return exp.DistributedByProperty( kind=exp.Var(this=kind_upper), @@ -2951,9 +3046,16 @@ def _build_distributed_by_exp( if distributed_info: raw_kind = distributed_info.get("kind") - expressions = distributed_info.get("expressions") + expressions = ( + distributed_info.get("expressions") + if distributed_info.get("expressions") is not None + else distributed_info.get("columns") + ) buckets = distributed_info.get("buckets", top_level_buckets) + if raw_kind is None and expressions is not None: + raw_kind = "HASH" + if raw_kind: kind_upper = str(raw_kind).strip().upper() if kind_upper not in ("HASH", "RANDOM"): @@ -2967,6 +3069,15 @@ def _build_distributed_by_exp( col_names = self._extract_column_names(expressions) expressions_list = [exp.to_column(c) for c in col_names] + if kind_upper == "RANDOM" and expressions_list: + raise SQLMeshError( + "Doris DISTRIBUTED BY RANDOM does not accept column expressions." + ) + if kind_upper == "HASH" and not expressions_list: + raise SQLMeshError( + "Doris DISTRIBUTED BY HASH requires at least one column expression." + ) + return exp.DistributedByProperty( kind=exp.Var(this=kind_upper), expressions=expressions_list, diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 0795782c4d..59c826f165 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -2282,3 +2282,101 @@ def test_state_sync_or_key_delete_and_boolean_column_and_alter_key_coercion_guar "ALTER TABLE `db`.`current_tbl` ADD COLUMN `id` INT", "ALTER TABLE `db`.`current_tbl` MODIFY COLUMN `notes` STRING", ] + + +def test_unannotated_view_column_comments_and_starrocks_property_parity( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Unannotated VIEW and MATERIALIZED VIEW with column_descriptions (as called by ViewStrategy.create + # where target_columns_to_types=None) infer projection column names to emit column comments. + adapter.create_view( + "db.unannotated_view", + parse_one("SELECT id, name FROM db.users", read="doris"), + replace=False, + table_description="Users view", + column_descriptions={"id": "Primary user ID"}, + ) + adapter.create_view( + "db.unannotated_mv", + parse_one("SELECT id, name FROM db.users", read="doris"), + replace=False, + materialized=True, + table_description="Users MV", + column_descriptions={"id": "Primary user ID"}, + ) + + # 2. Parenthesized string forms for unique_key and order_by, structured distributed_by with `columns=...`, + # and tuple generic properties (e.g. excluded_trigger_tables = ('db.t1', 'db.t2')). + adapter.create_table( + "db.paren_keys_tbl", + target_columns_to_types={ + "name": exp.DataType.build("VARCHAR(50)"), + "id": exp.DataType.build("INT"), + "dt": exp.DataType.build("DATE"), + }, + table_properties={ + "unique_key": exp.Literal.string("(id, dt)"), + "order_by": exp.Literal.string("'(dt, id)'"), + "distributed_by": exp.Tuple( + expressions=[ + exp.EQ(this=exp.column("columns"), expression=exp.Literal.string("(id, dt)")), + exp.EQ(this=exp.column("buckets"), expression=exp.Literal.number(12)), + ] + ), + "excluded_trigger_tables": exp.Tuple( + expressions=[exp.Literal.string("db.t1"), exp.Literal.string("db.t2")] + ), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE VIEW `db`.`unannotated_view` (`id` COMMENT 'Primary user ID', `name`) COMMENT 'Users view' AS SELECT `id`, `name` FROM `db`.`users`", + "CREATE MATERIALIZED VIEW `db`.`unannotated_mv` (`id` COMMENT 'Primary user ID', `name`) COMMENT 'Users MV' AS SELECT `id`, `name` FROM `db`.`users`", + "CREATE TABLE IF NOT EXISTS `db`.`paren_keys_tbl` (`id` INT, `dt` DATE, `name` VARCHAR(50)) UNIQUE KEY (`id`, `dt`) ORDER BY (`dt`, `id`) DISTRIBUTED BY HASH (`id`, `dt`) BUCKETS 12 PROPERTIES ('excluded_trigger_tables'='db.t1,db.t2')", + ] + + # 3. Fail-fast validation for unsupported primary_key, typo `partition`, conflicting partition definitions, + # and invalid RANDOM(col) / column-less HASH distribution. + cols = {"id": exp.DataType.build("INT"), "dt": exp.DataType.build("DATE")} + + with pytest.raises(SQLMeshError, match="does not support 'primary_key'"): + adapter.create_table( + "db.bad_pk", + target_columns_to_types=cols, + table_properties={"primary_key": exp.to_column("id")}, + ) + + with pytest.raises(SQLMeshError, match="Invalid property 'partition'"): + adapter.create_table( + "db.bad_partition_typo", + target_columns_to_types=cols, + table_properties={"partition": exp.to_column("dt")}, + ) + + with pytest.raises(SQLMeshError, match="Multiple partition definitions"): + adapter.create_table( + "db.conflict_part", + target_columns_to_types=cols, + partitioned_by=[exp.to_column("dt")], + table_properties={"partition_by": exp.to_column("dt")}, + ) + + with pytest.raises( + SQLMeshError, match="DISTRIBUTED BY RANDOM does not accept column expressions" + ): + adapter.create_table( + "db.bad_random_cols", + target_columns_to_types=cols, + table_properties={"distributed_by": exp.Literal.string("RANDOM(id)")}, + ) + + with pytest.raises( + SQLMeshError, match="DISTRIBUTED BY HASH requires at least one column expression" + ): + adapter.create_table( + "db.bad_hash_no_cols", + target_columns_to_types=cols, + table_properties={"distributed_by": exp.Literal.string("HASH BUCKETS 10")}, + ) From 32310c9a8554a7f43036c34f51d0803ecb0a1db7 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Fri, 23 Oct 2026 07:46:30 -0700 Subject: [PATCH 20/27] fix(doris): handle IS TRUE/FALSE, negative OR keys, UNION deletes, and TableDiff CTAS Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 163 ++++++++++++++++++++---- tests/core/engine_adapter/test_doris.py | 80 ++++++++++++ 2 files changed, 216 insertions(+), 27 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index c1a9aeaa99..5fc9c6f6db 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -1322,27 +1322,47 @@ def _extract_literal_tuple_in_from_where( return normalized_keys, tuple_rows, extra_conditions + @staticmethod + def _is_literal_key_value(expr: exp.Expr) -> bool: + return isinstance(expr, (exp.Literal, exp.Boolean, exp.Null)) or ( + isinstance(expr, exp.Neg) and isinstance(expr.this, exp.Literal) + ) + + @classmethod + def _normalize_single_boolean_conjunct(cls, cond: exp.Expr) -> t.Optional[exp.Expr]: + inner = cond.unnest() + if isinstance(inner, exp.Column): + return inner.copy().eq(exp.true()) + if isinstance(inner, exp.Is) and isinstance(inner.expression, exp.Boolean): + return inner.this.copy().eq(inner.expression.copy()) + if isinstance(inner, exp.Not): + neg_inner = inner.this.unnest() + if isinstance(neg_inner, exp.Column): + return neg_inner.copy().eq(exp.false()) + if isinstance(neg_inner, exp.Is) and isinstance(neg_inner.expression, exp.Boolean): + return neg_inner.this.copy().eq( + exp.Boolean(this=not bool(neg_inner.expression.this)) + ) + return None + @classmethod def _normalize_boolean_column_conjuncts(cls, where: exp.Expr) -> exp.Expr: - """Rewrite top-level bare boolean column conjuncts (`AND is_dev`, `AND NOT is_dev`) to - explicit binary comparisons (`is_dev = TRUE`, `is_dev = FALSE`) required by Doris DeleteHandler. + """Rewrite top-level bare boolean column conjuncts (`AND is_dev`, `AND NOT is_dev`) and + `IS [NOT] TRUE/FALSE` predicates to explicit binary comparisons (`is_dev = TRUE`, + `is_dev = FALSE`) required by Doris DeleteHandler. """ + single = cls._normalize_single_boolean_conjunct(where) + if single is not None: + return single unwrapped = where.unnest() - if isinstance(unwrapped, exp.Column): - return unwrapped.copy().eq(exp.true()) - if isinstance(unwrapped, exp.Not) and isinstance(unwrapped.this.unnest(), exp.Column): - return unwrapped.this.unnest().copy().eq(exp.false()) if isinstance(unwrapped, exp.And): conjuncts = list(unwrapped.flatten()) changed = False new_conjuncts: t.List[exp.Expr] = [] for cond in conjuncts: - inner = cond.unnest() - if isinstance(inner, exp.Column): - new_conjuncts.append(inner.copy().eq(exp.true())) - changed = True - elif isinstance(inner, exp.Not) and isinstance(inner.this.unnest(), exp.Column): - new_conjuncts.append(inner.this.unnest().copy().eq(exp.false())) + normalized = cls._normalize_single_boolean_conjunct(cond) + if normalized is not None: + new_conjuncts.append(normalized) changed = True else: new_conjuncts.append(cond) @@ -1391,21 +1411,22 @@ def _extract_or_key_rows_from_where( col_to_val: t.Dict[str, exp.Expr] = {} col_to_expr: t.Dict[str, exp.Expr] = {} for item in branch_conjuncts: - item_unwrapped = item.unnest() + normalized_bool = self._normalize_single_boolean_conjunct(item) + item_unwrapped = normalized_bool if normalized_bool is not None else item.unnest() if isinstance(item_unwrapped, exp.EQ): left, right = item_unwrapped.this, item_unwrapped.expression - if isinstance(left, (exp.Column, exp.Identifier)) and isinstance( - right, (exp.Literal, exp.Boolean, exp.Null) - ): + if isinstance( + left, (exp.Column, exp.Identifier) + ) and self._is_literal_key_value(right): col_expr = ( left.copy() if isinstance(left, exp.Column) else exp.Column(this=left.copy()) ) val_expr = right.copy() - elif isinstance(right, (exp.Column, exp.Identifier)) and isinstance( - left, (exp.Literal, exp.Boolean, exp.Null) - ): + elif isinstance( + right, (exp.Column, exp.Identifier) + ) and self._is_literal_key_value(left): col_expr = ( right.copy() if isinstance(right, exp.Column) @@ -1468,13 +1489,13 @@ def _get_in_subquery(in_expr: exp.In) -> t.Optional[exp.Subquery]: query = in_expr.args.get("query") if isinstance(query, exp.Subquery): return query.copy() - if isinstance(query, exp.Select): + if isinstance(query, (exp.Select, exp.SetOperation)): return exp.Subquery(this=query.copy()) if len(in_expr.expressions) == 1: first = in_expr.expressions[0] if isinstance(first, exp.Subquery): return first.copy() - if isinstance(first, exp.Select): + if isinstance(first, (exp.Select, exp.SetOperation)): return exp.Subquery(this=first.copy()) return None @@ -1492,6 +1513,7 @@ def _execute_delete_with_subquery( return join_conditions: t.List[exp.Expr] = [] + subquery = subquery.copy() first_select = subquery.this while isinstance(first_select, exp.SetOperation): first_select = first_select.this @@ -1505,15 +1527,12 @@ def _execute_delete_with_subquery( candidate_name = sel_item.alias_or_name if candidate_name and candidate_name.lower() not in used_t2_aliases: t2_col_name = candidate_name - elif isinstance(subquery.this, exp.Select): + else: alias_name = f"_col_{idx}" - subquery = subquery.copy() - sub_select = t.cast(exp.Select, subquery.this) - new_exprs = list(sub_select.expressions) + new_exprs = list(first_select.expressions) inner_expr = sel_item.this if isinstance(sel_item, exp.Alias) else sel_item new_exprs[idx] = exp.alias_(inner_expr.copy(), alias_name, quoted=True) - sub_select.set("expressions", new_exprs) - first_select = sub_select + first_select.set("expressions", new_exprs) t2_col_name = alias_name used_t2_aliases.add(t2_col_name.lower()) t2_col = exp.column(t2_col_name, table="_t2") @@ -1788,6 +1807,13 @@ def _create_table_from_source_queries( # `EngineAdapter.replace_query` creates the empty table via `_create_table_from_columns` # first and then calls `_create_table_from_source_queries(..., replace=False)`. Executing # `CREATE TABLE IF NOT EXISTS ... AS SELECT` would be a no-op, so we insert the source queries. + if is_internal_iceberg and not target_columns_to_types_known: + raise SQLMeshError( + f"Doris does not support `CREATE TABLE ... AS SELECT` with `ENGINE=iceberg` on table '{table_name}' " + "without known column types. Either qualify the table with an external Iceberg catalog " + "(`..
`) or specify explicit column types." + ) + if ( is_internal_iceberg or (self_referencing and not replace) ) and target_columns_to_types_known: @@ -2030,6 +2056,33 @@ def _build_create_table_exp( **kwargs, ) + @staticmethod + def _is_integer_indicator_projection(proj: exp.Expr) -> bool: + inner = proj.this if isinstance(proj, exp.Alias) else proj + if isinstance(inner, exp.Literal) and inner.is_int: + return True + if isinstance(inner, exp.If): + t_val, f_val = inner.args.get("true"), inner.args.get("false") + return bool( + isinstance(t_val, exp.Literal) + and t_val.is_int + and isinstance(f_val, exp.Literal) + and f_val.is_int + ) + if isinstance(inner, exp.Case): + ifs = inner.args.get("ifs") or [] + default = inner.args.get("default") + return bool( + ifs + and isinstance(default, exp.Literal) + and default.is_int + and all( + isinstance(i.args.get("true"), exp.Literal) and i.args["true"].is_int + for i in ifs + ) + ) + return False + def _reorder_unannotated_ctas_query( self, expression: exp.Expr, kwargs: t.Dict[str, t.Any] ) -> exp.Expr: @@ -2058,6 +2111,47 @@ def _reorder_unannotated_ctas_query( "expressions", [proj_by_name[c.lower()] for c in reordered_select_cols], ) + elif any(isinstance(proj, exp.Star) for proj in expression.expressions): + table_props = kwargs.get("table_properties") or {} + has_explicit_key = bool( + table_props.get("unique_key") + or kwargs.get("unique_key") + or table_props.get("duplicate_key") + or table_props.get("aggregate_key") + ) + with_clause = expression.args.get("with_") + from_clause = expression.args.get("from_") + if ( + not has_explicit_key + and isinstance(with_clause, exp.With) + and isinstance(from_clause, exp.From) + and isinstance(from_clause.this, exp.Table) + ): + target_cte_name = from_clause.this.name.lower() + for cte_idx, cte in enumerate(with_clause.expressions): + if ( + isinstance(cte, exp.CTE) + and cte.alias.lower() == target_cte_name + and isinstance(cte.this, exp.Select) + and len(cte.this.expressions) >= 2 + and not self._is_integer_indicator_projection(cte.this.expressions[0]) + ): + indicator_idx = next( + ( + idx + for idx, p in enumerate(cte.this.expressions) + if idx > 0 and self._is_integer_indicator_projection(p) + ), + None, + ) + if indicator_idx is not None: + expression = expression.copy() + copied_cte = expression.args["with_"].expressions[cte_idx] + copied_select = copied_cte.this + projs = list(copied_select.expressions) + indicator_proj = projs.pop(indicator_idx) + copied_select.set("expressions", [indicator_proj, *projs]) + break return expression @classmethod @@ -2863,6 +2957,21 @@ def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> if select_node.args.get("locks"): select_node.set("locks", None) logger.debug("[Doris] Removed unsupported FOR UPDATE lock from SELECT statement") + if any( + isinstance(is_node.expression, exp.Boolean) for is_node in expression.find_all(exp.Is) + ): + expression = expression.copy() + for is_node in list(expression.find_all(exp.Is)): + if isinstance(is_node.expression, exp.Boolean): + eq_expr: exp.Expr = exp.EQ( + this=is_node.this.copy(), expression=is_node.expression.copy() + ) + if isinstance(is_node.parent, exp.Not): + eq_expr = exp.Paren(this=eq_expr) + if is_node is expression: + expression = eq_expr + else: + is_node.replace(eq_expr) if expression.find(exp.DataType): expression = expression.copy() for dtype in reversed(list(expression.find_all(exp.DataType))): diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 59c826f165..099368cf50 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -2380,3 +2380,83 @@ def test_unannotated_view_column_comments_and_starrocks_property_parity( target_columns_to_types=cols, table_properties={"distributed_by": exp.Literal.string("HASH BUCKETS 10")}, ) + + +def test_is_bool_rewrites_negative_or_keys_union_delete_and_table_diff_ctas( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. `IS TRUE` / `IS FALSE` / `IS NOT TRUE` in UPDATE (e.g. SnapshotState.unpause_snapshots) and DELETE + adapter.update_table( + "sqlmesh._snapshots", + {"unpaused_ts": 1700000000}, + where=exp.and_( + exp.column("forward_only").is_(exp.true()), + exp.column("name").eq(exp.Literal.string("db.m1")), + ), + ) + adapter.delete_from( + "sqlmesh._intervals", + where=exp.and_( + exp.column("is_dev").is_(exp.true()).not_(), + exp.column("name").eq(exp.Literal.string("db.m1")), + ), + ) + + # 2. Top-level OR of key rows with negative number literals and bare boolean / IS TRUE conjuncts + adapter.delete_from( + "sqlmesh._intervals", + where=parse_one( + "(id = -1 AND is_dev) OR (id = -2 AND NOT is_dev)", + read="doris", + ), + ) + + # 3. DELETE ... USING with a UNION ALL subquery containing duplicate/unaliased projection names + union_where = exp.Tuple(expressions=[exp.column("k1"), exp.column("k2")]).isin( + query=parse_one( + "SELECT a.id, b.id FROM db.a AS a JOIN db.b AS b ON a.id = b.id " + "UNION ALL " + "SELECT c.id, d.id FROM db.c AS c JOIN db.d AS d ON c.id = d.id", + read="doris", + ) + ) + adapter.delete_from("db.target_tbl", where=union_where) + + # 4. TableDiff-style unannotated CTAS (`SELECT *, CASE ... AS row_full_match FROM __stats`) + # promotes the integer indicator column (`s_exists`) to index 0 of `__stats` so string/float + # first columns do not fail Doris's implicit DUPLICATE KEY rule. + table_diff_query = parse_one( + "WITH `__source` AS (SELECT `s`.`name` AS `name`, `s`.`id` AS `__sqlmesh_join_key` FROM `db`.`src` AS `s`), " + "`__target` AS (SELECT `t`.`name` AS `name`, `t`.`id` AS `__sqlmesh_join_key` FROM `db`.`tgt` AS `t`), " + "`__stats` AS (" + "SELECT `s`.`name` AS `s__name`, `t`.`name` AS `t__name`, " + "IF(NOT `s`.`__sqlmesh_join_key` IS NULL, 1, 0) AS `s_exists`, " + "IF(NOT `t`.`__sqlmesh_join_key` IS NULL, 1, 0) AS `t_exists` " + "FROM `__source` AS `s` LEFT JOIN `__target` AS `t` ON `s`.`__sqlmesh_join_key` = `t`.`__sqlmesh_join_key`" + ") " + "SELECT *, CASE WHEN `s__name` = `t__name` THEN 1 ELSE 0 END AS `row_full_match` FROM `__stats`", + read="doris", + ) + adapter.ctas("sqlmesh_temp.diff", table_diff_query) + + assert to_sql_calls(adapter) == [ + "UPDATE `sqlmesh`.`_snapshots` SET `unpaused_ts` = 1700000000 WHERE `forward_only` = TRUE AND `name` = 'db.m1'", + "DELETE FROM `sqlmesh`.`_intervals` WHERE `is_dev` = FALSE AND `name` = 'db.m1'", + "DELETE FROM `sqlmesh`.`_intervals` WHERE `id` = -1 AND `is_dev` IN (TRUE)", + "DELETE FROM `sqlmesh`.`_intervals` WHERE `id` = -2 AND `is_dev` IN (FALSE)", + "DELETE FROM `db`.`target_tbl` AS `_t1` USING (SELECT `a`.`id`, `b`.`id` AS `_col_1` FROM `db`.`a` AS `a` JOIN `db`.`b` AS `b` ON `a`.`id` = `b`.`id` UNION ALL SELECT `c`.`id`, `d`.`id` FROM `db`.`c` AS `c` JOIN `db`.`d` AS `d` ON `c`.`id` = `d`.`id`) AS `_t2` WHERE `_t1`.`k1` = `_t2`.`id` AND `_t1`.`k2` = `_t2`.`_col_1`", + "CREATE TABLE IF NOT EXISTS `sqlmesh_temp`.`diff` AS WITH `__source` AS (SELECT `s`.`name` AS `name`, `s`.`id` AS `__sqlmesh_join_key` FROM `db`.`src` AS `s`), `__target` AS (SELECT `t`.`name` AS `name`, `t`.`id` AS `__sqlmesh_join_key` FROM `db`.`tgt` AS `t`), `__stats` AS (SELECT CASE WHEN NOT `s`.`__sqlmesh_join_key` IS NULL THEN 1 ELSE 0 END AS `s_exists`, `s`.`name` AS `s__name`, `t`.`name` AS `t__name`, CASE WHEN NOT `t`.`__sqlmesh_join_key` IS NULL THEN 1 ELSE 0 END AS `t_exists` FROM `__source` AS `s` LEFT JOIN `__target` AS `t` ON `s`.`__sqlmesh_join_key` = `t`.`__sqlmesh_join_key`) SELECT *, CASE WHEN `s__name` = `t__name` THEN 1 ELSE 0 END AS `row_full_match` FROM `__stats`", + ] + + # 5. Unannotated 2-part/internal Iceberg CTAS raises a clear SQLMeshError instead of silently dropping ENGINE=iceberg + with pytest.raises( + SQLMeshError, + match="Doris does not support `CREATE TABLE ... AS SELECT` with `ENGINE=iceberg`", + ): + adapter.ctas( + "silver.unannotated_iceberg", + parse_one("SELECT id, ds FROM bronze.events", read="doris"), + table_format="iceberg", + ) From eb1ad0f9b4f38bdfe81401e62808a45a0dd4665f Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Fri, 23 Oct 2026 08:24:17 -0700 Subject: [PATCH 21/27] docs(doris): document StarRocks and dbt-doris physical_properties aliases Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/integrations/engines/doris.md | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index 8617491142..db6c471428 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -181,7 +181,7 @@ Doris requires a `DISTRIBUTED BY` clause on every OLAP table. SQLMesh supports t -- or 'RANDOM BUCKETS AUTO' ) ``` -3. **Structured dictionary `distributed_by`**: +3. **Structured dictionary or column tuple `distributed_by`**: ```sql physical_properties ( distributed_by = ( @@ -189,6 +189,8 @@ Doris requires a `DISTRIBUTED BY` clause on every OLAP table. SQLMesh supports t expressions = 'user_id', buckets = 'AUTO' ) + -- Also supports `columns = ('user_id', 'event_time')` or `distributed_by = ('user_id', 'event_time')`, + -- or `random = 'true'` for `DISTRIBUTED BY RANDOM` on DUPLICATE KEY tables. ) ``` @@ -200,7 +202,7 @@ If none of `clustered_by` or `distributed_by` is provided: ## Partitioning -Doris OLAP tables support **Multi-Range Batch Partitioning**, **Explicit Partition Definitions**, **Dynamic Partitioning**, and **Auto Partitioning** (Doris 2.1+). +Doris OLAP tables support **Multi-Range Batch Partitioning**, **Explicit Partition Definitions**, **Dynamic Partitioning**, and **Auto Partitioning** (Doris 2.1+). You can specify the partition expression via the top-level `partitioned_by` model property or via `partitioned_by` / `partition_by` inside `physical_properties`. !!! note "OLAP Partitioning Requirements" - On Doris OLAP tables, `PARTITION BY` requires at least one of `partitions`, `auto_partition = 'true'`, or `'dynamic_partition.enable' = 'true'` in `physical_properties`. If none is provided, SQLMesh omits the `PARTITION BY` clause so Doris does not create a zero-partition table that rejects inserts. @@ -467,7 +469,7 @@ GROUP BY user_id, CAST(event_time AS DATE); !!! note "Automatic Physical Table Resolution & Non-Destructive Evaluation" - Doris materialized views are only created on initial deployment or when a rebuild is required (`RECREATE_MATERIALIZED_VIEW_ON_EVALUATION = False`), preserving MTMV refresh state across routine `sqlmesh run` invocations. - - When `excluded_trigger_tables` (on async MTMVs) or `source_table` (on synchronous MVs) references a SQLMesh-managed model, SQLMesh automatically resolves the model name to its underlying physical snapshot table name. + - When `excluded_trigger_tables` (supports a single string, comma-separated string, or tuple `('db.t1', 'db.t2')` on async MTMVs) or `source_table` (on synchronous MVs) references a SQLMesh-managed model, SQLMesh automatically resolves the model name to its underlying physical snapshot table name. ### Synchronous Materialized Views (`materialized_type = 'SYNC'`) @@ -515,16 +517,18 @@ FROM raw.logs; --- -## Migrating from `dbt-doris` +## Migrating from `dbt-doris` or StarRocks -If you are migrating an existing `dbt-doris` project to native SQLMesh models, map your `dbt-doris` config keys as follows: +If you are migrating an existing `dbt-doris` project or StarRocks SQLMesh models to Doris, map your configuration keys as follows: -| `dbt-doris` Config | SQLMesh Model Equivalent | +| `dbt-doris` / StarRocks Config | SQLMesh Doris Equivalent | | --- | --- | | `materialized='table'` | `kind FULL` | | `materialized='incremental'` (with `unique_key`) | `kind INCREMENTAL_BY_UNIQUE_KEY (unique_key ...)` | | `materialized='incremental'` (partition overwrite) | `kind INCREMENTAL_BY_TIME_RANGE (time_column ...)` or `kind INCREMENTAL_BY_PARTITION` | | `duplicate_key=['c1', 'c2']` | `physical_properties (duplicate_key = ('c1', 'c2'))` | | `unique_key=['c1']` | `kind INCREMENTAL_BY_UNIQUE_KEY (unique_key c1)` or `physical_properties (unique_key = ('c1'))` | -| `distributed_by=['c1']`, `buckets=10` | `clustered_by (c1), physical_properties (buckets = 10)` | +| `table_type='PRIMARY' \| 'UNIQUE' \| 'DUPLICATE' \| 'AGGREGATE'`, `keys=['c1', 'c2']` | Supported directly in `physical_properties` (`PRIMARY` maps to Doris `UNIQUE KEY`) | +| `distributed_by=['c1']`, `buckets=10` | `clustered_by (c1), physical_properties (buckets = 10)` or `physical_properties (distributed_by = ('c1'), buckets = 10)` | +| `partition_by` | Supported in `physical_properties (partition_by = ...)` or top-level `partitioned_by` | | `properties={...}` | Key-value entries inside `physical_properties (...)` | From f40cbb34232af8555f744ecb71c45977e079381c Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:44:26 -0700 Subject: [PATCH 22/27] fix(doris): support table_type/keys, random dist flag, MV exists, and tuple ref resolution Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/doris.py | 121 +++++++++++++---- sqlmesh/core/model/definition.py | 36 +++++- tests/core/engine_adapter/test_doris.py | 164 ++++++++++++++++++++++++ 3 files changed, 293 insertions(+), 28 deletions(-) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 5fc9c6f6db..ab48299c41 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -872,6 +872,7 @@ def _create_materialized_view( this=schema, kind="VIEW", replace=False, + exists=bool(create_kwargs.get("exists", False)) if not is_sync_mv else False, expression=query, properties=properties_exp, ), @@ -1703,6 +1704,59 @@ def _to_expr(val: t.Any) -> exp.Expr: preds = [*prefix_preds, batch_col.is_(exp.null()), *extra_conditions] self.execute(exp.delete(table_name, where=exp.and_(*preds))) + @classmethod + def _resolve_table_type_and_keys( + cls, table_props: t.Dict[str, t.Any], pop: bool = False + ) -> t.Tuple[t.Any, t.Any, t.Any]: + get_fn = table_props.pop if pop else table_props.get + unique_key = get_fn("unique_key", None) + duplicate_key = get_fn("duplicate_key", None) + aggregate_key = get_fn("aggregate_key", None) + raw_table_type = get_fn("table_type", None) + raw_keys = get_fn("keys", None) + + if raw_table_type is not None and raw_keys is None: + raise SQLMeshError( + "Doris property 'table_type' requires 'keys' to be specified in physical_properties." + ) + if raw_keys is not None: + tt_str = ( + (cls._extract_str_value(raw_table_type) or "DUPLICATE") + .strip() + .strip("'\"`") + .upper() + ) + if tt_str in ( + "PRIMARY", + "PRIMARY_KEY", + "PRIMARY KEY", + "UNIQUE", + "UNIQUE_KEY", + "UNIQUE KEY", + ): + if unique_key is not None: + raise SQLMeshError( + "Cannot specify both 'unique_key' and 'table_type'/'keys' in physical_properties." + ) + unique_key = raw_keys + elif tt_str in ("DUPLICATE", "DUPLICATE_KEY", "DUPLICATE KEY"): + if duplicate_key is not None: + raise SQLMeshError( + "Cannot specify both 'duplicate_key' and 'table_type'/'keys' in physical_properties." + ) + duplicate_key = raw_keys + elif tt_str in ("AGGREGATE", "AGGREGATE_KEY", "AGGREGATE KEY"): + if aggregate_key is not None: + raise SQLMeshError( + "Cannot specify both 'aggregate_key' and 'table_type'/'keys' in physical_properties." + ) + aggregate_key = raw_keys + else: + raise SQLMeshError( + f"Invalid Doris table_type '{tt_str}'. Expected PRIMARY, UNIQUE, DUPLICATE, or AGGREGATE." + ) + return unique_key, duplicate_key, aggregate_key + def adjust_physical_properties_for_incremental( self, physical_properties: t.Dict[str, t.Any], @@ -1714,11 +1768,14 @@ def adjust_physical_properties_for_incremental( """Promote ``unique_key`` from ``INCREMENTAL_BY_UNIQUE_KEY`` models into ``physical_properties`` when no explicit OLAP key model (``unique_key``, ``duplicate_key``, ``aggregate_key``) is set. """ + unique_key_prop, duplicate_key_prop, aggregate_key_prop = self._resolve_table_type_and_keys( + physical_properties + ) if ( unique_key - and "unique_key" not in physical_properties - and "duplicate_key" not in physical_properties - and "aggregate_key" not in physical_properties + and not unique_key_prop + and not duplicate_key_prop + and not aggregate_key_prop and all(isinstance(k, (exp.Column, exp.Identifier)) for k in unique_key) ): cleaned_keys = [exp.column(k.name) for k in unique_key] @@ -1739,12 +1796,15 @@ def _create_table_from_columns( **kwargs: t.Any, ) -> None: table_properties = dict(kwargs.get("table_properties") or {}) + unique_key_prop, duplicate_key_prop, aggregate_key_prop = self._resolve_table_type_and_keys( + table_properties + ) if ( primary_key - and "unique_key" not in table_properties - and "duplicate_key" not in table_properties - and "aggregate_key" not in table_properties + and not unique_key_prop + and not duplicate_key_prop + and not aggregate_key_prop ): table_properties["unique_key"] = exp.Tuple( expressions=[exp.to_column(col) for col in primary_key] @@ -1983,11 +2043,8 @@ def _build_create_table_exp( for expr in parsed_part_exprs for col in expr.find_all(exp.Column) } - strict_key = ( - table_props.get("unique_key") - or kwargs.get("unique_key") - or table_props.get("aggregate_key") - ) + resolved_uk, _, resolved_ak = self._resolve_table_type_and_keys(table_props) + strict_key = resolved_uk or kwargs.get("unique_key") or resolved_ak strict_key_set = {k.lower() for k in self._extract_column_names(strict_key)} if strict_key_set and not auto_range_cols.issubset(strict_key_set): auto_range_cols = set() @@ -2113,11 +2170,11 @@ def _reorder_unannotated_ctas_query( ) elif any(isinstance(proj, exp.Star) for proj in expression.expressions): table_props = kwargs.get("table_properties") or {} + resolved_uk, resolved_dk, resolved_ak = self._resolve_table_type_and_keys( + table_props + ) has_explicit_key = bool( - table_props.get("unique_key") - or kwargs.get("unique_key") - or table_props.get("duplicate_key") - or table_props.get("aggregate_key") + resolved_uk or kwargs.get("unique_key") or resolved_dk or resolved_ak ) with_clause = expression.args.get("with_") from_clause = expression.args.get("from_") @@ -2191,9 +2248,8 @@ def _reorder_olap_columns( raise SQLMeshError( "Multiple partition definitions found: specify only one of model 'partitioned_by' or physical_properties 'partitioned_by' / 'partition_by'." ) - unique_key = table_props.get("unique_key") or kwargs.get("unique_key") - duplicate_key = table_props.get("duplicate_key") - aggregate_key = table_props.get("aggregate_key") + resolved_uk, duplicate_key, aggregate_key = self._resolve_table_type_and_keys(table_props) + unique_key = resolved_uk or kwargs.get("unique_key") specified_key_types = [ name @@ -2601,9 +2657,10 @@ def _build_table_properties_exp( raise SQLMeshError( "Invalid property 'partition' in physical_properties; did you mean 'partition_by', 'partitioned_by', or 'partitions'?" ) - unique_key = table_properties_copy.pop("unique_key", None) or kwargs.get("unique_key") - duplicate_key = table_properties_copy.pop("duplicate_key", None) - aggregate_key = table_properties_copy.pop("aggregate_key", None) + resolved_uk, duplicate_key, aggregate_key = self._resolve_table_type_and_keys( + table_properties_copy, pop=True + ) + unique_key = resolved_uk or kwargs.get("unique_key") order_by = table_properties_copy.pop("order_by", None) unique_key_cols = self._extract_column_names(unique_key) @@ -3155,6 +3212,7 @@ def _build_distributed_by_exp( if distributed_info: raw_kind = distributed_info.get("kind") + raw_random = distributed_info.get("random") expressions = ( distributed_info.get("expressions") if distributed_info.get("expressions") is not None @@ -3162,6 +3220,16 @@ def _build_distributed_by_exp( ) buckets = distributed_info.get("buckets", top_level_buckets) + if raw_random is not None: + rand_str = ( + str(self._extract_str_value(raw_random) or raw_random).strip().lower() + ) + if rand_str in ("true", "1"): + if raw_kind is not None and str(raw_kind).strip().upper() != "RANDOM": + raise SQLMeshError( + f"Conflicting distributed_by configuration: kind='{raw_kind}' with random=true." + ) + raw_kind = "RANDOM" if raw_kind is None and expressions is not None: raw_kind = "HASH" @@ -3279,6 +3347,8 @@ def _properties_to_expressions(self, properties: t.Dict[str, t.Any]) -> t.List[e "unique_key", "duplicate_key", "aggregate_key", + "table_type", + "keys", "order_by", "engine", "properties", @@ -3289,6 +3359,13 @@ def _properties_to_expressions(self, properties: t.Dict[str, t.Any]) -> t.List[e if key in skip_keys or value is None: continue raw_str = self._extract_str_value(value) - str_literal = exp.Literal.string(raw_str if raw_str is not None else str(value)) + val_str = raw_str if raw_str is not None else str(value) + if key == "excluded_trigger_tables" and val_str: + val_str = ",".join( + re.sub(r"(?i)^`?internal`?\.", "", part.strip()) + for part in val_str.split(",") + if part.strip() + ) + str_literal = exp.Literal.string(val_str) expressions.append(exp.Property(this=exp.Literal.string(str(key)), value=str_literal)) return expressions diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py index e8e122dece..6eedc65fd1 100644 --- a/sqlmesh/core/model/definition.py +++ b/sqlmesh/core/model/definition.py @@ -737,7 +737,10 @@ def render_physical_properties(self, **render_kwargs: t.Any) -> t.Dict[str, t.An ) for key in keys_to_resolve: rendered[key] = _resolve_model_refs_to_physical_tables( - rendered[key], table_mapping, self.dialect + rendered[key], + table_mapping, + self.dialect, + default_catalog=self.default_catalog, ) return rendered @@ -2830,23 +2833,44 @@ def _split_sql_model_statements( def _resolve_model_refs_to_physical_tables( - value: exp.Expr, table_mapping: t.Dict[str, str], dialect: DialectType + value: exp.Expr, + table_mapping: t.Dict[str, str], + dialect: DialectType, + default_catalog: t.Optional[str] = None, ) -> exp.Literal: """Resolve managed-model references in a property value to their physical table names. - The value is a single table reference or a comma-separated list of them. Each reference that - matches a managed model (via ``table_mapping``) is swapped for its physical ``db.table`` name; - anything else (e.g. a raw source) is kept as written. Returns a single string literal so the - property renders just like a hand-written value. + The value is a single table reference, a comma-separated string of references, or a tuple/array + of references. Each reference that matches a managed model (via ``table_mapping``) is swapped + for its physical ``db.table`` name; anything else (e.g. a raw source) is kept as written. + Returns a single string literal so the property renders just like a hand-written value. """ + if isinstance(value, exp.Paren): + value = value.this + if isinstance(value, exp.Literal) and value.is_string: refs = value.this.split(",") + elif isinstance(value, (exp.Tuple, exp.Array)): + refs = [ + part + for elem in value.expressions + for part in ( + elem.this.split(",") + if isinstance(elem, exp.Literal) and elem.is_string + else [elem.sql(dialect=dialect)] + ) + ] else: refs = [value.sql(dialect=dialect)] def resolve(ref: str) -> str: table = exp.to_table(ref.strip(), dialect=dialect) physical = table_mapping.get(exp.table_name(table, identify=True)) + if not physical and default_catalog and not table.catalog: + normalized_name = d.normalize_model_name( + table, default_catalog=default_catalog, dialect=dialect + ) + physical = table_mapping.get(normalized_name) # Managed model -> physical table; otherwise keep the reference (just unquoted/normalized). return exp.table_name( exp.to_table(physical, dialect=dialect) if physical else table, identify=False diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 099368cf50..bca25f1bbf 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -2460,3 +2460,167 @@ def test_is_bool_rewrites_negative_or_keys_union_delete_and_table_diff_ctas( parse_one("SELECT id, ds FROM bronze.events", read="doris"), table_format="iceberg", ) + + +def test_table_type_keys_random_distribution_mv_exists_and_tuple_ref_resolution( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + from sqlmesh.core.dialect import parse + from sqlmesh.core.model import SqlModel, load_sql_based_model + from sqlmesh.core.snapshot import Snapshot, SnapshotChangeCategory + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + cols = { + "payload": exp.DataType.build("VARCHAR(50)"), + "id": exp.DataType.build("INT"), + "dt": exp.DataType.build("DATE"), + } + + # 1. StarRocks-style `table_type` ('PRIMARY' | 'UNIQUE' | 'DUPLICATE' | 'AGGREGATE') + `keys` + adapter.create_table( + "db.sr_primary", + target_columns_to_types=cols, + table_properties={ + "table_type": exp.Literal.string("PRIMARY"), + "keys": exp.Tuple(expressions=[exp.to_column("id"), exp.to_column("dt")]), + "buckets": exp.Literal.number(8), + }, + ) + adapter.create_table( + "db.sr_duplicate_random", + target_columns_to_types=cols, + table_properties={ + "table_type": exp.Literal.string("DUPLICATE"), + "keys": exp.Tuple(expressions=[exp.to_column("id")]), + "distributed_by": exp.Tuple( + expressions=[ + exp.EQ(this=exp.column("random"), expression=exp.Literal.string("true")), + exp.EQ(this=exp.column("buckets"), expression=exp.Literal.number(6)), + ] + ), + }, + ) + adapter.create_table( + "db.sr_aggregate", + target_columns_to_types={"id": exp.DataType.build("INT"), "dt": exp.DataType.build("DATE")}, + table_properties={ + "table_type": exp.Literal.string("AGGREGATE"), + "keys": exp.Literal.string("(id, dt)"), + "buckets": exp.Literal.number(4), + }, + ) + + # 2. `create_view(..., materialized=True, replace=False, exists=True)` emits IF NOT EXISTS + adapter.create_view( + "db.mv_if_not_exists", + parse_one("SELECT id FROM db.sr_primary", read="doris"), + replace=False, + materialized=True, + exists=True, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`sr_primary` (`id` INT, `dt` DATE, `payload` VARCHAR(50)) UNIQUE KEY (`id`, `dt`) DISTRIBUTED BY HASH (`id`) BUCKETS 8", + "CREATE TABLE IF NOT EXISTS `db`.`sr_duplicate_random` (`id` INT, `payload` VARCHAR(50), `dt` DATE) DUPLICATE KEY (`id`) DISTRIBUTED BY RANDOM BUCKETS 6", + "CREATE TABLE IF NOT EXISTS `db`.`sr_aggregate` (`id` INT, `dt` DATE) AGGREGATE KEY (`id`, `dt`) DISTRIBUTED BY HASH (`id`) BUCKETS 4", + "CREATE MATERIALIZED VIEW IF NOT EXISTS `db`.`mv_if_not_exists` AS SELECT `id` FROM `db`.`sr_primary`", + ] + + # 3. Validation errors for invalid `table_type` / missing `keys` / conflicting `random=true` + with pytest.raises(SQLMeshError, match="requires 'keys' to be specified"): + adapter.create_table( + "db.bad_tt_no_keys", + target_columns_to_types=cols, + table_properties={"table_type": exp.Literal.string("UNIQUE")}, + ) + + with pytest.raises(SQLMeshError, match="Invalid Doris table_type"): + adapter.create_table( + "db.bad_tt_val", + target_columns_to_types=cols, + table_properties={ + "table_type": exp.Literal.string("UNKNOWN_MODEL"), + "keys": exp.to_column("id"), + }, + ) + + with pytest.raises(SQLMeshError, match="Conflicting distributed_by configuration"): + adapter.create_table( + "db.bad_dist_conflict", + target_columns_to_types=cols, + table_properties={ + "distributed_by": exp.Tuple( + expressions=[ + exp.EQ(this=exp.column("kind"), expression=exp.Literal.string("HASH")), + exp.EQ(this=exp.column("random"), expression=exp.Literal.string("true")), + ] + ) + }, + ) + + # 4. Tuple `excluded_trigger_tables` + `default_catalog="internal"` resolves 2-part model refs + # and strips the leading `internal.` catalog prefix from the physical table name in Doris MTMV DDL. + adapter.cursor.execute.reset_mock() + base_model = t.cast( + SqlModel, + load_sql_based_model( + parse( + """ + MODEL ( + name doris_db.orders, + kind FULL, + dialect doris, + columns (id INT, amount INT) + ); + SELECT 1 AS id, 10 AS amount; + """, + default_dialect="doris", + ), + default_catalog="internal", + ), + ) + base_snapshot = Snapshot.from_node(base_model, nodes={}, ttl="in 1 week") + base_snapshot.categorize_as(SnapshotChangeCategory.BREAKING) + physical_tbl = exp.to_table(base_snapshot.table_name()) + expected_physical = f"{physical_tbl.db}.{physical_tbl.name}" + snapshots = {base_snapshot.name: base_snapshot} + + tuple_mv_model = t.cast( + SqlModel, + load_sql_based_model( + parse( + """ + MODEL ( + name doris_db.orders_tuple_mtmv, + kind VIEW (materialized true), + dialect doris, + columns (id INT, amount INT), + physical_properties ( + build = IMMEDIATE, + refresh = AUTO, + excluded_trigger_tables = ('doris_db.orders', 'ext_cat.ext_db.raw_events') + ) + ); + SELECT id, amount FROM doris_db.orders; + """, + default_dialect="doris", + ), + default_catalog="internal", + ), + ) + rendered_tuple_props = tuple_mv_model.render_physical_properties( + snapshots=snapshots, engine_adapter=adapter + ) + adapter.create_view( + tuple_mv_model.name, + tuple_mv_model.render_query() or parse_one("SELECT id, amount FROM doris_db.orders"), + replace=False, + materialized=True, + target_columns_to_types=tuple_mv_model.columns_to_types, + view_properties=rendered_tuple_props, + ) + tuple_mv_ddl = to_sql_calls(adapter)[-1] + assert ( + f"'excluded_trigger_tables'='{expected_physical},ext_cat.ext_db.raw_events'" in tuple_mv_ddl + ) From c20b8908caf0781aff4a7bf68b6da08cc1dd85f2 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Fri, 23 Oct 2026 18:58:26 -0700 Subject: [PATCH 23/27] fix(doris): time-range delete cast unwrap, OLAP key prefix, SCD2 self-ref & catalog reconnect Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/engine_adapter/base.py | 28 ++-- sqlmesh/core/engine_adapter/doris.py | 149 +++++++++++++++--- .../integration/test_integration_doris.py | 122 ++++++++++++++ tests/core/engine_adapter/test_doris.py | 138 ++++++++++++++++ 4 files changed, 402 insertions(+), 35 deletions(-) diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index 1ac852f4f7..8592972556 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -525,10 +525,23 @@ def replace_query( if not target_columns_to_types and table_exists: target_columns_to_types = self.columns(target_table) query = source_queries[0].query_factory() - self_referencing = any( - quote_identifiers(table) == quote_identifiers(target_table) - for table in query.find_all(exp.Table) - ) + + def _is_target_table(node: exp.Expr) -> bool: + if not isinstance(node, exp.Table): + return False + if quote_identifiers(node.copy()) == quote_identifiers(target_table.copy()): + return True + if self.default_catalog: + norm_node = node.copy() + if not norm_node.catalog and norm_node.db: + norm_node.set("catalog", exp.parse_identifier(self.default_catalog)) + norm_target = target_table.copy() + if not norm_target.catalog and norm_target.db: + norm_target.set("catalog", exp.parse_identifier(self.default_catalog)) + return quote_identifiers(norm_node) == quote_identifiers(norm_target) + return False + + self_referencing = any(_is_target_table(table) for table in query.find_all(exp.Table)) # If a query references itself then it must have a table created regardless of approach used. if self_referencing: if not target_columns_to_types: @@ -572,10 +585,7 @@ def replace_query( for source_query in source_queries: source_query.add_transform( lambda node: ( # type: ignore - temp_table # type: ignore - if isinstance(node, exp.Table) - and quote_identifiers(node) == quote_identifiers(target_table) - else node + temp_table if _is_target_table(node) else node # type: ignore ) ) return self._insert_overwrite_by_condition( @@ -2718,7 +2728,7 @@ def temp_table( Yields: The table expression """ - name = exp.to_table(name) + name = exp.to_table(name).copy() # ensure that we use default catalog if none is not specified if isinstance(name, exp.Table) and not name.catalog and name.db and self.default_catalog: name.set("catalog", exp.parse_identifier(self.default_catalog)) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index ab48299c41..2a5a594fe7 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -178,6 +178,22 @@ def get_current_catalog(self) -> t.Optional[str]: def set_current_catalog(self, catalog: str) -> None: catalog_sql = self._to_sql(exp.to_identifier(catalog, quoted=True)) self.execute(exp.Command(this="SWITCH", expression=catalog_sql)) + self._active_catalog: t.Optional[str] = catalog + self._needs_catalog_restore: bool = False + + def _restore_active_catalog_on_reconnect(self, sql: t.Optional[str] = None) -> None: + active_cat = getattr(self, "_active_catalog", None) + self._needs_catalog_restore = False + if not active_cat: + return + if sql and re.match( + r"\A(?:\s|/\*.*?\*/|--[^\r\n]*(?:[\r\n]|\Z))*SWITCH\b", + sql, + flags=re.IGNORECASE | re.DOTALL, + ): + return + catalog_sql = self._to_sql(exp.to_identifier(active_cat, quoted=True)) + self.cursor.execute(f"SWITCH {catalog_sql}") def _is_iceberg_table( self, @@ -294,7 +310,10 @@ def _fetch_native_df( except Exception: logger.info("Connection to the database was lost. Reconnecting...") self._connection_pool.close() + self._needs_catalog_restore = True self._connection_pool.get_cursor() + if getattr(self, "_needs_catalog_restore", False): + self._restore_active_catalog_on_reconnect() try: return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) except Exception as e: @@ -316,12 +335,15 @@ def _fetch_native_df( self._connection_pool.close() try: self._connection_pool.get_cursor() + self._restore_active_catalog_on_reconnect() return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) except Exception: self._connection_pool.close() + self._needs_catalog_restore = True raise if is_disconnect and not self._connection_pool.is_transaction_active: self._connection_pool.close() + self._needs_catalog_restore = True raise _SAFE_RETRY_SQL_PREFIX_RE = re.compile( @@ -344,6 +366,9 @@ def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any except Exception: logger.info("Connection to the database was lost. Reconnecting...") self._connection_pool.close() + self._needs_catalog_restore = True + if getattr(self, "_needs_catalog_restore", False): + self._restore_active_catalog_on_reconnect(sql) try: super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) except Exception as e: @@ -370,14 +395,20 @@ def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any self._connection_pool.close_cursor() try: conn = self._connection_pool.get() - conn.ping(reconnect=True) + try: + conn.ping(reconnect=True) + except Exception: + self._connection_pool.close() + self._restore_active_catalog_on_reconnect(sql) super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) except Exception: self._connection_pool.close() + self._needs_catalog_restore = True raise else: if is_disconnect and not self._connection_pool.is_transaction_active: self._connection_pool.close() + self._needs_catalog_restore = True raise def columns( @@ -1119,6 +1150,7 @@ def delete_from( if where == exp.false(): return + where = self._unwrap_literal_casts_in_where(where) where = self._rewrite_between_in_where(where) where = self._normalize_boolean_column_conjuncts(where) @@ -1189,6 +1221,22 @@ def _simplify_boolean_where(cls, expr: exp.Expr) -> exp.Expr: return expr.this return expr + @classmethod + def _unwrap_literal_casts_in_where(cls, where: exp.Expr) -> exp.Expr: + """Unwrap top-level `CAST( AS ...)` expressions in `DELETE WHERE` predicates + (e.g. from `to_time_column` in `insert_overwrite_by_time_partition`) because Doris's OLAP + `DeleteHandler` only accepts plain literals on the RHS of comparison and `IN` predicates. + """ + for cast_node in list(where.find_all(exp.Cast, exp.TryCast)): + if cast_node.find_ancestor(exp.Subquery, exp.Select): + continue + inner = cast_node.this + while isinstance(inner, (exp.Cast, exp.TryCast)): + inner = inner.this + if inner is not None and cls._is_literal_key_value(inner): + cast_node.replace(inner.copy()) + return where + @staticmethod def _rewrite_between_in_where(where: exp.Expr) -> exp.Expr: if isinstance(where, exp.Between): @@ -1853,12 +1901,23 @@ def _create_table_from_source_queries( target_columns_to_types and columns_to_types_all_known(target_columns_to_types) ) first_query = source_queries[0].query_factory() if source_queries else None + + def _matches_target_table(t_node: exp.Table) -> bool: + if quote_identifiers(t_node.copy()) == quote_identifiers(table.copy()): + return True + if self.default_catalog: + norm_node = t_node.copy() + if not norm_node.catalog and norm_node.db: + norm_node.set("catalog", exp.parse_identifier(self.default_catalog)) + norm_target = table.copy() + if not norm_target.catalog and norm_target.db: + norm_target.set("catalog", exp.parse_identifier(self.default_catalog)) + return quote_identifiers(norm_node) == quote_identifiers(norm_target) + return False + self_referencing = bool( first_query - and any( - quote_identifiers(t_node.copy()) == quote_identifiers(table.copy()) - for t_node in first_query.find_all(exp.Table) - ) + and any(_matches_target_table(t_node) for t_node in first_query.find_all(exp.Table)) ) # 1. Doris's CTAS grammar forbids `ENGINE=iceberg`, so internal/2-part Iceberg tables with known @@ -2299,6 +2358,28 @@ def _reorder_olap_columns( return {c: updated_types[c] for c in new_order} updated_types = dict(target_columns_to_types) + raw_part_by = ( + kwargs.get("partitioned_by") + or table_props.get("partitioned_by") + or table_props.get("partition_by") + ) + part_col_names: t.List[str] = [] + if raw_part_by: + part_list = raw_part_by if isinstance(raw_part_by, list) else [raw_part_by] + parsed_part_exprs, _ = self._parse_partition_expressions(part_list) + col_by_lower = {c.lower(): c for c in col_names} + for expr in parsed_part_exprs: + for col in expr.find_all(exp.Column): + canonical = col_by_lower.get(col.name.lower()) + if canonical and canonical not in part_col_names: + part_col_names.append(canonical) + if not self._is_valid_olap_key_type( + updated_types[canonical] + ) and self._is_coercible_string_key_type(updated_types[canonical]): + updated_types[canonical] = exp.DataType.build( + "VARCHAR(255)", dialect=self.dialect + ) + first_dtype = updated_types[col_names[0]] if not self._is_valid_olap_key_type(first_dtype): valid_col = next( @@ -2319,23 +2400,15 @@ def _reorder_olap_columns( if col_names[0] != coercible_col: col_names = [coercible_col] + [c for c in col_names if c != coercible_col] - # Ensure any OLAP partition columns with valid key types precede non-keyable columns - raw_part_by = ( - kwargs.get("partitioned_by") - or table_props.get("partitioned_by") - or table_props.get("partition_by") - ) - if raw_part_by: - part_list = raw_part_by if isinstance(raw_part_by, list) else [raw_part_by] - parsed_part_exprs, _ = self._parse_partition_expressions(part_list) - col_by_lower = {c.lower(): c for c in col_names} + # Ensure OLAP partition/time columns are included in Doris's implicit DUPLICATE KEY prefix: + # Doris stops implicit key inference at the first non-keyable column, immediately after the + # first VARCHAR column, or after 3 columns (`shortkey_max_column_count`) when no physical + # OLAP PARTITION BY clause is emitted. + if part_col_names: part_cols = [ - col_by_lower[col.name.lower()] - for expr in parsed_part_exprs - for col in expr.find_all(exp.Column) - if col.name.lower() in col_by_lower - and self._is_valid_olap_key_type(updated_types[col_by_lower[col.name.lower()]]) + c for c in part_col_names if self._is_valid_olap_key_type(updated_types[c]) ] + part_cols_set = set(part_cols) first_non_key_idx = next( ( i @@ -2344,16 +2417,40 @@ def _reorder_olap_columns( ), None, ) + first_varchar_idx = next( + ( + i + for i, c in enumerate(col_names) + if c not in part_cols_set and updated_types[c].is_type("varchar") + ), + None, + ) + auto_str = self._extract_str_value(table_props.get("auto_partition")) + dyn_str = self._extract_str_value(table_props.get("dynamic_partition.enable")) + has_physical_olap_partition = bool( + table_props.get("partitions") + or (auto_str and str(auto_str).strip().lower() in ("true", "1")) + or (dyn_str and str(dyn_str).strip().lower() in ("true", "1")) + ) + shortkey_limit_idx = ( + max(1, 3 - len(part_cols)) if not has_physical_olap_partition else None + ) + cutoff_candidates = [ + idx + for idx in (first_non_key_idx, first_varchar_idx, shortkey_limit_idx) + if idx is not None + ] + cutoff_idx = min(cutoff_candidates) if cutoff_candidates else None if ( part_cols - and first_non_key_idx is not None - and any(col_names.index(p) > first_non_key_idx for p in part_cols) + and cutoff_idx is not None + and any(col_names.index(p) >= cutoff_idx for p in part_cols) ): - moved_parts = [p for p in part_cols if col_names.index(p) > first_non_key_idx] + moved_parts = [p for p in part_cols if col_names.index(p) >= cutoff_idx] moved_set = set(moved_parts) - before_non_key = [c for c in col_names[:first_non_key_idx] if c not in moved_set] - after_non_key = [c for c in col_names[first_non_key_idx:] if c not in moved_set] - col_names = before_non_key + moved_parts + after_non_key + before_cutoff = [c for c in col_names[:cutoff_idx] if c not in moved_set] + after_cutoff = [c for c in col_names[cutoff_idx:] if c not in moved_set] + col_names = before_cutoff + moved_parts + after_cutoff if col_names != list(target_columns_to_types) or updated_types != target_columns_to_types: return {c: updated_types[c] for c in col_names} diff --git a/tests/core/engine_adapter/integration/test_integration_doris.py b/tests/core/engine_adapter/integration/test_integration_doris.py index cae796ec0a..13ac651ce0 100644 --- a/tests/core/engine_adapter/integration/test_integration_doris.py +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -363,3 +363,125 @@ def test_duplicate_key_delete_with_subquery_and_between(ctx: TestContext): ] ), ) + + +def test_insert_overwrite_by_time_partition_and_scd_type_2(ctx: TestContext): + from datetime import datetime + import pandas as pd # noqa: TID253 + from sqlmesh.utils.date import to_time_column + + # 1. insert_overwrite_by_time_partition with CAST(... AS DATE) bounds and partition column after VARCHAR + time_table = ctx.table("test_time_range_overwrite") + time_cols = { + "code": exp.DataType.build("VARCHAR(64)", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + "amount": exp.DataType.build("DOUBLE", dialect="doris"), + } + ctx.engine_adapter.create_table( + time_table, + time_cols, + partitioned_by=[exp.to_column("ds")], + ) + initial_df = pd.DataFrame( + [ + {"code": "a", "ds": "2022-01-01", "amount": 10.0}, + {"code": "b", "ds": "2022-01-02", "amount": 20.0}, + {"code": "c", "ds": "2022-01-03", "amount": 30.0}, + ] + ) + ctx.engine_adapter.insert_append( + time_table, + ctx.input_data(initial_df, time_cols), + target_columns_to_types=time_cols, + ) + replacement_df = pd.DataFrame( + [ + {"code": "b_new", "ds": "2022-01-02", "amount": 25.0}, + ] + ) + ctx.engine_adapter.insert_overwrite_by_time_partition( + time_table, + ctx.input_data(replacement_df, time_cols), + start="2022-01-02", + end="2022-01-02", + time_column="ds", + time_formatter=lambda v, _: to_time_column( + v, exp.DataType.build("DATE", dialect="doris"), "doris" + ), + target_columns_to_types=time_cols, + ) + rows = ctx.engine_adapter.fetchall( + exp.select("code", exp.cast(exp.column("ds"), "VARCHAR").as_("ds"), "amount") + .from_(time_table) + .order_by("ds") + ) + assert list(rows) == [ + ("a", "2022-01-01", 10.0), + ("b_new", "2022-01-02", 25.0), + ("c", "2022-01-03", 30.0), + ] + + # 2. scd_type_2_by_time across initial load and incremental update + scd_table = ctx.table("test_scd_type_2") + scd_src_table = ctx.table("test_scd_type_2_src") + scd_src_cols = { + "id": exp.DataType.build("INT", dialect="doris"), + "name": exp.DataType.build("VARCHAR(64)", dialect="doris"), + "updated_at": exp.DataType.build("DATETIME", dialect="doris"), + } + scd_target_cols = { + **scd_src_cols, + "valid_from": exp.DataType.build("DATETIME", dialect="doris"), + "valid_to": exp.DataType.build("DATETIME", dialect="doris"), + } + ctx.engine_adapter.create_table(scd_table, scd_target_cols) + batch1_df = pd.DataFrame( + [ + {"id": 1, "name": "alice_v1", "updated_at": "2022-01-01 00:00:00"}, + {"id": 2, "name": "bob_v1", "updated_at": "2022-01-01 00:00:00"}, + ] + ) + ctx.engine_adapter.ctas(scd_src_table, ctx.input_data(batch1_df, scd_src_cols), scd_src_cols) + ctx.engine_adapter.scd_type_2_by_time( + target_table=scd_table, + source_table=exp.select("id", "name", "updated_at").from_(scd_src_table), + unique_key=[exp.column("id")], + valid_from_col=exp.column("valid_from", quoted=True), + valid_to_col=exp.column("valid_to", quoted=True), + updated_at_col=exp.column("updated_at", quoted=True), + target_columns_to_types=scd_target_cols, + execution_time=datetime(2022, 1, 2, 0, 0, 0), + truncate=True, + ) + batch2_df = pd.DataFrame( + [ + {"id": 1, "name": "alice_v2", "updated_at": "2022-01-05 00:00:00"}, + {"id": 2, "name": "bob_v1", "updated_at": "2022-01-01 00:00:00"}, + ] + ) + ctx.engine_adapter.replace_query( + scd_src_table, + ctx.input_data(batch2_df, scd_src_cols), + target_columns_to_types=scd_src_cols, + ) + ctx.engine_adapter.scd_type_2_by_time( + target_table=scd_table, + source_table=exp.select("id", "name", "updated_at").from_(scd_src_table), + unique_key=[exp.column("id")], + valid_from_col=exp.column("valid_from", quoted=True), + valid_to_col=exp.column("valid_to", quoted=True), + updated_at_col=exp.column("updated_at", quoted=True), + target_columns_to_types=scd_target_cols, + execution_time=datetime(2022, 1, 6, 0, 0, 0), + truncate=False, + ) + scd_rows = ctx.engine_adapter.fetchall( + exp.select("id", "name", exp.column("valid_to").is_(exp.null()).as_("is_current")) + .from_(scd_table) + .order_by("id", "valid_from") + ) + assert list(scd_rows) == [ + (1, "alice_v1", 0), + (1, "alice_v2", 1), + (2, "bob_v1", 1), + ] diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index bca25f1bbf..00847ca54a 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -2624,3 +2624,141 @@ def test_table_type_keys_random_distribution_mv_exists_and_tuple_ref_resolution( assert ( f"'excluded_trigger_tables'='{expected_physical},ext_cat.ext_db.raw_events'" in tuple_mv_ddl ) + + +def test_time_range_delete_cast_unwrap_partition_key_prefix_scd_type_2_and_reconnect_catalog( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + from datetime import datetime + from sqlmesh.utils.date import to_time_column + + adapter = make_mocked_engine_adapter(DorisEngineAdapter, default_catalog="internal") + + # 1. insert_overwrite_by_time_partition with `to_time_column` (`CAST('...' AS DATE)`) unwraps + # CAST literals in top-level DELETE WHERE while preserving CAST in INSERT SELECT WHERE. + adapter.insert_overwrite_by_time_partition( + "db.events", + parse_one("SELECT id, ds FROM db.src_events", read="doris"), + start="2024-01-01", + end="2024-01-31", + time_column="ds", + time_formatter=lambda v, _: to_time_column( + v, exp.DataType.build("DATE", dialect="doris"), "doris" + ), + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + }, + ) + + time_overwrite_calls = to_sql_calls(adapter) + assert time_overwrite_calls == [ + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-01' AND `ds` <= '2024-01-31'", + "INSERT INTO `db`.`events` (`id`, `ds`) SELECT `id`, `ds` FROM (SELECT `id`, `ds` FROM `db`.`src_events`) AS `_subquery` WHERE `ds` BETWEEN CAST('2024-01-01' AS DATE) AND CAST('2024-01-31' AS DATE)", + ] + + # 2. Implicit DUPLICATE KEY prefix ordering for partitioned_by / time_column: + # (a) TEXT partition column coerced to VARCHAR(255) and moved before DOUBLE; + # (b) DATE partition column moved before VARCHAR column (which cuts off Doris implicit key); + # (c) DATE partition column moved within the first 3 columns (`shortkey_max_column_count = 3`). + adapter.cursor.execute.reset_mock() + adapter.create_table( + "db.part_text_col", + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "amount": exp.DataType.build("DOUBLE", dialect="doris"), + "ds": exp.DataType.build("TEXT", dialect="doris"), + }, + partitioned_by=[exp.to_column("ds")], + ) + adapter.create_table( + "db.part_after_varchar", + target_columns_to_types={ + "code": exp.DataType.build("VARCHAR(100)", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + "amount": exp.DataType.build("DOUBLE", dialect="doris"), + }, + partitioned_by=[exp.to_column("ds")], + ) + adapter.create_table( + "db.part_fourth_col", + target_columns_to_types={ + "c1": exp.DataType.build("INT", dialect="doris"), + "c2": exp.DataType.build("INT", dialect="doris"), + "c3": exp.DataType.build("INT", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + }, + partitioned_by=[exp.to_column("ds")], + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`part_text_col` (`id` INT, `ds` VARCHAR(255), `amount` DOUBLE)", + "CREATE TABLE IF NOT EXISTS `db`.`part_after_varchar` (`ds` DATE, `code` VARCHAR(100), `amount` DOUBLE)", + "CREATE TABLE IF NOT EXISTS `db`.`part_fourth_col` (`c1` INT, `c2` INT, `ds` DATE, `c3` INT)", + ] + + # 3. SCD Type 2 on an existing table when `default_catalog="internal"` does not mutate `target_table` + # in `temp_table` and properly rewrites `db.scd_tbl` references to `__temp_scd_tbl_...`. + adapter.cursor.execute.reset_mock() + mocker.patch("sqlmesh.core.engine_adapter.doris.random_id", return_value="abcd1234") + mocker.patch.object( + adapter, + "get_data_object", + return_value=DataObject( + catalog="internal", schema="db", name="scd_tbl", type=DataObjectType.TABLE + ), + ) + adapter.scd_type_2_by_time( + target_table="db.scd_tbl", + source_table=t.cast( + exp.Select, + parse_one("SELECT id, name, updated_at FROM db.src_tbl", read="doris"), + ), + unique_key=[exp.column("id")], + valid_from_col=exp.column("valid_from", quoted=True), + valid_to_col=exp.column("valid_to", quoted=True), + updated_at_col=exp.column("updated_at", quoted=True), + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "name": exp.DataType.build("VARCHAR(50)", dialect="doris"), + "updated_at": exp.DataType.build("DATETIME", dialect="doris"), + "valid_from": exp.DataType.build("DATETIME", dialect="doris"), + "valid_to": exp.DataType.build("DATETIME", dialect="doris"), + }, + execution_time=datetime(2024, 1, 2, 0, 0, 0), + ) + scd_calls = to_sql_calls(adapter) + insert_scd_sql = next(c for c in scd_calls if c.startswith("INSERT INTO `db`.`scd_tbl`")) + assert "FROM `internal`.`db`.`__temp_scd_tbl_abcd1234`" in insert_scd_sql + assert "FROM `db`.`scd_tbl`" not in insert_scd_sql + + # 4. Reconnect in `_execute` restores active catalog set via `set_current_catalog` and falls back + # to `self._connection_pool.close()` when `conn.ping(reconnect=True)` raises on a closed socket. + class InterfaceError(Exception): + pass + + adapter.cursor.execute.reset_mock() + adapter.set_current_catalog("iceberg_cat") + adapter.cursor.execute.reset_mock() + + conn_mock = mocker.MagicMock() + conn_mock.ping.side_effect = InterfaceError(0, "Already closed") + mocker.patch.object(adapter._connection_pool, "get", return_value=conn_mock) + close_spy = mocker.spy(adapter._connection_pool, "close") + + exec_attempt = 0 + + def reconnect_side_effect(sql, *args, **kwargs): + nonlocal exec_attempt + exec_attempt += 1 + if exec_attempt == 1: + raise InterfaceError(0, "Already closed") + + adapter.cursor.execute.side_effect = reconnect_side_effect + adapter._execute("SELECT 1") + + assert close_spy.call_count == 1 + assert exec_attempt == 1 + assert to_sql_calls(adapter) == [ + "SWITCH `iceberg_cat`", + "SELECT 1", + ] From 1724764772a25c75d97bc7a954be5a195220a5a1 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sat, 24 Oct 2026 19:00:56 -0700 Subject: [PATCH 24/27] fix(doris): thread-local state, partition range delete, TableDiff shortkey & init config Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/integrations/engines/doris.md | 16 + sqlmesh/cli/project_init.py | 7 + sqlmesh/core/engine_adapter/doris.py | 539 ++++++++++++++++-- .../integration/test_integration_doris.py | 45 ++ tests/core/engine_adapter/test_doris.py | 238 +++++++- 5 files changed, 793 insertions(+), 52 deletions(-) diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index db6c471428..13835b4e92 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -517,6 +517,22 @@ FROM raw.logs; --- +## Supported Model Kinds & Catalog Matrix + +| SQLMesh Model Kind | Doris Internal OLAP Catalog (`internal`) | External Apache Iceberg Catalog (`table_format iceberg`) | +| --- | :---: | :---: | +| `FULL` | Supported (`CREATE TABLE ... AS SELECT` / `DELETE` + `INSERT`) | Supported (`CREATE TABLE ... ENGINE=iceberg` / `INSERT OVERWRITE TABLE`) | +| `INCREMENTAL_BY_TIME_RANGE` | Supported (`DELETE WHERE time_col >= low AND time_col <= high` + `INSERT`) | Supported (`DELETE` + `INSERT`) | +| `INCREMENTAL_BY_UNIQUE_KEY` | Supported (Auto-promotes `unique_key` to `UNIQUE KEY`, `DELETE ... USING` + `INSERT`) | Supported (Batched key `DELETE` + `INSERT`) | +| `INCREMENTAL_BY_PARTITION` | Supported (Column & time-truncation range `DELETE` + `INSERT`) | Supported (`INSERT OVERWRITE` / partition `DELETE` + `INSERT`) | +| `SCD_TYPE_2_BY_TIME` / `SCD_TYPE_2_BY_COLUMN` | Supported | Supported | +| `VIEW` | Supported (`CREATE OR REPLACE VIEW`) | Requires `internal` catalog (`environment_catalog_mapping`) | +| `VIEW (materialized true)` | Supported (Async MTMV & Sync MV with `materialized_type = 'SYNC'`) | Not supported by Doris inside external catalogs | +| `SEED` / Python DataFrame Models | Supported (Automatic `VARCHAR(255)` OLAP key coercion) | Supported | +| `CUSTOM ('doris_ddl_only')` | Supported (DDL + `ALTER TABLE` schema evolution only) | Supported | + +--- + ## Migrating from `dbt-doris` or StarRocks If you are migrating an existing `dbt-doris` project or StarRocks SQLMesh models to Doris, map your configuration keys as follows: diff --git a/sqlmesh/cli/project_init.py b/sqlmesh/cli/project_init.py index e3132a6de3..b149b44bb9 100644 --- a/sqlmesh/cli/project_init.py +++ b/sqlmesh/cli/project_init.py @@ -88,6 +88,13 @@ def _gen_config( " # https://sqlmesh.readthedocs.io/en/stable/reference/configuration/#connection\n" f" # {doc_link.format(engine_link=engine_link)}\n{connection_settings}" ) + if engine_type == "doris": + connection_settings += ( + " # Doris cannot be used as a state_connection; use DuckDB for local dev or Postgres/MySQL in production:\n" + " state_connection:\n" + " type: duckdb\n" + " database: state.db\n" + ) default_configs = { ProjectTemplate.DEFAULT: f"""# --- Gateway Connection --- diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 2a5a594fe7..d5fa52834d 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -1,9 +1,12 @@ from __future__ import annotations from collections import defaultdict +import datetime +from decimal import Decimal import itertools import logging import re +import threading import typing as t import pandas as pd # noqa: TID253 @@ -39,6 +42,7 @@ from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF from sqlmesh.core.node import IntervalUnit from sqlmesh.core.schema_diff import TableAlterOperation + from sqlmesh.utils.date import TimeLike logger = logging.getLogger(__name__) @@ -66,6 +70,7 @@ class DorisEngineAdapter( SUPPORTS_MATERIALIZED_VIEW_SCHEMA = True RECREATE_MATERIALIZED_VIEW_ON_EVALUATION = False SUPPORTS_CREATE_DROP_CATALOG = False + SUPPORTS_QUERY_EXECUTION_TRACKING = True INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.DELETE_INSERT CURRENT_CATALOG_EXPRESSION = exp.func("current_catalog") RESOLVE_TABLE_REFS_IN_PHYSICAL_PROPERTIES: t.FrozenSet[str] = frozenset( @@ -158,6 +163,47 @@ class DorisEngineAdapter( "AGG_STATE", } + def _get_thread_state(self) -> threading.local: + state = getattr(self._connection_pool, "_doris_thread_state", None) + if not isinstance(state, threading.local): + state = self._extra_config.get("_doris_thread_state") + if not isinstance(state, threading.local): + state = getattr(self, "_doris_thread_state", None) + if not isinstance(state, threading.local): + state = threading.local() + self._doris_thread_state = state + if isinstance(getattr(self, "_extra_config", None), dict): + self._extra_config["_doris_thread_state"] = state + try: + setattr(self._connection_pool, "_doris_thread_state", state) + except Exception: + pass + return state + + @property + def _active_catalog(self) -> t.Optional[str]: + return getattr(self._get_thread_state(), "active_catalog", None) + + @_active_catalog.setter + def _active_catalog(self, value: t.Optional[str]) -> None: + self._get_thread_state().active_catalog = value + + @property + def _needs_catalog_restore(self) -> bool: + return bool(getattr(self._get_thread_state(), "needs_catalog_restore", False)) + + @_needs_catalog_restore.setter + def _needs_catalog_restore(self, value: bool) -> None: + self._get_thread_state().needs_catalog_restore = bool(value) + + @property + def _last_described_key_cols(self) -> t.Set[str]: + return getattr(self._get_thread_state(), "last_described_key_cols", set()) + + @_last_described_key_cols.setter + def _last_described_key_cols(self, value: t.Set[str]) -> None: + self._get_thread_state().last_described_key_cols = set(value) + @property def catalog_support(self) -> CatalogSupport: return CatalogSupport.FULL_SUPPORT @@ -178,11 +224,11 @@ def get_current_catalog(self) -> t.Optional[str]: def set_current_catalog(self, catalog: str) -> None: catalog_sql = self._to_sql(exp.to_identifier(catalog, quoted=True)) self.execute(exp.Command(this="SWITCH", expression=catalog_sql)) - self._active_catalog: t.Optional[str] = catalog - self._needs_catalog_restore: bool = False + self._active_catalog = catalog + self._needs_catalog_restore = False def _restore_active_catalog_on_reconnect(self, sql: t.Optional[str] = None) -> None: - active_cat = getattr(self, "_active_catalog", None) + active_cat = self._active_catalog self._needs_catalog_restore = False if not active_cat: return @@ -194,6 +240,11 @@ def _restore_active_catalog_on_reconnect(self, sql: t.Optional[str] = None) -> N return catalog_sql = self._to_sql(exp.to_identifier(active_cat, quoted=True)) self.cursor.execute(f"SWITCH {catalog_sql}") + database = self._extra_config.get("database") + default_cat = self.default_catalog or "internal" + if database and active_cat.lower() == str(default_cat).lower(): + db_sql = self._to_sql(exp.to_identifier(str(database), quoted=True)) + self.cursor.execute(f"USE {db_sql}") def _is_iceberg_table( self, @@ -1091,6 +1142,104 @@ def _insert_overwrite_by_condition( **kwargs, ) + _TIME_UNIT_PRECEDENCE: t.ClassVar[t.Dict[str, int]] = { + "second": 0, + "minute": 1, + "hour": 2, + "day": 3, + "month": 4, + "year": 5, + } + + @classmethod + def _extract_time_partition_transform( + cls, expr: exp.Expr + ) -> t.Optional[t.Tuple[exp.Column, str]]: + """Extract `(column, unit)` from time-truncation partition expressions (`DATE_TRUNC`, + `TO_DATE`, `DATE`, `YEAR`/`YEARS`, `MONTH`/`MONTHS`, `DAY`/`DAYS`, `HOUR`/`HOURS`). + """ + curr = expr.unnest() + col_node: t.Optional[exp.Column] = None + raw_unit: t.Optional[str] = None + + if isinstance(curr, (exp.TimestampTrunc, exp.DateTrunc)): + unit_arg = curr.args.get("unit") + if isinstance(curr.this, exp.Column): + col_node = curr.this + raw_unit = cls._extract_str_value(unit_arg) + elif isinstance(unit_arg, exp.Column): + col_node = unit_arg + raw_unit = cls._extract_str_value(curr.this) + elif isinstance(curr, (exp.TsOrDsToDate, exp.Date)): + if isinstance(curr.this, exp.Column): + col_node = curr.this + raw_unit = "day" + elif isinstance(curr, (exp.Year, exp.Month, exp.Day, exp.Hour)): + inner_arg = curr.this.unnest() if isinstance(curr.this, exp.Expr) else None + if isinstance(inner_arg, (exp.TsOrDsToDate, exp.Date)) and isinstance( + inner_arg.this, exp.Expr + ): + inner_arg = inner_arg.this.unnest() + if isinstance(inner_arg, exp.Column): + col_node = inner_arg + if isinstance(curr, exp.Year): + raw_unit = "year" + elif isinstance(curr, exp.Month): + raw_unit = "month" + elif isinstance(curr, exp.Day): + raw_unit = "day" + else: + raw_unit = "hour" + elif isinstance(curr, exp.Anonymous): + fn_name = curr.name.upper() + if fn_name == "DATE_TRUNC" and len(curr.expressions) == 2: + arg0, arg1 = curr.expressions + if isinstance(arg0, exp.Column): + col_node = arg0 + raw_unit = cls._extract_str_value(arg1) + elif isinstance(arg1, exp.Column): + col_node = arg1 + raw_unit = cls._extract_str_value(arg0) + elif ( + fn_name in ("TO_DATE", "DATE") + and len(curr.expressions) == 1 + and isinstance(curr.expressions[0], exp.Column) + ): + col_node = curr.expressions[0] + raw_unit = "day" + elif ( + fn_name + in ( + "YEAR", + "YEARS", + "MONTH", + "MONTHS", + "DAY", + "DAYS", + "HOUR", + "HOURS", + "MINUTE", + "MINUTES", + "SECOND", + "SECONDS", + ) + and len(curr.expressions) == 1 + and isinstance(curr.expressions[0], exp.Column) + ): + col_node = curr.expressions[0] + raw_unit = fn_name.lower() + + if col_node is not None and raw_unit: + unit_clean = str(raw_unit).strip().strip("'\"").lower() + if unit_clean.endswith("s") and unit_clean[:-1] in cls._TIME_UNIT_PRECEDENCE: + unit_clean = unit_clean[:-1] + if unit_clean in cls._TIME_UNIT_PRECEDENCE: + clean_col = exp.column( + col_node.name, quoted=bool(getattr(col_node.this, "quoted", False)) + ) + return clean_col, unit_clean + return None + def insert_overwrite_by_partition( self, table_name: TableName, @@ -1100,28 +1249,72 @@ def insert_overwrite_by_partition( source_columns: t.Optional[t.List[str]] = None, ) -> None: parsed_exprs, _ = self._parse_partition_expressions(partitioned_by) - raw_key_exprs: t.List[exp.Expr] = [] + col_order: t.List[str] = [] + col_to_expr: t.Dict[str, exp.Column] = {} + col_to_unit: t.Dict[str, t.Optional[str]] = {} + for expr in parsed_exprs or partitioned_by: if isinstance(expr, exp.Column): - raw_key_exprs.append(expr) + col_lower = expr.name.lower() + if col_lower not in col_to_expr: + col_order.append(col_lower) + col_to_expr[col_lower] = exp.column( + expr.name, quoted=bool(getattr(expr.this, "quoted", False)) + ) + col_to_unit[col_lower] = None + elif isinstance(expr, exp.Identifier): + col_lower = expr.name.lower() + if col_lower not in col_to_expr: + col_order.append(col_lower) + col_to_expr[col_lower] = exp.column(expr.name, quoted=bool(expr.quoted)) + col_to_unit[col_lower] = None + elif (transform := self._extract_time_partition_transform(expr)) is not None: + col_expr, unit = transform + col_lower = col_expr.name.lower() + if col_lower not in col_to_expr: + col_order.append(col_lower) + col_to_expr[col_lower] = col_expr + col_to_unit[col_lower] = unit + elif col_to_unit[col_lower] is not None: + prev_unit = col_to_unit[col_lower] + assert prev_unit is not None + if self._TIME_UNIT_PRECEDENCE[unit] < self._TIME_UNIT_PRECEDENCE[prev_unit]: + col_to_unit[col_lower] = unit else: - cols = list(expr.find_all(exp.Column)) - if cols: - raw_key_exprs.extend(cols) - else: - raw_key_exprs.append(expr) + raise SQLMeshError( + f"Doris insert_overwrite_by_partition does not support partition expression '{expr.sql(dialect=self.dialect)}'. " + "Only column names and time-truncation transforms (DATE_TRUNC, TO_DATE, DATE, YEAR, MONTH, DAY, HOUR) are supported." + ) + key_exprs: t.List[exp.Expr] = [] - seen_cols: t.Set[str] = set() - for expr in raw_key_exprs: - if isinstance(expr, exp.Column): - col_lower = expr.name.lower() - if col_lower not in seen_cols: - seen_cols.add(col_lower) - key_exprs.append( - exp.column(expr.name, quoted=bool(getattr(expr.this, "quoted", False))) - ) + for col_lower in col_order: + col_expr = col_to_expr[col_lower] + col_unit = col_to_unit[col_lower] + if col_unit is None: + key_exprs.append(col_expr) else: - key_exprs.append(expr) + key_exprs.append( + exp.Anonymous( + this="DATE_TRUNC", + expressions=[col_expr, exp.Literal.string(col_unit)], + ) + ) + + target_table = exp.to_table(table_name) + if ( + target_columns_to_types + and columns_to_types_all_known(target_columns_to_types) + and not self._is_iceberg_table(catalog_name=target_table.catalog) + ): + target_columns_to_types = self._reorder_olap_columns( + target_columns_to_types, {"partitioned_by": partitioned_by} + ) + if source_columns: + source_col_set = set(source_columns) + source_columns = [c for c in target_columns_to_types if c in source_col_set] + [ + c for c in source_columns if c not in target_columns_to_types + ] + super().insert_overwrite_by_partition( table_name, query_or_df, @@ -1311,6 +1504,14 @@ def _extract_subquery_from_where( normalized_keys.append(k) elif isinstance(k, exp.Identifier): normalized_keys.append(exp.Column(this=k.copy())) + elif (transform := self._extract_time_partition_transform(k)) is not None: + col_expr, unit = transform + normalized_keys.append( + exp.Anonymous( + this="DATE_TRUNC", + expressions=[col_expr, exp.Literal.string(unit)], + ) + ) else: return None @@ -1556,6 +1757,10 @@ def _execute_delete_with_subquery( extra_conditions: t.List[exp.Expr], ) -> None: """Execute DELETE FROM ... USING (SELECT ...) in Doris, with literal batch fallback for DUPLICATE KEY/Iceberg.""" + if any(self._extract_time_partition_transform(k) is not None for k in key_exprs): + self._delete_by_time_transform_keys(table_name, key_exprs, subquery, extra_conditions) + return + target_table = exp.to_table(table_name) if self._is_iceberg_table(catalog_name=target_table.catalog): self._delete_by_literal_keys(table_name, key_exprs, subquery, extra_conditions) @@ -1641,6 +1846,108 @@ def _sanitize_key_value(val: t.Any) -> t.Any: pass return val + @staticmethod + def _to_delete_literal_expr(val: t.Any, *, is_ast_values: bool = False) -> exp.Expr: + if is_ast_values and isinstance(val, exp.Expr): + return val.copy() + if isinstance(val, (datetime.date, datetime.datetime, datetime.time)): + return exp.Literal.string(str(val)) + if isinstance(val, Decimal): + return exp.Literal.number(str(val)) + return exp.convert(val) + + def _delete_by_time_transform_keys( + self, + table_name: TableName, + key_exprs: t.List[exp.Expr], + subquery: t.Union[exp.Subquery, exp.Select], + extra_conditions: t.Optional[t.List[exp.Expr]] = None, + ) -> None: + """Delete partition ranges for time-truncation partition keys (`DATE_TRUNC`, `DAY`, `MONTH`, `YEAR`, `HOUR`, `TO_DATE`) + using Doris-compatible half-open bare-column range predicates (`col >= AND col < `). + """ + select_query = ( + subquery.this.copy() if isinstance(subquery, exp.Subquery) else subquery.copy() + ) + if not isinstance(select_query, exp.Select): + return + + projections: t.List[exp.Expr] = [] + for idx, k in enumerate(key_exprs): + transform = self._extract_time_partition_transform(k) + if transform is not None: + col_expr, unit = transform + trunc_expr = exp.Anonymous( + this="DATE_TRUNC", + expressions=[col_expr.copy(), exp.Literal.string(unit)], + ) + upper_expr = exp.DateAdd( + this=trunc_expr.copy(), + expression=exp.Literal.number(1), + unit=exp.var(unit.upper()), + ) + projections.append(exp.alias_(trunc_expr, f"_low_{idx}", quoted=True)) + projections.append(exp.alias_(upper_expr, f"_high_{idx}", quoted=True)) + else: + col_expr = exp.column(k.name, quoted=bool(getattr(k.this, "quoted", False))) + projections.append(exp.alias_(col_expr, f"_val_{idx}", quoted=True)) + + select_query.set("expressions", projections) + select_query = select_query.distinct() + rows = self.fetchall(select_query, quote_identifiers=True) + if not rows or not isinstance(rows, (list, tuple)): + return + + target_table = exp.to_table(table_name) + cleaned_extra_conditions: t.List[exp.Expr] = [] + for extra in extra_conditions or []: + cleaned_extra = extra.copy() + for col in cleaned_extra.find_all(exp.Column): + if not col.find_ancestor(exp.Subquery, exp.Select) and ( + col.table and col.table.lower() == target_table.name.lower() + ): + col.set("table", None) + col.set("db", None) + col.set("catalog", None) + cleaned_extra_conditions.append(cleaned_extra) + + for row in rows: + predicates: t.List[exp.Expr] = [] + pos = 0 + for k in key_exprs: + transform = self._extract_time_partition_transform(k) + if transform is not None: + col_expr, _ = transform + target_col = exp.column(col_expr.name, quoted=True) + low_val = self._sanitize_key_value(row[pos]) + high_val = self._sanitize_key_value(row[pos + 1]) + pos += 2 + if low_val is None or high_val is None: + predicates.append(target_col.is_(exp.null())) + else: + predicates.append( + exp.GTE( + this=target_col.copy(), + expression=self._to_delete_literal_expr(low_val), + ) + ) + predicates.append( + exp.LT( + this=target_col.copy(), + expression=self._to_delete_literal_expr(high_val), + ) + ) + else: + target_col = exp.column(k.name, quoted=True) + val = self._sanitize_key_value(row[pos]) + pos += 1 + if val is None: + predicates.append(target_col.is_(exp.null())) + else: + predicates.append(target_col.eq(self._to_delete_literal_expr(val))) + predicates.extend(cleaned_extra_conditions) + self.execute(exp.delete(table_name, where=exp.and_(*predicates))) + def _delete_by_literal_keys( self, table_name: TableName, @@ -1702,9 +2009,7 @@ def _is_null_val(val: t.Any) -> bool: return False def _to_expr(val: t.Any) -> exp.Expr: - if is_ast_values and isinstance(val, exp.Expr): - return val.copy() - return exp.convert(val) + return self._to_delete_literal_expr(val, is_ast_values=is_ast_values) if len(cols) == 1: col = cols[0] @@ -1870,6 +2175,102 @@ def _create_table_from_columns( **kwargs, ) + def replace_query( + self, + table_name: TableName, + query_or_df: QueryOrDF, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + source_columns: t.Optional[t.List[str]] = None, + supports_replace_table_override: t.Optional[bool] = None, + **kwargs: t.Any, + ) -> None: + target_table = exp.to_table(table_name) + is_iceberg = self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=target_table.catalog, + table_properties=kwargs.get("table_properties"), + ) + if ( + not is_iceberg + and target_columns_to_types + and columns_to_types_all_known(target_columns_to_types) + ): + target_columns_to_types = self._reorder_olap_columns(target_columns_to_types, kwargs) + if source_columns: + source_col_set = set(source_columns) + source_columns = [c for c in target_columns_to_types if c in source_col_set] + [ + c for c in source_columns if c not in target_columns_to_types + ] + super().replace_query( + table_name, + query_or_df, + target_columns_to_types=target_columns_to_types, + table_description=table_description, + column_descriptions=column_descriptions, + source_columns=source_columns, + supports_replace_table_override=supports_replace_table_override, + **kwargs, + ) + + def _scd_type_2( + self, + target_table: TableName, + source_table: QueryOrDF, + unique_key: t.Sequence[exp.Expr], + valid_from_col: exp.Column, + valid_to_col: exp.Column, + execution_time: t.Union[TimeLike, exp.Column], + invalidate_hard_deletes: bool = True, + updated_at_col: t.Optional[exp.Column] = None, + check_columns: t.Optional[t.Union[exp.Star, t.Sequence[exp.Expr]]] = None, + updated_at_as_valid_from: bool = False, + execution_time_as_valid_from: bool = False, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + truncate: bool = False, + source_columns: t.Optional[t.List[str]] = None, + **kwargs: t.Any, + ) -> None: + table = exp.to_table(target_table) + is_iceberg = self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=table.catalog, + table_properties=kwargs.get("table_properties"), + ) + if ( + not is_iceberg + and target_columns_to_types + and columns_to_types_all_known(target_columns_to_types) + ): + target_columns_to_types = self._reorder_olap_columns(target_columns_to_types, kwargs) + if source_columns: + source_col_set = set(source_columns) + source_columns = [c for c in target_columns_to_types if c in source_col_set] + [ + c for c in source_columns if c not in target_columns_to_types + ] + super()._scd_type_2( + target_table=target_table, + source_table=source_table, + unique_key=unique_key, + valid_from_col=valid_from_col, + valid_to_col=valid_to_col, + execution_time=execution_time, + invalidate_hard_deletes=invalidate_hard_deletes, + updated_at_col=updated_at_col, + check_columns=check_columns, + updated_at_as_valid_from=updated_at_as_valid_from, + execution_time_as_valid_from=execution_time_as_valid_from, + target_columns_to_types=target_columns_to_types, + table_description=table_description, + column_descriptions=column_descriptions, + truncate=truncate, + source_columns=source_columns, + **kwargs, + ) + def _create_table_from_source_queries( self, table_name: TableName, @@ -2199,6 +2600,24 @@ def _is_integer_indicator_projection(proj: exp.Expr) -> bool: ) return False + def _infer_projection_type(self, proj: exp.Expr) -> exp.DataType: + inner = proj.this if isinstance(proj, exp.Alias) else proj + inner = inner.unnest() + if isinstance(inner, (exp.Cast, exp.TryCast)) and isinstance(inner.to, exp.DataType): + return inner.to.copy() + if self._is_integer_indicator_projection(proj): + return exp.DataType.build("INT", dialect=self.dialect) + if isinstance(inner, exp.Literal): + if inner.is_int: + return exp.DataType.build("INT", dialect=self.dialect) + if inner.is_number: + return exp.DataType.build("DOUBLE", dialect=self.dialect) + if inner.is_string: + return exp.DataType.build("VARCHAR(255)", dialect=self.dialect) + if isinstance(inner, exp.Boolean): + return exp.DataType.build("BOOLEAN", dialect=self.dialect) + return exp.DataType.build("UNKNOWN") + def _reorder_unannotated_ctas_query( self, expression: exp.Expr, kwargs: t.Dict[str, t.Any] ) -> exp.Expr: @@ -2217,16 +2636,41 @@ def _reorder_unannotated_ctas_query( } if proj_by_name and len(proj_by_name) == len(expression.expressions): select_cols_to_types = { - proj.alias_or_name: exp.DataType.build("UNKNOWN") + proj.alias_or_name: self._infer_projection_type(proj) for proj in expression.expressions } reordered_select_cols = self._reorder_olap_columns(select_cols_to_types, kwargs) - if list(reordered_select_cols) != list(select_cols_to_types): + order_changed = list(reordered_select_cols) != list(select_cols_to_types) + types_changed = any( + reordered_select_cols[c] != select_cols_to_types[c] + for c in reordered_select_cols + ) + if order_changed or types_changed: expression = expression.copy() - expression.set( - "expressions", - [proj_by_name[c.lower()] for c in reordered_select_cols], - ) + updated_projs: t.List[exp.Expr] = [] + for c, new_dtype in reordered_select_cols.items(): + proj = proj_by_name[c.lower()].copy() + if types_changed and select_cols_to_types[c] != new_dtype: + if isinstance(proj, exp.Alias) and isinstance( + proj.this, (exp.Cast, exp.TryCast) + ): + proj.this.set("to", new_dtype.copy()) + elif isinstance(proj, (exp.Cast, exp.TryCast)): + proj = exp.alias_( + exp.cast(proj.this.copy(), to=new_dtype.copy()), + c, + quoted=True, + ) + elif isinstance(proj, exp.Alias): + proj.set("this", exp.cast(proj.this.copy(), to=new_dtype.copy())) + elif isinstance(proj, exp.Column): + proj = exp.alias_( + exp.cast(proj.copy(), to=new_dtype.copy()), + proj.name, + quoted=bool(getattr(proj.this, "quoted", True)), + ) + updated_projs.append(proj) + expression.set("expressions", updated_projs) elif any(isinstance(proj, exp.Star) for proj in expression.expressions): table_props = kwargs.get("table_properties") or {} resolved_uk, resolved_dk, resolved_ak = self._resolve_table_type_and_keys( @@ -2250,23 +2694,26 @@ def _reorder_unannotated_ctas_query( and cte.alias.lower() == target_cte_name and isinstance(cte.this, exp.Select) and len(cte.this.expressions) >= 2 - and not self._is_integer_indicator_projection(cte.this.expressions[0]) ): - indicator_idx = next( - ( - idx - for idx, p in enumerate(cte.this.expressions) - if idx > 0 and self._is_integer_indicator_projection(p) - ), - None, - ) - if indicator_idx is not None: + indicator_projs = [ + p + for p in cte.this.expressions + if self._is_integer_indicator_projection(p) + ] + other_projs = [ + p + for p in cte.this.expressions + if not self._is_integer_indicator_projection(p) + ] + if ( + indicator_projs + and list(cte.this.expressions[: len(indicator_projs)]) + != indicator_projs + ): expression = expression.copy() copied_cte = expression.args["with_"].expressions[cte_idx] copied_select = copied_cte.this - projs = list(copied_select.expressions) - indicator_proj = projs.pop(indicator_idx) - copied_select.set("expressions", [indicator_proj, *projs]) + copied_select.set("expressions", indicator_projs + other_projs) break return expression @@ -2383,6 +2830,14 @@ def _reorder_olap_columns( first_dtype = updated_types[col_names[0]] if not self._is_valid_olap_key_type(first_dtype): valid_col = next( + ( + c + for c in col_names + if self._is_valid_olap_key_type(updated_types[c]) + and not updated_types[c].is_type("unknown", "null") + ), + None, + ) or next( (c for c in col_names if self._is_valid_olap_key_type(updated_types[c])), None, ) diff --git a/tests/core/engine_adapter/integration/test_integration_doris.py b/tests/core/engine_adapter/integration/test_integration_doris.py index 13ac651ce0..aa60c60563 100644 --- a/tests/core/engine_adapter/integration/test_integration_doris.py +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -277,6 +277,51 @@ def test_auto_partition_and_insert_overwrite_by_partition(ctx: TestContext): (20, "2022-01-02", "b_updated"), ] + # Verify DATETIME partition overwrite deletes the entire day range even when timestamps differ + ts_table = ctx.table("test_auto_partition_ts") + ts_cols = { + "payload": exp.DataType.build("STRING", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + "ts": exp.DataType.build("DATETIME", dialect="doris"), + } + ts_partitioned_by = [parse_one("RANGE(DATE_TRUNC(ts, 'day'))", read="doris")] + ctx.engine_adapter.create_table( + ts_table, + ts_cols, + partitioned_by=ts_partitioned_by, + table_properties={"auto_partition": exp.Literal.string("true")}, + ) + ts_initial = pd.DataFrame( + [ + {"payload": "d1", "id": 1, "ts": "2022-01-01 09:00:00"}, + {"payload": "d2_morning", "id": 2, "ts": "2022-01-02 08:00:00"}, + {"payload": "d2_evening", "id": 3, "ts": "2022-01-02 18:00:00"}, + ] + ) + ctx.engine_adapter.insert_append( + ts_table, + ctx.input_data(ts_initial, ts_cols), + target_columns_to_types=ts_cols, + ) + ts_overwrite = pd.DataFrame( + [ + {"payload": "d2_replaced", "id": 20, "ts": "2022-01-02 14:30:00"}, + ] + ) + ctx.engine_adapter.insert_overwrite_by_partition( + ts_table, + ctx.input_data(ts_overwrite, ts_cols), + partitioned_by=ts_partitioned_by, + target_columns_to_types=ts_cols, + ) + ts_rows = ctx.engine_adapter.fetchall( + exp.select("id", "payload").from_(ts_table).order_by("id") + ) + assert list(ts_rows) == [ + (1, "d1"), + (20, "d2_replaced"), + ] + def test_unique_key_order_by_and_key_reordering(ctx: TestContext): import pandas as pd # noqa: TID253 diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 00847ca54a..220126e932 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -1125,6 +1125,11 @@ def test_insert_overwrite_by_partition_unwraps_wrappers_and_transforms( "_get_temp_table", return_value=exp.table_("__temp_tbl", db="db"), ) + mocker.patch.object( + adapter.cursor, + "fetchall", + return_value=[("2024-01-15 00:00:00", "2024-01-16 00:00:00")], + ) adapter.insert_overwrite_by_partition( "db.events", @@ -1143,12 +1148,30 @@ def test_insert_overwrite_by_partition_unwraps_wrappers_and_transforms( "DELETE FROM `db`.`events` WHERE DATE_TRUNC(`ds`, 'day') IN (SELECT DISTINCT DATE_TRUNC(`ds`, 'day') FROM `db`.`__temp_tbl`)" not in sql_calls ) - assert any( - "SELECT DISTINCT `ds` FROM `db`.`__temp_tbl`" in call - or "USING (SELECT DISTINCT `ds` FROM `db`.`__temp_tbl`)" in call - for call in sql_calls + assert ( + "SELECT DISTINCT DATE_TRUNC(`ds`, 'day') AS `_low_0`, DATE_ADD(DATE_TRUNC(`ds`, 'day'), INTERVAL 1 DAY) AS `_high_0` FROM `db`.`__temp_tbl`" + in sql_calls + ) + assert ( + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-15 00:00:00' AND `ds` < '2024-01-16 00:00:00'" + in sql_calls ) + # Unsupported non-range partition transforms (e.g. BUCKET) raise a clear SQLMeshError + with pytest.raises( + SQLMeshError, + match="Doris insert_overwrite_by_partition does not support partition expression", + ): + adapter.insert_overwrite_by_partition( + "db.events", + parse_one("SELECT id, ds FROM db.src", read="doris"), + partitioned_by=[parse_one("BUCKET(16, id)", read="doris")], + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATETIME"), + }, + ) + def test_sync_materialized_view_drop_and_replace( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], @@ -1940,13 +1963,18 @@ def test_delete_tautology_and_union_subquery_and_dedup_partition_overwrite( "DELETE FROM `db`.`tbl` AS `_t1` USING (SELECT `other_id` + 1 AS `_col_0` FROM `db`.`a`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`_col_0`", ] - # 4. insert_overwrite_by_partition deduplicates columns and strips table qualifiers + # 4. insert_overwrite_by_partition deduplicates time transforms to finest unit (MONTH) and strips table qualifiers adapter.cursor.execute.reset_mock() mocker.patch.object( adapter, "_get_temp_table", return_value=exp.table_("__temp_tbl", db="db"), ) + mocker.patch.object( + adapter.cursor, + "fetchall", + return_value=[("2024-01-01", "2024-02-01")], + ) adapter.insert_overwrite_by_partition( "db.events", parse_one("SELECT id, ds FROM db.src", read="doris"), @@ -1960,10 +1988,13 @@ def test_delete_tautology_and_union_subquery_and_dedup_partition_overwrite( }, ) overwrite_calls = to_sql_calls(adapter) - assert any( - "DELETE FROM `db`.`events` AS `_t1` USING (SELECT DISTINCT `ds` FROM `db`.`__temp_tbl`) AS `_t2` WHERE `_t1`.`ds` = `_t2`.`ds`" - == call - for call in overwrite_calls + assert ( + "SELECT DISTINCT DATE_TRUNC(`ds`, 'month') AS `_low_0`, DATE_ADD(DATE_TRUNC(`ds`, 'month'), INTERVAL 1 MONTH) AS `_high_0` FROM `db`.`__temp_tbl`" + in overwrite_calls + ) + assert ( + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-01' AND `ds` < '2024-02-01'" + in overwrite_calls ) @@ -2447,7 +2478,7 @@ def test_is_bool_rewrites_negative_or_keys_union_delete_and_table_diff_ctas( "DELETE FROM `sqlmesh`.`_intervals` WHERE `id` = -1 AND `is_dev` IN (TRUE)", "DELETE FROM `sqlmesh`.`_intervals` WHERE `id` = -2 AND `is_dev` IN (FALSE)", "DELETE FROM `db`.`target_tbl` AS `_t1` USING (SELECT `a`.`id`, `b`.`id` AS `_col_1` FROM `db`.`a` AS `a` JOIN `db`.`b` AS `b` ON `a`.`id` = `b`.`id` UNION ALL SELECT `c`.`id`, `d`.`id` FROM `db`.`c` AS `c` JOIN `db`.`d` AS `d` ON `c`.`id` = `d`.`id`) AS `_t2` WHERE `_t1`.`k1` = `_t2`.`id` AND `_t1`.`k2` = `_t2`.`_col_1`", - "CREATE TABLE IF NOT EXISTS `sqlmesh_temp`.`diff` AS WITH `__source` AS (SELECT `s`.`name` AS `name`, `s`.`id` AS `__sqlmesh_join_key` FROM `db`.`src` AS `s`), `__target` AS (SELECT `t`.`name` AS `name`, `t`.`id` AS `__sqlmesh_join_key` FROM `db`.`tgt` AS `t`), `__stats` AS (SELECT CASE WHEN NOT `s`.`__sqlmesh_join_key` IS NULL THEN 1 ELSE 0 END AS `s_exists`, `s`.`name` AS `s__name`, `t`.`name` AS `t__name`, CASE WHEN NOT `t`.`__sqlmesh_join_key` IS NULL THEN 1 ELSE 0 END AS `t_exists` FROM `__source` AS `s` LEFT JOIN `__target` AS `t` ON `s`.`__sqlmesh_join_key` = `t`.`__sqlmesh_join_key`) SELECT *, CASE WHEN `s__name` = `t__name` THEN 1 ELSE 0 END AS `row_full_match` FROM `__stats`", + "CREATE TABLE IF NOT EXISTS `sqlmesh_temp`.`diff` AS WITH `__source` AS (SELECT `s`.`name` AS `name`, `s`.`id` AS `__sqlmesh_join_key` FROM `db`.`src` AS `s`), `__target` AS (SELECT `t`.`name` AS `name`, `t`.`id` AS `__sqlmesh_join_key` FROM `db`.`tgt` AS `t`), `__stats` AS (SELECT CASE WHEN NOT `s`.`__sqlmesh_join_key` IS NULL THEN 1 ELSE 0 END AS `s_exists`, CASE WHEN NOT `t`.`__sqlmesh_join_key` IS NULL THEN 1 ELSE 0 END AS `t_exists`, `s`.`name` AS `s__name`, `t`.`name` AS `t__name` FROM `__source` AS `s` LEFT JOIN `__target` AS `t` ON `s`.`__sqlmesh_join_key` = `t`.`__sqlmesh_join_key`) SELECT *, CASE WHEN `s__name` = `t__name` THEN 1 ELSE 0 END AS `row_full_match` FROM `__stats`", ] # 5. Unannotated 2-part/internal Iceberg CTAS raises a clear SQLMeshError instead of silently dropping ENGINE=iceberg @@ -2762,3 +2793,190 @@ def reconnect_side_effect(sql, *args, **kwargs): "SWITCH `iceberg_cat`", "SELECT 1", ] + + +def test_thread_local_state_isolation_time_transform_partition_delete_table_diff_and_pm_features( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + from datetime import date, datetime + from decimal import Decimal + import threading + import pandas as pd + from sqlmesh.cli.project_init import InitCliMode, ProjectTemplate, _gen_config + + adapter = make_mocked_engine_adapter( + DorisEngineAdapter, default_catalog="internal", database="analytics" + ) + assert adapter.SUPPORTS_QUERY_EXECUTION_TRACKING is True + + # 1. Thread-local state isolation for `_active_catalog`, `_needs_catalog_restore`, and + # `_last_described_key_cols` across worker threads, `with_settings()`, and `pool.close()`. + adapter._active_catalog = "internal" + adapter._needs_catalog_restore = False + adapter._last_described_key_cols = {"main_key"} + + cloned = t.cast(DorisEngineAdapter, adapter.with_settings(correlation_id="test_corr_id")) + assert cloned._active_catalog == "internal" + assert cloned._last_described_key_cols == {"main_key"} + + worker_results: t.Dict[str, t.Any] = {} + + def worker_fn() -> None: + worker_results["initial_cat"] = adapter._active_catalog + worker_results["initial_keys"] = set(adapter._last_described_key_cols) + adapter._active_catalog = "iceberg_worker" + adapter._needs_catalog_restore = True + adapter._last_described_key_cols = {"worker_key"} + worker_results["updated_cat"] = cloned._active_catalog + worker_results["updated_keys"] = set(cloned._last_described_key_cols) + + thread = threading.Thread(target=worker_fn) + thread.start() + thread.join() + + assert worker_results == { + "initial_cat": None, + "initial_keys": set(), + "updated_cat": "iceberg_worker", + "updated_keys": {"worker_key"}, + } + # Main thread state was not overwritten by worker thread and survives connection pool close + adapter._connection_pool.close() + assert adapter._active_catalog == "internal" + assert adapter._needs_catalog_restore is False + assert adapter._last_described_key_cols == {"main_key"} + + # 2. Reconnect in `_restore_active_catalog_on_reconnect` restores both `SWITCH ` + # and `USE ` when `database` is configured on the default catalog. + adapter.cursor.execute.reset_mock() + adapter._needs_catalog_restore = True + adapter._restore_active_catalog_on_reconnect("SELECT 1") + assert to_sql_calls(adapter) == [ + "SWITCH `internal`", + "USE `analytics`", + ] + + # 3. TableDiff-style query with 4 integer indicators and DOUBLE/STRING user columns promotes all + # indicators to the front of `__stats` (covering all 3 slots of `shortkey_max_column_count = 3`). + adapter.cursor.execute.reset_mock() + four_indicator_diff = parse_one( + "WITH `__stats` AS (" + "SELECT `s`.`score` AS `s__score`, `t`.`score` AS `t__score`, `s`.`bio` AS `s__bio`, " + "IF(NOT `s`.`k` IS NULL, 1, 0) AS `s_exists`, " + "IF(NOT `t`.`k` IS NULL, 1, 0) AS `t_exists`, " + "IF(NOT `s`.`k` IS NULL AND NOT `t`.`k` IS NULL, 1, 0) AS `row_joined`, " + "IF(`s`.`id` IS NULL AND `t`.`id` IS NULL, 1, 0) AS `null_grain` " + "FROM `db`.`s` LEFT JOIN `db`.`t` ON `s`.`k` = `t`.`k`" + ") SELECT *, CASE WHEN `s__score` = `t__score` THEN 1 ELSE 0 END AS `row_full_match` FROM `__stats`", + read="doris", + ) + adapter.ctas("sqlmesh_temp.diff_4", four_indicator_diff) + diff_sql = to_sql_calls(adapter)[0] + assert ( + "`__stats` AS (SELECT CASE WHEN NOT `s`.`k` IS NULL THEN 1 ELSE 0 END AS `s_exists`, " + "CASE WHEN NOT `t`.`k` IS NULL THEN 1 ELSE 0 END AS `t_exists`, " + "CASE WHEN NOT `s`.`k` IS NULL AND NOT `t`.`k` IS NULL THEN 1 ELSE 0 END AS `row_joined`, " + "CASE WHEN `s`.`id` IS NULL AND `t`.`id` IS NULL THEN 1 ELSE 0 END AS `null_grain`, " + "`s`.`score` AS `s__score`" in diff_sql + ) + + # 4. Non-star unannotated CTAS infers explicit CAST types (`DOUBLE`, `STRING`, `INT`) and + # prefers known valid key columns (`id` INT) or coerces `STRING` to `VARCHAR(255)` at column 0. + adapter.cursor.execute.reset_mock() + adapter.ctas( + "db.unannotated_cast_reordered", + parse_one( + "SELECT CAST(score AS DOUBLE) AS score, raw_col, CAST(id AS INT) AS id FROM db.src", + read="doris", + ), + ) + adapter.ctas( + "db.unannotated_cast_string", + parse_one( + "SELECT CAST(score AS DOUBLE) AS score, CAST(label AS STRING) AS label FROM db.src", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`unannotated_cast_reordered` AS SELECT CAST(`id` AS INT) AS `id`, CAST(`score` AS DOUBLE) AS `score`, `raw_col` FROM `db`.`src`", + "CREATE TABLE IF NOT EXISTS `db`.`unannotated_cast_string` AS SELECT CAST(`label` AS VARCHAR(255)) AS `label`, CAST(`score` AS DOUBLE) AS `score` FROM `db`.`src`", + ] + + # 5. `replace_query` on an existing table and `scd_type_2_by_column` when column 0 is `DOUBLE` + # or `TEXT` reorder/coerce `target_columns_to_types` (and `source_columns`). + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter, + "get_data_object", + return_value=DataObject( + catalog="internal", schema="db", name="existing_seed", type=DataObjectType.TABLE + ), + ) + seed_df = pd.DataFrame([{"score": 9.5, "code": "us"}]) + adapter.replace_query( + "db.existing_seed", + seed_df, + target_columns_to_types={ + "score": exp.DataType.build("DOUBLE", dialect="doris"), + "code": exp.DataType.build("TEXT", dialect="doris"), + }, + source_columns=["score", "code"], + ) + replace_calls = to_sql_calls(adapter) + assert replace_calls == [ + "TRUNCATE TABLE `db`.`existing_seed`", + "INSERT INTO `db`.`existing_seed` (`code`, `score`) SELECT CAST(`code` AS VARCHAR(255)) AS `code`, CAST(`score` AS DOUBLE) AS `score` FROM (SELECT 'us' AS `code`, 9.5 AS `score`) AS `t`", + ] + + adapter.cursor.execute.reset_mock() + mocker.patch("sqlmesh.core.engine_adapter.doris.random_id", return_value="efgh5678") + adapter.scd_type_2_by_column( + target_table="db.scd_col_tbl", + source_table=t.cast( + exp.Select, + parse_one("SELECT score, id FROM db.src_tbl", read="doris"), + ), + unique_key=[exp.column("id")], + valid_from_col=exp.column("valid_from", quoted=True), + valid_to_col=exp.column("valid_to", quoted=True), + check_columns=[exp.column("score")], + target_columns_to_types={ + "score": exp.DataType.build("DOUBLE", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + "valid_from": exp.DataType.build("DATETIME", dialect="doris"), + "valid_to": exp.DataType.build("DATETIME", dialect="doris"), + }, + execution_time=datetime(2024, 1, 2, 0, 0, 0), + ) + scd_col_calls = to_sql_calls(adapter) + assert scd_col_calls[0].startswith( + "CREATE TABLE IF NOT EXISTS `db`.`scd_col_tbl` (`id` INT, `score` DOUBLE, `valid_from` DATETIME, `valid_to` DATETIME)" + ) + + # 6. `_delete_by_literal_keys` with Python `date`, `datetime`, and `Decimal` values from PyMySQL + # emits plain string/number literals (`'2024-01-01'`, `12.5`) instead of `CAST('...' AS DATE)`. + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter, + "fetchall", + return_value=[(date(2024, 1, 1), datetime(2024, 1, 1, 12, 30, 0), Decimal("12.50"))], + ) + adapter._delete_by_literal_keys( + "db.typed_keys", + [exp.column("dt"), exp.column("ts"), exp.column("amt")], + parse_one("SELECT dt, ts, amt FROM db.staging", read="doris", into=exp.Select), + ) + assert to_sql_calls(adapter) == [ + "DELETE FROM `db`.`typed_keys` WHERE `dt` = '2024-01-01' AND `ts` = '2024-01-01 12:30:00' AND `amt` IN (12.50)" + ] + + # 7. `sqlmesh init doris` generates a working `state_connection` block (`type: duckdb`). + doris_init_cfg = _gen_config( + engine_type="doris", + settings=None, + start="2024-01-01", + template=ProjectTemplate.DEFAULT, + cli_mode=InitCliMode.DEFAULT, + ) + assert "state_connection:\n type: duckdb\n database: state.db" in doris_init_cfg + assert "dialect: doris" in doris_init_cfg From 87b75306e1300f117fc0caca038bdf5343e6279e Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sat, 24 Oct 2026 19:45:20 -0700 Subject: [PATCH 25/27] feat(doris): harden UNION/time-transform deletes, aggregate CTAS key inference, LIMIT folding, dbt-doris target, and docs Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/integrations/engines/doris.md | 140 ++++++++--- sqlmesh/core/engine_adapter/doris.py | 312 ++++++++++++++++++++---- sqlmesh/dbt/model.py | 74 ++++++ sqlmesh/dbt/target.py | 66 +++++ tests/core/engine_adapter/test_doris.py | 181 ++++++++++++++ 5 files changed, 699 insertions(+), 74 deletions(-) diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index 13835b4e92..389b8a8075 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -18,6 +18,26 @@ This page describes SQLMesh support for the [Apache Doris](https://doris.apache. pip install "sqlmesh[doris]" ``` +### 5-Minute Quickstart + +1. **Initialize a new Doris project** (generates `config.yaml` with `dialect: doris` and a local DuckDB `state_connection`): + ```bash + mkdir doris_demo && cd doris_demo + sqlmesh init doris + ``` +2. **Single-BE Docker / Local Dev Tip**: Doris defaults to `replication_num = 3`. When running against a single-BE Docker container, set `replication_num = 1` in `model_defaults.physical_properties` inside `config.yaml`: + ```yaml + model_defaults: + dialect: doris + start: 2024-01-01 + physical_properties: + replication_num: 1 + ``` +3. **Plan and apply**: + ```bash + sqlmesh plan + ``` + ### Connection options | Option | Description | Type | Required | @@ -48,34 +68,60 @@ pip install "sqlmesh[doris]" ### Connection configuration example -```yaml -gateways: - doris: - connection: - type: doris - host: fe.doris.cluster - port: 9030 - user: doris_user - password: {{ env_var('DORIS_PASSWORD') }} - database: analytics - catalog: internal - connect_timeout: 30 - session_variables: - workload_group: etl_group - query_timeout: 3600 - state_connection: - type: postgres - host: state-db.internal - port: 5432 - user: sqlmesh - password: {{ env_var('POSTGRES_PASSWORD') }} - database: sqlmesh_state - -default_gateway: doris - -model_defaults: - dialect: doris -``` +=== "Local Development (DuckDB State)" + + ```yaml + gateways: + doris: + connection: + type: doris + host: 127.0.0.1 + port: 9030 + user: root + password: "" + database: analytics + state_connection: + type: duckdb + database: state.db + + default_gateway: doris + + model_defaults: + dialect: doris + physical_properties: + replication_num: 1 + ``` + +=== "Production (Postgres / MySQL State)" + + ```yaml + gateways: + doris: + connection: + type: doris + host: fe.doris.cluster + port: 9030 + user: doris_user + password: {{ env_var('DORIS_PASSWORD') }} + database: analytics + catalog: internal + connect_timeout: 30 + session_variables: + workload_group: etl_group + query_timeout: 3600 + state_connection: + type: postgres + host: state-db.internal + port: 5432 + user: sqlmesh + password: {{ env_var('POSTGRES_PASSWORD') }} + database: sqlmesh_state + + default_gateway: doris + + model_defaults: + dialect: doris + ``` --- @@ -225,7 +271,7 @@ MODEL ( ### Auto partitioning (`auto_partition = true`) -In Doris 2.1+, `AUTO PARTITION` automatically creates partitions on demand as new data arrives: +In Doris 2.1+, `AUTO PARTITION` automatically creates partitions on demand as new data arrives (and SQLMesh automatically marks the range partition column `NOT NULL` and wraps CTAS projections in `NON_NULLABLE(...)`): ```sql MODEL ( @@ -240,6 +286,28 @@ MODEL ( ); ``` +### Dynamic partitioning (`dynamic_partition.enable = 'true'`) + +For rolling retention windows managed by Doris FE's dynamic partition scheduler: + +```sql +MODEL ( + name analytics.rolling_daily_events, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column event_date + ), + partitioned_by RANGE(event_date), + physical_properties ( + 'dynamic_partition.enable' = 'true', + 'dynamic_partition.time_unit' = 'DAY', + 'dynamic_partition.start' = '-30', + 'dynamic_partition.end' = '3', + 'dynamic_partition.prefix' = 'p', + 'dynamic_partition.buckets' = '16' + ) +); +``` + --- ## Apache Iceberg & Multi-Catalog Lakehouse Integration @@ -469,6 +537,7 @@ GROUP BY user_id, CAST(event_time AS DATE); !!! note "Automatic Physical Table Resolution & Non-Destructive Evaluation" - Doris materialized views are only created on initial deployment or when a rebuild is required (`RECREATE_MATERIALIZED_VIEW_ON_EVALUATION = False`), preserving MTMV refresh state across routine `sqlmesh run` invocations. + - When `build = 'DEFERRED'` is used on an asynchronous MTMV that also defines `audits (...)`, SQLMesh automatically executes `REFRESH MATERIALIZED VIEW AUTO` after creation so audits evaluate against populated data rather than an empty view. - When `excluded_trigger_tables` (supports a single string, comma-separated string, or tuple `('db.t1', 'db.t2')` on async MTMVs) or `source_table` (on synchronous MVs) references a SQLMesh-managed model, SQLMesh automatically resolves the model name to its underlying physical snapshot table name. ### Synchronous Materialized Views (`materialized_type = 'SYNC'`) @@ -535,7 +604,7 @@ FROM raw.logs; ## Migrating from `dbt-doris` or StarRocks -If you are migrating an existing `dbt-doris` project or StarRocks SQLMesh models to Doris, map your configuration keys as follows: +SQLMesh supports loading existing `dbt-doris` projects directly (`type: doris` in `profiles.yml`) as well as native SQLMesh Doris models. If you are migrating a `dbt-doris` project or StarRocks SQLMesh models to native SQLMesh Doris models, map your configuration keys as follows: | `dbt-doris` / StarRocks Config | SQLMesh Doris Equivalent | | --- | --- | @@ -548,3 +617,14 @@ If you are migrating an existing `dbt-doris` project or StarRocks SQLMesh models | `distributed_by=['c1']`, `buckets=10` | `clustered_by (c1), physical_properties (buckets = 10)` or `physical_properties (distributed_by = ('c1'), buckets = 10)` | | `partition_by` | Supported in `physical_properties (partition_by = ...)` or top-level `partitioned_by` | | `properties={...}` | Key-value entries inside `physical_properties (...)` | + +--- + +## Troubleshooting & Doris Gotchas + +| Symptom / Doris Error | Cause | Resolution | +| --- | --- | --- | +| `errCode = 2, replication num 3 is larger than available backends 1` | Doris defaults to `replication_num = 3`, which fails on single-BE local/CI clusters. | Set `replication_num: 1` under `model_defaults.physical_properties` in `config.yaml` for local dev. | +| `PARTITION BY` omitted on OLAP table (logged as a warning) | Doris OLAP tables with `PARTITION BY (...) ()` and no partition creation strategy have zero partitions and reject all `INSERT`s. | Add `auto_partition = 'true'`, `'dynamic_partition.enable' = 'true'`, or explicit `partitions = '...'` to `physical_properties`. | +| `Doris only supports creating views in the 'internal' catalog` | External Iceberg catalogs do not support Doris `VIEW` or `MATERIALIZED VIEW` objects. | Set `virtual_environment_mode: dev_only` and `environment_catalog_mapping: {"^(?!prod$).*": "internal"}` in `config.yaml`. | +| `Key column cannot be of type STRING / DOUBLE / ARRAY` | Doris forbids `STRING`/`TEXT`, floating-point (`FLOAT`/`DOUBLE`), and complex types as OLAP key columns. | SQLMesh automatically reorders valid key columns first and coerces `STRING`/`TEXT` key columns to `VARCHAR(255)`. For explicit keys on `DOUBLE`/`FLOAT`, use `DECIMAL` or an integer/date key column. | diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index d5fa52834d..e76e226bb5 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -536,7 +536,9 @@ def get_alter_operations( ignore_additive: bool = False, ) -> t.List[TableAlterOperation]: from sqlmesh.core.schema_diff import ( + TableAlterAddColumnOperation, TableAlterChangeColumnTypeOperation, + TableAlterDropColumnOperation, TableAlterOperation, ) @@ -556,6 +558,11 @@ def get_alter_operations( if current_cols and not current_key_cols: current_key_cols = {next(iter(current_cols)).lower()} target_cols = self.columns(target_table_name) + invalid_target_key_cols = { + c.lower() + for c, dtype in target_cols.items() + if c.lower() in current_key_cols and not self._is_valid_olap_key_type(dtype) + } operations = t.cast( t.List[TableAlterOperation], self.schema_differ.compare_columns( @@ -570,9 +577,15 @@ def get_alter_operations( op for op in operations if not ( - isinstance(op, TableAlterChangeColumnTypeOperation) - and op.column.name.lower() in current_key_cols - and not self._is_valid_olap_key_type(op.column_type) + ( + isinstance(op, TableAlterChangeColumnTypeOperation) + and op.column.name.lower() in current_key_cols + and not self._is_valid_olap_key_type(op.column_type) + ) + or ( + isinstance(op, (TableAlterDropColumnOperation, TableAlterAddColumnOperation)) + and op.column.name.lower() in invalid_target_key_cols + ) ) ] @@ -798,9 +811,6 @@ def _create_materialized_view( # handling to _get_source_queries_and_columns_to_types, which would otherwise materialize # a temporary table that gets dropped as soon as the SourceQuery context manager exits. if isinstance(query_or_df, pd.DataFrame): - values: t.List[t.Tuple[t.Any, ...]] = list( - query_or_df.itertuples(index=False, name=None) - ) target_columns_to_types, source_columns = self._columns_to_types( query_or_df, target_columns_to_types, source_columns ) @@ -809,6 +819,10 @@ def _create_materialized_view( source_columns_to_types = get_source_columns_to_types( target_columns_to_types, source_columns ) + ordered_df = query_or_df[list(source_columns or target_columns_to_types)] + values: t.List[t.Tuple[t.Any, ...]] = list( + ordered_df.itertuples(index=False, name=None) + ) query_or_df = self._values_to_sql( values, source_columns_to_types, @@ -1482,17 +1496,13 @@ def _extract_subquery_from_where( lhs_concat = self._unwrap_concat_ws(lhs) if lhs_concat is not None: - if not (isinstance(select_expr, exp.Select) and len(select_expr.expressions) == 1): - return None - first_sel = select_expr.expressions[0] - raw_sel = first_sel.this if isinstance(first_sel, exp.Alias) else first_sel - sel_concat = self._unwrap_concat_ws(raw_sel) - if sel_concat is None or len(sel_concat.expressions) != len(lhs_concat.expressions): + unwrapped_select_expr = self._unwrap_concat_ws_in_query( + select_expr, len(lhs_concat.expressions) + ) + if unwrapped_select_expr is None: return None key_exprs = [e.copy() for e in lhs_concat.expressions[1:]] - select_expr = select_expr.copy() - select_expr.set("expressions", [e.copy() for e in sel_concat.expressions[1:]]) - subquery = exp.Subquery(this=select_expr) + subquery = exp.Subquery(this=unwrapped_select_expr) elif isinstance(lhs, exp.Tuple): key_exprs = [e.copy() for e in lhs.expressions] else: @@ -1734,6 +1744,30 @@ def _unwrap_concat_ws(expr: exp.Expr) -> t.Optional[exp.Expr]: return curr return None + @classmethod + def _unwrap_concat_ws_in_query( + cls, query_expr: exp.Expr, expected_arity: int + ) -> t.Optional[exp.Expr]: + if isinstance(query_expr, exp.SetOperation): + left = cls._unwrap_concat_ws_in_query(query_expr.this, expected_arity) + right = cls._unwrap_concat_ws_in_query(query_expr.expression, expected_arity) + if left is None or right is None: + return None + unwrapped_set = query_expr.copy() + unwrapped_set.set("this", left) + unwrapped_set.set("expression", right) + return unwrapped_set + if isinstance(query_expr, exp.Select) and len(query_expr.expressions) == 1: + first_sel = query_expr.expressions[0] + raw_sel = first_sel.this if isinstance(first_sel, exp.Alias) else first_sel + sel_concat = cls._unwrap_concat_ws(raw_sel) + if sel_concat is None or len(sel_concat.expressions) != expected_arity: + return None + unwrapped_sel = query_expr.copy() + unwrapped_sel.set("expressions", [e.copy() for e in sel_concat.expressions[1:]]) + return unwrapped_sel + return None + @staticmethod def _get_in_subquery(in_expr: exp.In) -> t.Optional[exp.Subquery]: query = in_expr.args.get("query") @@ -1767,6 +1801,7 @@ def _execute_delete_with_subquery( return join_conditions: t.List[exp.Expr] = [] + orig_subquery = subquery.copy() subquery = subquery.copy() first_select = subquery.this while isinstance(first_select, exp.SetOperation): @@ -1778,7 +1813,11 @@ def _execute_delete_with_subquery( t2_col_name = col_name if isinstance(first_select, exp.Select) and idx < len(first_select.expressions): sel_item = first_select.expressions[idx] - candidate_name = sel_item.alias_or_name + candidate_name = ( + sel_item.alias_or_name + if isinstance(sel_item, (exp.Column, exp.Alias, exp.Identifier)) + else "" + ) if candidate_name and candidate_name.lower() not in used_t2_aliases: t2_col_name = candidate_name else: @@ -1831,7 +1870,7 @@ def _execute_delete_with_subquery( table_name, e, ) - self._delete_by_literal_keys(table_name, key_exprs, subquery, extra_conditions) + self._delete_by_literal_keys(table_name, key_exprs, orig_subquery, extra_conditions) else: raise @@ -1860,7 +1899,7 @@ def _delete_by_time_transform_keys( self, table_name: TableName, key_exprs: t.List[exp.Expr], - subquery: t.Union[exp.Subquery, exp.Select], + subquery: t.Union[exp.Subquery, exp.Select, exp.SetOperation], extra_conditions: t.Optional[t.List[exp.Expr]] = None, ) -> None: """Delete partition ranges for time-truncation partition keys (`DATE_TRUNC`, `DAY`, `MONTH`, `YEAR`, `HOUR`, `TO_DATE`) @@ -1869,32 +1908,102 @@ def _delete_by_time_transform_keys( select_query = ( subquery.this.copy() if isinstance(subquery, exp.Subquery) else subquery.copy() ) - if not isinstance(select_query, exp.Select): + first_select = select_query + while isinstance(first_select, exp.SetOperation): + first_select = first_select.this + if not isinstance(first_select, exp.Select): return projections: t.List[exp.Expr] = [] - for idx, k in enumerate(key_exprs): - transform = self._extract_time_partition_transform(k) - if transform is not None: - col_expr, unit = transform - trunc_expr = exp.Anonymous( - this="DATE_TRUNC", - expressions=[col_expr.copy(), exp.Literal.string(unit)], - ) - upper_expr = exp.DateAdd( - this=trunc_expr.copy(), - expression=exp.Literal.number(1), - unit=exp.var(unit.upper()), - ) - projections.append(exp.alias_(trunc_expr, f"_low_{idx}", quoted=True)) - projections.append(exp.alias_(upper_expr, f"_high_{idx}", quoted=True)) - else: - col_expr = exp.column(k.name, quoted=bool(getattr(k.this, "quoted", False))) - projections.append(exp.alias_(col_expr, f"_val_{idx}", quoted=True)) + if ( + isinstance(select_query, exp.Select) + and len(select_query.expressions) == len(key_exprs) + and not select_query.args.get("group") + and not select_query.args.get("having") + and not select_query.args.get("limit") + ): + for idx, k in enumerate(key_exprs): + sel_item = select_query.expressions[idx] + raw_sel = sel_item.this if isinstance(sel_item, exp.Alias) else sel_item + transform = self._extract_time_partition_transform(k) + if transform is not None: + _, unit = transform + src_transform = self._extract_time_partition_transform(raw_sel) + if src_transform is not None: + src_arg: exp.Expr = src_transform[0].copy() + else: + src_arg = raw_sel.copy() + trunc_expr = exp.Anonymous( + this="DATE_TRUNC", + expressions=[src_arg, exp.Literal.string(unit)], + ) + upper_expr = exp.DateAdd( + this=trunc_expr.copy(), + expression=exp.Literal.number(1), + unit=exp.var(unit.upper()), + ) + projections.append(exp.alias_(trunc_expr, f"_low_{idx}", quoted=True)) + projections.append(exp.alias_(upper_expr, f"_high_{idx}", quoted=True)) + else: + projections.append(exp.alias_(raw_sel.copy(), f"_val_{idx}", quoted=True)) + select_query.set("expressions", projections) + fetch_query: exp.Select = select_query.distinct() + elif len(first_select.expressions) == len(key_exprs): + all_branches: t.List[exp.Select] = [] + + def _collect_select_branches(node: exp.Expr) -> None: + if isinstance(node, exp.SetOperation): + _collect_select_branches(node.this) + _collect_select_branches(node.expression) + elif isinstance(node, exp.Select): + all_branches.append(node) + + _collect_select_branches(select_query) + for branch in all_branches: + if len(branch.expressions) == len(key_exprs) and not branch.args.get("group"): + branch_exprs = list(branch.expressions) + for idx, k in enumerate(key_exprs): + if self._extract_time_partition_transform(k) is not None: + sel_item = branch_exprs[idx] + raw_sel = sel_item.this if isinstance(sel_item, exp.Alias) else sel_item + src_transform = self._extract_time_partition_transform(raw_sel) + if src_transform is not None: + branch_exprs[idx] = src_transform[0].copy() + branch.set("expressions", branch_exprs) + + new_first_exprs = list(first_select.expressions) + for idx, sel_item in enumerate(new_first_exprs): + inner_expr = sel_item.this if isinstance(sel_item, exp.Alias) else sel_item + new_first_exprs[idx] = exp.alias_(inner_expr.copy(), f"_col_{idx}", quoted=True) + first_select.set("expressions", new_first_exprs) + + for idx, k in enumerate(key_exprs): + sq_col = exp.column(f"_col_{idx}", quoted=True) + transform = self._extract_time_partition_transform(k) + if transform is not None: + _, unit = transform + trunc_expr = exp.Anonymous( + this="DATE_TRUNC", + expressions=[sq_col, exp.Literal.string(unit)], + ) + upper_expr = exp.DateAdd( + this=trunc_expr.copy(), + expression=exp.Literal.number(1), + unit=exp.var(unit.upper()), + ) + projections.append(exp.alias_(trunc_expr, f"_low_{idx}", quoted=True)) + projections.append(exp.alias_(upper_expr, f"_high_{idx}", quoted=True)) + else: + projections.append(exp.alias_(sq_col, f"_val_{idx}", quoted=True)) + fetch_query = ( + exp.select(*projections) + .from_(exp.alias_(exp.Subquery(this=select_query), "_sq", table=True)) + .distinct() + ) + else: + return - select_query.set("expressions", projections) - select_query = select_query.distinct() - rows = self.fetchall(select_query, quote_identifiers=True) + rows = self.fetchall(fetch_query, quote_identifiers=True) if not rows or not isinstance(rows, (list, tuple)): return @@ -2011,9 +2120,19 @@ def _is_null_val(val: t.Any) -> bool: def _to_expr(val: t.Any) -> exp.Expr: return self._to_delete_literal_expr(val, is_ast_values=is_ast_values) + def _dedup_exprs(exprs: t.List[exp.Expr]) -> t.List[exp.Expr]: + seen: t.Set[str] = set() + unique_exprs: t.List[exp.Expr] = [] + for e in exprs: + key_sql = e.sql(dialect=self.dialect) + if key_sql not in seen: + seen.add(key_sql) + unique_exprs.append(e) + return unique_exprs + if len(cols) == 1: col = cols[0] - non_null_vals = [_to_expr(r[0]) for r in rows if not _is_null_val(r[0])] + non_null_vals = _dedup_exprs([_to_expr(r[0]) for r in rows if not _is_null_val(r[0])]) has_null = any(_is_null_val(r[0]) for r in rows) for i in range(0, len(non_null_vals), batch_size): @@ -2045,7 +2164,7 @@ def _to_expr(val: t.Any) -> exp.Expr: p_col.eq(_to_expr(p_val)) if not _is_null_val(p_val) else p_col.is_(exp.null()) for p_col, p_val in zip(prefix_cols, prefix_vals) ] - non_null_batch = [_to_expr(v) for v in batch_vals if not _is_null_val(v)] + non_null_batch = _dedup_exprs([_to_expr(v) for v in batch_vals if not _is_null_val(v)]) has_null_batch = any(_is_null_val(v) for v in batch_vals) for i in range(0, len(non_null_batch), batch_size): @@ -2616,6 +2735,75 @@ def _infer_projection_type(self, proj: exp.Expr) -> exp.DataType: return exp.DataType.build("VARCHAR(255)", dialect=self.dialect) if isinstance(inner, exp.Boolean): return exp.DataType.build("BOOLEAN", dialect=self.dialect) + if isinstance( + inner, + ( + exp.Avg, + exp.Stddev, + exp.StddevPop, + exp.StddevSamp, + exp.Variance, + exp.VariancePop, + ), + ): + return exp.DataType.build("DOUBLE", dialect=self.dialect) + if isinstance(inner, (exp.Array, exp.ArrayAgg)): + return exp.DataType.build("ARRAY", dialect=self.dialect) + if isinstance(inner, (exp.Map, exp.VarMap)): + return exp.DataType.build("MAP", dialect=self.dialect) + if isinstance(inner, exp.Struct): + return exp.DataType.build("STRUCT", dialect=self.dialect, udt=True) + if isinstance(inner, (exp.JSONObject, exp.JSONArray)): + return exp.DataType.build("JSON", dialect=self.dialect) + if isinstance(inner, exp.GroupConcat): + return exp.DataType.build("TEXT", dialect=self.dialect) + if isinstance(inner, (exp.Count, exp.RowNumber, exp.Rank, exp.DenseRank)): + return exp.DataType.build("BIGINT", dialect=self.dialect) + if isinstance( + inner, + ( + exp.Year, + exp.Month, + exp.Day, + exp.Hour, + exp.Minute, + exp.Second, + exp.DayOfWeek, + exp.DayOfYear, + exp.WeekOfYear, + ), + ): + return exp.DataType.build("INT", dialect=self.dialect) + if isinstance(inner, (exp.CurrentDate, exp.Date, exp.TsOrDsToDate)): + return exp.DataType.build("DATE", dialect=self.dialect) + if isinstance(inner, (exp.CurrentTimestamp, exp.DateTrunc, exp.TimestampTrunc)): + return exp.DataType.build("DATETIME", dialect=self.dialect) + if isinstance(inner, exp.Anonymous): + fn_upper = inner.name.upper() + if fn_upper in ( + "AVG", + "STDDEV", + "STDDEV_POP", + "STDDEV_SAMP", + "VARIANCE", + "VAR_POP", + "VAR_SAMP", + ): + return exp.DataType.build("DOUBLE", dialect=self.dialect) + if fn_upper in ("ARRAY", "ARRAY_AGG", "COLLECT_LIST", "COLLECT_SET"): + return exp.DataType.build("ARRAY", dialect=self.dialect) + if fn_upper in ("MAP", "MAP_AGG"): + return exp.DataType.build("MAP", dialect=self.dialect) + if fn_upper in ("STRUCT", "NAMED_STRUCT"): + return exp.DataType.build("STRUCT", dialect=self.dialect, udt=True) + if fn_upper in ("JSON_OBJECT", "JSON_ARRAY", "TO_JSON"): + return exp.DataType.build("JSON", dialect=self.dialect) + if fn_upper == "GROUP_CONCAT": + return exp.DataType.build("TEXT", dialect=self.dialect) + if fn_upper == "DATE_TRUNC": + return exp.DataType.build("DATETIME", dialect=self.dialect) + if fn_upper == "TO_DATE": + return exp.DataType.build("DATE", dialect=self.dialect) return exp.DataType.build("UNKNOWN") def _reorder_unannotated_ctas_query( @@ -2794,7 +2982,8 @@ def _reorder_olap_columns( updated_types[k] = exp.DataType.build("VARCHAR(255)", dialect=self.dialect) else: raise SQLMeshError( - f"Doris key column '{k}' has type '{k_dtype.sql(dialect=self.dialect)}' which is not supported as an OLAP key column." + f"Doris key column '{k}' has type '{k_dtype.sql(dialect=self.dialect)}' which is not supported as an OLAP key column. " + "Use an integer, DATE/DATETIME, DECIMAL, CHAR, or VARCHAR(n) column instead." ) key_cols_set = set(key_cols) new_order = ( @@ -3050,6 +3239,10 @@ def _build_partitioned_by_exp( # On Doris OLAP tables, if no explicit partitions, auto_partition, or dynamic_partition.enable # are configured, omit PARTITION BY so Doris does not create a 0-partition table that rejects inserts. if not partitions and not is_auto and not dynamic_partition_enabled: + logger.warning( + "[Doris] Omitting PARTITION BY on OLAP table because none of 'auto_partition = true', " + "'dynamic_partition.enable = true', or explicit 'partitions' were configured in physical_properties." + ) return None create_expressions: t.Optional[t.List[exp.Expr]] = None @@ -3271,7 +3464,8 @@ def _build_table_properties_exp( if order_by_cols and not is_materialized_view: if not unique_key_cols: raise SQLMeshError( - "Doris table property 'order_by' is only supported on UNIQUE KEY tables." + "Doris table property 'order_by' is only supported on UNIQUE KEY tables. " + "For DUPLICATE KEY or AGGREGATE KEY tables, sort order is determined by the key columns." ) if col_by_lower: missing_order_cols = [k for k in order_by_cols if k.lower() not in col_by_lower] @@ -3559,6 +3753,27 @@ def _resolve_buckets(self, raw_buckets: t.Any) -> t.Optional[exp.Expr]: f"Invalid buckets value for Doris DISTRIBUTED BY: {raw_buckets}. Expected a positive integer or 'AUTO'." ) + @classmethod + def _eval_int_limit_expr(cls, node: t.Optional[exp.Expr]) -> t.Optional[int]: + if node is None: + return None + curr = node.unnest() + if isinstance(curr, exp.Literal) and not curr.is_string: + try: + return int(str(curr.this).strip()) + except ValueError: + return None + if isinstance(curr, (exp.Add, exp.Sub, exp.Mul)): + left = cls._eval_int_limit_expr(curr.this) + right = cls._eval_int_limit_expr(curr.expression) + if left is not None and right is not None: + if isinstance(curr, exp.Add): + return left + right + if isinstance(curr, exp.Sub): + return max(0, left - right) + return left * right + return None + def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> str: if expression.find(exp.Lock): expression = expression.copy() @@ -3566,6 +3781,15 @@ def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> if select_node.args.get("locks"): select_node.set("locks", None) logger.debug("[Doris] Removed unsupported FOR UPDATE lock from SELECT statement") + if any( + not isinstance(lim.expression, exp.Literal) for lim in expression.find_all(exp.Limit) + ): + expression = expression.copy() + for lim in expression.find_all(exp.Limit): + if not isinstance(lim.expression, exp.Literal): + folded = self._eval_int_limit_expr(lim.expression) + if folded is not None: + lim.set("expression", exp.Literal.number(folded)) if any( isinstance(is_node.expression, exp.Boolean) for is_node in expression.find_all(exp.Is) ): diff --git a/sqlmesh/dbt/model.py b/sqlmesh/dbt/model.py index 55994abf85..0920a23fad 100644 --- a/sqlmesh/dbt/model.py +++ b/sqlmesh/dbt/model.py @@ -151,6 +151,15 @@ class ModelConfig(BaseModelConfig): inserts_only: t.Optional[bool] = None incremental_predicates: t.Optional[t.List[str]] = None + # doris + distributed_by: t.Optional[t.Union[t.List[str], str, t.Dict[str, t.Any]]] = None + buckets: t.Optional[t.Union[int, str]] = None + duplicate_key: t.Optional[t.Union[t.List[str], str]] = None + aggregate_key: t.Optional[t.Union[t.List[str], str]] = None + keys: t.Optional[t.Union[t.List[str], str]] = None + table_type: t.Optional[str] = None + properties: t.Optional[t.Dict[str, t.Any]] = None + _sql_validator = sql_str_validator _on_destructive_change_validator = on_destructive_change_validator _on_additive_change_validator = on_additive_change_validator @@ -715,6 +724,71 @@ def to_sqlmesh( if physical_properties: model_kwargs["physical_properties"] = physical_properties + if context.target.dialect == "doris": + for prop_name, raw_cols in ( + ("unique_key", self.unique_key), + ("duplicate_key", self.duplicate_key), + ("aggregate_key", self.aggregate_key), + ("keys", self.keys), + ("order_by", self.order_by), + ): + if raw_cols: + if prop_name == "unique_key" and ( + self.duplicate_key or self.aggregate_key or self.keys + ): + continue + col_list = raw_cols if isinstance(raw_cols, list) else [raw_cols] + parsed_cols = [d.parse_one(c, dialect=model_dialect) for c in col_list] + physical_properties[prop_name] = ( + parsed_cols[0] + if len(parsed_cols) == 1 + else exp.Tuple(expressions=parsed_cols) + ) + + if self.table_type: + physical_properties["table_type"] = exp.Literal.string(self.table_type) + + if self.distributed_by is not None: + if isinstance(self.distributed_by, list): + dist_cols = [d.parse_one(c, dialect=model_dialect) for c in self.distributed_by] + physical_properties["distributed_by"] = ( + dist_cols[0] if len(dist_cols) == 1 else exp.Tuple(expressions=dist_cols) + ) + elif isinstance(self.distributed_by, dict): + physical_properties["distributed_by"] = exp.Tuple( + expressions=[ + exp.EQ( + this=exp.column(str(k)), + expression=exp.Literal.number(v) + if isinstance(v, int) and not isinstance(v, bool) + else exp.Literal.string(str(v)), + ) + for k, v in self.distributed_by.items() + ] + ) + else: + physical_properties["distributed_by"] = exp.Literal.string( + str(self.distributed_by) + ) + + if self.buckets is not None: + physical_properties["buckets"] = ( + exp.Literal.number(self.buckets) + if isinstance(self.buckets, int) and not isinstance(self.buckets, bool) + else exp.Literal.string(str(self.buckets)) + ) + + if self.properties: + for k, v in self.properties.items(): + physical_properties[str(k)] = ( + exp.Literal.number(v) + if isinstance(v, int) and not isinstance(v, bool) + else exp.Literal.string(str(v)) + ) + + if physical_properties: + model_kwargs["physical_properties"] = physical_properties + kind = self.model_kind(context) # A falsy grants config (None or {}) is considered as unmanaged per dbt semantics diff --git a/sqlmesh/dbt/target.py b/sqlmesh/dbt/target.py index 62683ecfac..413108f89e 100644 --- a/sqlmesh/dbt/target.py +++ b/sqlmesh/dbt/target.py @@ -16,6 +16,7 @@ ClickhouseConnectionConfig, ConnectionConfig, DatabricksConnectionConfig, + DorisConnectionConfig, DuckDBConnectionConfig, MSSQLConnectionConfig, PostgresConnectionConfig, @@ -137,6 +138,8 @@ def load(cls, data: t.Dict[str, t.Any]) -> TargetConfig: return ClickhouseConfig(**data) if db_type == "athena": return AthenaConfig(**data) + if db_type == "doris": + return DorisConfig(**data) raise ConfigError(f"{db_type} not supported.") @@ -1098,6 +1101,68 @@ def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig: ) +class DorisConfig(TargetConfig): + """ + Project connection and operational configuration for the Apache Doris target (`dbt-doris`). + """ + + type: t.Literal["doris"] = "doris" + host: str = "localhost" + port: int = 9030 + user: str = Field(default="root", validation_alias=AliasChoices("user", "username")) + password: t.Optional[str] = Field( + default=None, validation_alias=AliasChoices("password", "pass") + ) + schema_: str = Field(default="default", alias="schema") + database: str = "" + catalog: t.Optional[str] = None + charset: t.Optional[str] = None + connect_timeout: int = Field( + default=10, validation_alias=AliasChoices("connect_timeout", "check_timeout") + ) + session_variables: t.Optional[t.Dict[str, t.Any]] = None + + def default_incremental_strategy(self, kind: IncrementalKind) -> str: + return "merge" if kind is IncrementalByUniqueKeyKind else "insert_overwrite" + + @classproperty + def relation_class(cls) -> t.Type[BaseRelation]: + try: + from dbt.adapters.doris.relation import DorisRelation # type: ignore + + return DorisRelation + except ImportError: + return BaseRelation + + @classproperty + def column_class(cls) -> t.Type[Column]: + try: + from dbt.adapters.doris.column import DorisColumn # type: ignore + + return DorisColumn + except ImportError: + return Column + + @with_schema_differ_overrides + def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig: + catalog = self.catalog or ( + self.database if self.database and self.database != self.schema_ else None + ) + return DorisConnectionConfig( + host=self.host, + port=self.port, + user=self.user, + password=self.password, + database=self.schema_, + catalog=catalog, + charset=self.charset, + connect_timeout=self.connect_timeout, + session_variables=self.session_variables, + concurrent_tasks=self.threads, + **kwargs, + ) + + TARGET_TYPE_TO_CONFIG_CLASS: t.Dict[str, t.Type[TargetConfig]] = { "databricks": DatabricksConfig, "duckdb": DuckDbConfig, @@ -1110,4 +1175,5 @@ def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig: "trino": TrinoConfig, "athena": AthenaConfig, "clickhouse": ClickhouseConfig, + "doris": DorisConfig, } diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 220126e932..fcbec0a7d4 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -2980,3 +2980,184 @@ def worker_fn() -> None: ) assert "state_connection:\n type: duckdb\n database: state.db" in doris_init_cfg assert "dialect: doris" in doris_init_cfg + + +def test_union_time_transform_delete_agg_ctas_mv_df_order_limit_fold_and_dbt_target( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + import pandas as pd + from sqlmesh.core.config.connection import DorisConnectionConfig + from sqlmesh.dbt.context import DbtContext + from sqlmesh.dbt.model import ModelConfig + from sqlmesh.dbt.target import DorisConfig, TargetConfig + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. `_delete_by_time_transform_keys` uses the source subquery's projection column (`src_ds`) + # when it differs from the target column (`target_ds`), and supports `UNION ALL` subqueries. + mocker.patch.object( + adapter.cursor, + "fetchall", + side_effect=[ + [("2024-01-15", "2024-01-16")], + [("2024-01-01", "2025-01-01")], + ], + ) + adapter.delete_from( + "db.events", + where=parse_one( + "DATE_TRUNC(target_ds, 'day') IN (SELECT DATE_TRUNC(src_ds, 'day') FROM db.staging)", + read="doris", + ), + ) + adapter.delete_from( + "db.events", + where=parse_one( + "YEAR(ds) IN (SELECT YEAR(ds) FROM db.s1 UNION ALL SELECT YEAR(ds) FROM db.s2)", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "SELECT DISTINCT DATE_TRUNC(`src_ds`, 'day') AS `_low_0`, DATE_ADD(DATE_TRUNC(`src_ds`, 'day'), INTERVAL 1 DAY) AS `_high_0` FROM `db`.`staging`", + "DELETE FROM `db`.`events` WHERE `target_ds` >= '2024-01-15' AND `target_ds` < '2024-01-16'", + "SELECT DISTINCT DATE_TRUNC(`_col_0`, 'year') AS `_low_0`, DATE_ADD(DATE_TRUNC(`_col_0`, 'year'), INTERVAL 1 YEAR) AS `_high_0` FROM (SELECT `ds` AS `_col_0` FROM `db`.`s1` UNION ALL SELECT `ds` FROM `db`.`s2`) AS `_sq`", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-01' AND `ds` < '2025-01-01'", + ] + + # 2. `CONCAT_WS` multi-column unwrap in `delete_from` supports `UNION ALL` subqueries, + # `_execute_delete_with_subquery` aliases unaliased string literals (`'us' AS _col_0`), + # and `_delete_key_rows` deduplicates repeated literal keys. + adapter.cursor.execute.reset_mock() + key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", exp.column("id"), exp.column("ds")) + union_subq = exp.union( + exp.select(key_exp).from_("db.s1"), + exp.select(key_exp).from_("db.s2"), + distinct=False, + ) + adapter.delete_from("db.target", where=key_exp.isin(query=union_subq)) + adapter.delete_from( + "db.target", + where=parse_one("(region, id) IN (SELECT 'us', id FROM db.staging)", read="doris"), + ) + adapter.delete_from( + "db.target", + where=parse_one( + "(region, id) IN (('us', 1), ('us', 1), ('us', 2))", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "DELETE FROM `db`.`target` AS `_t1` USING (SELECT `id`, `ds` FROM `db`.`s1` UNION ALL SELECT `id`, `ds` FROM `db`.`s2`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`ds` = `_t2`.`ds`", + "DELETE FROM `db`.`target` AS `_t1` USING (SELECT 'us' AS `_col_0`, `id` FROM `db`.`staging`) AS `_t2` WHERE `_t1`.`region` = `_t2`.`_col_0` AND `_t1`.`id` = `_t2`.`id`", + "DELETE FROM `db`.`target` WHERE `region` = 'us' AND `id` IN (1, 2)", + ] + + # 3. Unannotated CTAS with aggregate/complex expressions at column 0 (`AVG(...)`, `ARRAY(...)`) + # infers non-key types (`DOUBLE`, `ARRAY`) and promotes valid key columns (`store_id`, `cnt`) to column 0. + adapter.cursor.execute.reset_mock() + adapter.ctas( + "db.agg_ctas", + parse_one( + "SELECT AVG(price) AS avg_price, CAST(store_id AS INT) AS store_id FROM db.sales GROUP BY store_id", + read="doris", + ), + ) + adapter.ctas( + "db.array_ctas", + parse_one("SELECT ARRAY(1, 2) AS arr, COUNT(*) AS cnt FROM db.src", read="doris"), + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`agg_ctas` AS SELECT CAST(`store_id` AS INT) AS `store_id`, AVG(`price`) AS `avg_price` FROM `db`.`sales` GROUP BY `store_id`", + "CREATE TABLE IF NOT EXISTS `db`.`array_ctas` AS SELECT COUNT(*) AS `cnt`, ARRAY(1, 2) AS `arr` FROM `db`.`src`", + ] + + # 4. `_create_materialized_view` with a DataFrame whose column order (`b`, `a`) differs from + # `target_columns_to_types` (`a`, `b`) reorders the DataFrame columns before `itertuples()`. + adapter.cursor.execute.reset_mock() + mv_df = pd.DataFrame([{"b": "hello", "a": 42}]) + adapter.create_view( + "db.df_mv", + mv_df, + replace=False, + materialized=True, + target_columns_to_types={ + "a": exp.DataType.build("INT", dialect="doris"), + "b": exp.DataType.build("VARCHAR(20)", dialect="doris"), + }, + ) + assert to_sql_calls(adapter) == [ + "CREATE MATERIALIZED VIEW `db`.`df_mv` (`a`, `b`) AS SELECT CAST(`a` AS INT) AS `a`, CAST(`b` AS VARCHAR(20)) AS `b` FROM (SELECT 42 AS `a`, 'hello' AS `b`) AS `t`", + ] + + # 5. `get_alter_operations` suppresses destructive `DROP` + `ADD` pairs on OLAP key columns + # when the target column type (`TEXT`) is an invalid OLAP key type. + adapter.cursor.reset_mock() + mocker.patch.object( + adapter.cursor, + "fetchall", + side_effect=[ + [("code", "VARCHAR(255)", "Yes", "true", None, "")], + [("code", "TEXT", "Yes", "false", None, "")], + ], + ) + ops_destructive = adapter.get_alter_operations( + "db.current_tbl", "db.target_tbl", ignore_destructive=False + ) + assert ops_destructive == [] + + # 6. `_to_sql` folds constant integer arithmetic inside `exp.Limit` (`LIMIT 10 + 1` -> `LIMIT 11`) + # so built-in `number_of_rows` audit queries execute cleanly on Doris. + audit_query = parse_one( + "SELECT COUNT(*) FROM (SELECT 1 FROM db.events WHERE TRUE LIMIT 10 + 1) AS _q HAVING COUNT(*) <= 10", + read="doris", + ) + assert ( + adapter._to_sql(audit_query) + == "SELECT COUNT(*) FROM (SELECT 1 FROM `db`.`events` WHERE TRUE LIMIT 11) AS `_q` HAVING COUNT(*) <= 10" + ) + + # 7. `dbt-doris` TargetConfig (`DorisConfig`) and `ModelConfig` integration. + target = TargetConfig.load( + { + "type": "doris", + "name": "dev", + "host": "fe.internal", + "port": 9030, + "username": "doris_admin", + "pass": "secret", + "schema": "analytics", + "catalog": "iceberg_cat", + "threads": 8, + } + ) + assert isinstance(target, DorisConfig) + conn_cfg = target.to_sqlmesh() + assert isinstance(conn_cfg, DorisConnectionConfig) + assert conn_cfg.host == "fe.internal" + assert conn_cfg.user == "doris_admin" + assert conn_cfg.password == "secret" + assert conn_cfg.database == "analytics" + assert conn_cfg.catalog == "iceberg_cat" + assert conn_cfg.concurrent_tasks == 8 + + dbt_ctx = DbtContext() + dbt_ctx.project_name = "test" + dbt_ctx.target = target + dbt_model = ModelConfig( + name="orders_doris", + alias="orders_doris", + package_name="pkg", + schema="analytics", + materialized="incremental", + unique_key=["order_id"], + distributed_by=["order_id"], + buckets=16, + properties={"replication_num": "1"}, + sql="SELECT 1 AS order_id", + ) + sqlmesh_model = dbt_model.to_sqlmesh(dbt_ctx) + rendered_props = sqlmesh_model.render_physical_properties() + assert rendered_props["unique_key"] == exp.column("order_id", quoted=True) + assert rendered_props["distributed_by"] == exp.column("order_id", quoted=True) + assert rendered_props["buckets"] == exp.Literal.number(16) + assert rendered_props["replication_num"] == exp.Literal.string("1") From b1344ff0933ca424e1f007062a1b20b91ae50bf4 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sat, 24 Oct 2026 21:15:03 -0700 Subject: [PATCH 26/27] fix(doris): harden UNION CTAS reordering, DUPLICATE KEY partitioning, TLS config, and native docs Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/concepts/models/model_kinds.md | 5 + docs/integrations/engines/doris.md | 48 ++- sqlmesh/core/config/connection.py | 8 +- sqlmesh/core/engine_adapter/doris.py | 402 ++++++++++++++++++++---- sqlmesh/dbt/model.py | 74 ----- sqlmesh/dbt/target.py | 66 ---- tests/core/engine_adapter/test_doris.py | 146 ++++++--- 7 files changed, 483 insertions(+), 266 deletions(-) diff --git a/docs/concepts/models/model_kinds.md b/docs/concepts/models/model_kinds.md index cde104790a..6ef2cc47fe 100644 --- a/docs/concepts/models/model_kinds.md +++ b/docs/concepts/models/model_kinds.md @@ -444,6 +444,7 @@ Depending on the target engine, models of the `INCREMENTAL_BY_TIME_RANGE` kind a | Redshift | DELETE by time range, then INSERT | | Postgres | DELETE by time range, then INSERT | | DuckDB | DELETE by time range, then INSERT | +| Doris | DELETE by time range, then INSERT | ## INCREMENTAL_BY_UNIQUE_KEY @@ -669,6 +670,7 @@ Depending on the target engine, models of the `INCREMENTAL_BY_UNIQUE_KEY` kind a | Redshift | MERGE ON unique key | | Postgres | MERGE ON unique key | | DuckDB | DELETE ON matched + INSERT new rows | +| Doris | DELETE ON matched + INSERT new rows | ## FULL Models of the `FULL` kind cause the dataset associated with a model to be fully refreshed (rewritten) upon each model evaluation. @@ -761,6 +763,7 @@ Depending on the target engine, models of the `FULL` kind are materialized using | Redshift | DROP TABLE, CREATE TABLE, INSERT | | Postgres | DROP TABLE, CREATE TABLE, INSERT | | DuckDB | CREATE OR REPLACE TABLE | +| Doris | TRUNCATE TABLE, INSERT (OLAP) / INSERT OVERWRITE (Iceberg) | ## VIEW The model kinds described so far cause the output of a model query to be materialized and stored in a physical table. @@ -837,6 +840,7 @@ MODEL ( * BigQuery * Databricks +* Doris * Snowflake During the evaluation of a model of this kind, the view will be replaced or recreated only if the model's query rendered during evaluation does not match the query used during the previous view creation for this model, or if the target view does not exist. Thus, views are recreated only when necessary in order to realize all the benefits provided by materialized views. @@ -1629,6 +1633,7 @@ Depending on the target engine, models of the `INCREMENTAL_BY_PARTITION` kind ar | Redshift | DELETE by partitioning key, then INSERT | | Postgres | DELETE by partitioning key, then INSERT | | DuckDB | DELETE by partitioning key, then INSERT | +| Doris | DELETE by partitioning key, then INSERT | ## INCREMENTAL_UNMANAGED diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index 389b8a8075..e1167ac5ca 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -586,6 +586,31 @@ FROM raw.logs; --- +## Audits & Data Diff (`sqlmesh audit` & `sqlmesh table_diff`) + +All built-in and custom SQLMesh [audits](../../concepts/audits.md) (`not_null`, `unique_values`, `number_of_rows`, `accepted_values`, `forall`, etc.) and [`sqlmesh table_diff`](../../guides/tablediff.md) work out of the box on Doris: + +```sql +MODEL ( + name analytics.daily_orders, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column order_date, + partition_by_time_column false + ), + audits ( + not_null(columns := (order_id, customer_id)), + unique_values(columns := (order_id)), + number_of_rows(threshold := 1) + ) +); +``` + +- **Constant folding in `LIMIT` clauses**: Doris requires `LIMIT` arguments to be integer literals rather than arithmetic expressions (such as `LIMIT 10 + 1` emitted by `number_of_rows`). `DorisEngineAdapter` automatically folds constant arithmetic in `LIMIT` clauses before execution. +- **Deferred Materialized Views (`build = 'DEFERRED'`)**: When a materialized view is created with `build = 'DEFERRED'` and has `audits (...)` attached, SQLMesh automatically triggers `REFRESH MATERIALIZED VIEW ... AUTO` before running audits so they evaluate against populated data. +- **`sqlmesh table_diff` on floating-point and string columns**: Temporary diff tables created by `sqlmesh table_diff` automatically promote integer indicator columns into the Doris OLAP short-key prefix, so comparing tables whose leading user columns are `DOUBLE`, `FLOAT`, `STRING`, or `JSON` succeeds without manual key configuration. + +--- + ## Supported Model Kinds & Catalog Matrix | SQLMesh Model Kind | Doris Internal OLAP Catalog (`internal`) | External Apache Iceberg Catalog (`table_format iceberg`) | @@ -602,21 +627,18 @@ FROM raw.logs; --- -## Migrating from `dbt-doris` or StarRocks +## Property Aliases & StarRocks Compatibility -SQLMesh supports loading existing `dbt-doris` projects directly (`type: doris` in `profiles.yml`) as well as native SQLMesh Doris models. If you are migrating a `dbt-doris` project or StarRocks SQLMesh models to native SQLMesh Doris models, map your configuration keys as follows: +If you are migrating StarRocks SQLMesh models or existing Doris DDL conventions, `DorisEngineAdapter` supports these property shorthands in `physical_properties`: -| `dbt-doris` / StarRocks Config | SQLMesh Doris Equivalent | +| Property / Shorthand | SQLMesh Doris Behavior | | --- | --- | -| `materialized='table'` | `kind FULL` | -| `materialized='incremental'` (with `unique_key`) | `kind INCREMENTAL_BY_UNIQUE_KEY (unique_key ...)` | -| `materialized='incremental'` (partition overwrite) | `kind INCREMENTAL_BY_TIME_RANGE (time_column ...)` or `kind INCREMENTAL_BY_PARTITION` | -| `duplicate_key=['c1', 'c2']` | `physical_properties (duplicate_key = ('c1', 'c2'))` | -| `unique_key=['c1']` | `kind INCREMENTAL_BY_UNIQUE_KEY (unique_key c1)` or `physical_properties (unique_key = ('c1'))` | -| `table_type='PRIMARY' \| 'UNIQUE' \| 'DUPLICATE' \| 'AGGREGATE'`, `keys=['c1', 'c2']` | Supported directly in `physical_properties` (`PRIMARY` maps to Doris `UNIQUE KEY`) | -| `distributed_by=['c1']`, `buckets=10` | `clustered_by (c1), physical_properties (buckets = 10)` or `physical_properties (distributed_by = ('c1'), buckets = 10)` | -| `partition_by` | Supported in `physical_properties (partition_by = ...)` or top-level `partitioned_by` | -| `properties={...}` | Key-value entries inside `physical_properties (...)` | +| `duplicate_key = ('c1', 'c2')` | Emits `DUPLICATE KEY (c1, c2)` (and automatically appends active OLAP `partitioned_by` columns if omitted from `duplicate_key`) | +| `unique_key = ('c1')` | Emits `UNIQUE KEY (c1)` (automatic for `INCREMENTAL_BY_UNIQUE_KEY`) | +| `aggregate_key = ('c1', 'c2')` | Emits `AGGREGATE KEY (c1, c2)` | +| `table_type = 'PRIMARY' \| 'UNIQUE' \| 'DUPLICATE' \| 'AGGREGATE'`, `keys = ('c1', 'c2')` | Maps directly to the corresponding Doris key model (`PRIMARY` maps to Doris `UNIQUE KEY`) | +| `clustered_by (c1)` + `physical_properties (buckets = 10)` | Emits `DISTRIBUTED BY HASH(c1) BUCKETS 10` | +| `physical_properties (partition_by = ...)` | Supported alias for top-level `partitioned_by` | --- @@ -625,6 +647,6 @@ SQLMesh supports loading existing `dbt-doris` projects directly (`type: doris` i | Symptom / Doris Error | Cause | Resolution | | --- | --- | --- | | `errCode = 2, replication num 3 is larger than available backends 1` | Doris defaults to `replication_num = 3`, which fails on single-BE local/CI clusters. | Set `replication_num: 1` under `model_defaults.physical_properties` in `config.yaml` for local dev. | -| `PARTITION BY` omitted on OLAP table (logged as a warning) | Doris OLAP tables with `PARTITION BY (...) ()` and no partition creation strategy have zero partitions and reject all `INSERT`s. | Add `auto_partition = 'true'`, `'dynamic_partition.enable' = 'true'`, or explicit `partitions = '...'` to `physical_properties`. | +| `PARTITION BY` omitted on OLAP table (logged as a warning) | Doris OLAP tables with `PARTITION BY (...) ()` and no partition creation strategy have zero partitions and reject all `INSERT`s. | Add `auto_partition = 'true'`, `'dynamic_partition.enable' = 'true'`, or explicit `partitions = '...'` to `physical_properties`, or set `partition_by_time_column false` in `INCREMENTAL_BY_TIME_RANGE` if you do not want OLAP partitioning. | | `Doris only supports creating views in the 'internal' catalog` | External Iceberg catalogs do not support Doris `VIEW` or `MATERIALIZED VIEW` objects. | Set `virtual_environment_mode: dev_only` and `environment_catalog_mapping: {"^(?!prod$).*": "internal"}` in `config.yaml`. | | `Key column cannot be of type STRING / DOUBLE / ARRAY` | Doris forbids `STRING`/`TEXT`, floating-point (`FLOAT`/`DOUBLE`), and complex types as OLAP key columns. | SQLMesh automatically reorders valid key columns first and coerces `STRING`/`TEXT` key columns to `VARCHAR(255)`. For explicit keys on `DOUBLE`/`FLOAT`, use `DECIMAL` or an integer/date key column. | diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 1334810a1c..b3f517b7f2 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -2786,7 +2786,13 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]: val = getattr(self, f"ssl_{key}", None) if val is not None: ssl_config[key] = val - if ssl_config or (self.ssl is not None and self.ssl is not False): + if not ssl_config and ( + (self.ssl is not None and self.ssl is not False) + or self.ssl_verify_cert is not None + or self.ssl_verify_identity is not None + ): + ssl_config["check_hostname"] = bool(self.ssl_verify_identity) + if ssl_config: kwargs["ssl"] = ssl_config for key in ( "ssl_ca", diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index e76e226bb5..2670fb9a9d 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -105,6 +105,7 @@ class DorisEngineAdapter( exp.DataType.build("INT", dialect=DIALECT), exp.DataType.build("BIGINT", dialect=DIALECT), exp.DataType.build("INT128", dialect=DIALECT), + exp.DataType.build("TEXT", dialect=DIALECT), }, exp.DataType.build("TINYINT", dialect=DIALECT): { exp.DataType.build("SMALLINT", dialect=DIALECT), @@ -112,26 +113,46 @@ class DorisEngineAdapter( exp.DataType.build("BIGINT", dialect=DIALECT), exp.DataType.build("INT128", dialect=DIALECT), exp.DataType.build("DOUBLE", dialect=DIALECT), + exp.DataType.build("TEXT", dialect=DIALECT), }, exp.DataType.build("SMALLINT", dialect=DIALECT): { exp.DataType.build("INT", dialect=DIALECT), exp.DataType.build("BIGINT", dialect=DIALECT), exp.DataType.build("INT128", dialect=DIALECT), exp.DataType.build("DOUBLE", dialect=DIALECT), + exp.DataType.build("TEXT", dialect=DIALECT), }, exp.DataType.build("INT", dialect=DIALECT): { exp.DataType.build("BIGINT", dialect=DIALECT), exp.DataType.build("INT128", dialect=DIALECT), exp.DataType.build("DOUBLE", dialect=DIALECT), + exp.DataType.build("TEXT", dialect=DIALECT), }, exp.DataType.build("BIGINT", dialect=DIALECT): { exp.DataType.build("INT128", dialect=DIALECT), + exp.DataType.build("TEXT", dialect=DIALECT), + }, + exp.DataType.build("INT128", dialect=DIALECT): { + exp.DataType.build("LARGEINT", dialect=DIALECT, udt=True), + exp.DataType.build("TEXT", dialect=DIALECT), + }, + exp.DataType.build("LARGEINT", dialect=DIALECT, udt=True): { + exp.DataType.build("INT128", dialect=DIALECT), + exp.DataType.build("TEXT", dialect=DIALECT), }, exp.DataType.build("FLOAT", dialect=DIALECT): { exp.DataType.build("DOUBLE", dialect=DIALECT), + exp.DataType.build("TEXT", dialect=DIALECT), + }, + exp.DataType.build("DOUBLE", dialect=DIALECT): { + exp.DataType.build("TEXT", dialect=DIALECT), }, exp.DataType.build("DATE", dialect=DIALECT): { exp.DataType.build("DATETIME", dialect=DIALECT), + exp.DataType.build("TEXT", dialect=DIALECT), + }, + exp.DataType.build("DATETIME", dialect=DIALECT): { + exp.DataType.build("TEXT", dialect=DIALECT), }, }, } @@ -573,6 +594,13 @@ def get_alter_operations( ignore_additive=ignore_additive, ), ) + dropped_key_cols = { + op.column.name.lower() + for op in operations + if isinstance(op, TableAlterDropColumnOperation) + and op.column.name.lower() in current_key_cols + } + protected_key_cols = invalid_target_key_cols | dropped_key_cols return [ op for op in operations @@ -584,7 +612,7 @@ def get_alter_operations( ) or ( isinstance(op, (TableAlterDropColumnOperation, TableAlterAddColumnOperation)) - and op.column.name.lower() in invalid_target_key_cols + and op.column.name.lower() in protected_key_cols ) ) ] @@ -1262,6 +1290,14 @@ def insert_overwrite_by_partition( target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, source_columns: t.Optional[t.List[str]] = None, ) -> None: + if not partitioned_by: + self.replace_query( + table_name, + query_or_df, + target_columns_to_types=target_columns_to_types, + source_columns=source_columns, + ) + return parsed_exprs, _ = self._parse_partition_expressions(partitioned_by) col_order: t.List[str] = [] col_to_expr: t.Dict[str, exp.Column] = {} @@ -1514,14 +1550,8 @@ def _extract_subquery_from_where( normalized_keys.append(k) elif isinstance(k, exp.Identifier): normalized_keys.append(exp.Column(this=k.copy())) - elif (transform := self._extract_time_partition_transform(k)) is not None: - col_expr, unit = transform - normalized_keys.append( - exp.Anonymous( - this="DATE_TRUNC", - expressions=[col_expr, exp.Literal.string(unit)], - ) - ) + elif self._extract_time_partition_transform(k) is not None: + normalized_keys.append(k.copy()) else: return None @@ -1915,6 +1945,7 @@ def _delete_by_time_transform_keys( return projections: t.List[exp.Expr] = [] + raw_year_indices: t.Set[int] = set() if ( isinstance(select_query, exp.Select) and len(select_query.expressions) == len(key_exprs) @@ -1928,9 +1959,22 @@ def _delete_by_time_transform_keys( transform = self._extract_time_partition_transform(k) if transform is not None: _, unit = transform + k_unnested = k.unnest() + is_year_fn = ( + unit == "year" + and not isinstance(k_unnested, (exp.TimestampTrunc, exp.DateTrunc)) + and not ( + isinstance(k_unnested, exp.Anonymous) + and k_unnested.name.upper() == "DATE_TRUNC" + ) + ) src_transform = self._extract_time_partition_transform(raw_sel) if src_transform is not None: src_arg: exp.Expr = src_transform[0].copy() + elif is_year_fn: + raw_year_indices.add(idx) + projections.append(exp.alias_(raw_sel.copy(), f"_val_{idx}", quoted=True)) + continue else: src_arg = raw_sel.copy() trunc_expr = exp.Anonymous( @@ -1959,16 +2003,24 @@ def _collect_select_branches(node: exp.Expr) -> None: all_branches.append(node) _collect_select_branches(select_query) + unwrapped_indices: t.Set[int] = set(range(len(key_exprs))) for branch in all_branches: - if len(branch.expressions) == len(key_exprs) and not branch.args.get("group"): + if len(branch.expressions) == len(key_exprs): + can_unwrap_group = not branch.args.get("having") and not branch.args.get( + "limit" + ) branch_exprs = list(branch.expressions) for idx, k in enumerate(key_exprs): if self._extract_time_partition_transform(k) is not None: sel_item = branch_exprs[idx] raw_sel = sel_item.this if isinstance(sel_item, exp.Alias) else sel_item src_transform = self._extract_time_partition_transform(raw_sel) - if src_transform is not None: + if src_transform is not None and can_unwrap_group: branch_exprs[idx] = src_transform[0].copy() + if branch.args.get("group"): + branch.set("group", None) + else: + unwrapped_indices.discard(idx) branch.set("expressions", branch_exprs) new_first_exprs = list(first_select.expressions) @@ -1982,6 +2034,19 @@ def _collect_select_branches(node: exp.Expr) -> None: transform = self._extract_time_partition_transform(k) if transform is not None: _, unit = transform + k_unnested = k.unnest() + is_year_fn = ( + unit == "year" + and not isinstance(k_unnested, (exp.TimestampTrunc, exp.DateTrunc)) + and not ( + isinstance(k_unnested, exp.Anonymous) + and k_unnested.name.upper() == "DATE_TRUNC" + ) + ) + if idx not in unwrapped_indices and is_year_fn: + raw_year_indices.add(idx) + projections.append(exp.alias_(sq_col, f"_val_{idx}", quoted=True)) + continue trunc_expr = exp.Anonymous( this="DATE_TRUNC", expressions=[sq_col, exp.Literal.string(unit)], @@ -2023,11 +2088,33 @@ def _collect_select_branches(node: exp.Expr) -> None: for row in rows: predicates: t.List[exp.Expr] = [] pos = 0 - for k in key_exprs: + for idx, k in enumerate(key_exprs): transform = self._extract_time_partition_transform(k) if transform is not None: col_expr, _ = transform target_col = exp.column(col_expr.name, quoted=True) + if idx in raw_year_indices: + yr_val = self._sanitize_key_value(row[pos]) + pos += 1 + if yr_val is None: + predicates.append(target_col.is_(exp.null())) + else: + yr_int = int(yr_val) + low_val: t.Any = f"{yr_int:04d}-01-01" + high_val: t.Any = f"{yr_int + 1:04d}-01-01" + predicates.append( + exp.GTE( + this=target_col.copy(), + expression=self._to_delete_literal_expr(low_val), + ) + ) + predicates.append( + exp.LT( + this=target_col.copy(), + expression=self._to_delete_literal_expr(high_val), + ) + ) + continue low_val = self._sanitize_key_value(row[pos]) high_val = self._sanitize_key_value(row[pos + 1]) pos += 2 @@ -2652,34 +2739,67 @@ def _build_create_table_exp( ) updated_exprs.append(expr) table_name_or_schema.set("expressions", updated_exprs) - elif auto_range_cols and isinstance(expression, exp.Select): + elif auto_range_cols and isinstance(expression, (exp.Select, exp.SetOperation)): expression = expression.copy() - new_select_exprs: t.List[exp.Expr] = [] - for proj in expression.expressions: - if proj.alias_or_name.lower() in auto_range_cols: - if isinstance(proj, exp.Alias): - inner = proj.this - if not ( - isinstance(inner, exp.Anonymous) - and inner.name.upper() == "NON_NULLABLE" + auto_select_branches: t.List[exp.Select] = [] + + def _collect_auto_branches(node: exp.Expr) -> None: + if isinstance(node, exp.SetOperation): + _collect_auto_branches(node.this) + _collect_auto_branches(node.expression) + elif isinstance(node, exp.Select): + auto_select_branches.append(node) + + _collect_auto_branches(expression) + if auto_select_branches: + target_indices = { + idx + for idx, proj in enumerate(auto_select_branches[0].expressions) + if proj.alias_or_name.lower() in auto_range_cols + } + for branch_idx, branch in enumerate(auto_select_branches): + new_select_exprs: t.List[exp.Expr] = [] + for idx, proj in enumerate(branch.expressions): + if ( + proj.alias_or_name.lower() in auto_range_cols + or idx in target_indices ): - proj = proj.copy() - proj.set( - "this", - exp.Anonymous( - this="NON_NULLABLE", expressions=[inner.copy()] - ), - ) - elif isinstance(proj, exp.Column): - proj = exp.alias_( - exp.Anonymous( - this="NON_NULLABLE", expressions=[proj.copy()] - ), - proj.name, - quoted=bool(getattr(proj.this, "quoted", True)), - ) - new_select_exprs.append(proj) - expression.set("expressions", new_select_exprs) + if isinstance(proj, exp.Alias): + inner = proj.this + if not ( + isinstance(inner, exp.Anonymous) + and inner.name.upper() == "NON_NULLABLE" + ): + proj = proj.copy() + proj.set( + "this", + exp.Anonymous( + this="NON_NULLABLE", + expressions=[inner.copy()], + ), + ) + elif isinstance(proj, exp.Column): + wrapped = exp.Anonymous( + this="NON_NULLABLE", expressions=[proj.copy()] + ) + proj = ( + exp.alias_( + wrapped, + proj.name, + quoted=bool(getattr(proj.this, "quoted", True)), + ) + if branch_idx == 0 + else wrapped + ) + elif not ( + isinstance(proj, exp.Anonymous) + and proj.name.upper() == "NON_NULLABLE" + ): + proj = exp.Anonymous( + this="NON_NULLABLE", expressions=[proj.copy()] + ) + new_select_exprs.append(proj) + branch.set("expressions", new_select_exprs) return super()._build_create_table_exp( table_name_or_schema, @@ -2811,6 +2931,91 @@ def _reorder_unannotated_ctas_query( ) -> exp.Expr: if isinstance(expression, exp.SetOperation): expression = expression.copy() + branches: t.List[exp.Select] = [] + + def _collect_branches(node: exp.Expr) -> None: + if isinstance(node, exp.SetOperation): + _collect_branches(node.this) + _collect_branches(node.expression) + elif isinstance(node, exp.Select): + branches.append(node) + + _collect_branches(expression) + if not branches: + return expression + first_exprs = branches[0].expressions + num_cols = len(first_exprs) + if ( + num_cols > 0 + and all( + len(b.expressions) == num_cols + and not any(isinstance(p, exp.Star) for p in b.expressions) + for b in branches + ) + and len( + { + p.alias_or_name.lower() + for p in first_exprs + if p.alias_or_name and not isinstance(p, exp.Star) + } + ) + == num_cols + ): + orig_names = [p.alias_or_name for p in first_exprs] + merged_cols_to_types: t.Dict[str, exp.DataType] = {} + for col_idx, col_name in enumerate(orig_names): + inferred = exp.DataType.build("UNKNOWN") + for b in branches: + cand = self._infer_projection_type(b.expressions[col_idx]) + if cand.this != exp.DataType.Type.UNKNOWN: + inferred = cand + break + merged_cols_to_types[col_name] = inferred + reordered_select_cols = self._reorder_olap_columns(merged_cols_to_types, kwargs) + perm = [orig_names.index(c) for c in reordered_select_cols] + coercions = { + orig_names.index(c): reordered_select_cols[c] + for c in reordered_select_cols + if reordered_select_cols[c] != merged_cols_to_types[c] + } + if perm != list(range(num_cols)) or coercions: + for b_idx, branch in enumerate(branches): + old_projs = list(branch.expressions) + updated_branch_projs: t.List[exp.Expr] = [] + for orig_idx in perm: + proj = old_projs[orig_idx].copy() + new_dtype = coercions.get(orig_idx) + if new_dtype is not None: + col_name = orig_names[orig_idx] + if isinstance(proj, exp.Alias) and isinstance( + proj.this, (exp.Cast, exp.TryCast) + ): + proj.this.set("to", new_dtype.copy()) + elif isinstance(proj, (exp.Cast, exp.TryCast)): + cast_expr = exp.cast(proj.this.copy(), to=new_dtype.copy()) + proj = ( + exp.alias_(cast_expr, col_name, quoted=True) + if b_idx == 0 + else cast_expr + ) + elif isinstance(proj, exp.Alias): + proj.set( + "this", exp.cast(proj.this.copy(), to=new_dtype.copy()) + ) + elif isinstance(proj, exp.Column): + cast_expr = exp.cast(proj.copy(), to=new_dtype.copy()) + proj = ( + exp.alias_( + cast_expr, + proj.name, + quoted=bool(getattr(proj.this, "quoted", True)), + ) + if b_idx == 0 + else cast_expr + ) + updated_branch_projs.append(proj) + branch.set("expressions", updated_branch_projs) + return expression expression.set("this", self._reorder_unannotated_ctas_query(expression.this, kwargs)) expression.set( "expression", self._reorder_unannotated_ctas_query(expression.expression, kwargs) @@ -2974,6 +3179,42 @@ def _reorder_olap_columns( f"Doris key columns {missing_cols} not found in table schema columns {col_names}." ) key_cols = [col_by_lower[k.lower()] for k in raw_key_cols if k.lower() in col_by_lower] + if duplicate_key and not unique_key and not aggregate_key: + raw_part_by_dup = ( + kwargs.get("partitioned_by") + or table_props.get("partitioned_by") + or table_props.get("partition_by") + ) + has_active_olap_part = bool( + raw_part_by_dup + and ( + table_props.get("partitions") + or kwargs.get("partitions") + or str(self._extract_str_value(table_props.get("auto_partition")) or "") + .strip() + .lower() + in ("true", "1") + or str( + self._extract_str_value(table_props.get("dynamic_partition.enable")) + or "" + ) + .strip() + .lower() + in ("true", "1") + ) + ) + if has_active_olap_part: + part_list_dup = ( + raw_part_by_dup if isinstance(raw_part_by_dup, list) else [raw_part_by_dup] + ) + parsed_part_dup, _ = self._parse_partition_expressions(part_list_dup) + key_cols_lower = {k.lower() for k in key_cols} + for expr in parsed_part_dup: + for col_node in expr.find_all(exp.Column): + canonical = col_by_lower.get(col_node.name.lower()) + if canonical and canonical.lower() not in key_cols_lower: + key_cols.append(canonical) + key_cols_lower.add(canonical.lower()) updated_types = dict(target_columns_to_types) for k in key_cols: k_dtype = updated_types[k] @@ -3241,7 +3482,9 @@ def _build_partitioned_by_exp( if not partitions and not is_auto and not dynamic_partition_enabled: logger.warning( "[Doris] Omitting PARTITION BY on OLAP table because none of 'auto_partition = true', " - "'dynamic_partition.enable = true', or explicit 'partitions' were configured in physical_properties." + "'dynamic_partition.enable = true', or explicit 'partitions' were configured in physical_properties. " + "Add one of those properties to physical_properties, or set 'partition_by_time_column false' in your " + "INCREMENTAL_BY_TIME_RANGE model kind to silence this warning." ) return None @@ -3416,6 +3659,33 @@ def _build_table_properties_exp( {c.lower(): c for c in target_columns_to_types} if target_columns_to_types else {} ) + # Pop partitioning-related properties before building key properties so explicit duplicate_key + # can automatically include active OLAP partition columns (required by Doris FE validation). + partitions = table_properties_copy.pop("partitions", None) + auto_partition = table_properties_copy.pop("auto_partition", None) + raw_part_by = table_properties_copy.pop("partitioned_by", None) + raw_part_by_alias = table_properties_copy.pop("partition_by", None) + if ( + sum( + x is not None + for x in ( + partitioned_by or None, + raw_part_by, + raw_part_by_alias, + ) + ) + > 1 + ): + raise SQLMeshError( + "Multiple partition definitions found: specify only one of model 'partitioned_by' or physical_properties 'partitioned_by' / 'partition_by'." + ) + if not partitioned_by: + chosen_part_by = raw_part_by if raw_part_by is not None else raw_part_by_alias + if chosen_part_by is not None: + partitioned_by = ( + chosen_part_by if isinstance(chosen_part_by, list) else [chosen_part_by] + ) + if not is_iceberg: specified_key_types = [ name @@ -3445,6 +3715,34 @@ def _build_table_properties_exp( elif aggregate_key_cols: aggregate_key_cols = [col_by_lower[k.lower()] for k in aggregate_key_cols] + if ( + duplicate_key_cols + and partitioned_by + and not is_materialized_view + and ( + partitions + or str(self._extract_str_value(auto_partition) or "").strip().lower() + in ("true", "1") + or str( + self._extract_str_value( + table_properties_copy.get("dynamic_partition.enable") + ) + or "" + ) + .strip() + .lower() + in ("true", "1") + ) + ): + parsed_dup_part_exprs, _ = self._parse_partition_expressions(partitioned_by) + dup_lower_set = {k.lower() for k in duplicate_key_cols} + for expr in parsed_dup_part_exprs: + for col_node in expr.find_all(exp.Column): + part_col_name = col_by_lower.get(col_node.name.lower(), col_node.name) + if part_col_name.lower() not in dup_lower_set: + duplicate_key_cols.append(part_col_name) + dup_lower_set.add(part_col_name.lower()) + if unique_key_cols: properties.append( exp.UniqueKeyProperty(expressions=[exp.to_column(k) for k in unique_key_cols]) @@ -3486,32 +3784,6 @@ def _build_table_properties_exp( ) ) - # Pop partitioning-related properties - partitions = table_properties_copy.pop("partitions", None) - auto_partition = table_properties_copy.pop("auto_partition", None) - raw_part_by = table_properties_copy.pop("partitioned_by", None) - raw_part_by_alias = table_properties_copy.pop("partition_by", None) - if ( - sum( - x is not None - for x in ( - partitioned_by or None, - raw_part_by, - raw_part_by_alias, - ) - ) - > 1 - ): - raise SQLMeshError( - "Multiple partition definitions found: specify only one of model 'partitioned_by' or physical_properties 'partitioned_by' / 'partition_by'." - ) - if not partitioned_by: - chosen_part_by = raw_part_by if raw_part_by is not None else raw_part_by_alias - if chosen_part_by is not None: - partitioned_by = ( - chosen_part_by if isinstance(chosen_part_by, list) else [chosen_part_by] - ) - add_partition = True if partitioned_by: parsed_part_exprs, _ = self._parse_partition_expressions(partitioned_by) diff --git a/sqlmesh/dbt/model.py b/sqlmesh/dbt/model.py index 0920a23fad..55994abf85 100644 --- a/sqlmesh/dbt/model.py +++ b/sqlmesh/dbt/model.py @@ -151,15 +151,6 @@ class ModelConfig(BaseModelConfig): inserts_only: t.Optional[bool] = None incremental_predicates: t.Optional[t.List[str]] = None - # doris - distributed_by: t.Optional[t.Union[t.List[str], str, t.Dict[str, t.Any]]] = None - buckets: t.Optional[t.Union[int, str]] = None - duplicate_key: t.Optional[t.Union[t.List[str], str]] = None - aggregate_key: t.Optional[t.Union[t.List[str], str]] = None - keys: t.Optional[t.Union[t.List[str], str]] = None - table_type: t.Optional[str] = None - properties: t.Optional[t.Dict[str, t.Any]] = None - _sql_validator = sql_str_validator _on_destructive_change_validator = on_destructive_change_validator _on_additive_change_validator = on_additive_change_validator @@ -724,71 +715,6 @@ def to_sqlmesh( if physical_properties: model_kwargs["physical_properties"] = physical_properties - if context.target.dialect == "doris": - for prop_name, raw_cols in ( - ("unique_key", self.unique_key), - ("duplicate_key", self.duplicate_key), - ("aggregate_key", self.aggregate_key), - ("keys", self.keys), - ("order_by", self.order_by), - ): - if raw_cols: - if prop_name == "unique_key" and ( - self.duplicate_key or self.aggregate_key or self.keys - ): - continue - col_list = raw_cols if isinstance(raw_cols, list) else [raw_cols] - parsed_cols = [d.parse_one(c, dialect=model_dialect) for c in col_list] - physical_properties[prop_name] = ( - parsed_cols[0] - if len(parsed_cols) == 1 - else exp.Tuple(expressions=parsed_cols) - ) - - if self.table_type: - physical_properties["table_type"] = exp.Literal.string(self.table_type) - - if self.distributed_by is not None: - if isinstance(self.distributed_by, list): - dist_cols = [d.parse_one(c, dialect=model_dialect) for c in self.distributed_by] - physical_properties["distributed_by"] = ( - dist_cols[0] if len(dist_cols) == 1 else exp.Tuple(expressions=dist_cols) - ) - elif isinstance(self.distributed_by, dict): - physical_properties["distributed_by"] = exp.Tuple( - expressions=[ - exp.EQ( - this=exp.column(str(k)), - expression=exp.Literal.number(v) - if isinstance(v, int) and not isinstance(v, bool) - else exp.Literal.string(str(v)), - ) - for k, v in self.distributed_by.items() - ] - ) - else: - physical_properties["distributed_by"] = exp.Literal.string( - str(self.distributed_by) - ) - - if self.buckets is not None: - physical_properties["buckets"] = ( - exp.Literal.number(self.buckets) - if isinstance(self.buckets, int) and not isinstance(self.buckets, bool) - else exp.Literal.string(str(self.buckets)) - ) - - if self.properties: - for k, v in self.properties.items(): - physical_properties[str(k)] = ( - exp.Literal.number(v) - if isinstance(v, int) and not isinstance(v, bool) - else exp.Literal.string(str(v)) - ) - - if physical_properties: - model_kwargs["physical_properties"] = physical_properties - kind = self.model_kind(context) # A falsy grants config (None or {}) is considered as unmanaged per dbt semantics diff --git a/sqlmesh/dbt/target.py b/sqlmesh/dbt/target.py index 413108f89e..62683ecfac 100644 --- a/sqlmesh/dbt/target.py +++ b/sqlmesh/dbt/target.py @@ -16,7 +16,6 @@ ClickhouseConnectionConfig, ConnectionConfig, DatabricksConnectionConfig, - DorisConnectionConfig, DuckDBConnectionConfig, MSSQLConnectionConfig, PostgresConnectionConfig, @@ -138,8 +137,6 @@ def load(cls, data: t.Dict[str, t.Any]) -> TargetConfig: return ClickhouseConfig(**data) if db_type == "athena": return AthenaConfig(**data) - if db_type == "doris": - return DorisConfig(**data) raise ConfigError(f"{db_type} not supported.") @@ -1101,68 +1098,6 @@ def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig: ) -class DorisConfig(TargetConfig): - """ - Project connection and operational configuration for the Apache Doris target (`dbt-doris`). - """ - - type: t.Literal["doris"] = "doris" - host: str = "localhost" - port: int = 9030 - user: str = Field(default="root", validation_alias=AliasChoices("user", "username")) - password: t.Optional[str] = Field( - default=None, validation_alias=AliasChoices("password", "pass") - ) - schema_: str = Field(default="default", alias="schema") - database: str = "" - catalog: t.Optional[str] = None - charset: t.Optional[str] = None - connect_timeout: int = Field( - default=10, validation_alias=AliasChoices("connect_timeout", "check_timeout") - ) - session_variables: t.Optional[t.Dict[str, t.Any]] = None - - def default_incremental_strategy(self, kind: IncrementalKind) -> str: - return "merge" if kind is IncrementalByUniqueKeyKind else "insert_overwrite" - - @classproperty - def relation_class(cls) -> t.Type[BaseRelation]: - try: - from dbt.adapters.doris.relation import DorisRelation # type: ignore - - return DorisRelation - except ImportError: - return BaseRelation - - @classproperty - def column_class(cls) -> t.Type[Column]: - try: - from dbt.adapters.doris.column import DorisColumn # type: ignore - - return DorisColumn - except ImportError: - return Column - - @with_schema_differ_overrides - def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig: - catalog = self.catalog or ( - self.database if self.database and self.database != self.schema_ else None - ) - return DorisConnectionConfig( - host=self.host, - port=self.port, - user=self.user, - password=self.password, - database=self.schema_, - catalog=catalog, - charset=self.charset, - connect_timeout=self.connect_timeout, - session_variables=self.session_variables, - concurrent_tasks=self.threads, - **kwargs, - ) - - TARGET_TYPE_TO_CONFIG_CLASS: t.Dict[str, t.Type[TargetConfig]] = { "databricks": DatabricksConfig, "duckdb": DuckDbConfig, @@ -1175,5 +1110,4 @@ def to_sqlmesh(self, **kwargs: t.Any) -> ConnectionConfig: "trino": TrinoConfig, "athena": AthenaConfig, "clickhouse": ClickhouseConfig, - "doris": DorisConfig, } diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index fcbec0a7d4..91e720c3c9 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -2982,14 +2982,11 @@ def worker_fn() -> None: assert "dialect: doris" in doris_init_cfg -def test_union_time_transform_delete_agg_ctas_mv_df_order_limit_fold_and_dbt_target( +def test_union_time_transform_delete_agg_ctas_mv_df_order_and_limit_fold( make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture ): import pandas as pd from sqlmesh.core.config.connection import DorisConnectionConfig - from sqlmesh.dbt.context import DbtContext - from sqlmesh.dbt.model import ModelConfig - from sqlmesh.dbt.target import DorisConfig, TargetConfig adapter = make_mocked_engine_adapter(DorisEngineAdapter) @@ -3116,48 +3113,103 @@ def test_union_time_transform_delete_agg_ctas_mv_df_order_limit_fold_and_dbt_tar == "SELECT COUNT(*) FROM (SELECT 1 FROM `db`.`events` WHERE TRUE LIMIT 11) AS `_q` HAVING COUNT(*) <= 10" ) - # 7. `dbt-doris` TargetConfig (`DorisConfig`) and `ModelConfig` integration. - target = TargetConfig.load( + # 7. Lockstep `UNION ALL` CTAS reordering when only one branch has explicit types (`AVG(score)`, `CAST(id AS INT)`) + # plus `NON_NULLABLE(...)` wrapping across all `UNION ALL` branches when `auto_partition = true`. + adapter.cursor.execute.reset_mock() + adapter.ctas( + "db.union_lockstep_ctas", + parse_one( + "SELECT AVG(score) AS score, CAST(id AS INT) AS id FROM db.s1 GROUP BY id " + "UNION ALL SELECT score, id FROM db.s2", + read="doris", + ), + ) + adapter.ctas( + "db.union_auto_part_ctas", + parse_one( + "SELECT id, ds FROM db.s1 UNION ALL SELECT id, ds FROM db.s2", + read="doris", + ), + partitioned_by=[parse_one("RANGE(ds)", read="doris")], + table_properties={"auto_partition": exp.Literal.string("true")}, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`union_lockstep_ctas` AS SELECT CAST(`id` AS INT) AS `id`, AVG(`score`) AS `score` FROM `db`.`s1` GROUP BY `id` UNION ALL SELECT `id`, `score` FROM `db`.`s2`", + "CREATE TABLE IF NOT EXISTS `db`.`union_auto_part_ctas` AUTO PARTITION BY RANGE (`ds`) () AS SELECT `id`, NON_NULLABLE(`ds`) AS `ds` FROM `db`.`s1` UNION ALL SELECT `id`, NON_NULLABLE(`ds`) FROM `db`.`s2`", + ] + + # 8. Explicit `duplicate_key` with active OLAP `partitioned_by` automatically includes the + # partition column in `DUPLICATE KEY` and the leading schema prefix (required by Doris FE). + adapter.cursor.execute.reset_mock() + adapter.create_table( + "db.dup_part_tbl", { - "type": "doris", - "name": "dev", - "host": "fe.internal", - "port": 9030, - "username": "doris_admin", - "pass": "secret", - "schema": "analytics", - "catalog": "iceberg_cat", - "threads": 8, - } + "score": exp.DataType.build("DOUBLE", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + }, + partitioned_by=[parse_one("RANGE(ds)", read="doris")], + table_properties={ + "duplicate_key": exp.Tuple(expressions=[exp.Literal.string("id")]), + "auto_partition": exp.Literal.string("true"), + "buckets": exp.Literal.number(8), + }, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`dup_part_tbl` (`id` INT, `ds` DATE NOT NULL, `score` DOUBLE) DUPLICATE KEY (`id`, `ds`) AUTO PARTITION BY RANGE (`ds`) () DISTRIBUTED BY HASH (`id`) BUCKETS 8" + ] + + # 9. `_delete_by_time_transform_keys` handles pre-extracted integer `YEAR` columns (`SELECT yr`), + # `GROUP BY` inside `UNION ALL` branches, and `NULL` partition values (`IS NULL`). + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter.cursor, + "fetchall", + side_effect=[ + [(2024,), (None,)], + [("2024-01-01", "2025-01-01")], + ], + ) + adapter.delete_from( + "db.events", + where=parse_one("YEAR(ds) IN (SELECT yr FROM db.staging)", read="doris"), + ) + adapter.delete_from( + "db.events", + where=parse_one( + "YEAR(ds) IN (SELECT YEAR(ds) FROM db.s1 GROUP BY YEAR(ds) UNION ALL SELECT YEAR(ds) FROM db.s2)", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "SELECT DISTINCT `yr` AS `_val_0` FROM `db`.`staging`", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-01' AND `ds` < '2025-01-01'", + "DELETE FROM `db`.`events` WHERE `ds` IS NULL", + "SELECT DISTINCT DATE_TRUNC(`_col_0`, 'year') AS `_low_0`, DATE_ADD(DATE_TRUNC(`_col_0`, 'year'), INTERVAL 1 YEAR) AS `_high_0` FROM (SELECT `ds` AS `_col_0` FROM `db`.`s1` UNION ALL SELECT `ds` FROM `db`.`s2`) AS `_sq`", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-01' AND `ds` < '2025-01-01'", + ] + + # 10. `insert_overwrite_by_partition` with empty `partitioned_by=[]` delegates cleanly to `replace_query`. + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter, + "get_data_object", + return_value=DataObject( + catalog="internal", schema="db", name="unpart_tbl", type=DataObjectType.TABLE + ), + ) + adapter.insert_overwrite_by_partition( + "db.unpart_tbl", + parse_one("SELECT 1 AS id", read="doris"), + partitioned_by=[], + target_columns_to_types={"id": exp.DataType.build("INT", dialect="doris")}, ) - assert isinstance(target, DorisConfig) - conn_cfg = target.to_sqlmesh() - assert isinstance(conn_cfg, DorisConnectionConfig) - assert conn_cfg.host == "fe.internal" - assert conn_cfg.user == "doris_admin" - assert conn_cfg.password == "secret" - assert conn_cfg.database == "analytics" - assert conn_cfg.catalog == "iceberg_cat" - assert conn_cfg.concurrent_tasks == 8 - - dbt_ctx = DbtContext() - dbt_ctx.project_name = "test" - dbt_ctx.target = target - dbt_model = ModelConfig( - name="orders_doris", - alias="orders_doris", - package_name="pkg", - schema="analytics", - materialized="incremental", - unique_key=["order_id"], - distributed_by=["order_id"], - buckets=16, - properties={"replication_num": "1"}, - sql="SELECT 1 AS order_id", - ) - sqlmesh_model = dbt_model.to_sqlmesh(dbt_ctx) - rendered_props = sqlmesh_model.render_physical_properties() - assert rendered_props["unique_key"] == exp.column("order_id", quoted=True) - assert rendered_props["distributed_by"] == exp.column("order_id", quoted=True) - assert rendered_props["buckets"] == exp.Literal.number(16) - assert rendered_props["replication_num"] == exp.Literal.string("1") + assert to_sql_calls(adapter) == [ + "TRUNCATE TABLE `db`.`unpart_tbl`", + "INSERT INTO `db`.`unpart_tbl` (`id`) SELECT 1 AS `id`", + ] + + # 11. `DorisConnectionConfig` with `ssl=True` populates a non-empty `ssl` dict so PyMySQL enables TLS. + ssl_conn_cfg = DorisConnectionConfig(host="fe.internal", user="root", ssl=True) + assert ssl_conn_cfg._static_connection_kwargs["ssl"] == {"check_hostname": False} + assert bool(ssl_conn_cfg._static_connection_kwargs["ssl"]) is True From 2388534e91932b1ca975940d57486167564d2e2e Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sat, 24 Oct 2026 22:33:43 -0700 Subject: [PATCH 27/27] test(doris): add schema differ, UNION CTAS, TLS matrix, and SnapshotEvaluator tests Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/config/connection.py | 10 +- sqlmesh/core/engine_adapter/doris.py | 10 +- tests/core/engine_adapter/test_doris.py | 365 ++++++++++++++++++++++++ 3 files changed, 378 insertions(+), 7 deletions(-) diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index b3f517b7f2..be49001dd8 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -2786,11 +2786,11 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]: val = getattr(self, f"ssl_{key}", None) if val is not None: ssl_config[key] = val - if not ssl_config and ( - (self.ssl is not None and self.ssl is not False) - or self.ssl_verify_cert is not None - or self.ssl_verify_identity is not None - ): + if self.ssl_verify_cert is not None: + ssl_config["verify_mode"] = self.ssl_verify_cert + if self.ssl_verify_identity is not None: + ssl_config["check_hostname"] = bool(self.ssl_verify_identity) + if not ssl_config and (self.ssl is not None and self.ssl is not False): ssl_config["check_hostname"] = bool(self.ssl_verify_identity) if ssl_config: kwargs["ssl"] = ssl_config diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 2670fb9a9d..3376e7ffc8 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -607,8 +607,14 @@ def get_alter_operations( if not ( ( isinstance(op, TableAlterChangeColumnTypeOperation) - and op.column.name.lower() in current_key_cols - and not self._is_valid_olap_key_type(op.column_type) + and ( + self._to_sql(op.current_type).upper() + == self._to_sql(op.column_type).upper() + or ( + op.column.name.lower() in current_key_cols + and not self._is_valid_olap_key_type(op.column_type) + ) + ) ) or ( isinstance(op, (TableAlterDropColumnOperation, TableAlterAddColumnOperation)) diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 91e720c3c9..0879d9918b 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -3213,3 +3213,368 @@ def test_union_time_transform_delete_agg_ctas_mv_df_order_and_limit_fold( ssl_conn_cfg = DorisConnectionConfig(host="fe.internal", user="root", ssl=True) assert ssl_conn_cfg._static_connection_kwargs["ssl"] == {"check_hostname": False} assert bool(ssl_conn_cfg._static_connection_kwargs["ssl"]) is True + + +def test_schema_differ_text_widening_largeint_and_key_protection( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Widening non-key value columns to TEXT (`INT -> TEXT`, `DATE -> TEXT`) emits in-place + # `MODIFY COLUMN ... TEXT` operations instead of destructive `DROP` + `ADD`. + mocker.patch.object( + adapter.cursor, + "fetchall", + side_effect=[ + [ + ("id", "INT", "Yes", "true", None, ""), + ("val", "INT", "Yes", "false", None, ""), + ("dt", "DATE", "Yes", "false", None, ""), + ], + [ + ("id", "INT", "Yes", "true", None, ""), + ("val", "TEXT", "Yes", "false", None, ""), + ("dt", "TEXT", "Yes", "false", None, ""), + ], + ], + ) + ops_widen = adapter.get_alter_operations( + "db.current_tbl", "db.target_tbl", ignore_destructive=True + ) + assert len(ops_widen) == 2 + adapter.cursor.execute.reset_mock() + adapter.alter_table(ops_widen) + assert to_sql_calls(adapter) == [ + "ALTER TABLE `db`.`current_tbl` MODIFY COLUMN `val` STRING", + "ALTER TABLE `db`.`current_tbl` MODIFY COLUMN `dt` STRING", + ] + + # 2. `LARGEINT` vs `INT128` is treated as mutually compatible and produces zero alter operations. + mocker.patch.object( + adapter, + "columns", + side_effect=[ + {"id": exp.DataType.build("LARGEINT", udt=True)}, + {"id": exp.DataType.build("INT128")}, + ], + ) + assert adapter.get_alter_operations("db.current_tbl", "db.target_tbl") == [] + + # 3. Changing an OLAP key column to an incompatible valid key type (`DATETIME -> INT`) + # is blocked from emitting a destructive `DROP COLUMN` + `ADD COLUMN` on a key column, + # while non-key column destructive replacements (`DATETIME -> INT`) are still emitted when `ignore_destructive=False`. + mocker.stopall() + mocker.patch.object( + adapter.cursor, + "fetchall", + side_effect=[ + [ + ("k_ts", "DATETIME", "Yes", "true", None, ""), + ("v_ts", "DATETIME", "Yes", "false", None, ""), + ], + [ + ("k_ts", "INT", "Yes", "true", None, ""), + ("v_ts", "INT", "Yes", "false", None, ""), + ], + ], + ) + ops_key_protected = adapter.get_alter_operations( + "db.current_tbl", "db.target_tbl", ignore_destructive=False + ) + adapter.cursor.execute.reset_mock() + adapter.alter_table(ops_key_protected) + assert to_sql_calls(adapter) == [ + "ALTER TABLE `db`.`current_tbl` DROP COLUMN `v_ts`", + "ALTER TABLE `db`.`current_tbl` ADD COLUMN `v_ts` INT", + ] + + +def test_union_ctas_string_coercion_and_duplicate_key_list_partition_coercion( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Lockstep `UNION ALL` unannotated CTAS where column 0 is non-key (`AVG(score)`) and column 1 + # is `CAST(label AS STRING)`: both branches are reordered in lockstep AND cast `label` to `VARCHAR(255)`. + adapter.ctas( + "db.union_str_ctas", + parse_one( + "SELECT AVG(score) AS score, CAST(label AS STRING) AS label FROM db.s1 GROUP BY label " + "UNION ALL SELECT score, label FROM db.s2", + read="doris", + ), + ) + + # 2. Explicit `duplicate_key = ('id')` combined with `PARTITION BY LIST(region)` where `region` is `TEXT`: + # `region` is automatically appended to `DUPLICATE KEY`, moved into the key prefix, AND coerced to `VARCHAR(255)`. + adapter.create_table( + "db.dup_list_part_tbl", + { + "score": exp.DataType.build("DOUBLE", dialect="doris"), + "region": exp.DataType.build("TEXT", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + }, + partitioned_by=[parse_one("LIST(region)", read="doris")], + table_properties={ + "duplicate_key": exp.Tuple(expressions=[exp.Literal.string("id")]), + "partitions": exp.Literal.string("PARTITION p_us VALUES IN ('us')"), + "buckets": exp.Literal.number(4), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`union_str_ctas` AS SELECT CAST(`label` AS VARCHAR(255)) AS `label`, AVG(`score`) AS `score` FROM `db`.`s1` GROUP BY `label` UNION ALL SELECT CAST(`label` AS VARCHAR(255)), `score` FROM `db`.`s2`", + "CREATE TABLE IF NOT EXISTS `db`.`dup_list_part_tbl` (`id` INT, `region` VARCHAR(255), `score` DOUBLE) DUPLICATE KEY (`id`, `region`) PARTITION BY LIST (`region`) (PARTITION p_us VALUES IN ('us')) DISTRIBUTED BY HASH (`id`) BUCKETS 4", + ] + + +def test_doris_connection_config_tls_matrix(): + from sqlmesh.core.config.connection import DorisConnectionConfig + + # 1. `ssl_verify_identity=True` without explicit `ssl=True` still populates `kwargs["ssl"]` with `check_hostname=True`. + cfg_identity = DorisConnectionConfig(host="fe.internal", user="root", ssl_verify_identity=True) + assert cfg_identity._static_connection_kwargs["ssl"] == {"check_hostname": True} + + # 2. `ssl_verify_cert=True` populates `verify_mode=True` in `kwargs["ssl"]`. + cfg_cert = DorisConnectionConfig(host="fe.internal", user="root", ssl_verify_cert=True) + assert cfg_cert._static_connection_kwargs["ssl"] == {"verify_mode": True} + + # 3. `ssl_disabled=True` overrides `ssl=True` and omits the `ssl` dictionary. + cfg_disabled = DorisConnectionConfig( + host="fe.internal", user="root", ssl=True, ssl_disabled=True + ) + assert "ssl" not in cfg_disabled._static_connection_kwargs + + +def test_doris_snapshot_evaluator_end_to_end_model_kinds( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], + make_snapshot: t.Callable, + mocker: MockerFixture, +): + from sqlmesh.core.model import load_sql_based_model + from sqlmesh.core.dialect import parse as d_parse + from sqlmesh.core.snapshot import SnapshotChangeCategory, SnapshotEvaluator + + adapter = make_mocked_engine_adapter(DorisEngineAdapter, default_catalog="internal") + adapter.cursor.fetchall.return_value = [] + mocker.patch.object(DorisEngineAdapter, "_get_data_objects", return_value=[]) + mocker.patch.object(DorisEngineAdapter, "table_exists", return_value=True) + mocker.patch.object( + DorisEngineAdapter, + "get_data_object", + side_effect=lambda t: DataObject( + catalog="internal", + schema=exp.to_table(t).db or "db", + name=exp.to_table(t).name, + type=DataObjectType.TABLE, + ), + ) + evaluator = SnapshotEvaluator(adapter) + + # 1. INCREMENTAL_BY_TIME_RANGE model evaluation through SnapshotEvaluator + inc_time_model = load_sql_based_model( + d_parse( + """ + MODEL ( + name db.inc_time, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column ds + ), + dialect doris, + physical_properties ( + duplicate_key = (id), + auto_partition = true, + buckets = 4 + ), + partitioned_by RANGE(ds) + ); + SELECT CAST(id AS INT) AS id, CAST(ds AS DATE) AS ds, CAST(amount AS DOUBLE) AS amount + FROM db.raw_events + WHERE ds BETWEEN @start_ds AND @end_ds + """, + default_dialect="doris", + ), + default_catalog="internal", + ) + inc_time_snap = make_snapshot(inc_time_model) + inc_time_snap.categorize_as(SnapshotChangeCategory.BREAKING) + inc_time_snap.add_interval("2023-12-31", "2023-12-31") + + adapter.cursor.execute.reset_mock() + evaluator.create([inc_time_snap], {}) + evaluator.evaluate( + inc_time_snap, + start="2024-01-01", + end="2024-01-02", + execution_time="2024-01-03", + snapshots={inc_time_snap.name: inc_time_snap}, + ) + time_sql = to_sql_calls(adapter) + assert any( + "DUPLICATE KEY (`id`, `ds`) AUTO PARTITION BY RANGE (`ds`) () DISTRIBUTED BY HASH (`id`) BUCKETS 4" + in s + for s in time_sql + ) + assert any( + "DELETE FROM `internal`.`sqlmesh__db`.`db__inc_time__" in s + and "WHERE `ds` >= '2024-01-01' AND `ds` <= '2024-01-02'" in s + for s in time_sql + ) + assert any("INSERT INTO `internal`.`sqlmesh__db`.`db__inc_time__" in s for s in time_sql) + + # 2. INCREMENTAL_BY_UNIQUE_KEY model evaluation through SnapshotEvaluator + inc_uk_model = load_sql_based_model( + d_parse( + """ + MODEL ( + name db.inc_uk, + kind INCREMENTAL_BY_UNIQUE_KEY ( + unique_key id + ), + dialect doris, + physical_properties ( + buckets = 8 + ) + ); + SELECT CAST(id AS INT) AS id, CAST(status AS VARCHAR(32)) AS status + FROM db.raw_users + """, + default_dialect="doris", + ), + default_catalog="internal", + ) + inc_uk_snap = make_snapshot(inc_uk_model) + inc_uk_snap.categorize_as(SnapshotChangeCategory.BREAKING) + inc_uk_snap.add_interval("2023-12-31", "2023-12-31") + + adapter.cursor.execute.reset_mock() + adapter.cursor.fetchall.return_value = [ + ("id", "INT", "Yes", "true", None, ""), + ("status", "VARCHAR(32)", "Yes", "false", None, ""), + ] + evaluator.create([inc_uk_snap], {}) + evaluator.evaluate( + inc_uk_snap, + start="2024-01-01", + end="2024-01-02", + execution_time="2024-01-03", + snapshots={inc_uk_snap.name: inc_uk_snap}, + ) + uk_sql = to_sql_calls(adapter) + assert any("UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 8" in s for s in uk_sql) + assert any("INSERT INTO `internal`.`sqlmesh__db`.`db__inc_uk__" in s for s in uk_sql) + assert not any("MERGE INTO" in s for s in uk_sql) + + # 3. SCD_TYPE_2_BY_TIME model evaluation through SnapshotEvaluator + scd_model = load_sql_based_model( + d_parse( + """ + MODEL ( + name db.scd_users, + kind SCD_TYPE_2_BY_TIME ( + unique_key id, + updated_at_name updated_at, + time_data_type DATETIME + ), + columns ( + id INT, + status VARCHAR(32), + updated_at DATETIME, + valid_from DATETIME, + valid_to DATETIME + ), + dialect doris, + physical_properties ( + buckets = 4 + ) + ); + SELECT CAST(id AS INT) AS id, CAST(status AS VARCHAR(32)) AS status, CAST(updated_at AS DATETIME) AS updated_at + FROM db.raw_users + """, + default_dialect="doris", + ), + default_catalog="internal", + ) + scd_snap = make_snapshot(scd_model) + scd_snap.categorize_as(SnapshotChangeCategory.BREAKING) + scd_snap.add_interval("2023-12-31", "2023-12-31") + + adapter.cursor.execute.reset_mock() + adapter.cursor.fetchall.return_value = [ + ("id", "INT", "Yes", "true", None, ""), + ("status", "VARCHAR(32)", "Yes", "true", None, ""), + ("updated_at", "DATETIME", "Yes", "false", None, ""), + ("valid_from", "DATETIME", "Yes", "false", None, ""), + ("valid_to", "DATETIME", "Yes", "false", None, ""), + ] + evaluator.create([scd_snap], {}) + evaluator.evaluate( + scd_snap, + start="2024-01-01", + end="2024-01-02", + execution_time="2024-01-03", + snapshots={scd_snap.name: scd_snap}, + ) + scd_sql = to_sql_calls(adapter) + assert any( + "CREATE TABLE IF NOT EXISTS `internal`.`sqlmesh__db`.`db__scd_users__" in s + and "DISTRIBUTED BY RANDOM BUCKETS 4" in s + for s in scd_sql + ) + assert any("TRUNCATE TABLE `internal`.`sqlmesh__db`.`db__scd_users__" in s for s in scd_sql) + assert any("INSERT INTO `internal`.`sqlmesh__db`.`db__scd_users__" in s for s in scd_sql) + + # 4. INCREMENTAL_BY_PARTITION model evaluation through SnapshotEvaluator + inc_part_model = load_sql_based_model( + d_parse( + """ + MODEL ( + name db.inc_part, + kind INCREMENTAL_BY_PARTITION, + dialect doris, + partitioned_by LIST(region), + physical_properties ( + duplicate_key = (id), + partitions = "PARTITION p_us VALUES IN ('us')", + buckets = 4 + ) + ); + SELECT CAST(id AS INT) AS id, CAST(region AS VARCHAR(16)) AS region, CAST(amount AS DOUBLE) AS amount + FROM db.raw_regional_sales + """, + default_dialect="doris", + ), + default_catalog="internal", + ) + inc_part_snap = make_snapshot(inc_part_model) + inc_part_snap.categorize_as(SnapshotChangeCategory.BREAKING) + inc_part_snap.add_interval("2023-12-31", "2023-12-31") + + adapter.cursor.execute.reset_mock() + adapter.cursor.fetchall.return_value = [ + ("id", "INT", "Yes", "true", None, ""), + ("region", "VARCHAR(16)", "Yes", "true", None, ""), + ("amount", "DOUBLE", "Yes", "false", None, ""), + ] + evaluator.create([inc_part_snap], {}) + evaluator.evaluate( + inc_part_snap, + start="2024-01-01", + end="2024-01-02", + execution_time="2024-01-03", + snapshots={inc_part_snap.name: inc_part_snap}, + ) + part_sql = to_sql_calls(adapter) + assert any( + "DUPLICATE KEY (`id`, `region`) PARTITION BY LIST (`region`) (PARTITION p_us VALUES IN ('us')) DISTRIBUTED BY HASH (`id`) BUCKETS 4" + in s + for s in part_sql + ) + assert any( + "DELETE FROM `internal`.`sqlmesh__db`.`db__inc_part__" in s + and "USING (SELECT DISTINCT `region` FROM `internal`.`sqlmesh__db`.`__temp_db__inc_part__" + in s + and "WHERE `_t1`.`region` = `_t2`.`region`" in s + for s in part_sql + ) + assert any("INSERT INTO `internal`.`sqlmesh__db`.`db__inc_part__" in s for s in part_sql)