diff --git a/.github/scripts/install-prerequisites.sh b/.github/scripts/install-prerequisites.sh index 6997633a31..09998836d0 100755 --- a/.github/scripts/install-prerequisites.sh +++ b/.github/scripts/install-prerequisites.sh @@ -28,6 +28,8 @@ elif [ "$ENGINE" == "fabric" ]; then rm packages-microsoft-prod.deb ENGINE_DEPENDENCIES="msodbcsql18" +elif [ "$ENGINE" == "doris" ]; then + sudo sysctl -w vm.max_map_count=2000000 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..489cdd4583 100755 --- a/.github/scripts/wait-for-db.sh +++ b/.github/scripts/wait-for-db.sh @@ -104,6 +104,46 @@ db2_ready() { done } +doris_ready() { + 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 *$" || 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..." 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/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/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 38f6814c8b..a34fd189f1 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) @@ -975,6 +976,7 @@ 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) 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..83778c05d0 --- /dev/null +++ b/docs/integrations/engines/doris.md @@ -0,0 +1,654 @@ +# 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 + 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]" +``` + +### 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 | +| --- | --- | :---: | :---: | +| `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 | +| `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 (`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 | +| `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 | + +### Connection configuration example + +=== "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 + ``` + +--- + +## 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` or `cluster_by` (emitting `CLUSTER 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`. 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_keys, + 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 +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 or column tuple `distributed_by`**: + ```sql + physical_properties ( + distributed_by = ( + kind = 'HASH', + 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. + ) + ``` + +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+). 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. + - 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 + +```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 (and SQLMesh automatically marks the range partition column `NOT NULL` and wraps CTAS projections in `NON_NULLABLE(...)`): + +```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' + ) +); +``` + +### 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 + +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`). + +!!! 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 raise an error when attempting to create views in an external catalog): + ```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', + excluded_trigger_tables = 'analytics.user_events' + ) +); + +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); +``` + +!!! 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'`) + +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 ( + 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'`: + +```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', + ALTER TABLE @this_model ADD INDEX idx_message (message) USING INVERTED COMMENT 'Full-text inverted index' +); +``` + +--- + +## 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`) | +| --- | :---: | :---: | +| `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 | + +--- + +## Property Aliases & StarRocks Compatibility + +If you are migrating StarRocks SQLMesh models or existing Doris DDL conventions, `DorisEngineAdapter` supports these property shorthands in `physical_properties`: + +| Property / Shorthand | SQLMesh Doris Behavior | +| --- | --- | +| `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` | + +--- + +## 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`, 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/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/quickstart/cli.md b/docs/quickstart/cli.md index 6aa4fc34ed..4b300c9018 100644 --- a/docs/quickstart/cli.md +++ b/docs/quickstart/cli.md @@ -118,6 +118,8 @@ Choose your SQL engine: [16] RisingWave [17] Fabric [18] StarRocks + [19] Db2 + [20] Apache Doris Enter a number: 1 ``` 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 161db263c3..73e16bbd0b 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/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/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..45ef2f4989 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,184 @@ 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 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] + ) + 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) + 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 or (isinstance(self.ssl, dict) and len(self.ssl) == 0): + 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") + 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 + + 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. tpe.all_field_infos()["type_"].default: tpe diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py index bcfb300ba1..b5cfac2a86 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, 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 +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,66 @@ 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.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, + 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 +1574,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/__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/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 new file mode 100644 index 0000000000..ae446e70ff --- /dev/null +++ b/sqlmesh/core/engine_adapter/doris.py @@ -0,0 +1,5638 @@ +from __future__ import annotations + +from collections import defaultdict +import datetime +from decimal import Decimal +import hashlib +import itertools +import logging +import re +import threading +import typing as t + +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 +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 columns_to_types_all_known, 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.model.kind import TimeColumn + from sqlmesh.core.node import IntervalUnit + from sqlmesh.core.schema_diff import TableAlterOperation + from sqlmesh.utils.date import TimeLike + +logger = logging.getLogger(__name__) + + +@set_catalog() +class DorisEngineAdapter( + EngineAdapterWithIndexSupport, + LogicalMergeMixin, + GetCurrentCatalogFromFunctionMixin, + PandasNativeFetchDFSupportMixin, + RowDiffMixin, +): + DIALECT = "doris" + 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 + 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 + 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( + {"excluded_trigger_tables", "source_table"} + ) + 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("TEXT", 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("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), + }, + }, + } + + INVALID_KEY_TYPES = { + "FLOAT", + "DOUBLE", + "STRING", + "TEXT", + "TINYTEXT", + "MEDIUMTEXT", + "LONGTEXT", + "BLOB", + "TINYBLOB", + "MEDIUMBLOB", + "LONGBLOB", + "BINARY", + "VARBINARY", + "JSON", + "JSONB", + "VARIANT", + "OBJECT", + "ARRAY", + "MAP", + "STRUCT", + "BITMAP", + "HLL", + "HLLSKETCH", + "QUANTILE_STATE", + "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 _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) + 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 + + @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 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)) + 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" + 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 + 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 + 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 + 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 (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, + 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.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): + 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, exp.Table)): + return ".".join(part.name for part in val.parts) + 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) + + @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: + 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 cls._split_cleaned_column_names(key_expr.name) + if isinstance(key_expr, (exp.Identifier, exp.Literal, exp.Var)): + return cls._split_cleaned_column_names(str(key_expr.this)) + if isinstance(key_expr, str): + 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: + 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.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 + + 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: + 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() + except Exception: + logger.info("Connection to the database was lost. Reconnecting...") + self._connection_pool.close() + self._needs_catalog_restore = True + 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 + 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() + 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( + 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 catalog/database switch 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) + + @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" + 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( + 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: + self.ping() + except Exception: + logger.info("Connection to the database was lost. Reconnecting...") + self._connection_pool.close() + self._needs_catalog_restore = True + 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 + 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 + and self._is_safe_to_retry_sql(sql) + ): + logger.warning( + "Doris connection lost (%s: %s) during read/session statement, reconnecting and retrying once...", + err_name, + err_code, + ) + self._connection_pool.close_cursor() + try: + conn = self._connection_pool.get() + 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( + 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] = {} + 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 + ) + 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) + 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 ( + TableAlterAddColumnOperation, + TableAlterChangeColumnTypeOperation, + TableAlterDropColumnOperation, + TableAlterOperation, + ) + + current_table = exp.to_table(current_table_name) + is_iceberg = self._is_iceberg_table(catalog_name=current_table.catalog) + if is_iceberg: + 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) + 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) + 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( + current_table_name, + current_cols, + target_cols, + ignore_destructive=ignore_destructive, + 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 + if not ( + ( + isinstance(op, TableAlterChangeColumnTypeOperation) + 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)) + and op.column.name.lower() in protected_key_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.""" + 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, + 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) if schema_name else exp.Table() + 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 schema_exp.db else exp.func("DATABASE") + ) + ) + ) + short_to_orig: t.Dict[str, str] = {} + if 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] = [] + 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" + 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() + 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 + + # 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_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_arg})", + 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( + 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: + """ + Doris view creation behavior: + - 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 (after validating the new query and properties). + """ + target_table = exp.to_table(view_name) + effective_catalog = target_table.catalog or self.default_catalog + if effective_catalog and effective_catalog.lower() != "internal": + 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." + ) + + if not materialized: + return super().create_view( + view_name, + query_or_df, + target_columns_to_types, + replace=replace, + 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, + replace=replace, + 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, + 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, + view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, + source_columns: t.Optional[t.List[str]] = None, + **create_kwargs: t.Any, + ) -> None: + 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): + 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 + ) + 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) + ) + query_or_df = self._values_to_sql( + values, + source_columns_to_types, + batch_start=0, + 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, + None if is_sync_mv else 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_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 schema_cols_to_types and not is_sync_mv: + schema = self._build_schema_exp( + target_table, + schema_cols_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") + + 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 not is_sync_mv + and 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", + ) + + 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) + + 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: + 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") + and sync_source_tbl is not None + ): + drop_view_props["source_table"] = sync_source_tbl + self.drop_view( + view_name, + ignore_if_not_exists=True, + materialized=True, + view_properties=drop_view_props, + ) + + # 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 + ) + 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 + + 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: + registry[f"{target_table.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: + 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 + 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 _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() + 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, + 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) + 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 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: + 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 `. + 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 "" + 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) + + 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, + 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 = 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 = 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}" + + def _truncate_table(self, table_name: TableName) -> None: + 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) + 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 + + # 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) + 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=self._get_table_properties(kwargs), + ) + ): + 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, + ) + + _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 + + @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, + 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, + **kwargs: t.Any, + ) -> 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, + **kwargs, + ) + return + parsed_exprs, _ = self._parse_partition_expressions(partitioned_by) + 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): + 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: + 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] = [] + 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( + exp.Anonymous( + this="DATE_TRUNC", + expressions=[col_expr, exp.Literal.string(col_unit)], + ) + ) + + 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=self._get_table_properties(kwargs), + ) + if ( + target_columns_to_types + and columns_to_types_all_known(target_columns_to_types) + 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, reorder_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 + ] + + 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=self._get_table_properties(kwargs), + ) + 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, + 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( + self, table_name: TableName, where: t.Optional[t.Union[str, exp.Expr]] = None + ) -> None: + """Delete from a table in Doris.""" + if not where: + self._truncate_table(table_name) + return + + if isinstance(where, str): + where = parse_one(where, read=self.dialect) + 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._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) + if subquery_info is not None: + self._execute_delete_with_subquery(table_name, *subquery_info) + 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 + + 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 + + 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 ( + 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: + 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.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): + 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 + + @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 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: + 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.NEQ, + exp.GTE, + exp.GT, + exp.LTE, + exp.LT, + exp.In, + ) + ): + if node.find_ancestor(exp.Subquery, exp.Select): + continue + 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 + ): + 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 + 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(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: + 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 + 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): + 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 = self._flatten_and_conjuncts(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: + 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:]] + subquery = exp.Subquery(this=unwrapped_select_expr) + elif isinstance(lhs, exp.Tuple): + key_exprs = [e.copy() for e in lhs.expressions] + else: + key_exprs = [lhs.copy()] + + 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())) + elif self._extract_time_partition_transform(k) is not None: + normalized_keys.append(k.copy()) + else: + return None + + if not normalized_keys: + return None + + 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.""" + conjuncts = self._flatten_and_conjuncts(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, self._normalize_extra_conditions(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`) 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.And): + 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) + if normalized is not None: + new_conjuncts.append(normalized) + 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. + """ + conjuncts = self._flatten_and_conjuncts(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 = self._flatten_or_disjuncts(target_or) + 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_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: + 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 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 self._is_literal_key_value(left): + 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, self._normalize_extra_conditions(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 + + @classmethod + 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) + 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") + if isinstance(query, exp.Subquery): + return query.copy() + 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, exp.SetOperation)): + 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.""" + 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) 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() + subq_branches: t.List[exp.Select] = [] + + def _collect_subq_branches(node: t.Optional[exp.Expr]) -> None: + 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): + 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 + 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) + and not isinstance(first_select.expressions[idx], exp.Star) + ): + sel_item = first_select.expressions[idx] + 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: + alias_name = f"_col_{idx}" + 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) + 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") + join_conditions.append(t1_col.eq(t2_col)) + + 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) + + try: + 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( + 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, orig_subquery, extra_conditions) + 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 + + @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, 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`) + 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() + ) + first_select = select_query + while isinstance(first_select, (exp.SetOperation, exp.Subquery)): + first_select = first_select.this + if not isinstance(first_select, exp.Select): + 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) + 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] + 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 + 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 + 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 src_transform is not None or self._is_date_part_extract_fn(raw_sel): + raise SQLMeshError( + 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() + 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("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.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): + all_branches.append(node) + + _collect_select_branches(select_query) + 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) + + # 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) + or branch.args.get("limit") + or branch.args.get("offset") + ): + 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): + 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 + 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 ( + 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 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) + + 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 + 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 + 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 + 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 + + rows = self.fetchall(fetch_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) + 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] = [] + pos = 0 + 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: + 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( + 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 + 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))) + 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, + 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): + 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 + 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, + 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 + + 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_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 + + 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: + 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 = _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): + batch = non_null_vals[i : i + batch_size] + 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: + 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 (...) + # 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 rows}), -idx), + ) + grouped: t.Dict[t.Tuple[t.Any, ...], t.List[t.Any]] = defaultdict(list) + for row in 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(_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 = _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): + batch = non_null_batch[i : i + batch_size] + 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: + 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( + 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], + *, + 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. + """ + unique_key_prop, duplicate_key_prop, aggregate_key_prop = self._resolve_table_type_and_keys( + physical_properties + ) + if ( + unique_key + 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] + physical_properties = dict(physical_properties) + physical_properties["unique_key"] = ( + cleaned_keys[0] if len(cleaned_keys) == 1 else exp.Tuple(expressions=cleaned_keys) + ) + 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(self._get_table_properties(kwargs) or {}) + unique_key_prop, duplicate_key_prop, aggregate_key_prop = self._resolve_table_type_and_keys( + table_properties + ) + + if ( + primary_key + 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] + ) + + 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 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=self._get_table_properties(kwargs), + ) + 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 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, + 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, + **kwargs: t.Any, + ) -> 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( + table_format=kwargs.get("table_format"), + catalog_name=table.catalog, + 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 + for k in key + ] + key_cols: t.List[exp.Column] = [] + 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: + 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): + 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 + ) + 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__'", *normalized_key, dialect=self.dialect) + if len(normalized_key) > 1 + else normalized_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, + table_format=kwargs.get("table_format"), + ) + + 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(*normalized_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: + 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)) + + self.execute(insert_statement, track_rows_processed=True) + finally: + self.drop_table(temp_table) + + 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=self._get_table_properties(kwargs), + ) + 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, + 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=self._get_table_properties(kwargs), + ) + is_internal_iceberg = is_iceberg and ( + not effective_catalog or effective_catalog.lower() == "internal" + ) + + 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: + 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(_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 + # 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 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: + 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, + ) + + @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 + + 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: + _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: + _remap_lit(item) + + 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: + 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 + 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=self._get_table_properties(kwargs), + ) + + 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) + 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 = { + col_def.name.lower(): 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.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", + new_col_defs + other_exprs, + ) + if ( + (order_changed or types_changed or query_order_differs) + and 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.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.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] = [] + 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( + 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) + 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)): + 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, exp.Subquery)): + 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: + table_props = self._get_table_properties(kwargs) 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") + 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) + 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) + } + 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() + 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: + 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) + elif auto_range_cols and isinstance(expression, (exp.Select, exp.SetOperation)): + expression = expression.copy() + auto_select_branches: t.List[exp.Select] = [] + + def _collect_auto_branches(node: exp.Expr) -> None: + 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): + 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 + ): + 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, + expression=expression, + exists=exists, + replace=replace, + target_columns_to_types=target_columns_to_types, + table_description=table_description, + table_kind=table_kind, + **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 _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() + 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): + 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) + 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, + ( + 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, exp.ParseJSON)): + 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, exp.Ntile)): + 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", + "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"): + 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": + 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, + 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] = [] + + def _collect_branches(node: exp.Expr) -> None: + 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): + branches.append(node) + + _collect_branches(expression) + if not branches: + 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( + len(b.expressions) == num_cols + and not any(isinstance(p, exp.Star) for p in b.expressions) + for b in branches + ) + and all(orig_names) + and len({n.lower() for n in orig_names}) == num_cols + ): + merged_cols_to_types: t.Dict[str, exp.DataType] = {} + for col_idx, col_name in enumerate(orig_names): + 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 = { + 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: + 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] = [] + 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] + is_synthetic = col_name.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)): + cast_expr = exp.cast(proj.this.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 + ) + 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 + ) + 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) + self._remap_positional_group_and_order(expression, old_to_new_one_based) + return expression + expression.set( + "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): + 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 = { + 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) + types_changed = any( + reordered_select_cols[c] != select_cols_to_types[c] + for c in reordered_select_cols + ) + if order_changed or types_changed: + old_pos_by_lower = { + k.lower(): idx + 1 for idx, (k, _) in enumerate(keyed_projs) + } + expression = expression.copy() + updated_projs: t.List[exp.Expr] = [] + 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: + 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)): + 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())) + 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)), + ) + 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) + elif any(isinstance(proj, exp.Star) for proj in expression.expressions): + table_props = self._get_table_properties(kwargs) or {} + resolved_uk, resolved_dk, resolved_ak = self._resolve_table_type_and_keys( + table_props + ) + has_explicit_key = bool( + 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_") + if ( + 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 + ): + 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 ( + not has_explicit_key + and 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 + 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 + + @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], + kwargs: t.Optional[t.Dict[str, t.Any]] = None, + ) -> t.Dict[str, exp.DataType]: + kwargs = kwargs 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." + ) + 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'." + ) + 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 + 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: + 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.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 = [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] + 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. " + "Use an integer, DATE/DATETIME, DECIMAL, CHAR, or VARCHAR(n) column instead." + ) + 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} + + 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): + known_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, + ) + 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 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 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 + # 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 = [ + 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 + for i, c in enumerate(col_names) + if not self._is_valid_olap_key_type(updated_types[c]) + ), + 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 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) >= cutoff_idx] + moved_set = set(moved_parts) + 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} + + 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 + + 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 + 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) + 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 + + @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], + *, + 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 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: + 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=iceberg_part_exprs, + 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: + 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. " + "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 + + create_expressions: t.Optional[t.List[exp.Expr]] = None + if partitions: + if isinstance(partitions, exp.Tuple): + create_expressions = [self._to_raw_partition_sql(e) for e in partitions.expressions] + else: + create_expressions = [self._to_raw_partition_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.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" + ) + 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 + 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'?" + ) + 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) 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) + 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 {} + ) + + # 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 + 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 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 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 ( + 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]) + ) + 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: + if not unique_key_cols: + raise SQLMeshError( + "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] + 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"CLUSTER BY ({order_cols_sql})")) + ) + + if table_description: + properties.append( + exp.SchemaCommentProperty( + this=exp.Literal.string(self._truncate_table_comment(table_description)) + ) + ) + + add_partition = True + if partitioned_by: + parsed_part_exprs, _ = self._parse_partition_expressions(partitioned_by) + 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] {key_label} 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: + 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 (...) + 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: + 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=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 strict_key_cols and dist_kind == "RANDOM": + key_kind_label = "UNIQUE KEY" if unique_key_cols else "AGGREGATE KEY" + raise SQLMeshError( + 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: + 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], + ) + 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 + # 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: + 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, + 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): + 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) + 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": + 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 + 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 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))) + 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'." + ) + + @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: + expression = self._normalize_schema_tmp_identifiers(expression) + 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 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) + ): + 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))): + if dtype.is_type("int128"): + replacement = exp.DataType.build("LARGEINT", udt=True) + if dtype is expression: + expression = replacement + else: + dtype.replace(replacement) + elif dtype.is_type("struct"): + 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, 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_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( + 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) or RANDOM() + if isinstance(distributed_by, (exp.Anonymous, exp.Func)) and not isinstance( + distributed_by, (exp.Tuple, exp.Array, exp.Paren) + ): + 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." + ) + 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, + buckets=default_buckets_expr or (exp.Literal.number(10) if cols else None), + order=None, + ) + + # 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 = ( + 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*\(([^)]*)\)|\b)\s*(?:BUCKETS\s+(\S+))?$", + 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) 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), + 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): + 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)): + 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.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 + + 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 + else distributed_info.get("columns") + ) + 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" + + 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] + + 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, + buckets=buckets_expr, + order=None, + ) + + # Fallback 1: Top-level clustered_by (or physical_properties clustered_by) + if clustered_by: + 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"), + 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", + "partition_by", + "partitioned_by", + "partitions", + "auto_partition", + "unique_key", + "duplicate_key", + "aggregate_key", + "table_type", + "keys", + "order_by", + "cluster_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) + 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/integration/__init__.py b/tests/core/engine_adapter/integration/__init__.py index fe1951add6..3aa2639a92 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""" @@ -754,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) @@ -765,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/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..2638848e9e --- /dev/null +++ b/tests/core/engine_adapter/integration/docker/compose.doris.yaml @@ -0,0 +1,58 @@ +services: + doris-fe: + image: apache/doris:fe-4.0.8 + 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" + volumes: + - iceberg-warehouse:/tmp/doris_iceberg_warehouse + networks: + doris_net: + ipv4_address: 172.30.80.2 + + doris-be: + image: apache/doris:be-4.0.8 + container_name: doris-be + hostname: doris-be + depends_on: + - 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 = 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 + 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..dbad929aa6 100644 --- a/tests/core/engine_adapter/integration/test_integration.py +++ b/tests/core/engine_adapter/integration/test_integration.py @@ -158,6 +158,10 @@ def create_objects_and_validate(schema_name: str): 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." @@ -439,6 +443,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"}, @@ -2684,6 +2694,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 +3257,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 +3269,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 +3294,12 @@ 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 + ( + # "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) + ): v for k, v in columns_to_types.items() } @@ -3941,6 +3957,11 @@ 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( + "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") mview_name = ctx.table("test_mview") @@ -4008,6 +4029,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 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" + ) 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..32cfde900a --- /dev/null +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -0,0 +1,611 @@ +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") + + +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=30, 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=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): + 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"}, + ] + ), + ) + + +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"), + ] + + # 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 + + 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"}, + ] + ), + ) + + +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), + ] + + +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 new file mode 100644 index 0000000000..effc026348 --- /dev/null +++ b/tests/core/engine_adapter/test_doris.py @@ -0,0 +1,4587 @@ +import typing as t + +import pandas as pd # noqa: TID253 +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 sqlmesh.utils.errors import SQLMeshError +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) == [ + "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`", + ] + + +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"), + 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 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`", + ] + + +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`) CLUSTER 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 (...) () and marks range partition column NOT NULL + adapter.create_table( + "olap_auto_partition", + 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 NOT NULL, `payload` STRING) 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) + close_spy = mocker.spy(adapter._connection_pool, "close") + + 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("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): + 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() + close_spy.assert_called_once() + + +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 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 also raises SQLMeshError + with pytest.raises( + SQLMeshError, + match="Doris only supports creating materialized views in the 'internal' catalog", + ): + adapter.create_view( + "iceberg_cat.silver.events_mv", + parse_one("SELECT 1 AS id"), + 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) == [] + + # 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], +): + 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"), + ) + 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", + 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 ( + "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], +): + 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) == [ + "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() + + 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`", + ] + + 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 +): + 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) + + 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=raw_where, + ) + + assert to_sql_calls(adapter) == [ + "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", + ] + + +@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", + "BITMAP", + "HLL", + "QUANTILE_STATE", + "AGG_STATE(sum(INT))", + ], +) +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", udt=True), + "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) + 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 (table qualifiers stripped) + props = adapter.adjust_physical_properties_for_incremental( + {}, + requires_delete_capable_table=True, + unique_key=[exp.to_column("src.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" + ] + + +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 + ) + + +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"), + }, + ) + + 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"), + "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 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", + }, + ) + + # 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], +): + 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")}, + ) + + # 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_unannotated_reordered` UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10 AS SELECT `id`, `name` FROM `db`.`src`", + ] + + +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`"] + + # 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'" + ] + + # 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 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] == "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( + 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 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"), + 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 ( + "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 + ) + + +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 + + +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 + + +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", + ] + + +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", + ] + + +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` 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`) 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, + # 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")}, + ) + + +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`, 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 + 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", + ) + + +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 + ) + + +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", + ] + + +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 + + +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 + + 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. 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", + { + "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 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 + + +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) + + +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) == [ + "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 + # 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, + } + + +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) == [ + "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 + # `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", + ] + + +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" + ) + + +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`", + ] diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py index c506d401d9..b529eeac38 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,174 @@ 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", + "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 + 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"] == {"check_hostname": False} + + ssl_false_config = make_config( + type="doris", + 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", + 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, + )