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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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/31] 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) From ad784f47e875d5d33e296e325f7c7a149bb9893a Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sun, 25 Oct 2026 15:43:28 -0700 Subject: [PATCH 28/31] fix(doris): resolve CTAS unknown types, sync MV, delete and TLS edge cases Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- sqlmesh/core/config/connection.py | 31 +- sqlmesh/core/engine_adapter/doris.py | 512 +++++++++++++++++++----- tests/core/engine_adapter/test_doris.py | 180 +++++++++ tests/core/test_connection_config.py | 12 +- 4 files changed, 613 insertions(+), 122 deletions(-) diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index be49001dd8..4490886065 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -2781,31 +2781,36 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]: "local_infile": False, } if not self.ssl_disabled: - ssl_config: t.Dict[str, t.Any] = dict(self.ssl) if isinstance(self.ssl, dict) else {} + ssl_config: t.Dict[str, t.Any] = ( + dict(self.ssl) if isinstance(self.ssl, dict) else {} # type: ignore[arg-type] + ) for key in ("ca", "cert", "key"): val = getattr(self, f"ssl_{key}", None) if val is not None: ssl_config[key] = val if self.ssl_verify_cert is not None: ssl_config["verify_mode"] = self.ssl_verify_cert + if ( + self.ssl_verify_cert is False + and self.ssl_verify_identity is None + and "check_hostname" not in ssl_config + ): + ssl_config["check_hostname"] = False 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 - 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 + elif self.ssl is True: + kwargs["ssl"] = {"check_hostname": False} return kwargs + @property + def _extra_engine_config(self) -> t.Dict[str, t.Any]: + return { + "database": self.database, + "catalog": self.catalog, + } + @property def _connection_kwargs_keys(self) -> t.Set[str]: is_external_catalog = bool(self.catalog and self.catalog.lower() != "internal") diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 3376e7ffc8..0015cd573a 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -3,6 +3,7 @@ from collections import defaultdict import datetime from decimal import Decimal +import hashlib import itertools import logging import re @@ -11,6 +12,7 @@ import pandas as pd # noqa: TID253 from sqlglot import exp, parse, parse_one +from sqlglot.helper import ensure_list from sqlglot.optimizer.qualify_columns import quote_identifiers from sqlmesh.core.dialect import to_schema @@ -431,6 +433,81 @@ def _is_safe_to_retry_sql(cls, sql: str) -> bool: def ping(self) -> None: self._connection_pool.get().ping(reconnect=False) + @classmethod + def _normalize_schema_tmp_identifiers(cls, expression: exp.Expr) -> exp.Expr: + max_len = cls.MAX_IDENTIFIER_LENGTH or 64 + suffix = "_schema_tmp" + changed = False + for tbl in expression.find_all(exp.Table): + if ( + isinstance(tbl.this, exp.Identifier) + and len(tbl.name) > max_len + and tbl.name.endswith(suffix) + ): + if not changed: + expression = expression.copy() + changed = True + break + if changed: + for tbl in expression.find_all(exp.Table): + if ( + isinstance(tbl.this, exp.Identifier) + and len(tbl.name) > max_len + and tbl.name.endswith(suffix) + ): + digest = hashlib.md5(tbl.name.encode("utf-8")).hexdigest()[:8] + prefix_len = max_len - len(suffix) - 9 + short_name = f"{tbl.name[:prefix_len]}_{digest}{suffix}" + tbl.set("this", exp.to_identifier(short_name, quoted=bool(tbl.this.quoted))) + return expression + + def execute( + self, + expressions: t.Union[str, exp.Expr, t.Sequence[exp.Expr]], + ignore_unsupported_errors: bool = False, + quote_identifiers: bool = True, + track_rows_processed: bool = False, + **kwargs: t.Any, + ) -> None: + if isinstance(expressions, str): + super().execute( + expressions, + ignore_unsupported_errors=ignore_unsupported_errors, + quote_identifiers=quote_identifiers, + track_rows_processed=track_rows_processed, + **kwargs, + ) + return + + processed_expressions: t.List[exp.Expr] = [] + for e in ensure_list(expressions): + if not isinstance(e, exp.Expr): + super().execute( + expressions, + ignore_unsupported_errors=ignore_unsupported_errors, + quote_identifiers=quote_identifiers, + track_rows_processed=track_rows_processed, + **kwargs, + ) + return + if isinstance(e, exp.Select) and e.args.get("locks"): + e = e.copy() + e.set("locks", None) + logger.warning( + "[Doris] Removed FOR UPDATE from SELECT statement: %s", + e.sql(dialect=self.dialect, identify=quote_identifiers), + ) + e = self._normalize_schema_tmp_identifiers(e) + processed_expressions.append(e) + + super().execute( + processed_expressions, + ignore_unsupported_errors=ignore_unsupported_errors, + quote_identifiers=quote_identifiers, + track_rows_processed=track_rows_processed, + **kwargs, + ) + def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any) -> None: if self._pre_ping: try: @@ -864,9 +941,17 @@ def _create_materialized_view( batch_end=len(values), ) + is_sync_mv = bool( + view_properties + and ( + (self._extract_str_value(view_properties.get("materialized_type")) or "").upper() + == "SYNC" + ) + ) + source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( query_or_df, - target_columns_to_types, + None if is_sync_mv else target_columns_to_types, batch_size=0, target_table=view_name, source_columns=source_columns, @@ -877,16 +962,6 @@ 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_cols_to_types = target_columns_to_types if ( not schema_cols_to_types @@ -927,7 +1002,9 @@ def _create_materialized_view( target_columns_to_types=target_columns_to_types, table_description=( table_description - if self.COMMENT_CREATION_VIEW.supports_schema_def and self.comments_enabled + if not is_sync_mv + and self.COMMENT_CREATION_VIEW.supports_schema_def + and self.comments_enabled else None ), partitioned_by=partitioned_by, @@ -960,6 +1037,15 @@ def _create_materialized_view( aliased_selects.append(sel_expr) query.set("expressions", aliased_selects) + sync_source_tbl: t.Optional[exp.Table] = None + if is_sync_mv: + if view_properties and view_properties.get("source_table"): + sync_source_tbl = self._to_source_table(view_properties["source_table"]) + else: + first_tbl = query.find(exp.Table) + if first_tbl: + sync_source_tbl = self._to_source_table(first_tbl) + # 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: @@ -971,10 +1057,9 @@ def _create_materialized_view( is_sync_mv and drop_view_props is not None and not drop_view_props.get("source_table") + and sync_source_tbl is not None ): - first_tbl = query.find(exp.Table) - if first_tbl: - drop_view_props["source_table"] = first_tbl + drop_view_props["source_table"] = sync_source_tbl self.drop_view( view_name, ignore_if_not_exists=True, @@ -986,13 +1071,7 @@ def _create_materialized_view( # (`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"): - 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: - source_db = first_tbl.db + source_db: t.Optional[str] = sync_source_tbl.db if sync_source_tbl else None if source_db: db_sql = self._to_sql(exp.to_identifier(source_db, quoted=True)) self.execute(f"USE {db_sql}") @@ -1008,6 +1087,20 @@ def _create_materialized_view( ), quote_identifiers=self.QUOTE_IDENTIFIERS_IN_VIEWS, ) + if is_sync_mv and sync_source_tbl is not None: + registry = getattr(self, "_sync_mv_source_tables", None) + if registry is None: + registry = {} + self._sync_mv_source_tables = registry + registry[target_table.name.lower()] = sync_source_tbl.copy() + if target_table.db: + registry[f"{target_table.db.lower()}.{target_table.name.lower()}"] = ( + sync_source_tbl.copy() + ) + elif sync_source_tbl.db: + registry[f"{sync_source_tbl.db.lower()}.{target_table.name.lower()}"] = ( + sync_source_tbl.copy() + ) build_mode = ( self._to_upper_str(view_properties["build"]) if view_properties and view_properties.get("build") is not None @@ -1051,24 +1144,35 @@ def drop_view( return view_properties = kwargs.pop("view_properties", None) + sync_source_table: t.Any = 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 = self._to_sql(target_table) - 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: - db_sql = self._to_sql(exp.to_identifier(source_tbl_exp.db, quoted=True)) - self.execute(f"USE {db_sql}") - 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}" - ) - self._clear_data_object_cache(view_name) - return + if mat_type and mat_type.upper() == "SYNC": + sync_source_table = view_properties.get("source_table") + elif materialized: + registry = getattr(self, "_sync_mv_source_tables", None) or {} + lookup_key = ( + f"{target_table.db.lower()}.{target_table.name.lower()}" + if target_table.db + else target_table.name.lower() + ) + sync_source_table = registry.get(lookup_key) or registry.get(target_table.name.lower()) + + if sync_source_table is not None: + view_table_sql = self._to_sql(target_table) + source_tbl_exp = self._to_source_table(sync_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: + db_sql = self._to_sql(exp.to_identifier(source_tbl_exp.db, quoted=True)) + self.execute(f"USE {db_sql}") + 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}" + ) + self._clear_data_object_cache(view_name) + return super().drop_view(view_name, ignore_if_not_exists, materialized, **kwargs) def create_index( @@ -1288,6 +1392,28 @@ def _extract_time_partition_transform( return clean_col, unit_clean return None + @staticmethod + def _is_date_part_extract_fn(expr: exp.Expr) -> bool: + curr = expr.unnest() + if isinstance(curr, (exp.Year, exp.Month, exp.Day, exp.Hour)): + return True + if isinstance(curr, exp.Anonymous) and curr.name.upper() in ( + "YEAR", + "YEARS", + "MONTH", + "MONTHS", + "DAY", + "DAYS", + "HOUR", + "HOURS", + "MINUTE", + "MINUTES", + "SECOND", + "SECONDS", + ): + return True + return False + def insert_overwrite_by_partition( self, table_name: TableName, @@ -1295,6 +1421,7 @@ def insert_overwrite_by_partition( 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, + **kwargs: t.Any, ) -> None: if not partitioned_by: self.replace_query( @@ -1302,6 +1429,7 @@ def insert_overwrite_by_partition( query_or_df, target_columns_to_types=target_columns_to_types, source_columns=source_columns, + **kwargs, ) return parsed_exprs, _ = self._parse_partition_expressions(partitioned_by) @@ -1357,13 +1485,20 @@ def insert_overwrite_by_partition( ) 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 ( 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) + and not is_iceberg ): + reorder_kwargs = dict(kwargs) + reorder_kwargs["partitioned_by"] = partitioned_by target_columns_to_types = self._reorder_olap_columns( - target_columns_to_types, {"partitioned_by": partitioned_by} + target_columns_to_types, reorder_kwargs ) if source_columns: source_col_set = set(source_columns) @@ -1470,6 +1605,45 @@ def _simplify_boolean_where(cls, expr: exp.Expr) -> exp.Expr: return expr.this return expr + @classmethod + def _flatten_and_conjuncts(cls, expr: exp.Expr) -> t.List[exp.Expr]: + unwrapped = expr.unnest() + if isinstance(unwrapped, exp.And): + result: t.List[exp.Expr] = [] + for child in unwrapped.flatten(): + result.extend(cls._flatten_and_conjuncts(child)) + return result + return [unwrapped] + + @classmethod + def _flatten_or_disjuncts(cls, expr: exp.Expr) -> t.List[exp.Expr]: + unwrapped = expr.unnest() + if isinstance(unwrapped, exp.Or): + result: t.List[exp.Expr] = [] + for child in unwrapped.flatten(): + result.extend(cls._flatten_or_disjuncts(child)) + return result + return [unwrapped] + + def _normalize_extra_conditions(self, extra_conditions: t.List[exp.Expr]) -> t.List[exp.Expr]: + normalized: t.List[exp.Expr] = [] + for cond in extra_conditions: + unwrapped = cond.unnest() + if isinstance(unwrapped, exp.Or): + or_info = self._extract_or_key_rows_from_where(unwrapped) + if ( + or_info is not None + and len(or_info[0]) == 1 + and not or_info[2] + and all(not isinstance(r[0], exp.Null) for r in or_info[1]) + ): + col_expr = exp.column(or_info[0][0].name, quoted=True) + vals = [r[0].copy() for r in or_info[1]] + normalized.append(col_expr.isin(*vals)) + continue + normalized.append(cond) + return normalized + @classmethod def _unwrap_literal_casts_in_where(cls, where: exp.Expr) -> exp.Expr: """Unwrap top-level `CAST( AS ...)` expressions in `DELETE WHERE` predicates @@ -1511,10 +1685,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] - ) + conjuncts = self._flatten_and_conjuncts(where) target_in: t.Optional[exp.In] = None subquery: t.Optional[exp.Subquery] = None extra_conditions: t.List[exp.Expr] = [] @@ -1564,16 +1735,13 @@ def _extract_subquery_from_where( if not normalized_keys: return None - return normalized_keys, subquery, extra_conditions + return normalized_keys, subquery, self._normalize_extra_conditions(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] - ) + conjuncts = self._flatten_and_conjuncts(where) target_in: t.Optional[exp.In] = None extra_conditions: t.List[exp.Expr] = [] @@ -1616,7 +1784,7 @@ def _extract_literal_tuple_in_from_where( if not tuple_rows: return None - return normalized_keys, tuple_rows, extra_conditions + return normalized_keys, tuple_rows, self._normalize_extra_conditions(extra_conditions) @staticmethod def _is_literal_key_value(expr: exp.Expr) -> bool: @@ -1652,8 +1820,8 @@ def _normalize_boolean_column_conjuncts(cls, where: exp.Expr) -> exp.Expr: return single unwrapped = where.unnest() if isinstance(unwrapped, exp.And): - conjuncts = list(unwrapped.flatten()) - changed = False + conjuncts = cls._flatten_and_conjuncts(unwrapped) + changed = len(conjuncts) != len(list(unwrapped.flatten())) new_conjuncts: t.List[exp.Expr] = [] for cond in conjuncts: normalized = cls._normalize_single_boolean_conjunct(cond) @@ -1672,10 +1840,7 @@ def _extract_or_key_rows_from_where( """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] - ) + conjuncts = self._flatten_and_conjuncts(where) target_or: t.Optional[exp.Or] = None extra_conditions: t.List[exp.Expr] = [] @@ -1689,7 +1854,7 @@ def _extract_or_key_rows_from_where( if target_or is None: return None - disjuncts = list(target_or.flatten()) + disjuncts = self._flatten_or_disjuncts(target_or) if len(disjuncts) < 2: return None @@ -1698,12 +1863,7 @@ def _extract_or_key_rows_from_where( 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] - ) + branch_conjuncts = self._flatten_and_conjuncts(branch) col_to_val: t.Dict[str, exp.Expr] = {} col_to_expr: t.Dict[str, exp.Expr] = {} for item in branch_conjuncts: @@ -1767,7 +1927,7 @@ def _extract_or_key_rows_from_where( if not canonical_cols or not tuple_rows: return None - return canonical_cols, tuple_rows, extra_conditions + return canonical_cols, tuple_rows, self._normalize_extra_conditions(extra_conditions) @staticmethod def _unwrap_concat_ws(expr: exp.Expr) -> t.Optional[exp.Expr]: @@ -1842,12 +2002,21 @@ def _execute_delete_with_subquery( first_select = subquery.this while isinstance(first_select, exp.SetOperation): first_select = first_select.this + if isinstance(first_select, exp.Select) and any( + isinstance(e, exp.Star) for e in first_select.expressions + ): + self._delete_by_literal_keys(table_name, key_exprs, orig_subquery, extra_conditions) + return 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): + if ( + isinstance(first_select, exp.Select) + and idx < len(first_select.expressions) + and not isinstance(first_select.expressions[idx], exp.Star) + ): sel_item = first_select.expressions[idx] candidate_name = ( sel_item.alias_or_name @@ -1965,15 +2134,8 @@ 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" - ) - ) + is_extract_fn = self._is_date_part_extract_fn(k) + is_year_fn = unit == "year" and is_extract_fn src_transform = self._extract_time_partition_transform(raw_sel) if src_transform is not None: src_arg: exp.Expr = src_transform[0].copy() @@ -1981,6 +2143,11 @@ def _delete_by_time_transform_keys( raw_year_indices.add(idx) projections.append(exp.alias_(raw_sel.copy(), f"_val_{idx}", quoted=True)) continue + elif is_extract_fn: + raise SQLMeshError( + f"Doris delete_from cannot rewrite date-part predicate '{k.sql(dialect=self.dialect)}' " + f"with non-date subquery projection '{raw_sel.sql(dialect=self.dialect)}' into a contiguous range." + ) else: src_arg = raw_sel.copy() trunc_expr = exp.Anonymous( @@ -2009,24 +2176,50 @@ 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))) + transform_indices = { + idx + for idx, k in enumerate(key_exprs) + if self._extract_time_partition_transform(k) is not None + } + unwrapped_indices: t.Set[int] = set(transform_indices) + + # Pass 1: Validate which transform indices can be unwrapped across ALL branches. for branch in all_branches: - if len(branch.expressions) == len(key_exprs): - can_unwrap_group = not branch.args.get("having") and not branch.args.get( - "limit" - ) + if len(branch.expressions) != len(key_exprs): + unwrapped_indices.clear() + break + has_group = bool(branch.args.get("group")) + can_unwrap_group = ( + not branch.args.get("having") + and not branch.args.get("limit") + and not any(e.find(exp.AggFunc) for e in branch.expressions) + ) + if has_group and not can_unwrap_group: + unwrapped_indices.clear() + break + for idx in list(unwrapped_indices): + sel_item = branch.expressions[idx] + raw_sel = sel_item.this if isinstance(sel_item, exp.Alias) else sel_item + if self._extract_time_partition_transform(raw_sel) is None: + unwrapped_indices.discard(idx) + + if any(branch.args.get("group") for branch in all_branches) and ( + unwrapped_indices != transform_indices + ): + unwrapped_indices.clear() + + # Pass 2: Mutate branch projections only for indices verified across all branches. + if unwrapped_indices: + for branch in all_branches: 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 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) + for idx in unwrapped_indices: + 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) + assert src_transform is not None + branch_exprs[idx] = src_transform[0].copy() + if branch.args.get("group"): + branch.set("group", None) branch.set("expressions", branch_exprs) new_first_exprs = list(first_select.expressions) @@ -2040,19 +2233,17 @@ 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" - ) - ) + is_extract_fn = self._is_date_part_extract_fn(k) + is_year_fn = unit == "year" and is_extract_fn 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 + if idx not in unwrapped_indices and is_extract_fn: + raise SQLMeshError( + f"Doris delete_from cannot rewrite date-part predicate '{k.sql(dialect=self.dialect)}' " + "into a contiguous range when subquery branches return non-date expressions." + ) trunc_expr = exp.Anonymous( this="DATE_TRUNC", expressions=[sq_col, exp.Literal.string(unit)], @@ -2105,7 +2296,11 @@ def _collect_select_branches(node: exp.Expr) -> None: if yr_val is None: predicates.append(target_col.is_(exp.null())) else: - yr_int = int(yr_val) + yr_int = ( + yr_val.year + if isinstance(yr_val, (datetime.date, datetime.datetime)) + else 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( @@ -2426,6 +2621,84 @@ def replace_query( **kwargs, ) + def _replace_by_key( + self, + target_table: TableName, + source_table: QueryOrDF, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], + key: t.Sequence[exp.Expr], + is_unique_key: bool, + source_columns: t.Optional[t.List[str]] = None, + ) -> None: + if target_columns_to_types is None: + target_columns_to_types = self.columns(target_table) + + table = exp.to_table(target_table) + is_iceberg = self._is_iceberg_table(catalog_name=table.catalog) + key_cols: t.List[exp.Column] = [] + if is_unique_key and all(isinstance(k, (exp.Column, exp.Identifier)) for k in key): + key_cols = [exp.column(k.name) for k in key] + + temp_table_props: t.Optional[t.Dict[str, t.Any]] = None + if not is_iceberg and key_cols: + temp_table_props = { + "unique_key": ( + key_cols[0] if len(key_cols) == 1 else exp.Tuple(expressions=key_cols) + ) + } + if 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, {"table_properties": temp_table_props} + ) + 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 + ] + + temp_table = self._get_temp_table(target_table) + key_exp = ( + exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", *key, dialect=self.dialect) + if len(key) > 1 + else key[0] + ) + column_names = list(target_columns_to_types or []) + + with self.transaction(): + self.ctas( + temp_table, + source_table, + target_columns_to_types=target_columns_to_types, + exists=False, + source_columns=source_columns, + table_properties=temp_table_props, + ) + + try: + delete_query = exp.select(key_exp).from_(temp_table) + insert_query = self._select_columns(target_columns_to_types).from_(temp_table) + if not is_unique_key: + delete_query = delete_query.distinct() + else: + insert_query = insert_query.distinct(*key) + + insert_statement = exp.insert( + insert_query, + target_table, + columns=column_names, + ) + delete_filter = key_exp.isin(query=delete_query) + + if not self.INSERT_OVERWRITE_STRATEGY.is_replace_where: + self.delete_from(target_table, delete_filter) + else: + insert_statement.set("where", delete_filter) + insert_statement.set("this", exp.to_table(target_table)) + + self.execute(insert_statement, track_rows_processed=True) + finally: + self.drop_table(temp_table) + def _scd_type_2( self, target_table: TableName, @@ -2507,12 +2780,11 @@ def _create_table_from_source_queries( 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) ) + if not is_iceberg and target_columns_to_types and target_columns_to_types_known: + target_columns_to_types = self._reorder_olap_columns(target_columns_to_types, kwargs) first_query = source_queries[0].query_factory() if source_queries else None def _matches_target_table(t_node: exp.Table) -> bool: @@ -2614,12 +2886,22 @@ def _build_create_table_exp( table_properties=kwargs.get("table_properties"), ) - if not is_iceberg and target_columns_to_types: + if ( + not is_iceberg + and target_columns_to_types + and columns_to_types_all_known(target_columns_to_types) + ): original_cols = dict(target_columns_to_types) reordered_cols = self._reorder_olap_columns(target_columns_to_types, kwargs) 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: + query_order_differs = bool( + isinstance(expression, exp.Select) + and len(expression.expressions) == len(reordered_cols) + and [p.alias_or_name.lower() for p in expression.expressions if p.alias_or_name] + != [c.lower() for c in reordered_cols] + ) + if order_changed or types_changed or query_order_differs: target_columns_to_types = reordered_cols if isinstance(table_name_or_schema, exp.Schema): col_def_map = { @@ -2648,7 +2930,7 @@ def _build_create_table_exp( new_col_defs + other_exprs, ) if ( - (order_changed or types_changed) + (order_changed or types_changed or query_order_differs) and expression is not None and isinstance(expression, exp.Query) ): @@ -2691,6 +2973,24 @@ def _build_create_table_exp( ) elif not is_iceberg and isinstance(expression, (exp.Select, exp.SetOperation)): expression = self._reorder_unannotated_ctas_query(expression, kwargs) + if target_columns_to_types: + first_select: t.Optional[exp.Select] = None + curr_node: t.Optional[exp.Expr] = expression + while isinstance(curr_node, exp.SetOperation): + curr_node = curr_node.this + if isinstance(curr_node, exp.Select): + first_select = curr_node + if first_select is not None: + col_by_lower = {c.lower(): c for c in target_columns_to_types} + reordered_names = [ + col_by_lower[p.alias_or_name.lower()] + for p in first_select.expressions + if p.alias_or_name and p.alias_or_name.lower() in col_by_lower + ] + if len(reordered_names) == len(target_columns_to_types): + target_columns_to_types = { + c: target_columns_to_types[c] for c in reordered_names + } # Doris AUTO PARTITION BY RANGE requires the partition column(s) to be defined as NOT NULL. if not is_iceberg: diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 0879d9918b..386635d1df 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -3578,3 +3578,183 @@ def test_doris_snapshot_evaluator_end_to_end_model_kinds( for s in part_sql ) assert any("INSERT INTO `internal`.`sqlmesh__db`.`db__inc_part__" in s for s in part_sql) + + +def test_unannotated_unknown_columns_ctas_and_sync_mv_and_schema_tmp_truncation( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Unannotated SqlModel passes `target_columns_to_types` with `UNKNOWN` types to `ctas`: + # must reorder the `SELECT` projections so `id` (the `unique_key`) is column 0. + adapter.ctas( + "db.unannotated_uk_model", + parse_one("SELECT name, id FROM db.src", read="doris"), + target_columns_to_types={ + "name": exp.DataType.build("UNKNOWN"), + "id": exp.DataType.build("UNKNOWN"), + }, + table_properties={"unique_key": exp.to_column("id")}, + ) + + # 2. Unannotated SqlModel with `UNKNOWN` column types and `AVG(score)` at column 0: + # falls through to `_reorder_unannotated_ctas_query` and moves `CAST(id AS INT)` to column 0. + adapter.ctas( + "db.unannotated_agg_model", + parse_one( + "SELECT AVG(score) AS score, CAST(id AS INT) AS id FROM db.src GROUP BY id", + read="doris", + ), + target_columns_to_types={ + "score": exp.DataType.build("UNKNOWN"), + "id": exp.DataType.build("UNKNOWN"), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`unannotated_uk_model` UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10 AS SELECT `id`, `name` FROM `db`.`src`", + "CREATE TABLE IF NOT EXISTS `db`.`unannotated_agg_model` AS SELECT CAST(`id` AS INT) AS `id`, AVG(`score`) AS `score` FROM `db`.`src` GROUP BY `id`", + ] + + # 3. Sync MV (`materialized_type='SYNC'`) with `target_columns_to_types` and `table_description`: + # must NOT wrap the query in `SELECT CAST(...) FROM (...) AS _subquery`, must NOT emit `COMMENT`, + # and `drop_view(view_name, materialized=True)` without `view_properties` still emits `ON `. + adapter.cursor.execute.reset_mock() + adapter.create_view( + "db.sync_mv_typed", + parse_one("SELECT id, SUM(amount) AS total FROM db.orders GROUP BY id", read="doris"), + replace=False, + materialized=True, + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "total": exp.DataType.build("BIGINT", dialect="doris"), + }, + table_description="Sync rollup MV", + view_properties={ + "materialized_type": exp.Literal.string("SYNC"), + "source_table": exp.Literal.string("db.orders"), + }, + ) + adapter.drop_view("db.sync_mv_typed", materialized=True) + + assert to_sql_calls(adapter) == [ + "CREATE MATERIALIZED VIEW `db`.`sync_mv_typed` AS SELECT `id` AS `mv_id`, SUM(`amount`) AS `total` FROM `db`.`orders` GROUP BY `id`", + "DROP MATERIALIZED VIEW IF EXISTS `db`.`sync_mv_typed` ON `db`.`orders`", + ] + + # 4. Overlong `_schema_tmp` table identifier (> 64 chars) is deterministically shortened to <= 64 chars + # across `CREATE TABLE`, `DESCRIBE`, and `DROP TABLE`. + adapter.cursor.execute.reset_mock() + long_schema_tmp = f"db.{'very_long_model_name_' * 3}_schema_tmp" + assert len(exp.to_table(long_schema_tmp).name) > 64 + adapter.cursor.fetchall.return_value = [("id", "INT", "Yes", "true", None, "")] + adapter.ctas(long_schema_tmp, parse_one("SELECT 1 AS id", read="doris")) + adapter.columns(long_schema_tmp) + adapter.drop_table(long_schema_tmp) + + schema_tmp_calls = to_sql_calls(adapter) + assert len(schema_tmp_calls) == 3 + created_tbl = parse_one(schema_tmp_calls[0], read="doris").this + assert len(created_tbl.name) <= 64 + assert created_tbl.name.endswith("_schema_tmp") + assert f"`{created_tbl.name}`" in schema_tmp_calls[1] + assert f"`{created_tbl.name}`" in schema_tmp_calls[2] + + +def test_nested_delete_and_mixed_union_time_transform_and_merge_text_preservation( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + from sqlmesh.core.config.connection import DorisConnectionConfig + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Nested parenthesized `AND`s with `IN (SELECT ...)` and extra `OR` condition: + # extracts the subquery and normalizes `(status = 'x' OR status = 'y')` into `status IN ('x', 'y')`. + adapter.delete_from( + "db.target", + where=parse_one( + "(tenant_id = 1 AND id IN (SELECT id FROM db.staging)) AND (status = 'x' OR status = 'y')", + read="doris", + ), + ) + + # 2. Single-column `id IN (SELECT * FROM db.staging_ids)` falls back to literal key batching + # instead of emitting invalid `SELECT * AS _col_0`. + mocker.patch.object(adapter, "fetchall", return_value=[(10,), (20,)]) + adapter.delete_from( + "db.target", + where=parse_one("id IN (SELECT * FROM db.staging_ids)", read="doris"), + ) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `db`.`target` AS `_t1` USING (SELECT `id` FROM `db`.`staging`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`tenant_id` = 1 AND `_t1`.`status` IN ('x', 'y')", + "DELETE FROM `db`.`target` WHERE `id` IN (10, 20)", + ] + + # 3. `_delete_by_time_transform_keys`: + # (a) Mixed `UNION ALL` branches (`SELECT YEAR(ds) FROM db.s1 UNION ALL SELECT yr FROM db.s2`) + # use `branch_is_raw_col` instead of wrapping integer `yr` in `DATE_TRUNC(yr, 'year')`. + # (b) `GROUP BY` with aggregate `MAX(region)` in `WHERE`/`HAVING` is preserved instead of stripping `GROUP BY`. + adapter.cursor.execute.reset_mock() + mocker.stopall() + mocker.patch.object( + adapter.cursor, + "fetchall", + side_effect=[ + [(2024,)], + [("2024-01-01", "2024-01-02")], + ], + ) + adapter.delete_from( + "db.events", + where=parse_one( + "YEAR(ds) IN (SELECT YEAR(ds) FROM db.s1 UNION ALL SELECT yr FROM db.s2)", + read="doris", + ), + ) + adapter.delete_from( + "db.events", + where=parse_one( + "DATE_TRUNC(ds, 'day') IN (SELECT DATE_TRUNC(ds, 'day') FROM db.s1 GROUP BY DATE_TRUNC(ds, 'day') HAVING COUNT(*) > 1)", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "SELECT DISTINCT `_col_0` AS `_val_0` FROM (SELECT YEAR(`ds`) AS `_col_0` FROM `db`.`s1` UNION ALL SELECT `yr` FROM `db`.`s2`) AS `_sq`", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-01' AND `ds` < '2025-01-01'", + "SELECT DISTINCT DATE_TRUNC(`_col_0`, 'day') AS `_low_0`, DATE_ADD(DATE_TRUNC(`_col_0`, 'day'), INTERVAL 1 DAY) AS `_high_0` FROM (SELECT DATE_TRUNC(`ds`, 'DAY') AS `_col_0` FROM `db`.`s1` GROUP BY DATE_TRUNC(`ds`, 'DAY') HAVING COUNT(*) > 1) AS `_sq`", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-01' AND `ds` < '2024-01-02'", + ] + + # 4. `merge` with `{"notes": TEXT, "code": TEXT}` and `unique_key=[code]`: + # `temp_table` uses `UNIQUE KEY (code)` so `code` is coerced to `VARCHAR(255)` while non-key `notes` stays `STRING` (`TEXT`). + adapter.cursor.execute.reset_mock() + mocker.patch("sqlmesh.core.engine_adapter.doris.random_id", return_value="abcd1234") + adapter.merge( + target_table="db.notes_tbl", + source_table=parse_one("SELECT notes, code FROM db.staging_notes", read="doris"), + target_columns_to_types={ + "notes": exp.DataType.build("TEXT", dialect="doris"), + "code": exp.DataType.build("TEXT", dialect="doris"), + }, + unique_key=[exp.to_column("code")], + ) + merge_calls = to_sql_calls(adapter) + assert ( + merge_calls[0] + == "CREATE TABLE `db`.`__temp_notes_tbl_abcd1234` UNIQUE KEY (`code`) DISTRIBUTED BY HASH (`code`) BUCKETS 10 AS SELECT CAST(`code` AS VARCHAR(255)) AS `code`, CAST(`notes` AS STRING) AS `notes` FROM (SELECT `notes`, `code` FROM `db`.`staging_notes`) AS `_subquery`" + ) + + # 5. `DorisConnectionConfig` TLS dictionary with `ssl_verify_cert=False` sets `check_hostname=False` + # (preventing Python 3.13 `SSLContext` `ValueError`) and preserves `ca` when `ssl` is a dict. + cfg_ca_no_verify = DorisConnectionConfig( + host="fe.internal", + user="root", + ssl={"ca": "/path/to/ca.pem"}, + ssl_verify_cert=False, + ) + assert cfg_ca_no_verify._static_connection_kwargs["ssl"] == { + "ca": "/path/to/ca.pem", + "verify_mode": False, + "check_hostname": False, + } diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py index 800fb5fb19..adb39540ec 100644 --- a/tests/core/test_connection_config.py +++ b/tests/core/test_connection_config.py @@ -2545,8 +2545,14 @@ def test_doris(make_config): 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 + assert advanced_config._static_connection_kwargs["ssl"] == { + "ca": "/path/to/ca.pem", + "verify_mode": True, + } + assert advanced_config._extra_engine_config == { + "database": "analytics", + "catalog": "iceberg_cat", + } mock_cursor = MagicMock() assert advanced_config._cursor_init is not None @@ -2565,7 +2571,7 @@ def test_doris(make_config): ssl=True, check_import=False, ) - assert ssl_bool_config._static_connection_kwargs["ssl"] == {} + assert ssl_bool_config._static_connection_kwargs["ssl"] == {"check_hostname": False} ssl_false_config = make_config( type="doris", From f1df8ba8dfaa386ae81da66658da002ab5254a7b Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sun, 25 Oct 2026 16:38:55 -0700 Subject: [PATCH 29/31] fix(doris): harden schema_tmp shortening, TLS config, time-transform delete, and CTAS/merge reordering 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 | 556 ++++++++++++++++++------ tests/core/engine_adapter/test_doris.py | 424 ++++++++++++++++++ tests/core/test_connection_config.py | 52 +++ 4 files changed, 907 insertions(+), 135 deletions(-) diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 4490886065..45ef2f4989 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -2780,7 +2780,7 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]: "charset": self.charset or "utf8mb4", "local_infile": False, } - if not self.ssl_disabled: + if not self.ssl_disabled and self.ssl is not False: ssl_config: t.Dict[str, t.Any] = ( dict(self.ssl) if isinstance(self.ssl, dict) else {} # type: ignore[arg-type] ) @@ -2798,9 +2798,15 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]: ssl_config["check_hostname"] = False if self.ssl_verify_identity is not None: ssl_config["check_hostname"] = bool(self.ssl_verify_identity) + verify_mode = ssl_config.get("verify_mode") + if verify_mode is False or ( + isinstance(verify_mode, str) + and verify_mode.strip().lower() in ("false", "0", "no", "none") + ): + ssl_config["check_hostname"] = False if ssl_config: kwargs["ssl"] = ssl_config - elif self.ssl is True: + elif self.ssl is True or (isinstance(self.ssl, dict) and len(self.ssl) == 0): kwargs["ssl"] = {"check_hostname": False} return kwargs diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 0015cd573a..302f0c3f23 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -42,6 +42,7 @@ if t.TYPE_CHECKING: from sqlmesh.core._typing import SchemaName, TableName from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF + from sqlmesh.core.model.kind import TimeColumn from sqlmesh.core.node import IntervalUnit from sqlmesh.core.schema_diff import TableAlterOperation from sqlmesh.utils.date import TimeLike @@ -227,6 +228,24 @@ def _last_described_key_cols(self) -> t.Set[str]: def _last_described_key_cols(self, value: t.Set[str]) -> None: self._get_thread_state().last_described_key_cols = set(value) + @property + def _sync_mv_source_tables(self) -> t.Dict[str, exp.Table]: + registry = getattr(self._connection_pool, "_doris_sync_mv_source_tables", None) + if not isinstance(registry, dict): + registry = self._extra_config.get("_doris_sync_mv_source_tables") + if not isinstance(registry, dict): + registry = getattr(self, "_doris_sync_mv_source_tables", None) + if not isinstance(registry, dict): + registry = {} + self._doris_sync_mv_source_tables = registry + if isinstance(getattr(self, "_extra_config", None), dict): + self._extra_config["_doris_sync_mv_source_tables"] = registry + try: + setattr(self._connection_pool, "_doris_sync_mv_source_tables", registry) + except Exception: + pass + return registry + @property def catalog_support(self) -> CatalogSupport: return CatalogSupport.FULL_SUPPORT @@ -249,25 +268,33 @@ def set_current_catalog(self, catalog: str) -> None: self.execute(exp.Command(this="SWITCH", expression=catalog_sql)) self._active_catalog = catalog self._needs_catalog_restore = False + database = self._extra_config.get("database") + default_cat = self.default_catalog or "internal" + if database and catalog.lower() == str(default_cat).lower(): + db_sql = self._to_sql(exp.to_identifier(str(database), quoted=True)) + self.execute(f"USE {db_sql}") def _restore_active_catalog_on_reconnect(self, sql: t.Optional[str] = None) -> None: - active_cat = self._active_catalog - self._needs_catalog_restore = False - if not active_cat: + active_cat = self._active_catalog or self.default_catalog + database = self._extra_config.get("database") + if not active_cat and not database: + self._needs_catalog_restore = False return if sql and re.match( r"\A(?:\s|/\*.*?\*/|--[^\r\n]*(?:[\r\n]|\Z))*SWITCH\b", sql, flags=re.IGNORECASE | re.DOTALL, ): + self._needs_catalog_restore = False 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") + if active_cat: + catalog_sql = self._to_sql(exp.to_identifier(active_cat, quoted=True)) + self.cursor.execute(f"SWITCH {catalog_sql}") default_cat = self.default_catalog or "internal" - if database and active_cat.lower() == str(default_cat).lower(): + if database and (not active_cat or 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}") + self._needs_catalog_restore = False def _is_iceberg_table( self, @@ -378,6 +405,8 @@ 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 isinstance(query, exp.Expr): + query = self._normalize_schema_tmp_identifiers(query) if self._pre_ping: try: self.ping() @@ -385,10 +414,10 @@ def _fetch_native_df( 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: + self._connection_pool.get_cursor() + if getattr(self, "_needs_catalog_restore", False): + self._restore_active_catalog_on_reconnect() 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 @@ -433,32 +462,29 @@ def _is_safe_to_retry_sql(cls, sql: str) -> bool: def ping(self) -> None: self._connection_pool.get().ping(reconnect=False) + @classmethod + def _shorten_schema_tmp_name(cls, name: str) -> str: + max_len = cls.MAX_IDENTIFIER_LENGTH or 64 + suffix = "_schema_tmp" + if len(name) > max_len and name.endswith(suffix): + digest = hashlib.md5(name.encode("utf-8")).hexdigest()[:8] + prefix_len = max_len - len(suffix) - 9 + return f"{name[:prefix_len]}_{digest}{suffix}" + return name + @classmethod def _normalize_schema_tmp_identifiers(cls, expression: exp.Expr) -> exp.Expr: max_len = cls.MAX_IDENTIFIER_LENGTH or 64 suffix = "_schema_tmp" - changed = False - for tbl in expression.find_all(exp.Table): - if ( - isinstance(tbl.this, exp.Identifier) - and len(tbl.name) > max_len - and tbl.name.endswith(suffix) - ): - if not changed: - expression = expression.copy() - changed = True - break - if changed: - for tbl in expression.find_all(exp.Table): - if ( - isinstance(tbl.this, exp.Identifier) - and len(tbl.name) > max_len - and tbl.name.endswith(suffix) - ): - digest = hashlib.md5(tbl.name.encode("utf-8")).hexdigest()[:8] - prefix_len = max_len - len(suffix) - 9 - short_name = f"{tbl.name[:prefix_len]}_{digest}{suffix}" - tbl.set("this", exp.to_identifier(short_name, quoted=bool(tbl.this.quoted))) + if not any( + len(ident.name) > max_len and ident.name.endswith(suffix) + for ident in expression.find_all(exp.Identifier) + ): + return expression + expression = expression.copy() + for ident in expression.find_all(exp.Identifier): + if len(ident.name) > max_len and ident.name.endswith(suffix): + ident.set("this", cls._shorten_schema_tmp_name(ident.name)) return expression def execute( @@ -516,9 +542,9 @@ def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any 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: + if getattr(self, "_needs_catalog_restore", False): + self._restore_active_catalog_on_reconnect(sql) 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 @@ -772,8 +798,14 @@ def _get_data_objects( ) ) ) + short_to_orig: t.Dict[str, str] = {} if object_names: - lowered_names = [name.lower() for name in object_names] + lowered_names: t.List[str] = [] + for name in object_names: + short_name = self._shorten_schema_tmp_name(name) + if short_name != name: + short_to_orig[short_name.lower()] = name + lowered_names.append(short_name.lower()) query = query.where(exp.func("LOWER", exp.column("table_name")).isin(*lowered_names)) result: t.List[DataObject] = [] @@ -781,6 +813,8 @@ def _get_data_objects( 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" + if name_str.lower() in short_to_orig: + name_str = short_to_orig[name_str.lower()] obj_type = _decoded_str(table_type) if table_type is not None else "table" upper_type = obj_type.upper() @@ -1087,20 +1121,23 @@ def _create_materialized_view( ), quote_identifiers=self.QUOTE_IDENTIFIERS_IN_VIEWS, ) + registry = self._sync_mv_source_tables if is_sync_mv and sync_source_tbl is not None: - registry = getattr(self, "_sync_mv_source_tables", None) - if registry is None: - registry = {} - self._sync_mv_source_tables = registry - registry[target_table.name.lower()] = sync_source_tbl.copy() if target_table.db: registry[f"{target_table.db.lower()}.{target_table.name.lower()}"] = ( sync_source_tbl.copy() ) - elif sync_source_tbl.db: - registry[f"{sync_source_tbl.db.lower()}.{target_table.name.lower()}"] = ( - sync_source_tbl.copy() - ) + else: + registry[target_table.name.lower()] = sync_source_tbl.copy() + if sync_source_tbl.db: + registry[f"{sync_source_tbl.db.lower()}.{target_table.name.lower()}"] = ( + sync_source_tbl.copy() + ) + elif not is_sync_mv: + if target_table.db: + registry.pop(f"{target_table.db.lower()}.{target_table.name.lower()}", None) + else: + registry.pop(target_table.name.lower(), None) build_mode = ( self._to_upper_str(view_properties["build"]) if view_properties and view_properties.get("build") is not None @@ -1145,20 +1182,28 @@ def drop_view( view_properties = kwargs.pop("view_properties", None) sync_source_table: t.Any = None + is_explicit_async = False if materialized and view_properties: mat_type = self._extract_str_value(view_properties.get("materialized_type")) if mat_type and mat_type.upper() == "SYNC": sync_source_table = view_properties.get("source_table") - elif materialized: - registry = getattr(self, "_sync_mv_source_tables", None) or {} - lookup_key = ( - f"{target_table.db.lower()}.{target_table.name.lower()}" - if target_table.db - else target_table.name.lower() - ) - sync_source_table = registry.get(lookup_key) or registry.get(target_table.name.lower()) + elif mat_type and mat_type.upper() != "SYNC": + is_explicit_async = True + if materialized and sync_source_table is None and not is_explicit_async: + registry = self._sync_mv_source_tables + if target_table.db: + sync_source_table = registry.get( + f"{target_table.db.lower()}.{target_table.name.lower()}" + ) + else: + sync_source_table = registry.get(target_table.name.lower()) if sync_source_table is not None: + registry = self._sync_mv_source_tables + if target_table.db: + registry.pop(f"{target_table.db.lower()}.{target_table.name.lower()}", None) + else: + registry.pop(target_table.name.lower(), None) view_table_sql = self._to_sql(target_table) source_tbl_exp = self._to_source_table(sync_source_table) # Unqualified synchronous materialized views require the source table's database @@ -1506,12 +1551,64 @@ def insert_overwrite_by_partition( c for c in source_columns if c not in target_columns_to_types ] - super().insert_overwrite_by_partition( + if self.INSERT_OVERWRITE_STRATEGY.is_insert_overwrite: + 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, + ) + else: + self._replace_by_key( + table_name, + query_or_df, + target_columns_to_types, + key_exprs or partitioned_by, + is_unique_key=False, + source_columns=source_columns, + **kwargs, + ) + + def insert_overwrite_by_time_partition( + self, + table_name: TableName, + query_or_df: QueryOrDF, + start: TimeLike, + end: TimeLike, + time_formatter: t.Callable[[TimeLike, t.Optional[t.Dict[str, exp.DataType]]], exp.Expr], + time_column: TimeColumn | exp.Expr | str, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + source_columns: t.Optional[t.List[str]] = 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 + ] + return super().insert_overwrite_by_time_partition( table_name, query_or_df, - key_exprs or partitioned_by, + start=start, + end=end, + time_formatter=time_formatter, + time_column=time_column, target_columns_to_types=target_columns_to_types, source_columns=source_columns, + **kwargs, ) def delete_from( @@ -1644,6 +1741,24 @@ def _normalize_extra_conditions(self, extra_conditions: t.List[exp.Expr]) -> t.L normalized.append(cond) return normalized + def _expand_or_free_extra_condition_sets( + self, extra_conditions: t.Sequence[exp.Expr] + ) -> t.List[t.List[exp.Expr]]: + condition_sets: t.List[t.List[exp.Expr]] = [[]] + for cond in extra_conditions: + disjuncts = [ + self._normalize_single_boolean_conjunct(d) or d + for d in self._flatten_or_disjuncts(cond) + ] + if len(disjuncts) <= 1: + for cset in condition_sets: + cset.extend(disjuncts) + else: + condition_sets = [ + [*cset, branch.copy()] for cset in condition_sets for branch in disjuncts + ] + return condition_sets + @classmethod def _unwrap_literal_casts_in_where(cls, where: exp.Expr) -> exp.Expr: """Unwrap top-level `CAST( AS ...)` expressions in `DELETE WHERE` predicates @@ -1992,21 +2107,33 @@ def _execute_delete_with_subquery( return target_table = exp.to_table(table_name) - if self._is_iceberg_table(catalog_name=target_table.catalog): + if self._is_iceberg_table(catalog_name=target_table.catalog) or getattr( + self, "_in_iceberg_replace_by_key", False + ): self._delete_by_literal_keys(table_name, key_exprs, subquery, extra_conditions) return join_conditions: t.List[exp.Expr] = [] orig_subquery = subquery.copy() subquery = subquery.copy() - first_select = subquery.this - while isinstance(first_select, exp.SetOperation): - first_select = first_select.this - if isinstance(first_select, exp.Select) and any( - isinstance(e, exp.Star) for e in first_select.expressions + subq_branches: t.List[exp.Select] = [] + + def _collect_subq_branches(node: t.Optional[exp.Expr]) -> None: + if isinstance(node, exp.SetOperation): + _collect_subq_branches(node.this) + _collect_subq_branches(node.expression) + elif isinstance(node, exp.Select): + subq_branches.append(node) + + _collect_subq_branches(subquery.this) + if not subq_branches or any( + len(b.expressions) < len(key_exprs) + or any(isinstance(e, exp.Star) for e in b.expressions) + for b in subq_branches ): self._delete_by_literal_keys(table_name, key_exprs, orig_subquery, extra_conditions) return + first_select = subq_branches[0] used_t2_aliases: t.Set[str] = set() for idx, col in enumerate(key_exprs): col_name = col.name @@ -2110,6 +2237,18 @@ def _delete_by_time_transform_keys( """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 < `). """ + for k in key_exprs: + k_transform = self._extract_time_partition_transform(k) + if ( + k_transform is not None + and self._is_date_part_extract_fn(k) + and k_transform[1] != "year" + ): + raise SQLMeshError( + f"Doris delete_from cannot rewrite non-year date-part predicate '{k.sql(dialect=self.dialect)}' " + "into a contiguous range. Use DATE_TRUNC, TO_DATE, DATE, or YEAR instead." + ) + select_query = ( subquery.this.copy() if isinstance(subquery, exp.Subquery) else subquery.copy() ) @@ -2137,16 +2276,20 @@ def _delete_by_time_transform_keys( is_extract_fn = self._is_date_part_extract_fn(k) is_year_fn = unit == "year" and is_extract_fn src_transform = self._extract_time_partition_transform(raw_sel) - if src_transform is not None: + if ( + src_transform is not None + and src_transform[1] == unit + and (not self._is_date_part_extract_fn(raw_sel) or unit == "year") + ): 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 - elif is_extract_fn: + elif src_transform is not None or self._is_date_part_extract_fn(raw_sel): raise SQLMeshError( - f"Doris delete_from cannot rewrite date-part predicate '{k.sql(dialect=self.dialect)}' " - f"with non-date subquery projection '{raw_sel.sql(dialect=self.dialect)}' into a contiguous range." + f"Doris delete_from cannot rewrite time-partition predicate '{k.sql(dialect=self.dialect)}' " + f"with mismatched subquery projection '{raw_sel.sql(dialect=self.dialect)}' into a contiguous range." ) else: src_arg = raw_sel.copy() @@ -2198,9 +2341,16 @@ def _collect_select_branches(node: exp.Expr) -> None: unwrapped_indices.clear() break for idx in list(unwrapped_indices): + expected_transform = self._extract_time_partition_transform(key_exprs[idx]) + expected_unit = expected_transform[1] if expected_transform else "" sel_item = branch.expressions[idx] raw_sel = sel_item.this if isinstance(sel_item, exp.Alias) else sel_item - if self._extract_time_partition_transform(raw_sel) is None: + src_transform = self._extract_time_partition_transform(raw_sel) + if ( + src_transform is None + or src_transform[1] != expected_unit + or (self._is_date_part_extract_fn(raw_sel) and expected_unit != "year") + ): unwrapped_indices.discard(idx) if any(branch.args.get("group") for branch in all_branches) and ( @@ -2222,6 +2372,25 @@ def _collect_select_branches(node: exp.Expr) -> None: branch.set("group", None) branch.set("expressions", branch_exprs) + for idx, k in enumerate(key_exprs): + transform = self._extract_time_partition_transform(k) + if transform is not None and idx not in unwrapped_indices: + _, unit = transform + is_year_fn = unit == "year" and self._is_date_part_extract_fn(k) + if not is_year_fn: + for branch in all_branches: + if idx < len(branch.expressions): + b_sel = branch.expressions[idx] + b_raw = b_sel.this if isinstance(b_sel, exp.Alias) else b_sel + b_trans = self._extract_time_partition_transform(b_raw) + if self._is_date_part_extract_fn(b_raw) or ( + b_trans is not None and b_trans[1] != unit + ): + raise SQLMeshError( + f"Doris delete_from cannot rewrite time-partition predicate '{k.sql(dialect=self.dialect)}' " + f"with mismatched subquery projection '{b_raw.sql(dialect=self.dialect)}' into a contiguous range." + ) + 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 @@ -2239,11 +2408,6 @@ def _collect_select_branches(node: exp.Expr) -> None: raw_year_indices.add(idx) projections.append(exp.alias_(sq_col, f"_val_{idx}", quoted=True)) continue - if idx not in unwrapped_indices and is_extract_fn: - raise SQLMeshError( - f"Doris delete_from cannot rewrite date-part predicate '{k.sql(dialect=self.dialect)}' " - "into a contiguous range when subquery branches return non-date expressions." - ) trunc_expr = exp.Anonymous( this="DATE_TRUNC", expressions=[sq_col, exp.Literal.string(unit)], @@ -2281,6 +2445,8 @@ def _collect_select_branches(node: exp.Expr) -> None: col.set("db", None) col.set("catalog", None) cleaned_extra_conditions.append(cleaned_extra) + extra_condition_sets = self._expand_or_free_extra_condition_sets(cleaned_extra_conditions) + executed_delete_sqls: t.Set[str] = set() for row in rows: predicates: t.List[exp.Expr] = [] @@ -2296,11 +2462,17 @@ def _collect_select_branches(node: exp.Expr) -> None: if yr_val is None: predicates.append(target_col.is_(exp.null())) else: - yr_int = ( - yr_val.year - if isinstance(yr_val, (datetime.date, datetime.datetime)) - else int(yr_val) - ) + if isinstance(yr_val, (datetime.date, datetime.datetime)): + yr_int = yr_val.year + elif isinstance(yr_val, str): + yr_str = yr_val.strip() + yr_int = ( + int(yr_str[:4]) + if len(yr_str) >= 4 and yr_str[:4].isdigit() + else int(yr_str) + ) + 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( @@ -2342,8 +2514,12 @@ def _collect_select_branches(node: exp.Expr) -> 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))) + for extra_set in extra_condition_sets: + delete_expr = exp.delete(table_name, where=exp.and_(*predicates, *extra_set)) + delete_sql = self._to_sql(delete_expr) + if delete_sql not in executed_delete_sqls: + executed_delete_sqls.add(delete_sql) + self.execute(delete_expr) def _delete_by_literal_keys( self, @@ -2361,7 +2537,12 @@ def _delete_by_literal_keys( 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] + num_keys = len(key_exprs) + sanitized_rows = [ + tuple(self._sanitize_key_value(v) for v in row[:num_keys]) + for row in rows + if len(row) >= num_keys + ] self._delete_key_rows( table_name, key_exprs, @@ -2394,7 +2575,7 @@ def _delete_key_rows( col.set("db", None) col.set("catalog", None) cleaned_extra_conditions.append(cleaned_extra) - extra_conditions = cleaned_extra_conditions + extra_condition_sets = self._expand_or_free_extra_condition_sets(cleaned_extra_conditions) cols = [exp.column(k.name, quoted=True) for k in key_exprs] batch_size = self.DEFAULT_BATCH_SIZE @@ -2425,12 +2606,14 @@ def _dedup_exprs(exprs: t.List[exp.Expr]) -> t.List[exp.Expr]: 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))) + for extra_set in extra_condition_sets: + predicates: t.List[exp.Expr] = [col.isin(*batch), *extra_set] + 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))) + for extra_set in extra_condition_sets: + predicates = [col.is_(exp.null()), *extra_set] + self.execute(exp.delete(table_name, where=exp.and_(*predicates))) return # Composite key (N >= 2): dynamically select the column index to batch via IN (...) @@ -2457,12 +2640,14 @@ def _dedup_exprs(exprs: t.List[exp.Expr]) -> t.List[exp.Expr]: 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))) + for extra_set in extra_condition_sets: + preds = [*prefix_preds, batch_col.isin(*batch), *extra_set] + 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))) + for extra_set in extra_condition_sets: + preds = [*prefix_preds, batch_col.is_(exp.null()), *extra_set] + self.execute(exp.delete(table_name, where=exp.and_(*preds))) @classmethod def _resolve_table_type_and_keys( @@ -2621,6 +2806,32 @@ def replace_query( **kwargs, ) + def merge( + self, + target_table: TableName, + source_table: QueryOrDF, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]], + unique_key: t.Sequence[exp.Expr], + when_matched: t.Optional[exp.Whens] = None, + merge_filter: t.Optional[exp.Expr] = None, + source_columns: t.Optional[t.List[str]] = None, + **kwargs: t.Any, + ) -> None: + if when_matched or merge_filter: + prop = "when_matched" if when_matched else "merge_filter" + raise SQLMeshError( + f"This engine does not support MERGE expressions and therefore `{prop}` is not supported." + ) + self._replace_by_key( + target_table, + source_table, + target_columns_to_types, + unique_key, + is_unique_key=True, + source_columns=source_columns, + **kwargs, + ) + def _replace_by_key( self, target_table: TableName, @@ -2629,26 +2840,41 @@ def _replace_by_key( key: t.Sequence[exp.Expr], is_unique_key: bool, source_columns: t.Optional[t.List[str]] = None, + **kwargs: t.Any, ) -> None: - if target_columns_to_types is None: + if not target_columns_to_types or not columns_to_types_all_known(target_columns_to_types): target_columns_to_types = self.columns(target_table) table = exp.to_table(target_table) - is_iceberg = self._is_iceberg_table(catalog_name=table.catalog) + is_iceberg = self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=table.catalog, + table_properties=kwargs.get("table_properties"), + ) + normalized_key: t.List[exp.Expr] = [ + exp.column(k.name, quoted=bool(k.quoted)) if isinstance(k, exp.Identifier) else k + for k in key + ] key_cols: t.List[exp.Column] = [] - if is_unique_key and all(isinstance(k, (exp.Column, exp.Identifier)) for k in key): - key_cols = [exp.column(k.name) for k in key] + if is_unique_key and all(isinstance(k, exp.Column) for k in normalized_key): + key_cols = [exp.column(k.name) for k in normalized_key] temp_table_props: t.Optional[t.Dict[str, t.Any]] = None - if not is_iceberg and key_cols: - temp_table_props = { - "unique_key": ( - key_cols[0] if len(key_cols) == 1 else exp.Tuple(expressions=key_cols) - ) - } + if not is_iceberg: + if key_cols: + temp_table_props = { + "unique_key": ( + key_cols[0] if len(key_cols) == 1 else exp.Tuple(expressions=key_cols) + ) + } if target_columns_to_types and columns_to_types_all_known(target_columns_to_types): + reorder_kwargs: t.Dict[str, t.Any] = ( + {"table_properties": temp_table_props} + if temp_table_props + else ({"partitioned_by": list(normalized_key)} if normalized_key else {}) + ) target_columns_to_types = self._reorder_olap_columns( - target_columns_to_types, {"table_properties": temp_table_props} + target_columns_to_types, reorder_kwargs ) if source_columns: source_col_set = set(source_columns) @@ -2658,9 +2884,9 @@ def _replace_by_key( temp_table = self._get_temp_table(target_table) key_exp = ( - exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", *key, dialect=self.dialect) - if len(key) > 1 - else key[0] + exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", *normalized_key, dialect=self.dialect) + if len(normalized_key) > 1 + else normalized_key[0] ) column_names = list(target_columns_to_types or []) @@ -2672,6 +2898,7 @@ def _replace_by_key( exists=False, source_columns=source_columns, table_properties=temp_table_props, + table_format=kwargs.get("table_format"), ) try: @@ -2680,7 +2907,7 @@ def _replace_by_key( if not is_unique_key: delete_query = delete_query.distinct() else: - insert_query = insert_query.distinct(*key) + insert_query = insert_query.distinct(*normalized_key) insert_statement = exp.insert( insert_query, @@ -2690,7 +2917,12 @@ def _replace_by_key( delete_filter = key_exp.isin(query=delete_query) if not self.INSERT_OVERWRITE_STRATEGY.is_replace_where: - self.delete_from(target_table, delete_filter) + prev_iceberg_ctx = getattr(self, "_in_iceberg_replace_by_key", False) + self._in_iceberg_replace_by_key = is_iceberg + try: + self.delete_from(target_table, delete_filter) + finally: + self._in_iceberg_replace_by_key = prev_iceberg_ctx else: insert_statement.set("where", delete_filter) insert_statement.set("this", exp.to_table(target_table)) @@ -3145,9 +3377,13 @@ def _is_integer_indicator_projection(proj: exp.Expr) -> bool: ) return False - def _infer_projection_type(self, proj: exp.Expr) -> exp.DataType: + def _infer_projection_type( + self, proj: exp.Expr, cte_types: t.Optional[t.Dict[str, exp.DataType]] = None + ) -> exp.DataType: inner = proj.this if isinstance(proj, exp.Alias) else proj inner = inner.unnest() + if isinstance(inner, exp.Column) and cte_types and inner.name.lower() in cte_types: + return cte_types[inner.name.lower()].copy() if isinstance(inner, (exp.Cast, exp.TryCast)) and isinstance(inner.to, exp.DataType): return inner.to.copy() if self._is_integer_indicator_projection(proj): @@ -3233,8 +3469,21 @@ def _infer_projection_type(self, proj: exp.Expr) -> exp.DataType: return exp.DataType.build("UNKNOWN") def _reorder_unannotated_ctas_query( - self, expression: exp.Expr, kwargs: t.Dict[str, t.Any] + self, + expression: exp.Expr, + kwargs: t.Dict[str, t.Any], + cte_types: t.Optional[t.Dict[str, exp.DataType]] = None, ) -> exp.Expr: + cte_col_types: t.Dict[str, exp.DataType] = dict(cte_types or {}) + with_node = expression.args.get("with_") + if isinstance(with_node, exp.With): + for cte_item in with_node.expressions: + if isinstance(cte_item, exp.CTE) and isinstance(cte_item.this, exp.Select): + for cte_proj in cte_item.this.expressions: + if cte_proj.alias_or_name and not isinstance(cte_proj, exp.Star): + inferred_cte_t = self._infer_projection_type(cte_proj, cte_col_types) + if inferred_cte_t.this != exp.DataType.Type.UNKNOWN: + cte_col_types[cte_proj.alias_or_name.lower()] = inferred_cte_t if isinstance(expression, exp.SetOperation): expression = expression.copy() branches: t.List[exp.Select] = [] @@ -3270,13 +3519,34 @@ def _collect_branches(node: exp.Expr) -> None: 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 + branch_cands = [ + cand + for b in branches + if ( + cand := self._infer_projection_type( + b.expressions[col_idx], cte_col_types + ) + ).this + != exp.DataType.Type.UNKNOWN + ] + non_coercible_invalid = next( + ( + c + for c in branch_cands + if not self._is_valid_olap_key_type(c) + and not self._is_coercible_string_key_type(c) + ), + None, + ) + coercible_str = next( + (c for c in branch_cands if self._is_coercible_string_key_type(c)), + None, + ) + merged_cols_to_types[col_name] = ( + non_coercible_invalid + or coercible_str + or (branch_cands[0] if branch_cands else exp.DataType.build("UNKNOWN")) + ) reordered_select_cols = self._reorder_olap_columns(merged_cols_to_types, kwargs) perm = [orig_names.index(c) for c in reordered_select_cols] coercions = { @@ -3322,9 +3592,13 @@ def _collect_branches(node: exp.Expr) -> None: 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) + "this", + self._reorder_unannotated_ctas_query(expression.this, kwargs, cte_col_types), + ) + expression.set( + "expression", + self._reorder_unannotated_ctas_query(expression.expression, kwargs, cte_col_types), ) return expression if isinstance(expression, exp.Select): @@ -3335,7 +3609,7 @@ def _collect_branches(node: exp.Expr) -> None: } if proj_by_name and len(proj_by_name) == len(expression.expressions): select_cols_to_types = { - proj.alias_or_name: self._infer_projection_type(proj) + proj.alias_or_name: self._infer_projection_type(proj, cte_col_types) for proj in expression.expressions } reordered_select_cols = self._reorder_olap_columns(select_cols_to_types, kwargs) @@ -3381,8 +3655,7 @@ def _collect_branches(node: exp.Expr) -> None: with_clause = expression.args.get("with_") from_clause = expression.args.get("from_") if ( - not has_explicit_key - and isinstance(with_clause, exp.With) + isinstance(with_clause, exp.With) and isinstance(from_clause, exp.From) and isinstance(from_clause.this, exp.Table) ): @@ -3405,7 +3678,8 @@ def _collect_branches(node: exp.Expr) -> None: if not self._is_integer_indicator_projection(p) ] if ( - indicator_projs + not has_explicit_key + and indicator_projs and list(cte.this.expressions[: len(indicator_projs)]) != indicator_projs ): @@ -3413,6 +3687,18 @@ def _collect_branches(node: exp.Expr) -> None: copied_cte = expression.args["with_"].expressions[cte_idx] copied_select = copied_cte.this copied_select.set("expressions", indicator_projs + other_projs) + elif ( + len(expression.expressions) == 1 + and isinstance(expression.expressions[0], exp.Star) + and not expression.args.get("joins") + ): + reordered_cte_sel = self._reorder_unannotated_ctas_query( + cte.this, kwargs, cte_col_types + ) + if reordered_cte_sel != cte.this: + expression = expression.copy() + copied_cte = expression.args["with_"].expressions[cte_idx] + copied_cte.set("this", reordered_cte_sel) break return expression @@ -3428,8 +3714,9 @@ def _is_coercible_string_key_type(cls, dtype: exp.DataType) -> bool: def _reorder_olap_columns( self, target_columns_to_types: t.Dict[str, exp.DataType], - kwargs: t.Dict[str, t.Any], + kwargs: t.Optional[t.Dict[str, t.Any]] = None, ) -> t.Dict[str, exp.DataType]: + kwargs = kwargs or {} 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( @@ -3565,7 +3852,7 @@ def _reorder_olap_columns( first_dtype = updated_types[col_names[0]] if not self._is_valid_olap_key_type(first_dtype): - valid_col = next( + known_valid_col = next( ( c for c in col_names @@ -3573,23 +3860,25 @@ def _reorder_olap_columns( and not updated_types[c].is_type("unknown", "null") ), None, - ) or next( + ) + coercible_col = next( + (c for c in col_names if self._is_coercible_string_key_type(updated_types[c])), + None, + ) + unknown_valid_col = next( (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 known_valid_col: + col_names = [known_valid_col] + [c for c in col_names if c != known_valid_col] + elif coercible_col: + updated_types[coercible_col] = exp.DataType.build( + "VARCHAR(255)", dialect=self.dialect ) - 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] + if col_names[0] != coercible_col: + col_names = [coercible_col] + [c for c in col_names if c != coercible_col] + elif unknown_valid_col: + col_names = [unknown_valid_col] + [c for c in col_names if c != unknown_valid_col] # 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 @@ -4353,6 +4642,7 @@ def _eval_int_limit_expr(cls, node: t.Optional[exp.Expr]) -> t.Optional[int]: return None def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> str: + expression = self._normalize_schema_tmp_identifiers(expression) if expression.find(exp.Lock): expression = expression.copy() for select_node in expression.find_all(exp.Select): diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index 386635d1df..c172d46f3e 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -1,5 +1,6 @@ import typing as t +import pandas as pd # noqa: TID253 import pytest from pytest_mock.plugin import MockerFixture from sqlglot import expressions as exp @@ -3758,3 +3759,426 @@ def test_nested_delete_and_mixed_union_time_transform_and_merge_text_preservatio "verify_mode": False, "check_hostname": False, } + + +def test_doris_schema_tmp_shortening_across_fetchall_fetchdf_and_data_objects( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter, patch_get_data_objects=False) + long_tmp_name = "temp_very_long_model_name_exceeding_sixty_four_characters_abcdefgh_schema_tmp" + assert len(long_tmp_name) > 64 + short_tmp_name = adapter._shorten_schema_tmp_name(long_tmp_name) + assert len(short_tmp_name) <= 64 + assert short_tmp_name.endswith("_schema_tmp") + + # 1. `fetchall` / `_to_sql` rewrites Table, TableAlias, and Column.table qualifiers for >64-char `_schema_tmp` + mocker.patch.object(adapter.cursor, "fetchall", return_value=[(1,)]) + query_with_qualifiers = parse_one( + f"SELECT `{long_tmp_name}`.`id` FROM `db`.`{long_tmp_name}` AS `{long_tmp_name}`", + read="doris", + ) + assert adapter.fetchall(query_with_qualifiers) == [(1,)] + assert to_sql_calls(adapter) == [ + f"SELECT `{short_tmp_name}`.`id` FROM `db`.`{short_tmp_name}` AS `{short_tmp_name}`" + ] + + # 2. `fetchdf` (`_fetch_native_df`) normalizes >64-char `_schema_tmp` identifiers before `pd.read_sql_query` + mock_read_sql = mocker.patch( + "pandas.io.sql.read_sql_query", + return_value=pd.DataFrame({"id": [1]}), + ) + df = adapter.fetchdf(query_with_qualifiers) + assert list(df["id"]) == [1] + assert ( + mock_read_sql.call_args[0][0] + == f"SELECT `{short_tmp_name}`.`id` FROM `db`.`{short_tmp_name}` AS `{short_tmp_name}`" + ) + + # 3. `_get_data_objects` queries `information_schema.tables` using the shortened `_schema_tmp` name + # and restores the caller's original unshortened name on the returned `DataObject`. + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter, + "fetchall", + side_effect=[ + [("db", short_tmp_name, "BASE TABLE")], + [], + ], + ) + objs = adapter._get_data_objects("db", object_names={long_tmp_name}) + assert len(objs) == 1 + assert objs[0].name == long_tmp_name + query_arg = adapter.fetchall.call_args_list[0][0][0] + rendered_info_schema_sql = adapter._to_sql(query_arg) + assert short_tmp_name.lower() in rendered_info_schema_sql.lower() + assert long_tmp_name.lower() not in rendered_info_schema_sql.lower() + + +def test_doris_delete_by_time_transform_validation_and_deduplication( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Non-YEAR date-part extract `MONTH(ds) IN (SELECT MONTH(src_ds) ...)` raises `SQLMeshError` + # instead of silently deleting an entire calendar month. + with pytest.raises( + SQLMeshError, + match="cannot rewrite non-year date-part predicate", + ): + adapter.delete_from( + "db.events", + where=parse_one( + "MONTH(ds) IN (SELECT MONTH(src_ds) FROM db.staging)", + read="doris", + ), + ) + + # 2. Mismatched subquery time-transform unit (`DATE_TRUNC(ds, 'day') IN (SELECT YEAR(src_ds) ...)`) + # raises `SQLMeshError`. + with pytest.raises( + SQLMeshError, + match="with mismatched subquery projection", + ): + adapter.delete_from( + "db.events", + where=parse_one( + "DATE_TRUNC(ds, 'day') IN (SELECT YEAR(src_ds) FROM db.staging)", + read="doris", + ), + ) + + # 3. `YEAR(ds) IN (SELECT ds FROM db.staging)` where `staging.ds` returns ISO date strings + # in the same year parses the year prefix and deduplicates identical `DELETE` statements. + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter.cursor, + "fetchall", + return_value=[("2024-05-10",), ("2024-08-20",)], + ) + adapter.delete_from( + "db.events", + where=parse_one( + "YEAR(ds) IN (SELECT ds FROM db.staging)", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "SELECT DISTINCT `ds` AS `_val_0` FROM `db`.`staging`", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-01' AND `ds` < '2025-01-01'", + ] + + +def test_doris_replace_by_key_merge_and_partition_overwrite_edge_cases( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch("sqlmesh.core.engine_adapter.doris.random_id", return_value="abcd1234") + + # 1. `insert_overwrite_by_partition` / `_replace_by_key` with `{"notes": TEXT, "region": VARCHAR(16)}` + # and `partitioned_by=[region]` moves `region` to column 0 in the temp table so `notes` stays `STRING` + # (never truncated to `VARCHAR(255)`). + adapter.insert_overwrite_by_partition( + table_name="db.notes_by_region", + query_or_df=parse_one("SELECT notes, region FROM db.src_notes", read="doris"), + partitioned_by=[exp.to_column("region")], + target_columns_to_types={ + "notes": exp.DataType.build("TEXT", dialect="doris"), + "region": exp.DataType.build("VARCHAR(16)", dialect="doris"), + }, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE `db`.`__temp_notes_by_region_abcd1234` AS SELECT CAST(`region` AS VARCHAR(16)) AS `region`, CAST(`notes` AS STRING) AS `notes` FROM (SELECT `notes`, `region` FROM `db`.`src_notes`) AS `_subquery`", + "DELETE FROM `db`.`notes_by_region` AS `_t1` USING (SELECT DISTINCT `region` FROM `db`.`__temp_notes_by_region_abcd1234`) AS `_t2` WHERE `_t1`.`region` = `_t2`.`region`", + "INSERT INTO `db`.`notes_by_region` (`region`, `notes`) SELECT `region`, `notes` FROM `db`.`__temp_notes_by_region_abcd1234`", + "DROP TABLE IF EXISTS `db`.`__temp_notes_by_region_abcd1234`", + ] + + # 2. `merge` when `target_columns_to_types` contains `UNKNOWN` falls back to `self.columns(target_table)`. + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter, + "columns", + return_value={ + "id": exp.DataType.build("INT", dialect="doris"), + "val": exp.DataType.build("VARCHAR(64)", dialect="doris"), + }, + ) + adapter.merge( + target_table="db.unknown_cols_target", + source_table=parse_one("SELECT id, val FROM db.staging", read="doris"), + target_columns_to_types={ + "id": exp.DataType.build("UNKNOWN"), + "val": exp.DataType.build("UNKNOWN"), + }, + unique_key=[exp.to_column("id")], + ) + assert ( + to_sql_calls(adapter)[0] + == "CREATE TABLE `db`.`__temp_unknown_cols_target_abcd1234` UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10 AS SELECT CAST(`id` AS INT) AS `id`, CAST(`val` AS VARCHAR(64)) AS `val` FROM (SELECT `id`, `val` FROM `db`.`staging`) AS `_subquery`" + ) + + # 3. `merge` with 2-part table name and `table_format="iceberg"` routes to Iceberg `_replace_by_key` + # (no OLAP `UNIQUE KEY` or `DISTRIBUTED BY` on the temp table, and literal-key DELETE instead of DELETE ... USING). + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter.cursor, + "fetchall", + return_value=[(10,)], + ) + adapter.merge( + target_table="iceberg_db.events", + source_table=parse_one("SELECT id, payload FROM iceberg_db.staging", read="doris"), + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "payload": exp.DataType.build("STRING", dialect="doris"), + }, + unique_key=[exp.to_column("id")], + table_format="iceberg", + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE `iceberg_db`.`__temp_events_abcd1234` (`id` INT, `payload` STRING) ENGINE=iceberg", + "INSERT INTO `iceberg_db`.`__temp_events_abcd1234` (`id`, `payload`) SELECT CAST(`id` AS INT) AS `id`, CAST(`payload` AS STRING) AS `payload` FROM (SELECT `id`, `payload` FROM `iceberg_db`.`staging`) AS `_subquery`", + "SELECT DISTINCT `id` FROM `iceberg_db`.`__temp_events_abcd1234`", + "DELETE FROM `iceberg_db`.`events` WHERE `id` IN (10)", + "INSERT INTO `iceberg_db`.`events` (`id`, `payload`) SELECT `id`, `payload` FROM (SELECT `id` AS `id`, `payload` AS `payload`, ROW_NUMBER() OVER (PARTITION BY `id` ORDER BY `id`) AS _row_number FROM `iceberg_db`.`__temp_events_abcd1234`) AS _t WHERE _row_number = 1", + "DROP TABLE IF EXISTS `iceberg_db`.`__temp_events_abcd1234`", + ] + + # 4. `insert_overwrite_by_time_partition` with DataFrame and `{"val": DOUBLE, "ds": DATE, "id": INT}` + # reorders `target_columns_to_types` and `source_columns` so the temp table created by `_values_to_sql` + # places a valid OLAP key column (`ds` / `id`) before `val DOUBLE`. + adapter.cursor.execute.reset_mock() + df = pd.DataFrame( + { + "val": [3.14], + "ds": ["2024-01-01"], + "id": [42], + } + ) + adapter.insert_overwrite_by_time_partition( + table_name="db.metrics", + query_or_df=df, + start="2024-01-01", + end="2024-01-01", + time_formatter=lambda x, _: exp.Literal.string(str(x)), + time_column="ds", + target_columns_to_types={ + "val": exp.DataType.build("DOUBLE", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + }, + ) + time_part_calls = to_sql_calls(adapter) + assert time_part_calls == [ + "DELETE FROM `db`.`metrics` WHERE `ds` >= '2024-01-01 00:00:00+00:00' AND `ds` <= '2024-01-01 23:59:59.999999+00:00'", + "INSERT INTO `db`.`metrics` (`ds`, `val`, `id`) SELECT `ds`, `val`, `id` FROM (SELECT CAST(`ds` AS DATE) AS `ds`, CAST(`val` AS DOUBLE) AS `val`, CAST(`id` AS INT) AS `id` FROM (SELECT '2024-01-01' AS `ds`, 3.14 AS `val`, 42 AS `id`) AS `t`) AS `_subquery` WHERE `ds` BETWEEN '2024-01-01 00:00:00+00:00' AND '2024-01-01 23:59:59.999999+00:00'", + ] + + +def test_doris_unannotated_ctas_union_all_and_cte_column_reordering( + make_mocked_engine_adapter: t.Callable, +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. `UNION ALL` where Branch 1 has `1 AS metric` (INT) and Branch 2 has `1.5 AS metric` (DOUBLE): + # Doris widens `metric` to DOUBLE across the UNION, so `id` (INT across all branches) is promoted to col 0. + adapter.ctas( + table_name="db.union_metrics", + query_or_df=parse_one( + "SELECT 1 AS metric, 10 AS id UNION ALL SELECT 1.5 AS metric, 20 AS id", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`union_metrics` AS SELECT 10 AS `id`, 1 AS `metric` UNION ALL SELECT 20 AS `id`, 1.5 AS `metric`" + ] + + # 2. CTE (`WITH agg AS (...) SELECT avg_score, id FROM agg`) propagates CTE column types to the outer SELECT + # and reorders `id` before `avg_score` (`AVG` -> `DOUBLE`). + adapter.cursor.execute.reset_mock() + adapter.ctas( + table_name="db.cte_explicit_cols", + query_or_df=parse_one( + "WITH agg AS (SELECT AVG(score) AS avg_score, CAST(id AS INT) AS id FROM db.src GROUP BY id) " + "SELECT avg_score, id FROM agg", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`cte_explicit_cols` AS WITH `agg` AS (SELECT AVG(`score`) AS `avg_score`, CAST(`id` AS INT) AS `id` FROM `db`.`src` GROUP BY `id`) SELECT `id`, `avg_score` FROM `agg`" + ] + + # 3. CTE with `SELECT * FROM agg` reorders the inner CTE projection when the CTE starts with `AVG(score)`. + adapter.cursor.execute.reset_mock() + adapter.ctas( + table_name="db.cte_star", + query_or_df=parse_one( + "WITH agg AS (SELECT AVG(score) AS avg_score, CAST(id AS INT) AS id FROM db.src GROUP BY id) " + "SELECT * FROM agg", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`cte_star` AS WITH `agg` AS (SELECT CAST(`id` AS INT) AS `id`, AVG(`score`) AS `avg_score` FROM `db`.`src` GROUP BY `id`) SELECT * FROM `agg`" + ] + + # 4. `_reorder_olap_columns` prefers coercible `STRING` over `UNKNOWN` when column 0 is `DOUBLE`. + reordered_cols = adapter._reorder_olap_columns( + { + "metric": exp.DataType.build("DOUBLE", dialect="doris"), + "mystery": exp.DataType.build("UNKNOWN"), + "label": exp.DataType.build("STRING", dialect="doris"), + } + ) + assert list(reordered_cols.keys()) == ["label", "metric", "mystery"] + assert reordered_cols["label"] == exp.DataType.build("VARCHAR(255)", dialect="doris") + + +def test_doris_sync_mv_source_tables_persistence_and_isolation( + make_mocked_engine_adapter: t.Callable, +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Creating a sync MV in `db1` and `db2` with the same view name does not collide, + # survives `adapter.with_settings()`, and `drop_view` with `materialized_type='SYNC'` + # (without `source_table` in `view_properties`) still finds the recorded base table. + adapter.create_view( + "db1.shared_mv", + parse_one("SELECT id, SUM(val) AS total FROM db1.orders GROUP BY id", read="doris"), + replace=False, + materialized=True, + view_properties={"materialized_type": exp.Literal.string("SYNC")}, + ) + adapter.create_view( + "db2.shared_mv", + parse_one("SELECT id, SUM(val) AS total FROM db2.events GROUP BY id", read="doris"), + replace=False, + materialized=True, + view_properties={"materialized_type": exp.Literal.string("SYNC")}, + ) + + cloned_adapter = adapter.with_settings() + adapter.cursor.execute.reset_mock() + cloned_adapter.drop_view( + "db1.shared_mv", + materialized=True, + view_properties={"materialized_type": exp.Literal.string("SYNC")}, + ) + cloned_adapter.drop_view( + "db2.shared_mv", + materialized=True, + view_properties={"materialized_type": exp.Literal.string("SYNC")}, + ) + assert to_sql_calls(adapter) == [ + "DROP MATERIALIZED VIEW IF EXISTS `db1`.`shared_mv` ON `db1`.`orders`", + "DROP MATERIALIZED VIEW IF EXISTS `db2`.`shared_mv` ON `db2`.`events`", + ] + + # 2. Recreating a view name as an async MTMV pops any stale sync MV entry so subsequent + # `drop_view(materialized=True)` emits `DROP MATERIALIZED VIEW IF EXISTS db1.shared_mv`. + adapter.create_view( + "db1.shared_mv", + parse_one("SELECT id, SUM(val) AS total FROM db1.orders GROUP BY id", read="doris"), + replace=False, + materialized=True, + view_properties={"materialized_type": exp.Literal.string("SYNC")}, + ) + adapter.create_view( + "db1.shared_mv", + parse_one("SELECT id, val FROM db1.orders", read="doris"), + replace=False, + materialized=True, + ) + adapter.cursor.execute.reset_mock() + adapter.drop_view("db1.shared_mv", materialized=True) + assert to_sql_calls(adapter) == [ + "DROP MATERIALIZED VIEW IF EXISTS `db1`.`shared_mv`", + ] + + +def test_doris_reconnect_catalog_db_restore_union_star_delete_and_extra_or_conditions( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + class OperationalError(Exception): + pass + + adapter = make_mocked_engine_adapter( + DorisEngineAdapter, + default_catalog=None, + database="analytics", + ) + + # 1. `set_current_catalog("internal")` restores `USE \`analytics\`` when `database` is configured. + adapter.set_current_catalog("iceberg_cat") + adapter.set_current_catalog("internal") + assert to_sql_calls(adapter) == [ + "SWITCH `iceberg_cat`", + "SWITCH `internal`", + "USE `analytics`", + ] + + # 2. When a query fails with an operational error and the reconnect's `SWITCH` also fails, + # `_needs_catalog_restore` stays `True` until restoration succeeds on a subsequent call. + adapter.set_current_catalog("iceberg_cat") + adapter.cursor.execute.reset_mock() + adapter.cursor.execute.side_effect = [ + OperationalError(2013, "Lost connection to MySQL server during query"), + OperationalError(2006, "MySQL server has gone away"), + None, + None, + ] + with pytest.raises(OperationalError, match="MySQL server has gone away"): + adapter.execute("SELECT 1") + assert adapter._needs_catalog_restore is True + + adapter.execute("SELECT 1") + assert adapter._needs_catalog_restore is False + assert to_sql_calls(adapter) == [ + "SELECT 1", + "SWITCH `iceberg_cat`", + "SWITCH `iceberg_cat`", + "SELECT 1", + ] + + # 3. `delete_from` with `UNION ALL` where the second branch uses `SELECT *` expands/projects key columns + # and slices extra trailing columns in `_delete_by_literal_keys`. + adapter.cursor.execute.reset_mock() + adapter.cursor.execute.side_effect = None + mocker.patch.object( + adapter.cursor, + "fetchall", + return_value=[(101, "ignored_extra_col"), (102, "ignored_extra_col")], + ) + adapter.delete_from( + "db.target", + where=parse_one( + "id IN (SELECT id FROM db.s1 UNION ALL SELECT * FROM db.s2)", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "SELECT `id` FROM `db`.`s1` UNION ALL SELECT * FROM `db`.`s2`", + "DELETE FROM `db`.`target` WHERE `id` IN (101, 102)", + ] + + # 4. Extra `OR` condition alongside `IN (SELECT ...)` (`id IN (SELECT id FROM db.s1) AND (status = 'a' OR status IS NULL)`) + # on an Iceberg or literal-key delete path is distributed into separate `OR`-free `DELETE` statements so Doris does not + # reject the `DELETE` with `WhereClause only supports compound predicate, binary predicate, is_null predicate or in predicate`. + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter.cursor, + "fetchall", + return_value=[(1,), (2,)], + ) + adapter.delete_from( + "iceberg_cat.db.target", + where=parse_one( + "id IN (SELECT id FROM db.s1) AND (status = 'a' OR status IS NULL)", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "SELECT DISTINCT `id` FROM `db`.`s1`", + "DELETE FROM `iceberg_cat`.`db`.`target` WHERE `id` IN (1, 2) AND `status` = 'a'", + "DELETE FROM `iceberg_cat`.`db`.`target` WHERE `id` IN (1, 2) AND `status` IS NULL", + ] diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py index adb39540ec..b529eeac38 100644 --- a/tests/core/test_connection_config.py +++ b/tests/core/test_connection_config.py @@ -2578,10 +2578,62 @@ def test_doris(make_config): host="doris-fe", user="admin", ssl=False, + ssl_ca="/path/to/ca.pem", check_import=False, ) assert "ssl" not in ssl_false_config._static_connection_kwargs + ssl_empty_dict_config = make_config( + type="doris", + host="doris-fe", + user="admin", + ssl={}, + check_import=False, + ) + assert ssl_empty_dict_config._static_connection_kwargs["ssl"] == {"check_hostname": False} + + ssl_verify_cert_false_identity_true = make_config( + type="doris", + host="doris-fe", + user="admin", + ssl_ca="/path/to/ca.pem", + ssl_verify_cert=False, + ssl_verify_identity=True, + check_import=False, + ) + assert ssl_verify_cert_false_identity_true._static_connection_kwargs["ssl"] == { + "ca": "/path/to/ca.pem", + "verify_mode": False, + "check_hostname": False, + } + + ssl_dict_verify_mode_false = make_config( + type="doris", + host="doris-fe", + user="admin", + ssl={"ca": "/path/to/ca.pem", "verify_mode": False, "check_hostname": True}, + check_import=False, + ) + assert ssl_dict_verify_mode_false._static_connection_kwargs["ssl"] == { + "ca": "/path/to/ca.pem", + "verify_mode": False, + "check_hostname": False, + } + + ssl_dict_verify_mode_none_str = make_config( + type="doris", + host="doris-fe", + user="admin", + ssl={"ca": "/path/to/ca.pem", "verify_mode": "none"}, + ssl_verify_identity=True, + check_import=False, + ) + assert ssl_dict_verify_mode_none_str._static_connection_kwargs["ssl"] == { + "ca": "/path/to/ca.pem", + "verify_mode": "none", + "check_hostname": False, + } + with pytest.raises(ConfigError, match="Invalid Doris session variable name"): make_config( type="doris", From 98d1883f2e1cf102bdf1b1f6af39a2f4af582df4 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sun, 25 Oct 2026 18:14:29 -0700 Subject: [PATCH 30/31] fix(doris): harden dialect types, CTAS/DELETE/MERGE edge cases, and RowDiff Co-authored-by: xinge.ji Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/integrations/engines/doris.md | 4 +- sqlmesh/core/dialect.py | 147 ++++++- sqlmesh/core/engine_adapter/doris.py | 494 ++++++++++++++++++++---- tests/core/engine_adapter/test_doris.py | 361 ++++++++++++++++- 4 files changed, 918 insertions(+), 88 deletions(-) diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index e1167ac5ca..26dc5d376e 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -63,6 +63,8 @@ pip install "sqlmesh[doris]" | `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 | +| `light_schema_change_wait_timeout` | Maximum time in seconds to poll `SHOW ALTER TABLE COLUMN` when an `ALTER TABLE` statement triggers an asynchronous Doris schema change job (Default: `300.0`) | float | N | +| `light_schema_change_poll_interval` | Interval in seconds between `SHOW ALTER TABLE COLUMN` status polls during an asynchronous Doris schema change job (Default: `0.5`) | float | 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 | @@ -542,7 +544,7 @@ 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`: +Doris also supports single-table synchronous materialized views (rollup indexes maintained atomically on a single `internal`-catalog OLAP base table; queries with `JOIN`s or external catalog tables must use asynchronous MTMVs instead) by setting `materialized_type = 'SYNC'` and `source_table` in `physical_properties`: ```sql MODEL ( diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py index bcfb300ba1..078e4b8636 100644 --- a/sqlmesh/core/dialect.py +++ b/sqlmesh/core/dialect.py @@ -10,12 +10,14 @@ from enum import Enum, auto from functools import lru_cache -from sqlglot import Dialect, Generator, ParseError, Parser, Tokenizer, TokenType, exp -from sqlglot.dialects.dialect import DialectType -from sqlglot.dialects import DuckDB, Snowflake, TSQL +from sqlglot import Dialect, Generator, ParseError, Parser, Tokenizer, TokenType, exp, transforms +from sqlglot.dialects.dialect import DialectType, var_map_sql +from sqlglot.dialects import Doris, DuckDB, Snowflake, TSQL import sqlglot.dialects.athena as athena import sqlglot.generators.athena as athena_generators +from sqlglot.parser import build_var_map from sqlglot.parsers.athena import AthenaTrinoParser +from sqlglot.parsers.hive import _build_named_struct from sqlglot.helper import seq_get from sqlglot.optimizer.normalize_identifiers import normalize_identifiers from sqlglot.optimizer.qualify_columns import quote_identifiers @@ -484,6 +486,38 @@ def _parse_types( with_collation: bool = False, ) -> t.Optional[exp.Expr]: start = self._curr + udt_name = ( + self._curr.text.upper() + if self._curr and self._curr.token_type == TokenType.USERDEFINED + else None + ) + if udt_name == "AGG_STATE" and not check_func: + self._advance() + if self._curr and self._curr.token_type in (TokenType.LT, TokenType.L_PAREN): + open_tok = self._curr.token_type + close_tok = TokenType.GT if open_tok == TokenType.LT else TokenType.R_PAREN + depth = 1 + self._advance() + while self._curr and depth > 0: + if self._curr.token_type == open_tok: + depth += 1 + elif self._curr.token_type == close_tok: + depth -= 1 + self._advance() + parsed_type: t.Optional[exp.Expr] = self.expression( + exp.DataType( + this=exp.DataType.Type.USERDEFINED, + kind=self._find_sql(start, self._prev), + ) + ) + else: + parsed_type = self.expression( + exp.DataType(this=exp.DataType.Type.USERDEFINED, kind="AGG_STATE") + ) + if schema and parsed_type: + parsed_type.meta["sql"] = self._find_sql(start, self._prev) + return parsed_type + parsed_type = self.__parse_types( # type: ignore check_func=check_func, schema=schema, @@ -491,6 +525,14 @@ def _parse_types( with_collation=with_collation, ) + if ( + udt_name + and isinstance(parsed_type, exp.DataType) + and parsed_type.this == exp.DataType.Type.USERDEFINED + and not parsed_type.args.get("kind") + ): + parsed_type.set("kind", udt_name) + if schema and parsed_type: parsed_type.meta["sql"] = self._find_sql(start, self._prev) @@ -581,6 +623,46 @@ def altercolumn_sql(self: Generator, expression: exp.AlterColumn) -> str: return sql +def columndef_sql(self: Generator, expression: exp.ColumnDef, sep: str = " ") -> str: + return self._columndef_sql( # type: ignore + expression, + sep=( + ": " + if isinstance(expression.parent, exp.DataType) and expression.parent.is_type("struct") + else sep + ), + ) + + +def _doris_map_sql(self: Generator, expression: exp.Map | exp.VarMap) -> str: + keys = expression.args.get("keys") + values = expression.args.get("values") + if isinstance(keys, exp.Array) and isinstance(values, exp.Array): + return var_map_sql(self, expression) + if keys is not None or values is not None: + return self.func("MAP", keys, values) + return self.func("MAP", *expression.expressions) + + +def _doris_struct_sql(self: Generator, expression: exp.Struct) -> str: + if any(isinstance(e, exp.PropertyEQ) for e in expression.expressions): + args: t.List[str] = [] + for e in expression.expressions: + if isinstance(e, exp.PropertyEQ): + key_sql = ( + self.sql(exp.Literal.string(e.this.name)) + if isinstance(e.this, (exp.Identifier, exp.Var)) + else self.sql(e.this) + ) + args.append(key_sql) + args.append(self.sql(e.expression)) + else: + args.append(self.sql(e)) + return self.func("NAMED_STRUCT", *args) + values = [e.this if isinstance(e, exp.Alias) else e for e in expression.expressions] + return self.func("STRUCT", *values) + + def _parse_if(self: Parser) -> t.Optional[exp.Expr]: # If we fail to parse an IF function with expressions as arguments, we then try # to parse a statement / command to support the macro @IF(condition, statement) @@ -1343,6 +1425,64 @@ def parse_virtual_statement( def extend_sqlglot() -> None: """Extend SQLGlot with SQLMesh's custom macro aware dialect.""" + Doris.SUPPORTS_USER_DEFINED_TYPES = True + if Doris.tokenizer_class.__name__ != "_DorisTokenizer": + + class _DorisTokenizer(Doris.tokenizer_class): # type: ignore + KEYWORDS = { + **Doris.tokenizer_class.KEYWORDS, + "AGG_STATE": TokenType.USERDEFINED, + "BITMAP": TokenType.USERDEFINED, + "DATETIMEV2": TokenType.DATETIME, + "DATEV2": TokenType.DATE, + "DECIMALV2": TokenType.DECIMAL, + "DECIMALV3": TokenType.DECIMAL, + "HLL": TokenType.HLLSKETCH, + "IPV4": TokenType.IPV4, + "IPV6": TokenType.IPV6, + "JSONB": TokenType.JSONB, + "LARGEINT": TokenType.INT128, + "QUANTILE_STATE": TokenType.USERDEFINED, + "VARIANT": TokenType.VARIANT, + } + + Doris.Tokenizer = _DorisTokenizer + Doris.tokenizer_class = _DorisTokenizer + + Doris.Parser.TABLE_ALIAS_TOKENS = Doris.Parser.TABLE_ALIAS_TOKENS - { + TokenType.ANTI, + TokenType.SEMI, + } + Doris.Parser.FUNCTIONS.update( + { + "MAP": build_var_map, + "NAMED_STRUCT": _build_named_struct, + } + ) + Doris.Generator.TYPE_MAPPING.update( + { + exp.DataType.Type.HLLSKETCH: "HLL", + exp.DataType.Type.INT128: "LARGEINT", + exp.DataType.Type.IPV4: "IPV4", + exp.DataType.Type.IPV6: "IPV6", + exp.DataType.Type.JSONB: "JSONB", + exp.DataType.Type.VARIANT: "VARIANT", + } + ) + Doris.Generator.TRANSFORMS.update( + { + exp.Map: _doris_map_sql, + exp.Select: transforms.preprocess( + [ + transforms.eliminate_distinct_on, + transforms.unnest_generate_date_array_using_recursive_cte, + ] + ), + exp.Struct: _doris_struct_sql, + exp.VarMap: _doris_map_sql, + } + ) + tokenizers = {Tokenizer} parsers = {Parser} generators = {Generator} @@ -1432,6 +1572,7 @@ def extend_sqlglot() -> None: _override(Snowflake.Parser, _parse_table_parts) _override(TSQL.Parser, _parse_alter_table_alter) _override(TSQL.Generator, altercolumn_sql) + _override(Doris.Generator, columndef_sql) # DuckDB's prefix absolute power operator `@` clashes with the macro syntax DuckDB.Parser.NO_PAREN_FUNCTION_PARSERS.pop("@", None) diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 302f0c3f23..4bf35dceab 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -183,6 +183,7 @@ class DorisEngineAdapter( "STRUCT", "BITMAP", "HLL", + "HLLSKETCH", "QUANTILE_STATE", "AGG_STATE", } @@ -228,6 +229,18 @@ def _last_described_key_cols(self) -> t.Set[str]: def _last_described_key_cols(self, value: t.Set[str]) -> None: self._get_thread_state().last_described_key_cols = set(value) + @property + def _in_iceberg_replace_by_key(self) -> bool: + return bool(getattr(self._get_thread_state(), "in_iceberg_replace_by_key", False)) + + @_in_iceberg_replace_by_key.setter + def _in_iceberg_replace_by_key(self, value: bool) -> None: + self._get_thread_state().in_iceberg_replace_by_key = bool(value) + + @staticmethod + def _get_table_properties(kwargs: t.Dict[str, t.Any]) -> t.Optional[t.Dict[str, t.Any]]: + return kwargs.get("table_properties") or kwargs.get("physical_properties") + @property def _sync_mv_source_tables(self) -> t.Dict[str, exp.Table]: registry = getattr(self._connection_pool, "_doris_sync_mv_source_tables", None) @@ -267,6 +280,13 @@ 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 = catalog + self._doris_last_switched_catalog = catalog + if isinstance(getattr(self, "_extra_config", None), dict): + self._extra_config["_doris_last_switched_catalog"] = catalog + try: + setattr(self._connection_pool, "_doris_last_switched_catalog", catalog) + except Exception: + pass self._needs_catalog_restore = False database = self._extra_config.get("database") default_cat = self.default_catalog or "internal" @@ -275,7 +295,13 @@ def set_current_catalog(self, catalog: str) -> None: self.execute(f"USE {db_sql}") def _restore_active_catalog_on_reconnect(self, sql: t.Optional[str] = None) -> None: - active_cat = self._active_catalog or self.default_catalog + active_cat = ( + self._active_catalog + or getattr(self._connection_pool, "_doris_last_switched_catalog", None) + or self._extra_config.get("_doris_last_switched_catalog") + or getattr(self, "_doris_last_switched_catalog", None) + or self.default_catalog + ) database = self._extra_config.get("database") if not active_cat and not database: self._needs_catalog_restore = False @@ -402,11 +428,26 @@ def _is_valid_olap_key_type(cls, dtype: exp.DataType) -> bool: return False return True + def _normalize_nested_value(self, expr: exp.Expr) -> exp.Expr: + return exp.cast(expr, exp.DataType.build("VARCHAR", dialect=self.dialect)) + def _fetch_native_df( self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False ) -> DF: if isinstance(query, exp.Expr): query = self._normalize_schema_tmp_identifiers(query) + elif isinstance(query, str) and ( + "_schema_tmp" in query + or any( + tok in query.upper() + for tok in ("IS TRUE", "IS FALSE", "IS NOT TRUE", "IS NOT FALSE") + ) + ): + try: + parsed_q = parse_one(query, read=self.dialect) + query = self._to_sql(parsed_q, quote=quote_identifiers) + except Exception: + pass if self._pre_ping: try: self.ping() @@ -751,6 +792,16 @@ def drop_schema( **drop_args: t.Dict[str, exp.Expr], ) -> None: """Drop schema in Doris. Note: Doris doesn't support CASCADE clause.""" + if cascade: + schema_exp = to_schema(schema_name) if schema_name else exp.Table() + effective_catalog = schema_exp.catalog or self.default_catalog + if effective_catalog and str(effective_catalog).strip().lower() != "internal": + try: + for obj in self._get_data_objects(schema_name): + if obj.type == DataObjectType.TABLE: + self.drop_table(obj.to_table(), exists=True) + except Exception: + logger.debug("Failed to enumerate tables before drop_schema on %s", schema_name) return self._drop_object( name=schema_name, exists=ignore_if_not_exists, @@ -964,7 +1015,7 @@ 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)] + ordered_df = query_or_df[list(source_columns_to_types)] values: t.List[t.Tuple[t.Any, ...]] = list( ordered_df.itertuples(index=False, name=None) ) @@ -1104,11 +1155,17 @@ def _create_materialized_view( # 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] = sync_source_tbl.db if sync_source_tbl else None - if source_db: - db_sql = self._to_sql(exp.to_identifier(source_db, quoted=True)) + if is_sync_mv: + sync_db: t.Optional[str] = target_table.db or ( + sync_source_tbl.db if sync_source_tbl else None + ) + if sync_db: + db_sql = self._to_sql(exp.to_identifier(sync_db, quoted=True)) self.execute(f"USE {db_sql}") + unqualified_mv = target_table.copy() + unqualified_mv.set("db", None) + unqualified_mv.set("catalog", None) + schema = unqualified_mv self.execute( exp.Create( @@ -1204,13 +1261,17 @@ def drop_view( registry.pop(f"{target_table.db.lower()}.{target_table.name.lower()}", None) else: registry.pop(target_table.name.lower(), None) - view_table_sql = self._to_sql(target_table) source_tbl_exp = self._to_source_table(sync_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: - db_sql = self._to_sql(exp.to_identifier(source_tbl_exp.db, quoted=True)) + # Synchronous materialized views require the target/source table's database + # to be active on the session and an unqualified `` in `DROP MATERIALIZED VIEW ON `. + sync_db = target_table.db or source_tbl_exp.db + if sync_db: + db_sql = self._to_sql(exp.to_identifier(sync_db, quoted=True)) self.execute(f"USE {db_sql}") + unqualified_mv = target_table.copy() + unqualified_mv.set("db", None) + unqualified_mv.set("catalog", None) + view_table_sql = self._to_sql(unqualified_mv) source_table_sql = self._to_sql(source_tbl_exp) exists_clause = "IF EXISTS " if ignore_if_not_exists else "" self.execute( @@ -1277,11 +1338,11 @@ def _build_create_comment_column_exp( 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) + table = t.cast(exp.Table, self._normalize_schema_tmp_identifiers(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) + super()._truncate_table(table) return except Exception as e: err_msg = str(e).lower() @@ -1293,7 +1354,7 @@ def _truncate_table(self, table_name: TableName) -> None: # 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) + cols_to_types = self.columns(table) if cols_to_types: valid_col_name = next( ( @@ -1326,7 +1387,7 @@ def _insert_overwrite_by_condition( and self._is_iceberg_table( table_format=kwargs.get("table_format"), catalog_name=table.catalog, - table_properties=kwargs.get("table_properties"), + table_properties=self._get_table_properties(kwargs), ) ): insert_overwrite_strategy_override = InsertOverwriteStrategy.INSERT_OVERWRITE @@ -1533,7 +1594,7 @@ def insert_overwrite_by_partition( is_iceberg = self._is_iceberg_table( table_format=kwargs.get("table_format"), catalog_name=target_table.catalog, - table_properties=kwargs.get("table_properties"), + table_properties=self._get_table_properties(kwargs), ) if ( target_columns_to_types @@ -1586,7 +1647,7 @@ def insert_overwrite_by_time_partition( is_iceberg = self._is_iceberg_table( table_format=kwargs.get("table_format"), catalog_name=target_table.catalog, - table_properties=kwargs.get("table_properties"), + table_properties=self._get_table_properties(kwargs), ) if ( not is_iceberg @@ -1633,6 +1694,7 @@ def delete_from( where = self._unwrap_literal_casts_in_where(where) where = self._rewrite_between_in_where(where) + where = self._rewrite_direct_time_transforms_in_where(where) where = self._normalize_boolean_column_conjuncts(where) subquery_info = self._extract_subquery_from_where(where) @@ -1650,6 +1712,13 @@ def delete_from( self._delete_key_rows(table_name, *or_key_info, is_ast_values=True) return + expanded_or_sets = self._expand_or_free_extra_condition_sets([where]) + if len(expanded_or_sets) > 1: + for branch_conjuncts in expanded_or_sets: + if branch_conjuncts: + self.delete_from(table_name, exp.and_(*branch_conjuncts, copy=False)) + 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 ( @@ -1724,41 +1793,201 @@ def _flatten_or_disjuncts(cls, expr: exp.Expr) -> t.List[exp.Expr]: def _normalize_extra_conditions(self, extra_conditions: t.List[exp.Expr]) -> t.List[exp.Expr]: normalized: t.List[exp.Expr] = [] - for cond in extra_conditions: - unwrapped = cond.unnest() - if isinstance(unwrapped, exp.Or): - or_info = self._extract_or_key_rows_from_where(unwrapped) - if ( - or_info is not None - and len(or_info[0]) == 1 - and not or_info[2] - and all(not isinstance(r[0], exp.Null) for r in or_info[1]) - ): - col_expr = exp.column(or_info[0][0].name, quoted=True) - vals = [r[0].copy() for r in or_info[1]] - normalized.append(col_expr.isin(*vals)) - continue - normalized.append(cond) + for raw_cond in extra_conditions: + for cond in self._flatten_and_conjuncts(raw_cond): + unwrapped = cond.unnest() + if isinstance(unwrapped, exp.Or): + or_info = self._extract_or_key_rows_from_where(unwrapped) + if ( + or_info is not None + and len(or_info[0]) == 1 + and not or_info[2] + and all(not isinstance(r[0], exp.Null) for r in or_info[1]) + ): + col_expr = exp.column(or_info[0][0].name, quoted=True) + vals = [r[0].copy() for r in or_info[1]] + normalized.append(col_expr.isin(*vals)) + continue + normalized.append(cond) return normalized + def _expand_or_to_conjunct_sets(self, expr: exp.Expr) -> t.List[t.List[exp.Expr]]: + conjuncts = self._flatten_and_conjuncts(expr) + if len(conjuncts) > 1: + sets: t.List[t.List[exp.Expr]] = [[]] + for c in conjuncts: + sub_sets = self._expand_or_to_conjunct_sets(c) + sets = [[*s, *b] for s in sets for b in sub_sets] + return sets + single = conjuncts[0] + disjuncts = self._flatten_or_disjuncts(single) + if len(disjuncts) > 1: + or_sets: t.List[t.List[exp.Expr]] = [] + for d in disjuncts: + or_sets.extend(self._expand_or_to_conjunct_sets(d)) + return or_sets + leaf = disjuncts[0] + norm = self._normalize_single_boolean_conjunct(leaf) + return [[norm if norm is not None else leaf.copy()]] + def _expand_or_free_extra_condition_sets( self, extra_conditions: t.Sequence[exp.Expr] ) -> t.List[t.List[exp.Expr]]: condition_sets: t.List[t.List[exp.Expr]] = [[]] for cond in extra_conditions: - disjuncts = [ - self._normalize_single_boolean_conjunct(d) or d - for d in self._flatten_or_disjuncts(cond) - ] - if len(disjuncts) <= 1: - for cset in condition_sets: - cset.extend(disjuncts) - else: - condition_sets = [ - [*cset, branch.copy()] for cset in condition_sets for branch in disjuncts - ] + sub_sets = self._expand_or_to_conjunct_sets(cond) + condition_sets = [[*cset, *branch] for cset in condition_sets for branch in sub_sets] return condition_sets + def _build_time_transform_range( + self, + t_expr: exp.Expr, + col_expr: exp.Column, + unit: str, + val_expr: exp.Expr, + ) -> t.Optional[exp.Expr]: + if self._is_date_part_extract_fn(t_expr) and unit != "year": + raise SQLMeshError( + f"Doris delete_from cannot rewrite non-year date-part predicate '{t_expr.sql(dialect=self.dialect)}' " + "into a contiguous range. Use DATE_TRUNC, TO_DATE, DATE, or YEAR instead." + ) + raw_val = self._extract_str_value(val_expr) + if raw_val is None: + return None + raw_str = raw_val.strip() + if self._is_date_part_extract_fn(t_expr) and unit == "year": + try: + yr_int = ( + int(raw_str[:4]) + if len(raw_str) >= 4 and raw_str[:4].isdigit() + else int(raw_str) + ) + except ValueError: + return None + return exp.and_( + exp.GTE( + this=col_expr.copy(), + expression=exp.Literal.string(f"{yr_int:04d}-01-01"), + ), + exp.LT( + this=col_expr.copy(), + expression=exp.Literal.string(f"{yr_int + 1:04d}-01-01"), + ), + copy=False, + ) + try: + dt = datetime.datetime.fromisoformat(raw_str) + except ValueError: + return None + has_time = " " in raw_str or "T" in raw_str or unit in ("hour", "minute", "second") + if unit == "year": + low_dt = dt.replace(month=1, day=1, hour=0, minute=0, second=0, microsecond=0) + high_dt = low_dt.replace(year=low_dt.year + 1) + elif unit == "quarter": + q_month = ((dt.month - 1) // 3) * 3 + 1 + low_dt = dt.replace(month=q_month, day=1, hour=0, minute=0, second=0, microsecond=0) + high_dt = low_dt.replace( + year=low_dt.year + (1 if q_month == 10 else 0), + month=(1 if q_month == 10 else q_month + 3), + ) + elif unit == "month": + low_dt = dt.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + high_dt = low_dt.replace( + year=low_dt.year + (1 if low_dt.month == 12 else 0), + month=(1 if low_dt.month == 12 else low_dt.month + 1), + ) + elif unit == "week": + low_dt = (dt - datetime.timedelta(days=dt.weekday())).replace( + hour=0, minute=0, second=0, microsecond=0 + ) + high_dt = low_dt + datetime.timedelta(days=7) + elif unit == "day": + low_dt = dt.replace(hour=0, minute=0, second=0, microsecond=0) + high_dt = low_dt + datetime.timedelta(days=1) + elif unit == "hour": + low_dt = dt.replace(minute=0, second=0, microsecond=0) + high_dt = low_dt + datetime.timedelta(hours=1) + elif unit == "minute": + low_dt = dt.replace(second=0, microsecond=0) + high_dt = low_dt + datetime.timedelta(minutes=1) + elif unit == "second": + low_dt = dt.replace(microsecond=0) + high_dt = low_dt + datetime.timedelta(seconds=1) + else: + return None + fmt = "%Y-%m-%d %H:%M:%S" if has_time else "%Y-%m-%d" + return exp.and_( + exp.GTE( + this=col_expr.copy(), + expression=exp.Literal.string(low_dt.strftime(fmt)), + ), + exp.LT( + this=col_expr.copy(), + expression=exp.Literal.string(high_dt.strftime(fmt)), + ), + copy=False, + ) + + def _rewrite_direct_time_transforms_in_where(self, where: exp.Expr) -> exp.Expr: + for node in list(where.find_all(exp.EQ, exp.In)): + if node.find_ancestor(exp.Subquery, exp.Select, exp.Not): + continue + if isinstance(node, exp.EQ): + left_trans = self._extract_time_partition_transform(node.this) + right_trans = self._extract_time_partition_transform(node.expression) + t_expr: t.Optional[exp.Expr] = None + col_expr: t.Optional[exp.Column] = None + unit: t.Optional[str] = None + val_expr: t.Optional[exp.Expr] = None + if left_trans is not None and self._is_literal_key_value(node.expression): + t_expr = node.this + col_expr, unit = left_trans + val_expr = node.expression + elif right_trans is not None and self._is_literal_key_value(node.this): + t_expr = node.expression + col_expr, unit = right_trans + val_expr = node.this + if ( + t_expr is not None + and col_expr is not None + and unit is not None + and val_expr is not None + ): + replacement = self._build_time_transform_range(t_expr, col_expr, unit, val_expr) + if replacement is not None: + if node is where: + where = replacement + else: + if not isinstance(node.parent, (exp.Where, exp.Paren)): + replacement = exp.Paren(this=replacement) + node.replace(replacement) + elif ( + isinstance(node, exp.In) + and self._get_in_subquery(node) is None + and node.expressions + ): + in_trans = self._extract_time_partition_transform(node.this) + if in_trans is not None and all( + self._is_literal_key_value(v) for v in node.expressions + ): + col_expr, unit = in_trans + ranges = [ + self._build_time_transform_range(node.this, col_expr, unit, v) + for v in node.expressions + ] + if ranges and all(r is not None for r in ranges): + or_expr: exp.Expr = exp.or_( + *[exp.Paren(this=r) for r in ranges if r is not None], + copy=False, + ) + if node is where: + where = or_expr + else: + if not isinstance(node.parent, (exp.Where, exp.Paren)): + or_expr = exp.Paren(this=or_expr) + node.replace(or_expr) + return where + @classmethod def _unwrap_literal_casts_in_where(cls, where: exp.Expr) -> exp.Expr: """Unwrap top-level `CAST( AS ...)` expressions in `DELETE WHERE` predicates @@ -2059,6 +2288,13 @@ def _unwrap_concat_ws(expr: exp.Expr) -> t.Optional[exp.Expr]: def _unwrap_concat_ws_in_query( cls, query_expr: exp.Expr, expected_arity: int ) -> t.Optional[exp.Expr]: + if isinstance(query_expr, exp.Subquery): + inner = cls._unwrap_concat_ws_in_query(query_expr.this, expected_arity) + if inner is None: + return None + unwrapped_sq = query_expr.copy() + unwrapped_sq.set("this", inner) + return unwrapped_sq 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) @@ -2119,7 +2355,9 @@ def _execute_delete_with_subquery( subq_branches: t.List[exp.Select] = [] def _collect_subq_branches(node: t.Optional[exp.Expr]) -> None: - if isinstance(node, exp.SetOperation): + if isinstance(node, exp.Subquery): + _collect_subq_branches(node.this) + elif isinstance(node, exp.SetOperation): _collect_subq_branches(node.this) _collect_subq_branches(node.expression) elif isinstance(node, exp.Select): @@ -2253,7 +2491,7 @@ def _delete_by_time_transform_keys( subquery.this.copy() if isinstance(subquery, exp.Subquery) else subquery.copy() ) first_select = select_query - while isinstance(first_select, exp.SetOperation): + while isinstance(first_select, (exp.SetOperation, exp.Subquery)): first_select = first_select.this if not isinstance(first_select, exp.Select): return @@ -2266,6 +2504,7 @@ def _delete_by_time_transform_keys( and not select_query.args.get("group") and not select_query.args.get("having") and not select_query.args.get("limit") + and not select_query.args.get("offset") ): for idx, k in enumerate(key_exprs): sel_item = select_query.expressions[idx] @@ -2306,13 +2545,16 @@ def _delete_by_time_transform_keys( 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("order", None) 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): + if isinstance(node, exp.Subquery): + _collect_select_branches(node.this) + elif isinstance(node, exp.SetOperation): _collect_select_branches(node.this) _collect_select_branches(node.expression) elif isinstance(node, exp.Select): @@ -2326,9 +2568,21 @@ def _collect_select_branches(node: exp.Expr) -> None: } unwrapped_indices: t.Set[int] = set(transform_indices) + # INTERSECT / EXCEPT compare full projection values across branches; unwrapping + # `DATE_TRUNC(ds, 'day')` into raw `ds` inside branches would change set semantics. + if any( + isinstance(op, (exp.Intersect, exp.Except)) + for op in select_query.find_all(exp.SetOperation) + ): + unwrapped_indices.clear() + # Pass 1: Validate which transform indices can be unwrapped across ALL branches. for branch in all_branches: - if len(branch.expressions) != len(key_exprs): + if ( + len(branch.expressions) != len(key_exprs) + or branch.args.get("limit") + or branch.args.get("offset") + ): unwrapped_indices.clear() break has_group = bool(branch.args.get("group")) @@ -2533,7 +2787,12 @@ def _delete_by_literal_keys( subquery.this.copy() if isinstance(subquery, exp.Subquery) else subquery.copy() ) if isinstance(select_query, exp.Select): - select_query = select_query.distinct() + if select_query.args.get("limit") or select_query.args.get("offset"): + if not select_query.args.get("order"): + select_query = select_query.distinct() + else: + select_query.set("order", None) + select_query = select_query.distinct() rows = self.fetchall(select_query, quote_identifiers=True) if not rows: return @@ -2740,7 +2999,7 @@ def _create_table_from_columns( column_descriptions: t.Optional[t.Dict[str, str]] = None, **kwargs: t.Any, ) -> None: - table_properties = dict(kwargs.get("table_properties") or {}) + table_properties = dict(self._get_table_properties(kwargs) or {}) unique_key_prop, duplicate_key_prop, aggregate_key_prop = self._resolve_table_type_and_keys( table_properties ) @@ -2782,7 +3041,7 @@ def replace_query( is_iceberg = self._is_iceberg_table( table_format=kwargs.get("table_format"), catalog_name=target_table.catalog, - table_properties=kwargs.get("table_properties"), + table_properties=self._get_table_properties(kwargs), ) if ( not is_iceberg @@ -2849,7 +3108,7 @@ def _replace_by_key( is_iceberg = self._is_iceberg_table( table_format=kwargs.get("table_format"), catalog_name=table.catalog, - table_properties=kwargs.get("table_properties"), + table_properties=self._get_table_properties(kwargs), ) normalized_key: t.List[exp.Expr] = [ exp.column(k.name, quoted=bool(k.quoted)) if isinstance(k, exp.Identifier) else k @@ -2861,18 +3120,42 @@ def _replace_by_key( temp_table_props: t.Optional[t.Dict[str, t.Any]] = None if not is_iceberg: - if key_cols: + col_types_by_lower = { + c.lower(): dtype for c, dtype in (target_columns_to_types or {}).items() + } + norm_key_names_lower = { + k.name.lower() for k in normalized_key if isinstance(k, exp.Column) + } + has_invalid_or_coercible_key_col = any( + k_low in col_types_by_lower + and ( + not self._is_valid_olap_key_type(col_types_by_lower[k_low]) + or self._is_coercible_string_key_type(col_types_by_lower[k_low]) + ) + for k_low in norm_key_names_lower + ) + has_other_valid_olap_key_col = any( + c_low not in norm_key_names_lower + and self._is_valid_olap_key_type(dtype) + and not self._is_coercible_string_key_type(dtype) + for c_low, dtype in col_types_by_lower.items() + ) + avoid_key_coercion = has_invalid_or_coercible_key_col and has_other_valid_olap_key_col + if key_cols and not avoid_key_coercion: temp_table_props = { "unique_key": ( key_cols[0] if len(key_cols) == 1 else exp.Tuple(expressions=key_cols) ) } if target_columns_to_types and columns_to_types_all_known(target_columns_to_types): - reorder_kwargs: t.Dict[str, t.Any] = ( - {"table_properties": temp_table_props} - if temp_table_props - else ({"partitioned_by": list(normalized_key)} if normalized_key else {}) - ) + if avoid_key_coercion: + reorder_kwargs: t.Dict[str, t.Any] = {} + elif temp_table_props: + reorder_kwargs = {"table_properties": temp_table_props} + elif normalized_key: + reorder_kwargs = {"partitioned_by": list(normalized_key)} + else: + reorder_kwargs = {} target_columns_to_types = self._reorder_olap_columns( target_columns_to_types, reorder_kwargs ) @@ -2955,7 +3238,7 @@ def _scd_type_2( is_iceberg = self._is_iceberg_table( table_format=kwargs.get("table_format"), catalog_name=table.catalog, - table_properties=kwargs.get("table_properties"), + table_properties=self._get_table_properties(kwargs), ) if ( not is_iceberg @@ -3006,7 +3289,7 @@ def _create_table_from_source_queries( is_iceberg = self._is_iceberg_table( table_format=kwargs.get("table_format"), catalog_name=table.catalog, - table_properties=kwargs.get("table_properties"), + table_properties=self._get_table_properties(kwargs), ) is_internal_iceberg = is_iceberg and ( not effective_catalog or effective_catalog.lower() == "internal" @@ -3091,6 +3374,28 @@ def _matches_target_table(t_node: exp.Table) -> bool: **kwargs, ) + @staticmethod + def _remap_positional_group_and_order( + node: exp.Expr, old_to_new_one_based: t.Dict[int, int] + ) -> None: + if all(k == v for k, v in old_to_new_one_based.items()): + return + group = node.args.get("group") + if isinstance(group, exp.Group): + for item in group.expressions: + if isinstance(item, exp.Literal) and item.is_int: + old_pos = int(item.this) + if old_pos in old_to_new_one_based: + item.set("this", str(old_to_new_one_based[old_pos])) + order = node.args.get("order") + if isinstance(order, exp.Order): + for item in order.expressions: + ord_inner = item.this if isinstance(item, exp.Ordered) else item + if isinstance(ord_inner, exp.Literal) and ord_inner.is_int: + old_pos = int(ord_inner.this) + if old_pos in old_to_new_one_based: + ord_inner.set("this", str(old_to_new_one_based[old_pos])) + def _build_create_table_exp( self, table_name_or_schema: t.Union[exp.Schema, TableName], @@ -3102,7 +3407,10 @@ def _build_create_table_exp( table_kind: t.Optional[str] = None, **kwargs: t.Any, ) -> exp.Create: - kwargs.setdefault("table_properties", None) + if kwargs.get("table_properties") is None and kwargs.get("physical_properties") is not None: + kwargs["table_properties"] = kwargs["physical_properties"] + else: + kwargs.setdefault("table_properties", None) kwargs["is_schema_ddl"] = isinstance(table_name_or_schema, exp.Schema) catalog_name = None @@ -3115,7 +3423,7 @@ def _build_create_table_exp( is_iceberg = self._is_iceberg_table( table_format=kwargs.get("table_format"), catalog_name=catalog_name, - table_properties=kwargs.get("table_properties"), + table_properties=self._get_table_properties(kwargs), ) if ( @@ -3176,9 +3484,18 @@ def _build_create_table_exp( if len(proj_by_name) == len(expression.expressions) == len( target_columns_to_types ) and all(c.lower() in proj_by_name for c in target_columns_to_types): + old_pos_by_lower = { + proj.alias_or_name.lower(): idx + 1 + for idx, proj in enumerate(expression.expressions) + if proj.alias_or_name + } expression = expression.copy() updated_projs: t.List[exp.Expr] = [] - for c, new_dtype in target_columns_to_types.items(): + old_to_new_one_based: t.Dict[int, int] = {} + for new_idx, (c, new_dtype) in enumerate( + target_columns_to_types.items(), start=1 + ): + old_to_new_one_based[old_pos_by_lower[c.lower()]] = new_idx proj = proj_by_name[c.lower()].copy() if types_changed and original_cols[c] != new_dtype: if isinstance(proj, exp.Alias) and isinstance( @@ -3198,17 +3515,29 @@ def _build_create_table_exp( ) updated_projs.append(proj) expression.set("expressions", updated_projs) + self._remap_positional_group_and_order(expression, old_to_new_one_based) reordered_in_place = True if not reordered_in_place: expression = self._order_projections_and_filter( 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) + known_col_types = ( + { + c.lower(): dtype.copy() + for c, dtype in target_columns_to_types.items() + if dtype.this != exp.DataType.Type.UNKNOWN + } + if target_columns_to_types + else None + ) + expression = self._reorder_unannotated_ctas_query( + expression, kwargs, cte_types=known_col_types + ) if target_columns_to_types: first_select: t.Optional[exp.Select] = None curr_node: t.Optional[exp.Expr] = expression - while isinstance(curr_node, exp.SetOperation): + while isinstance(curr_node, (exp.SetOperation, exp.Subquery)): curr_node = curr_node.this if isinstance(curr_node, exp.Select): first_select = curr_node @@ -3226,7 +3555,7 @@ def _build_create_table_exp( # Doris AUTO PARTITION BY RANGE requires the partition column(s) to be defined as NOT NULL. if not is_iceberg: - table_props = kwargs.get("table_properties") or {} + table_props = self._get_table_properties(kwargs) or {} auto_partition = table_props.get("auto_partition") is_auto = bool( auto_partition is not None @@ -3282,7 +3611,9 @@ def _build_create_table_exp( auto_select_branches: t.List[exp.Select] = [] def _collect_auto_branches(node: exp.Expr) -> None: - if isinstance(node, exp.SetOperation): + if isinstance(node, exp.Subquery): + _collect_auto_branches(node.this) + elif isinstance(node, exp.SetOperation): _collect_auto_branches(node.this) _collect_auto_branches(node.expression) elif isinstance(node, exp.Select): @@ -3382,6 +3713,8 @@ def _infer_projection_type( ) -> exp.DataType: inner = proj.this if isinstance(proj, exp.Alias) else proj inner = inner.unnest() + while isinstance(inner, exp.Window) and isinstance(inner.this, exp.Expr): + inner = inner.this.unnest() if isinstance(inner, exp.Column) and cte_types and inner.name.lower() in cte_types: return cte_types[inner.name.lower()].copy() if isinstance(inner, (exp.Cast, exp.TryCast)) and isinstance(inner.to, exp.DataType): @@ -3419,7 +3752,7 @@ def _infer_projection_type( 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)): + if isinstance(inner, (exp.Count, exp.RowNumber, exp.Rank, exp.DenseRank, exp.Ntile)): return exp.DataType.build("BIGINT", dialect=self.dialect) if isinstance( inner, @@ -3462,6 +3795,8 @@ def _infer_projection_type( return exp.DataType.build("JSON", dialect=self.dialect) if fn_upper == "GROUP_CONCAT": return exp.DataType.build("TEXT", dialect=self.dialect) + if fn_upper in ("COUNT", "ROW_NUMBER", "RANK", "DENSE_RANK", "NTILE"): + return exp.DataType.build("BIGINT", dialect=self.dialect) if fn_upper == "DATE_TRUNC": return exp.DataType.build("DATETIME", dialect=self.dialect) if fn_upper == "TO_DATE": @@ -3489,7 +3824,9 @@ def _reorder_unannotated_ctas_query( branches: t.List[exp.Select] = [] def _collect_branches(node: exp.Expr) -> None: - if isinstance(node, exp.SetOperation): + if isinstance(node, exp.Subquery): + _collect_branches(node.this) + elif isinstance(node, exp.SetOperation): _collect_branches(node.this) _collect_branches(node.expression) elif isinstance(node, exp.Select): @@ -3555,6 +3892,9 @@ def _collect_branches(node: exp.Expr) -> None: if reordered_select_cols[c] != merged_cols_to_types[c] } if perm != list(range(num_cols)) or coercions: + old_to_new_one_based = { + orig_idx + 1: new_idx + 1 for new_idx, orig_idx in enumerate(perm) + } for b_idx, branch in enumerate(branches): old_projs = list(branch.expressions) updated_branch_projs: t.List[exp.Expr] = [] @@ -3591,6 +3931,8 @@ def _collect_branches(node: exp.Expr) -> None: ) updated_branch_projs.append(proj) branch.set("expressions", updated_branch_projs) + self._remap_positional_group_and_order(branch, old_to_new_one_based) + self._remap_positional_group_and_order(expression, old_to_new_one_based) return expression expression.set( "this", @@ -3619,9 +3961,18 @@ def _collect_branches(node: exp.Expr) -> None: for c in reordered_select_cols ) if order_changed or types_changed: + old_pos_by_lower = { + proj.alias_or_name.lower(): idx + 1 + for idx, proj in enumerate(expression.expressions) + if proj.alias_or_name + } expression = expression.copy() updated_projs: t.List[exp.Expr] = [] - for c, new_dtype in reordered_select_cols.items(): + old_to_new_one_based = {} + for new_idx, (c, new_dtype) in enumerate( + reordered_select_cols.items(), start=1 + ): + old_to_new_one_based[old_pos_by_lower[c.lower()]] = new_idx 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( @@ -3644,8 +3995,9 @@ def _collect_branches(node: exp.Expr) -> None: ) updated_projs.append(proj) expression.set("expressions", updated_projs) + self._remap_positional_group_and_order(expression, old_to_new_one_based) elif any(isinstance(proj, exp.Star) for proj in expression.expressions): - table_props = kwargs.get("table_properties") or {} + table_props = self._get_table_properties(kwargs) or {} resolved_uk, resolved_dk, resolved_ak = self._resolve_table_type_and_keys( table_props ) @@ -3717,7 +4069,7 @@ def _reorder_olap_columns( kwargs: t.Optional[t.Dict[str, t.Any]] = None, ) -> t.Dict[str, exp.DataType]: kwargs = kwargs or {} - table_props = kwargs.get("table_properties") or {} + table_props = self._get_table_properties(kwargs) 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." diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index c172d46f3e..e45fe1f0d2 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -1193,8 +1193,10 @@ def test_sync_materialized_view_drop_and_replace( ) 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`", + "USE `db`", + "DROP MATERIALIZED VIEW IF EXISTS `sync_mv` ON `db`.`orders`", + "USE `db`", + "CREATE MATERIALIZED VIEW `sync_mv` AS SELECT `region` AS `mv_region`, SUM(`amount`) AS `total` FROM `db`.`orders` GROUP BY `region`", ] adapter.cursor.execute.reset_mock() @@ -1277,16 +1279,24 @@ def test_rewrite_between_in_where_preserves_not_between_and_wraps_or( ): adapter = make_mocked_engine_adapter(DorisEngineAdapter) + raw_where = parse_one( + "(ds BETWEEN '2024-01-01' AND '2024-01-07' OR status = 'expired') AND id NOT BETWEEN 10 AND 20", + read="doris", + ) + rewritten_where = adapter._rewrite_between_in_where(raw_where.copy()) + assert ( + adapter._to_sql(rewritten_where) + == "((`ds` >= '2024-01-01' AND `ds` <= '2024-01-07') OR `status` = 'expired') AND NOT `id` BETWEEN 10 AND 20" + ) + 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", - ), + where=raw_where, ) 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", + "DELETE FROM `target_table` WHERE `ds` >= '2024-01-01' AND `ds` <= '2024-01-07' AND NOT `id` BETWEEN 10 AND 20", + "DELETE FROM `target_table` WHERE `status` = 'expired' AND NOT `id` BETWEEN 10 AND 20", ] @@ -1795,9 +1805,11 @@ def test_resolve_table_refs_in_physical_properties_for_mv_and_sync_mv( 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}`" + assert sync_calls[0] == "USE `doris_db`" + assert sync_calls[1] == ( + f"DROP MATERIALIZED VIEW IF EXISTS `orders_sync_mv` ON `{physical_tbl.db}`.`{physical_tbl.name}`" ) + assert sync_calls[-1].startswith("CREATE MATERIALIZED VIEW `orders_sync_mv` AS ") def test_dataframe_seed_all_string_coercion_and_batch_size( @@ -3639,8 +3651,10 @@ def test_unannotated_unknown_columns_ctas_and_sync_mv_and_schema_tmp_truncation( adapter.drop_view("db.sync_mv_typed", materialized=True) assert to_sql_calls(adapter) == [ - "CREATE MATERIALIZED VIEW `db`.`sync_mv_typed` AS SELECT `id` AS `mv_id`, SUM(`amount`) AS `total` FROM `db`.`orders` GROUP BY `id`", - "DROP MATERIALIZED VIEW IF EXISTS `db`.`sync_mv_typed` ON `db`.`orders`", + "USE `db`", + "CREATE MATERIALIZED VIEW `sync_mv_typed` AS SELECT `id` AS `mv_id`, SUM(`amount`) AS `total` FROM `db`.`orders` GROUP BY `id`", + "USE `db`", + "DROP MATERIALIZED VIEW IF EXISTS `sync_mv_typed` ON `db`.`orders`", ] # 4. Overlong `_schema_tmp` table identifier (> 64 chars) is deterministically shortened to <= 64 chars @@ -4070,8 +4084,10 @@ def test_doris_sync_mv_source_tables_persistence_and_isolation( view_properties={"materialized_type": exp.Literal.string("SYNC")}, ) assert to_sql_calls(adapter) == [ - "DROP MATERIALIZED VIEW IF EXISTS `db1`.`shared_mv` ON `db1`.`orders`", - "DROP MATERIALIZED VIEW IF EXISTS `db2`.`shared_mv` ON `db2`.`events`", + "USE `db1`", + "DROP MATERIALIZED VIEW IF EXISTS `shared_mv` ON `db1`.`orders`", + "USE `db2`", + "DROP MATERIALIZED VIEW IF EXISTS `shared_mv` ON `db2`.`events`", ] # 2. Recreating a view name as an async MTMV pops any stale sync MV entry so subsequent @@ -4182,3 +4198,322 @@ class OperationalError(Exception): "DELETE FROM `iceberg_cat`.`db`.`target` WHERE `id` IN (1, 2) AND `status` = 'a'", "DELETE FROM `iceberg_cat`.`db`.`target` WHERE `id` IN (1, 2) AND `status` IS NULL", ] + + +def test_doris_dialect_native_types_and_complex_constructors_and_row_diff( + make_mocked_engine_adapter: t.Callable, +): + from sqlmesh.core.dialect import parse as d_parse + from sqlmesh.core.model import load_sql_based_model + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Doris-native data types parse cleanly in `MODEL (columns (...))` and round-trip in Doris SQL: + # `LARGEINT`, `IPV4`, `IPV6`, `JSONB`, `VARIANT`, `BITMAP`, `HLL`, `QUANTILE_STATE`, `AGG_STATE`, + # `DATEV2`, `DATETIMEV2(3)`, `DECIMALV3(18, 4)`, `ARRAY`, `MAP`, and `STRUCT`. + model = load_sql_based_model( + d_parse( + """ + MODEL ( + name db.doris_rich_types, + kind FULL, + dialect doris, + columns ( + id LARGEINT, + ip4 IPV4, + ip6 IPV6, + jb JSONB, + var_col VARIANT, + bm BITMAP, + hll_col HLL, + qs QUANTILE_STATE, + agg_s AGG_STATE, + dt2 DATEV2, + ts2 DATETIMEV2(3), + dec3 DECIMALV3(18, 4), + tags ARRAY, + attrs MAP, + profile STRUCT + ) + ); + SELECT + CAST(1 AS LARGEINT) AS id, + CAST('127.0.0.1' AS IPV4) AS ip4, + CAST('::1' AS IPV6) AS ip6, + CAST('{}' AS JSONB) AS jb, + CAST('{}' AS VARIANT) AS var_col, + BITMAP_EMPTY() AS bm, + HLL_EMPTY() AS hll_col, + TO_QUANTILE_STATE(1, 2048) AS qs, + max_by_state(1, 2) AS agg_s, + CAST('2024-01-01' AS DATEV2) AS dt2, + CAST('2024-01-01 00:00:00' AS DATETIMEV2(3)) AS ts2, + CAST(12.34 AS DECIMALV3(18, 4)) AS dec3, + ARRAY('a', 'b') AS tags, + MAP('k1', 1, 'k2', 2) AS attrs, + STRUCT(100, 'alice') AS profile + """, + default_dialect="doris", + ) + ) + cols = model.columns_to_types + assert cols is not None + adapter.create_table("db.doris_rich_types", cols) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`doris_rich_types` (" + "`id` LARGEINT, `ip4` IPV4, `ip6` IPV6, `jb` JSONB, `var_col` VARIANT, " + "`bm` BITMAP, `hll_col` HLL, `qs` QUANTILE_STATE, `agg_s` AGG_STATE, " + "`dt2` DATE, `ts2` DATETIME(3), `dec3` DECIMAL(18, 4), `tags` ARRAY, " + "`attrs` MAP, `profile` STRUCT<`uid`: BIGINT, `name`: VARCHAR(64)>)" + ] + + # 2. Complex literal constructors (`ARRAY`, `MAP`, `STRUCT`, `NAMED_STRUCT`), `LEFT SEMI JOIN`, + # `LEFT ANTI JOIN` on unaliased tables, and `DISTINCT ON` transpile cleanly in Doris dialect. + query_sql = parse_one( + "SELECT [1, 2], MAP('a', 1, 'b', 2), STRUCT(1, 'x'), NAMED_STRUCT('id', 1, 'name', 'x') " + "FROM db.t1 LEFT ANTI JOIN db.t2 ON db.t1.id = db.t2.id " + "LEFT SEMI JOIN db.t3 ON db.t1.id = db.t3.id", + read="doris", + ).sql("doris") + assert ( + query_sql + == "SELECT ARRAY(1, 2), MAP('a', 1, 'b', 2), STRUCT(1, 'x'), NAMED_STRUCT('id', 1, 'name', 'x') " + "FROM db.t1 LEFT ANTI JOIN db.t2 ON db.t1.id = db.t2.id " + "LEFT SEMI JOIN db.t3 ON db.t1.id = db.t3.id" + ) + distinct_on_sql = parse_one( + "SELECT DISTINCT ON (id) id, ts FROM db.events ORDER BY id, ts DESC", + read="doris", + ).sql("doris") + assert "ROW_NUMBER() OVER (PARTITION BY id ORDER BY id, ts DESC)" in distinct_on_sql + + # 3. `RowDiffMixin.concat_columns` (`_normalize_value`) casts complex & special Doris types + # to `VARCHAR` instead of `DataType.Type.STRUCT` / `ARRAY` / `USERDEFINED`. + concat_expr = adapter.concat_columns( + { + "id": cols["id"], + "tags": cols["tags"], + "attrs": cols["attrs"], + "profile": cols["profile"], + "bm": cols["bm"], + "var_col": cols["var_col"], + } + ) + rendered_concat = concat_expr.sql("doris") + assert "AS STRUCT" not in rendered_concat + assert "AS ARRAY" not in rendered_concat + assert "AS MAP" not in rendered_concat + assert "AS BITMAP" not in rendered_concat + assert "CAST(tags AS VARCHAR)" in rendered_concat + assert "CAST(`profile` AS VARCHAR)" in rendered_concat + assert ( + adapter._normalize_nested_value(exp.column("tags", "s")).sql("doris") + == "CAST(s.tags AS VARCHAR)" + ) + + +def test_doris_ctas_window_partially_annotated_positional_group_order_and_subquery_union( + make_mocked_engine_adapter: t.Callable, +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Unannotated CTAS with `ROW_NUMBER() OVER (...)` at col 1 and `AVG(score)` at col 0 + # recognizes `exp.Window` (`ROW_NUMBER` -> `BIGINT`) and promotes `rn` to col 0. + adapter.ctas( + "db.win_ctas", + parse_one( + "SELECT AVG(score) AS avg_score, ROW_NUMBER() OVER (ORDER BY id) AS rn FROM db.src GROUP BY id", + read="doris", + ), + ) + + # 2. Partially annotated CTAS (`score: DOUBLE`, `id: INT`, `extra: UNKNOWN`) uses known column + # types to promote `id` before `score` (`DOUBLE`) and remaps positional `GROUP BY 2 ORDER BY 2 DESC` -> `1`. + adapter.ctas( + "db.partial_annotated_ctas", + parse_one( + "SELECT AVG(score) AS score, id, extra FROM db.src GROUP BY 2, 3 ORDER BY 2 DESC", + read="doris", + ), + target_columns_to_types={ + "score": exp.DataType.build("DOUBLE", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + "extra": exp.DataType.build("UNKNOWN"), + }, + ) + + # 3. Parenthesized `UNION ALL` branches (`(SELECT ...) UNION ALL (SELECT ...)`) with `auto_partition = true` + # unwrap `exp.Subquery` nodes, reorder projections in lockstep, and wrap range partition column in `NON_NULLABLE`. + adapter.ctas( + "db.paren_union_auto_ctas", + parse_one( + "(SELECT AVG(score) AS score, CAST(id AS INT) AS id, CAST(ds AS DATE) AS ds FROM db.s1 GROUP BY 2, 3) " + "UNION ALL " + "(SELECT CAST(score AS DOUBLE) AS score, CAST(id AS INT) AS id, CAST(ds AS DATE) AS 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`.`win_ctas` AS SELECT ROW_NUMBER() OVER (ORDER BY `id`) AS `rn`, AVG(`score`) AS `avg_score` FROM `db`.`src` GROUP BY `id`", + "CREATE TABLE IF NOT EXISTS `db`.`partial_annotated_ctas` AS SELECT `id`, AVG(`score`) AS `score`, `extra` FROM `db`.`src` GROUP BY 1, 3 ORDER BY 1 DESC", + "CREATE TABLE IF NOT EXISTS `db`.`paren_union_auto_ctas` AUTO PARTITION BY RANGE (`ds`) () AS (SELECT CAST(`id` AS INT) AS `id`, NON_NULLABLE(CAST(`ds` AS DATE)) AS `ds`, AVG(`score`) AS `score` FROM `db`.`s1` GROUP BY 1, 2) UNION ALL (SELECT CAST(`id` AS INT) AS `id`, NON_NULLABLE(CAST(`ds` AS DATE)) AS `ds`, CAST(`score` AS DOUBLE) AS `score` FROM `db`.`s2`)", + ] + + +def test_doris_replace_by_key_text_preservation_intersect_delete_and_direct_time_delete( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch("sqlmesh.core.engine_adapter.doris.random_id", return_value="abcd1234") + + # 1. `merge` / `_replace_by_key` when `unique_key` is a `TEXT`/`STRING` column (`doc_id`) and + # another valid OLAP key column (`tenant_id INT`) exists: does NOT coerce `doc_id` to `VARCHAR(255)` + # on the temp table (preventing 255-byte truncation of long keys), while promoting `tenant_id` to col 0. + adapter.merge( + target_table="db.docs_tbl", + source_table=parse_one("SELECT doc_id, tenant_id FROM db.staging_docs", read="doris"), + target_columns_to_types={ + "doc_id": exp.DataType.build("TEXT", dialect="doris"), + "tenant_id": exp.DataType.build("INT", dialect="doris"), + }, + unique_key=[exp.to_column("doc_id")], + ) + merge_sqls = to_sql_calls(adapter) + assert ( + merge_sqls[0] + == "CREATE TABLE `db`.`__temp_docs_tbl_abcd1234` AS SELECT CAST(`tenant_id` AS INT) AS `tenant_id`, CAST(`doc_id` AS STRING) AS `doc_id` FROM (SELECT `doc_id`, `tenant_id` FROM `db`.`staging_docs`) AS `_subquery`" + ) + + # 2. `_delete_by_time_transform_keys` with `INTERSECT` or `EXCEPT` does NOT unwrap `DATE_TRUNC(ds, 'day')` + # into raw `ds` inside the branches (preserving day-level set intersection/difference semantics), + # and unwraps parenthesized `(SELECT ...)` branches in `DELETE ... USING`. + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter.cursor, + "fetchall", + return_value=[("2024-01-15", "2024-01-16")], + ) + adapter.delete_from( + "db.events", + where=parse_one( + "DATE_TRUNC(ds, 'day') IN ((SELECT DATE_TRUNC(ds, 'day') FROM db.s1) INTERSECT (SELECT DATE_TRUNC(ds, 'day') FROM db.s2))", + read="doris", + ), + ) + adapter.delete_from( + "db.events", + where=parse_one( + "id IN ((SELECT id FROM db.s1) UNION ALL (SELECT id FROM db.s2))", + read="doris", + ), + ) + assert to_sql_calls(adapter) == [ + "SELECT DISTINCT DATE_TRUNC(`_col_0`, 'day') AS `_low_0`, DATE_ADD(DATE_TRUNC(`_col_0`, 'day'), INTERVAL 1 DAY) AS `_high_0` FROM ((SELECT DATE_TRUNC(`ds`, 'DAY') AS `_col_0` FROM `db`.`s1`) INTERSECT (SELECT DATE_TRUNC(`ds`, 'DAY') FROM `db`.`s2`)) AS `_sq`", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-15' AND `ds` < '2024-01-16'", + "DELETE FROM `db`.`events` AS `_t1` USING ((SELECT `id` FROM `db`.`s1`) UNION ALL (SELECT `id` FROM `db`.`s2`)) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id`", + ] + + # 3. Direct time-transform predicates (`DATE_TRUNC(ds, 'day') = '2024-01-15'`, `YEAR(ds) = 2024`, + # `DATE_TRUNC(ds, 'day') IN ('2024-01-15', '2024-01-16')`) and heterogeneous `OR` predicates + # (`a = 1 OR b = 2`) in `delete_from` are rewritten into `OR`-free bare-column predicates. + adapter.cursor.execute.reset_mock() + adapter.delete_from("db.events", "DATE_TRUNC(ds, 'day') = '2024-01-15'") + adapter.delete_from("db.events", "YEAR(ds) = 2024") + adapter.delete_from("db.events", "DATE_TRUNC(ds, 'day') IN ('2024-01-15', '2024-01-16')") + adapter.delete_from("db.events", "a = 1 OR b = 2") + assert to_sql_calls(adapter) == [ + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-15' AND `ds` < '2024-01-16'", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-01' AND `ds` < '2025-01-01'", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-15' AND `ds` < '2024-01-16'", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-16' AND `ds` < '2024-01-17'", + "DELETE FROM `db`.`events` WHERE `a` = 1", + "DELETE FROM `db`.`events` WHERE `b` = 2", + ] + + +def test_doris_physical_properties_kwarg_drop_schema_cascade_and_mv_df_source_columns( + make_mocked_engine_adapter: t.Callable, mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. `physical_properties={'engine': exp.Literal.string('iceberg')}` (as passed by SnapshotEvaluator) + # is honored across `replace_query` / `_insert_overwrite_by_condition`. + mocker.patch.object( + adapter, + "get_data_object", + return_value=DataObject( + catalog="internal", schema="db", name="iceberg_by_prop", type=DataObjectType.TABLE + ), + ) + adapter.replace_query( + "db.iceberg_by_prop", + parse_one("SELECT 1 AS id", read="doris"), + target_columns_to_types={"id": exp.DataType.build("INT", dialect="doris")}, + physical_properties={"engine": exp.Literal.string("iceberg")}, + ) + assert to_sql_calls(adapter) == [ + "INSERT OVERWRITE TABLE `db`.`iceberg_by_prop` (`id`) SELECT 1 AS `id`" + ] + + # 2. `drop_schema(..., cascade=True)` on an external Iceberg catalog drops member tables first + # before issuing `DROP DATABASE`. + adapter.cursor.execute.reset_mock() + mocker.patch.object( + adapter, + "_get_data_objects", + return_value=[ + DataObject( + catalog="iceberg_cat", schema="silver", name="t1", type=DataObjectType.TABLE + ), + DataObject( + catalog="iceberg_cat", schema="silver", name="t2", type=DataObjectType.TABLE + ), + ], + ) + adapter.drop_schema("iceberg_cat.silver", ignore_if_not_exists=True, cascade=True) + assert to_sql_calls(adapter) == [ + "DROP TABLE IF EXISTS `iceberg_cat`.`silver`.`t1`", + "DROP TABLE IF EXISTS `iceberg_cat`.`silver`.`t2`", + "DROP DATABASE IF EXISTS `iceberg_cat`.`silver`", + ] + + # 3. `_create_materialized_view` with a DataFrame containing extra `source_columns` not in + # `target_columns_to_types` slices only `source_columns_to_types` so `_values_to_sql` never raises `ValueError`. + adapter.cursor.execute.reset_mock() + df_extra = pd.DataFrame([{"id": 1, "name": "alice", "extra_ignored": 999}]) + adapter.create_view( + "db.mv_from_df_subset", + df_extra, + replace=False, + materialized=True, + target_columns_to_types={ + "id": exp.DataType.build("INT", dialect="doris"), + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + source_columns=["id", "name", "extra_ignored"], + ) + assert to_sql_calls(adapter) == [ + "CREATE MATERIALIZED VIEW `db`.`mv_from_df_subset` (`id`, `name`) AS SELECT CAST(`id` AS INT) AS `id`, CAST(`name` AS VARCHAR(32)) AS `name` FROM (SELECT 1 AS `id`, 'alice' AS `name`) AS `t`" + ] + + # 4. `_truncate_table` normalizes >64-char `_schema_tmp` table names before `TRUNCATE TABLE` and fallback `DELETE`. + adapter.cursor.execute.reset_mock() + long_tmp = "db." + ("very_long_table_name_" * 3) + "_schema_tmp" + short_tmp = adapter._shorten_schema_tmp_name(exp.to_table(long_tmp).name) + adapter._truncate_table(long_tmp) + assert to_sql_calls(adapter) == [f"TRUNCATE TABLE `db`.`{short_tmp}`"] + + # 5. `_fetch_native_df` with a raw string query containing `IS TRUE` or >64-char `_schema_tmp` + # normalizes the SQL string before `pd.read_sql_query`. + mock_read_sql = mocker.patch( + "pandas.io.sql.read_sql_query", + return_value=pd.DataFrame({"id": [1]}), + ) + adapter.fetchdf(f"SELECT id FROM {long_tmp} WHERE is_active IS TRUE", quote_identifiers=True) + assert ( + mock_read_sql.call_args[0][0] + == f"SELECT `id` FROM `db`.`{short_tmp}` WHERE `is_active` = TRUE" + ) From 93eee586064e936fc39a92e2376a4f2bfe381d10 Mon Sep 17 00:00:00 2001 From: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> Date: Sun, 25 Oct 2026 22:03:41 -0700 Subject: [PATCH 31/31] fix(doris): pass full live Docker integration suite and harden DDL/DML/row-diff edge cases Signed-off-by: Bruce Arctor <5032356+brucearctor@users.noreply.github.com> --- docs/integrations/engines/doris.md | 2 +- sqlmesh/core/dialect.py | 4 +- sqlmesh/core/engine_adapter/doris.py | 491 ++++++++++++++---- .../integration/docker/compose.doris.yaml | 3 +- .../integration/test_integration_doris.py | 79 +++ tests/core/engine_adapter/test_doris.py | 74 ++- 6 files changed, 537 insertions(+), 116 deletions(-) diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md index 26dc5d376e..83778c05d0 100644 --- a/docs/integrations/engines/doris.md +++ b/docs/integrations/engines/doris.md @@ -181,7 +181,7 @@ 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` (emitting `ORDER BY (col1, col2)` 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` or `cluster_by` (emitting `CLUSTER BY (col1, col2)` on Doris 2.1+ Merge-on-Write tables). ### AGGREGATE KEY model diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py index 078e4b8636..b5cfac2a86 100644 --- a/sqlmesh/core/dialect.py +++ b/sqlmesh/core/dialect.py @@ -11,7 +11,7 @@ from functools import lru_cache from sqlglot import Dialect, Generator, ParseError, Parser, Tokenizer, TokenType, exp, transforms -from sqlglot.dialects.dialect import DialectType, var_map_sql +from sqlglot.dialects.dialect import DialectType, rename_func, var_map_sql from sqlglot.dialects import Doris, DuckDB, Snowflake, TSQL import sqlglot.dialects.athena as athena import sqlglot.generators.athena as athena_generators @@ -1471,7 +1471,9 @@ class _DorisTokenizer(Doris.tokenizer_class): # type: ignore ) Doris.Generator.TRANSFORMS.update( { + exp.Array: rename_func("ARRAY"), exp.Map: _doris_map_sql, + exp.ParseJSON: lambda self, e: f"CAST({self.sql(e, 'this')} AS JSON)", exp.Select: transforms.preprocess( [ transforms.eliminate_distinct_on, diff --git a/sqlmesh/core/engine_adapter/doris.py b/sqlmesh/core/engine_adapter/doris.py index 4bf35dceab..ae446e70ff 100644 --- a/sqlmesh/core/engine_adapter/doris.py +++ b/sqlmesh/core/engine_adapter/doris.py @@ -273,8 +273,17 @@ def get_current_catalog(self) -> t.Optional[str]: 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" + logger.debug( + "current_catalog() failed on Doris, falling back to active/default catalog" + ) + return ( + self._active_catalog + or getattr(self._connection_pool, "_doris_last_switched_catalog", None) + or self._extra_config.get("_doris_last_switched_catalog") + or getattr(self, "_doris_last_switched_catalog", None) + or self.default_catalog + or "internal" + ) def set_current_catalog(self, catalog: str) -> None: catalog_sql = self._to_sql(exp.to_identifier(catalog, quoted=True)) @@ -431,6 +440,40 @@ def _is_valid_olap_key_type(cls, dtype: exp.DataType) -> bool: def _normalize_nested_value(self, expr: exp.Expr) -> exp.Expr: return exp.cast(expr, exp.DataType.build("VARCHAR", dialect=self.dialect)) + def normalize_value( + self, + expr: exp.Expr, + type: exp.DataType, + decimal_precision: int = 3, + timestamp_precision: int = 6, + ) -> exp.Expr: + udt_kind = ( + str(type.args.get("kind") or "").upper().split("(", 1)[0].strip().strip("`") + if type.this == exp.DataType.Type.USERDEFINED + else "" + ) + if type.is_type("int128") or udt_kind == "LARGEINT": + return exp.cast( + exp.cast(expr, exp.DataType.build("LARGEINT", udt=True)), + to=exp.DataType.build("VARCHAR"), + ) + if udt_kind == "BITMAP": + return exp.cast( + exp.func("BITMAP_TO_STRING", expr), + to=exp.DataType.build("VARCHAR"), + ) + if type.is_type("hllsketch") or udt_kind in ("HLL", "HLLSKETCH"): + return exp.cast( + exp.func("HLL_TO_BASE64", expr), + to=exp.DataType.build("VARCHAR"), + ) + return super().normalize_value( + expr, + type, + decimal_precision=decimal_precision, + timestamp_precision=timestamp_precision, + ) + def _fetch_native_df( self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False ) -> DF: @@ -710,12 +753,21 @@ def get_alter_operations( 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( + iceberg_ops = super().get_alter_operations( current_table_name, target_table_name, ignore_destructive=ignore_destructive, ignore_additive=ignore_additive, ) + return [ + op + for op in iceberg_ops + if not ( + isinstance(op, TableAlterChangeColumnTypeOperation) + and self._to_sql(op.current_type).upper() + == self._to_sql(op.column_type).upper() + ) + ] self._last_described_key_cols = set() current_cols = self.columns(current_table_name) @@ -951,16 +1003,6 @@ 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, @@ -1155,6 +1197,7 @@ def _create_materialized_view( # 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. + sync_db_switched: t.Optional[str] = None if is_sync_mv: sync_db: t.Optional[str] = target_table.db or ( sync_source_tbl.db if sync_source_tbl else None @@ -1162,22 +1205,37 @@ def _create_materialized_view( if sync_db: db_sql = self._to_sql(exp.to_identifier(sync_db, quoted=True)) self.execute(f"USE {db_sql}") + sync_db_switched = sync_db unqualified_mv = target_table.copy() unqualified_mv.set("db", None) unqualified_mv.set("catalog", None) schema = unqualified_mv - self.execute( - exp.Create( - 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, - ), - quote_identifiers=self.QUOTE_IDENTIFIERS_IN_VIEWS, - ) + try: + self.execute( + exp.Create( + 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, + ), + quote_identifiers=self.QUOTE_IDENTIFIERS_IN_VIEWS, + ) + finally: + configured_db = self._extra_config.get("database") + if ( + sync_db_switched + and configured_db + and str(configured_db).lower() != sync_db_switched.lower() + ): + restore_db_sql = self._to_sql( + exp.to_identifier(str(configured_db), quoted=True) + ) + self.execute(f"USE {restore_db_sql}") registry = self._sync_mv_source_tables if is_sync_mv and sync_source_tbl is not None: if target_table.db: @@ -1191,10 +1249,7 @@ def _create_materialized_view( sync_source_tbl.copy() ) elif not is_sync_mv: - if target_table.db: - registry.pop(f"{target_table.db.lower()}.{target_table.name.lower()}", None) - else: - registry.pop(target_table.name.lower(), None) + self._pop_sync_mv_source_table(target_table) build_mode = ( self._to_upper_str(view_properties["build"]) if view_properties and view_properties.get("build") is not None @@ -1208,6 +1263,21 @@ def _create_materialized_view( self.execute(f"REFRESH MATERIALIZED VIEW {self._to_sql(target_table)} AUTO") self._clear_data_object_cache(view_name) + def _pop_sync_mv_source_table(self, target_table: exp.Table) -> t.Optional[exp.Table]: + registry = self._sync_mv_source_tables + name_lower = target_table.name.lower() + if target_table.db: + db_lower = target_table.db.lower() + popped = registry.pop(f"{db_lower}.{name_lower}", None) + unqual = registry.get(name_lower) + if unqual is not None and (not unqual.db or unqual.db.lower() == db_lower): + registry.pop(name_lower, None) + return popped + popped = registry.pop(name_lower, None) + if popped is not None and popped.db: + registry.pop(f"{popped.db.lower()}.{name_lower}", None) + return popped + def _to_source_table(self, raw_src: t.Any) -> exp.Table: if isinstance(raw_src, exp.Table): tbl = raw_src.copy() @@ -1256,11 +1326,7 @@ def drop_view( sync_source_table = registry.get(target_table.name.lower()) if sync_source_table is not None: - registry = self._sync_mv_source_tables - if target_table.db: - registry.pop(f"{target_table.db.lower()}.{target_table.name.lower()}", None) - else: - registry.pop(target_table.name.lower(), None) + self._pop_sync_mv_source_table(target_table) source_tbl_exp = self._to_source_table(sync_source_table) # Synchronous materialized views require the target/source table's database # to be active on the session and an unqualified `` in `DROP MATERIALIZED VIEW ON `. @@ -1274,9 +1340,17 @@ def drop_view( view_table_sql = self._to_sql(unqualified_mv) 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}" - ) + try: + self.execute( + f"DROP MATERIALIZED VIEW {exists_clause}{view_table_sql} ON {source_table_sql}" + ) + finally: + configured_db = self._extra_config.get("database") + if sync_db and configured_db and str(configured_db).lower() != sync_db.lower(): + restore_db_sql = self._to_sql( + exp.to_identifier(str(configured_db), quoted=True) + ) + self.execute(f"USE {restore_db_sql}") self._clear_data_object_cache(view_name) return super().drop_view(view_name, ignore_if_not_exists, materialized, **kwargs) @@ -1929,38 +2003,96 @@ def _build_time_transform_range( ) def _rewrite_direct_time_transforms_in_where(self, where: exp.Expr) -> exp.Expr: - for node in list(where.find_all(exp.EQ, exp.In)): - if node.find_ancestor(exp.Subquery, exp.Select, exp.Not): + for node in list( + where.find_all( + exp.EQ, + exp.NEQ, + exp.GTE, + exp.GT, + exp.LTE, + exp.LT, + exp.In, + ) + ): + if node.find_ancestor(exp.Subquery, exp.Select): continue - if isinstance(node, exp.EQ): + parent_not = node.parent if isinstance(node.parent, exp.Not) else None + if node.find_ancestor(exp.Not) and parent_not is None: + continue + target_node: exp.Expr = parent_not if parent_not is not None else node + if isinstance(node, (exp.EQ, exp.NEQ, exp.GTE, exp.GT, exp.LTE, exp.LT)): left_trans = self._extract_time_partition_transform(node.this) right_trans = self._extract_time_partition_transform(node.expression) t_expr: t.Optional[exp.Expr] = None col_expr: t.Optional[exp.Column] = None unit: t.Optional[str] = None val_expr: t.Optional[exp.Expr] = None + transform_on_left = True if left_trans is not None and self._is_literal_key_value(node.expression): t_expr = node.this col_expr, unit = left_trans val_expr = node.expression + transform_on_left = True elif right_trans is not None and self._is_literal_key_value(node.this): t_expr = node.expression col_expr, unit = right_trans val_expr = node.this + transform_on_left = False if ( t_expr is not None and col_expr is not None and unit is not None and val_expr is not None ): - replacement = self._build_time_transform_range(t_expr, col_expr, unit, val_expr) - if replacement is not None: - if node is where: - where = replacement - else: - if not isinstance(node.parent, (exp.Where, exp.Paren)): - replacement = exp.Paren(this=replacement) - node.replace(replacement) + range_expr = self._build_time_transform_range(t_expr, col_expr, unit, val_expr) + if range_expr is not None and isinstance(range_expr, exp.And): + low_lit = range_expr.this.expression + high_lit = range_expr.expression.expression + effective_op: t.Type[exp.Binary] = type(node) + if not transform_on_left: + flip_map: t.Dict[t.Type[exp.Binary], t.Type[exp.Binary]] = { + exp.GTE: exp.LTE, + exp.GT: exp.LT, + exp.LTE: exp.GTE, + exp.LT: exp.GT, + } + effective_op = flip_map.get(effective_op, effective_op) + if parent_not is not None: + neg_map: t.Dict[t.Type[exp.Binary], t.Type[exp.Binary]] = { + exp.EQ: exp.NEQ, + exp.NEQ: exp.EQ, + exp.GTE: exp.LT, + exp.GT: exp.LTE, + exp.LTE: exp.GT, + exp.LT: exp.GTE, + } + effective_op = neg_map.get(effective_op, effective_op) + replacement: t.Optional[exp.Expr] = None + if effective_op is exp.EQ: + replacement = range_expr + elif effective_op is exp.NEQ: + replacement = exp.or_( + exp.LT(this=col_expr.copy(), expression=low_lit.copy()), + exp.GTE(this=col_expr.copy(), expression=high_lit.copy()), + copy=False, + ) + elif effective_op is exp.GTE: + replacement = exp.GTE(this=col_expr.copy(), expression=low_lit.copy()) + elif effective_op is exp.GT: + replacement = exp.GTE(this=col_expr.copy(), expression=high_lit.copy()) + elif effective_op is exp.LTE: + replacement = exp.LT(this=col_expr.copy(), expression=high_lit.copy()) + elif effective_op is exp.LT: + replacement = exp.LT(this=col_expr.copy(), expression=low_lit.copy()) + if replacement is not None: + if target_node is where: + where = replacement + else: + if isinstance(replacement, (exp.And, exp.Or)) and not isinstance( + target_node.parent, (exp.Where, exp.Paren) + ): + replacement = exp.Paren(this=replacement) + target_node.replace(replacement) elif ( isinstance(node, exp.In) and self._get_in_subquery(node) is None @@ -1975,17 +2107,44 @@ def _rewrite_direct_time_transforms_in_where(self, where: exp.Expr) -> exp.Expr: self._build_time_transform_range(node.this, col_expr, unit, v) for v in node.expressions ] - if ranges and all(r is not None for r in ranges): - or_expr: exp.Expr = exp.or_( - *[exp.Paren(this=r) for r in ranges if r is not None], - copy=False, - ) - if node is where: - where = or_expr + if ranges and all(isinstance(r, exp.And) for r in ranges): + if parent_not is not None: + not_in_parts: t.List[exp.Expr] = [] + for r in ranges: + assert isinstance(r, exp.And) + not_in_parts.append( + exp.Paren( + this=exp.or_( + exp.LT( + this=col_expr.copy(), + expression=r.this.expression.copy(), + ), + exp.GTE( + this=col_expr.copy(), + expression=r.expression.expression.copy(), + ), + copy=False, + ) + ) + ) + and_expr: exp.Expr = exp.and_(*not_in_parts, copy=False) + if target_node is where: + where = and_expr + else: + if not isinstance(target_node.parent, (exp.Where, exp.Paren)): + and_expr = exp.Paren(this=and_expr) + target_node.replace(and_expr) else: - if not isinstance(node.parent, (exp.Where, exp.Paren)): - or_expr = exp.Paren(this=or_expr) - node.replace(or_expr) + or_expr: exp.Expr = exp.or_( + *[exp.Paren(this=r) for r in ranges if r is not None], + copy=False, + ) + if target_node is where: + where = or_expr + else: + if not isinstance(target_node.parent, (exp.Where, exp.Paren)): + or_expr = exp.Paren(this=or_expr) + target_node.replace(or_expr) return where @classmethod @@ -2401,27 +2560,28 @@ def _collect_subq_branches(node: t.Optional[exp.Expr]) -> None: 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 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) - + extra_condition_sets = self._expand_or_free_extra_condition_sets(extra_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) + for cond_set in extra_condition_sets: + branch_joins = [jc.copy() for jc in join_conditions] + for extra in cond_set: + 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 or col.table.lower() == target_table.name.lower() + ): + col.set("table", exp.to_identifier("_t1")) + col.set("db", None) + col.set("catalog", None) + branch_joins.append(qualified_extra) + where_expr = exp.and_(*branch_joins) + where_sql = self._to_sql(where_expr) + delete_sql = f"DELETE FROM {target_sql} USING {subquery_sql} WHERE {where_sql}" + self.execute(delete_sql) except Exception as e: err_msg = str(e).lower() if any( @@ -3380,21 +3540,34 @@ def _remap_positional_group_and_order( ) -> None: if all(k == v for k, v in old_to_new_one_based.items()): return + + def _remap_lit(item: exp.Expr) -> None: + inner = item.this if isinstance(item, exp.Ordered) else item + if isinstance(inner, exp.Literal) and inner.is_int: + old_pos = int(inner.this) + if old_pos in old_to_new_one_based: + inner.set("this", str(old_to_new_one_based[old_pos])) + elif isinstance(inner, (exp.Tuple, exp.Paren)): + for sub in inner.expressions or ([inner.this] if inner.this else []): + _remap_lit(sub) + group = node.args.get("group") if isinstance(group, exp.Group): for item in group.expressions: - if isinstance(item, exp.Literal) and item.is_int: - old_pos = int(item.this) - if old_pos in old_to_new_one_based: - item.set("this", str(old_to_new_one_based[old_pos])) + _remap_lit(item) + for prop_key in ("rollup", "cube", "grouping_sets"): + prop_val = group.args.get(prop_key) + if isinstance(prop_val, list): + for g_item in prop_val: + if isinstance(g_item, exp.Expr): + _remap_lit(g_item) + elif isinstance(prop_val, exp.Expr): + for g_item in prop_val.expressions: + _remap_lit(g_item) order = node.args.get("order") if isinstance(order, exp.Order): for item in order.expressions: - ord_inner = item.this if isinstance(item, exp.Ordered) else item - if isinstance(ord_inner, exp.Literal) and ord_inner.is_int: - old_pos = int(ord_inner.this) - if old_pos in old_to_new_one_based: - ord_inner.set("this", str(old_to_new_one_based[old_pos])) + _remap_lit(item) def _build_create_table_exp( self, @@ -3730,6 +3903,58 @@ def _infer_projection_type( 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.Div): + return exp.DataType.build("DOUBLE", dialect=self.dialect) + if isinstance(inner, (exp.Add, exp.Sub, exp.Mul)): + left_t = self._infer_projection_type(inner.this, cte_types) + right_t = self._infer_projection_type(inner.expression, cte_types) + if any(t_cand.is_type("float", "double") for t_cand in (left_t, right_t)): + return exp.DataType.build("DOUBLE", dialect=self.dialect) + if isinstance(inner, exp.If): + branch_types = [ + self._infer_projection_type(b_expr, cte_types) + for b_expr in (inner.args.get("true"), inner.args.get("false")) + if isinstance(b_expr, exp.Expr) + ] + known_b = [bt for bt in branch_types if bt.this != exp.DataType.Type.UNKNOWN] + if known_b: + non_key_b = next( + (bt for bt in known_b if not self._is_valid_olap_key_type(bt)), None + ) + return non_key_b or known_b[0] + if isinstance(inner, exp.Case): + case_exprs: t.List[exp.Expr] = [ + i.args["true"] + for i in (inner.args.get("ifs") or []) + if isinstance(i.args.get("true"), exp.Expr) + ] + if isinstance(inner.args.get("default"), exp.Expr): + case_exprs.append(inner.args["default"]) + known_c = [ + ct + for ce in case_exprs + if (ct := self._infer_projection_type(ce, cte_types)).this + != exp.DataType.Type.UNKNOWN + ] + if known_c: + non_key_c = next( + (ct for ct in known_c if not self._is_valid_olap_key_type(ct)), None + ) + return non_key_c or known_c[0] + if isinstance(inner, exp.Coalesce): + coal_exprs = [inner.this] + list(inner.expressions or []) + known_coal = [ + ct + for ce in coal_exprs + if isinstance(ce, exp.Expr) + and (ct := self._infer_projection_type(ce, cte_types)).this + != exp.DataType.Type.UNKNOWN + ] + if known_coal: + non_key_coal = next( + (ct for ct in known_coal if not self._is_valid_olap_key_type(ct)), None + ) + return non_key_coal or known_coal[0] if isinstance( inner, ( @@ -3748,7 +3973,7 @@ def _infer_projection_type( 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)): + if isinstance(inner, (exp.JSONObject, exp.JSONArray, exp.ParseJSON)): return exp.DataType.build("JSON", dialect=self.dialect) if isinstance(inner, exp.GroupConcat): return exp.DataType.build("TEXT", dialect=self.dialect) @@ -3791,8 +4016,28 @@ def _infer_projection_type( 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"): + if fn_upper in ( + "JSON_OBJECT", + "JSON_ARRAY", + "TO_JSON", + "JSON_PARSE", + "JSONB_PARSE", + "JSON_EXTRACT", + ): return exp.DataType.build("JSON", dialect=self.dialect) + if fn_upper in ( + "BITMAP_UNION", + "BITMAP_EMPTY", + "TO_BITMAP", + "BITMAP_HASH", + "BITMAP_HASH64", + "BITMAP_FROM_STRING", + ): + return exp.DataType.build("BITMAP", dialect=self.dialect, udt=True) + if fn_upper in ("HLL_UNION", "HLL_EMPTY", "HLL_HASH", "HLL_FROM_BASE64"): + return exp.DataType.build("HLL", dialect=self.dialect, udt=True) + if fn_upper in ("TO_QUANTILE_STATE", "QUANTILE_UNION"): + return exp.DataType.build("QUANTILE_STATE", dialect=self.dialect, udt=True) if fn_upper == "GROUP_CONCAT": return exp.DataType.build("TEXT", dialect=self.dialect) if fn_upper in ("COUNT", "ROW_NUMBER", "RANK", "DENSE_RANK", "NTILE"): @@ -3803,6 +4048,20 @@ def _infer_projection_type( return exp.DataType.build("DATE", dialect=self.dialect) return exp.DataType.build("UNKNOWN") + @staticmethod + def _projection_key_name(proj: exp.Expr, idx: int) -> str: + if isinstance(proj, exp.Star): + return "" + if proj.alias_or_name: + return proj.alias_or_name + if ( + isinstance(proj, (exp.Cast, exp.TryCast)) + and isinstance(proj.this, exp.Column) + and proj.this.name + ): + return proj.this.name + return f"__unaliased_col_{idx}" + def _reorder_unannotated_ctas_query( self, expression: exp.Expr, @@ -3837,6 +4096,7 @@ def _collect_branches(node: exp.Expr) -> None: return expression first_exprs = branches[0].expressions num_cols = len(first_exprs) + orig_names = [self._projection_key_name(p, i) for i, p in enumerate(first_exprs)] if ( num_cols > 0 and all( @@ -3844,16 +4104,9 @@ def _collect_branches(node: exp.Expr) -> None: 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 + and all(orig_names) + and len({n.lower() for n in orig_names}) == 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): branch_cands = [ @@ -3903,6 +4156,7 @@ def _collect_branches(node: exp.Expr) -> None: new_dtype = coercions.get(orig_idx) if new_dtype is not None: col_name = orig_names[orig_idx] + is_synthetic = col_name.startswith("__unaliased_col_") if isinstance(proj, exp.Alias) and isinstance( proj.this, (exp.Cast, exp.TryCast) ): @@ -3911,7 +4165,7 @@ def _collect_branches(node: exp.Expr) -> None: cast_expr = exp.cast(proj.this.copy(), to=new_dtype.copy()) proj = ( exp.alias_(cast_expr, col_name, quoted=True) - if b_idx == 0 + if b_idx == 0 and not is_synthetic else cast_expr ) elif isinstance(proj, exp.Alias): @@ -3929,6 +4183,13 @@ def _collect_branches(node: exp.Expr) -> None: if b_idx == 0 else cast_expr ) + else: + cast_expr = exp.cast(proj.copy(), to=new_dtype.copy()) + proj = ( + exp.alias_(cast_expr, col_name, quoted=True) + if b_idx == 0 and not is_synthetic + else cast_expr + ) updated_branch_projs.append(proj) branch.set("expressions", updated_branch_projs) self._remap_positional_group_and_order(branch, old_to_new_one_based) @@ -3944,15 +4205,15 @@ def _collect_branches(node: exp.Expr) -> None: ) 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) - } + keyed_projs = [ + (self._projection_key_name(proj, idx), proj) + for idx, proj in enumerate(expression.expressions) + if not isinstance(proj, exp.Star) + ] + proj_by_name = {k.lower(): proj for k, proj in keyed_projs if k} if proj_by_name and len(proj_by_name) == len(expression.expressions): select_cols_to_types = { - proj.alias_or_name: self._infer_projection_type(proj, cte_col_types) - for proj in expression.expressions + k: self._infer_projection_type(proj, cte_col_types) for k, proj in keyed_projs } reordered_select_cols = self._reorder_olap_columns(select_cols_to_types, kwargs) order_changed = list(reordered_select_cols) != list(select_cols_to_types) @@ -3962,9 +4223,7 @@ def _collect_branches(node: exp.Expr) -> None: ) if order_changed or types_changed: old_pos_by_lower = { - proj.alias_or_name.lower(): idx + 1 - for idx, proj in enumerate(expression.expressions) - if proj.alias_or_name + k.lower(): idx + 1 for idx, (k, _) in enumerate(keyed_projs) } expression = expression.copy() updated_projs: t.List[exp.Expr] = [] @@ -3975,15 +4234,17 @@ def _collect_branches(node: exp.Expr) -> None: old_to_new_one_based[old_pos_by_lower[c.lower()]] = new_idx proj = proj_by_name[c.lower()].copy() if types_changed and select_cols_to_types[c] != new_dtype: + is_synthetic = c.startswith("__unaliased_col_") 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, + cast_expr = exp.cast(proj.this.copy(), to=new_dtype.copy()) + proj = ( + cast_expr + if is_synthetic + else exp.alias_(cast_expr, c, quoted=True) ) elif isinstance(proj, exp.Alias): proj.set("this", exp.cast(proj.this.copy(), to=new_dtype.copy())) @@ -3993,6 +4254,13 @@ def _collect_branches(node: exp.Expr) -> None: proj.name, quoted=bool(getattr(proj.this, "quoted", True)), ) + else: + cast_expr = exp.cast(proj.copy(), to=new_dtype.copy()) + proj = ( + cast_expr + if is_synthetic + else exp.alias_(cast_expr, c, quoted=True) + ) updated_projs.append(proj) expression.set("expressions", updated_projs) self._remap_positional_group_and_order(expression, old_to_new_one_based) @@ -4596,7 +4864,9 @@ def _build_table_properties_exp( table_properties_copy, pop=True ) unique_key = resolved_uk or kwargs.get("unique_key") - order_by = table_properties_copy.pop("order_by", None) + order_by = table_properties_copy.pop("order_by", None) or table_properties_copy.pop( + "cluster_by", None + ) unique_key_cols = self._extract_column_names(unique_key) duplicate_key_cols = self._extract_column_names(duplicate_key) @@ -4721,7 +4991,7 @@ def _build_table_properties_exp( 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})")) + exp.SetConfigProperty(this=exp.var(f"CLUSTER BY ({order_cols_sql})")) ) if table_description: @@ -5346,6 +5616,7 @@ def _properties_to_expressions(self, properties: t.Dict[str, t.Any]) -> t.List[e "table_type", "keys", "order_by", + "cluster_by", "engine", "properties", "materialized_type", diff --git a/tests/core/engine_adapter/integration/docker/compose.doris.yaml b/tests/core/engine_adapter/integration/docker/compose.doris.yaml index 50376667c5..2638848e9e 100644 --- a/tests/core/engine_adapter/integration/docker/compose.doris.yaml +++ b/tests/core/engine_adapter/integration/docker/compose.doris.yaml @@ -34,7 +34,8 @@ services: 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 "mem_limit = 80%" >> /opt/apache-doris/be/conf/be.conf + echo "disable_memory_gc = false" >> /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 diff --git a/tests/core/engine_adapter/integration/test_integration_doris.py b/tests/core/engine_adapter/integration/test_integration_doris.py index aa60c60563..32cfde900a 100644 --- a/tests/core/engine_adapter/integration/test_integration_doris.py +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -530,3 +530,82 @@ def test_insert_overwrite_by_time_partition_and_scd_type_2(ctx: TestContext): (1, "alice_v2", 1), (2, "bob_v1", 1), ] + + +def test_doris_complex_types_and_alter_table_schema_evolution(ctx: TestContext): + from sqlglot import parse_one + + # 1. Create table with LARGEINT, ARRAY, MAP, STRUCT, and BITMAP, insert rows, and run row-diff normalization + rich_table = ctx.table("test_doris_complex_types") + rich_cols = { + "id": exp.DataType.build("INT128", dialect="doris"), + "tags": exp.DataType.build("ARRAY", dialect="doris"), + "attrs": exp.DataType.build("MAP", dialect="doris"), + "profile": exp.DataType.build( + "STRUCT", dialect="doris", udt=True + ), + "bm": exp.DataType.build("BITMAP", dialect="doris", udt=True), + } + ctx.engine_adapter.create_table( + rich_table, + rich_cols, + table_properties={"duplicate_key": exp.Tuple(expressions=[exp.to_column("id")])}, + ) + ctx.engine_adapter.execute( + f"INSERT INTO {ctx.engine_adapter._to_sql(rich_table)} (`id`, `tags`, `attrs`, `profile`, `bm`) " + f"VALUES (100, ARRAY('x', 'y'), MAP('k1', 10, 'k2', 20), STRUCT(7, 'neo'), TO_BITMAP(42))" + ) + described_cols = ctx.engine_adapter.columns(rich_table) + assert described_cols["id"] == exp.DataType.build("INT128") + concat_expr = t.cast(DorisEngineAdapter, ctx.engine_adapter).concat_columns(described_cols) + diff_rows = ctx.engine_adapter.fetchall(exp.select(concat_expr).from_(rich_table)) + assert len(diff_rows) == 1 + assert "100" in str(diff_rows[0][0]) + assert "42" in str(diff_rows[0][0]) + + # 2. Live schema evolution (ADD COLUMN + MODIFY COLUMN + suppressed key-column VARCHAR(255)->STRING no-op) + cur_tbl = ctx.table("test_alter_cur") + tgt_tbl = ctx.table("test_alter_tgt") + ctx.engine_adapter.create_table( + cur_tbl, + { + "code": exp.DataType.build("TEXT", dialect="doris"), + "notes": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + ) + ctx.engine_adapter.create_table( + tgt_tbl, + { + "code": exp.DataType.build("TEXT", dialect="doris"), + "notes": exp.DataType.build("TEXT", dialect="doris"), + "extra_cnt": exp.DataType.build("BIGINT", dialect="doris"), + }, + ) + alter_ops = ctx.engine_adapter.get_alter_operations(cur_tbl, tgt_tbl) + ctx.engine_adapter.alter_table(alter_ops) + evolved_cols = ctx.engine_adapter.columns(cur_tbl) + assert "extra_cnt" in evolved_cols + assert evolved_cols["code"] == exp.DataType.build("TEXT", dialect="doris") + + # 3. Direct time-transform inequality delete on OLAP table + ev_tbl = ctx.table("test_time_ineq_del") + ev_cols = { + "id": exp.DataType.build("INT", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + } + ctx.engine_adapter.create_table(ev_tbl, ev_cols) + ctx.engine_adapter.execute( + f"INSERT INTO {ctx.engine_adapter._to_sql(ev_tbl)} (`id`, `ds`) VALUES " + f"(1, '2024-01-14'), (2, '2024-01-15'), (3, '2024-01-16')" + ) + ctx.engine_adapter.delete_from( + ev_tbl, + parse_one( + "DATE_TRUNC(ds, 'day') >= '2024-01-15' AND DATE_TRUNC(ds, 'day') <= '2024-01-15'", + read="doris", + ), + ) + remaining_ids = [ + r[0] for r in ctx.engine_adapter.fetchall(exp.select("id").from_(ev_tbl).order_by("id")) + ] + assert remaining_ids == [1, 3] diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py index e45fe1f0d2..effc026348 100644 --- a/tests/core/engine_adapter/test_doris.py +++ b/tests/core/engine_adapter/test_doris.py @@ -686,7 +686,7 @@ def test_create_table_olap_key_prefix_reordering_and_order_by( ) 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", + "CREATE TABLE IF NOT EXISTS `events` (`id` BIGINT, `ds` DATE, `payload` VARCHAR(255), `updated_at` DATETIME) UNIQUE KEY (`id`, `ds`) CLUSTER BY (`updated_at`, `id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10", ] @@ -2376,9 +2376,9 @@ def test_unannotated_view_column_comments_and_starrocks_property_parity( ) 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 VIEW `db`.`unannotated_view` 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')", + "CREATE TABLE IF NOT EXISTS `db`.`paren_keys_tbl` (`id` INT, `dt` DATE, `name` VARCHAR(50)) UNIQUE KEY (`id`, `dt`) CLUSTER 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, @@ -4517,3 +4517,71 @@ def test_doris_physical_properties_kwarg_drop_schema_cascade_and_mv_df_source_co mock_read_sql.call_args[0][0] == f"SELECT `id` FROM `db`.`{short_tmp}` WHERE `is_active` = TRUE" ) + + +def test_doris_time_transform_inequalities_unaliased_ctas_and_sync_mv_db_restore( + make_mocked_engine_adapter: t.Callable, +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter, database="analytics") + + # 1. Direct time-transform inequalities (`>=`, `<=`, `!=`) in `delete_from` are rewritten into + # bare-column predicates (and `!=` is split into two OR-free DELETE statements). + adapter.delete_from( + "db.events", + "DATE_TRUNC(ds, 'day') >= '2024-01-15' AND DATE_TRUNC(ds, 'day') <= '2024-01-20'", + ) + adapter.delete_from("db.events", "DATE_TRUNC(ds, 'day') != '2024-01-15'") + assert to_sql_calls(adapter) == [ + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-15' AND `ds` < '2024-01-21'", + "DELETE FROM `db`.`events` WHERE `ds` < '2024-01-15'", + "DELETE FROM `db`.`events` WHERE `ds` >= '2024-01-16'", + ] + + # 2. `DELETE ... USING` with an extra `OR` condition (`id IN (SELECT id FROM db.src) AND (status = 'a' OR status IS NULL)`) + # expands `extra_conditions` into separate OR-free `DELETE ... USING` statements. + adapter.cursor.execute.reset_mock() + adapter.delete_from( + "db.target", + where=parse_one( + "id IN (SELECT id FROM db.src) AND (status = 'a' OR status IS NULL)", + read="doris", + ), + ) + 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`.`status` = 'a'", + "DELETE FROM `db`.`target` AS `_t1` USING (SELECT `id` FROM `db`.`src`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`status` IS NULL", + ] + + # 3. Unannotated CTAS with unaliased `AVG(score)` or division `a / b` at column 0 promotes `id` to column 0. + adapter.cursor.execute.reset_mock() + adapter.ctas( + "db.unaliased_ctas", + parse_one("SELECT AVG(score), id FROM db.src GROUP BY 2", read="doris"), + ) + adapter.ctas( + "db.div_ctas", + parse_one("SELECT a / b AS ratio, id FROM db.src", read="doris"), + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`unaliased_ctas` AS SELECT `id`, AVG(`score`) FROM `db`.`src` GROUP BY 1", + "CREATE TABLE IF NOT EXISTS `db`.`div_ctas` AS SELECT `id`, `a` / `b` AS `ratio` FROM `db`.`src`", + ] + + # 4. Sync MV creation and drop in another schema (`other_db`) restores the configured session database (`analytics`). + adapter.cursor.execute.reset_mock() + adapter.create_view( + "other_db.rollup_mv", + parse_one("SELECT id, SUM(val) AS s FROM other_db.fact GROUP BY id", read="doris"), + replace=False, + materialized=True, + view_properties={"materialized_type": exp.Literal.string("SYNC")}, + ) + adapter.drop_view("other_db.rollup_mv", materialized=True) + assert to_sql_calls(adapter) == [ + "USE `other_db`", + "CREATE MATERIALIZED VIEW `rollup_mv` AS SELECT `id` AS `mv_id`, SUM(`val`) AS `s` FROM `other_db`.`fact` GROUP BY `id`", + "USE `analytics`", + "USE `other_db`", + "DROP MATERIALIZED VIEW IF EXISTS `rollup_mv` ON `other_db`.`fact`", + "USE `analytics`", + ]