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..e1167ac5ca --- /dev/null +++ b/docs/integrations/engines/doris.md @@ -0,0 +1,652 @@ +# 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 | +| `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` (emitting `ORDER BY (col1, col2)` on Doris 2.1+ Merge-on-Write tables). + +### AGGREGATE KEY model + +For pre-aggregated rollup tables, specify the key columns using `aggregate_key` in `physical_properties`. 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 base table writes) by setting `materialized_type = 'SYNC'` and `source_table` in `physical_properties`: + +```sql +MODEL ( + name analytics.user_events_sync_mv, + kind VIEW ( + materialized true + ), + physical_properties ( + materialized_type = 'SYNC', + source_table = 'analytics.user_events' + ) +); + +SELECT + user_id, + COUNT(event_type) AS event_count +FROM analytics.user_events +GROUP BY user_id; +``` + +### Inverted Indexes (Doris 2.1–5.0) + +You can attach Doris inverted indexes (`USING INVERTED`) for full-text search and fast point filtering using post-statements gated on `@runtime_stage = 'creating'`: + +```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..be49001dd8 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,173 @@ def connect_db2(**kwargs: t.Any) -> t.Any: return connect_db2 +class DorisConnectionConfig(ConnectionConfig): + """Configuration for the Apache Doris connection. + + Apache Doris uses the MySQL network protocol and supports both internal OLAP tables + and external Multi-Catalog lakehouse catalogs (e.g., Apache Iceberg). + + Args: + host: The hostname of the Doris FE (Frontend) node. + user: The Doris username. + password: The optional Doris password. + port: The query port number of the Doris FE node. Default is 9030. + database: The optional database name. + catalog: The optional default catalog name (e.g., 'internal' or an external Iceberg catalog). + charset: The optional character set (defaults to 'utf8mb4'). + collation: The optional collation. + connect_timeout: Timeout in seconds for establishing a connection. Default is 10. + read_timeout: Timeout in seconds for reading from the connection. + write_timeout: Timeout in seconds for writing to the connection. + ssl: Optional dictionary of SSL arguments passed to PyMySQL. + ssl_ca: Path to the file that contains a PEM-formatted CA certificate. + ssl_cert: Path to the file that contains a PEM-formatted client certificate. + ssl_key: Path to the file that contains a PEM-formatted private key for the client certificate. + ssl_verify_cert: Whether to verify the server certificate against the CA certificate. + ssl_verify_identity: Whether to verify the server hostname against the certificate. + ssl_disabled: Whether to disable SSL connection. + session_variables: Optional dictionary of Doris session variables to set on each connection. + concurrent_tasks: The maximum number of tasks that can use this connection concurrently. + register_comments: Whether or not to register model comments with the SQL engine. + pre_ping: Whether or not to pre-ping the connection before starting a new statement. + """ + + host: str + user: str + password: t.Optional[str] = None + port: t.Optional[int] = 9030 + database: t.Optional[str] = None + catalog: t.Optional[str] = None + charset: t.Optional[str] = None + collation: t.Optional[str] = None + connect_timeout: t.Optional[int] = 10 + read_timeout: t.Optional[int] = None + write_timeout: t.Optional[int] = None + ssl: t.Optional[t.Union[bool, t.Dict[str, t.Any]]] = None + ssl_ca: t.Optional[str] = None + ssl_cert: t.Optional[str] = None + ssl_key: t.Optional[str] = None + ssl_verify_cert: t.Optional[bool] = None + ssl_verify_identity: t.Optional[bool] = None + ssl_disabled: t.Optional[bool] = None + session_variables: t.Optional[t.Dict[str, t.Any]] = None + + concurrent_tasks: int = 4 + register_comments: bool = True + pre_ping: bool = False + + type_: t.Literal["doris"] = Field(alias="type", default="doris") + DIALECT: t.ClassVar[t.Literal["doris"]] = "doris" + DISPLAY_NAME: t.ClassVar[t.Literal["Apache Doris"]] = "Apache Doris" + DISPLAY_ORDER: t.ClassVar[t.Literal[20]] = 20 + + _engine_import_validator = _get_engine_import_validator("pymysql", "doris") + + @field_validator("session_variables", mode="after") + @classmethod + def _validate_session_variables( + cls, v: t.Optional[t.Dict[str, t.Any]] + ) -> t.Optional[t.Dict[str, t.Any]]: + if v: + import math + + for k, val in v.items(): + if not isinstance(k, str) or not re.fullmatch(r"[a-zA-Z_][a-zA-Z0-9_]*\Z", k): + raise ConfigError(f"Invalid Doris session variable name: '{k}'") + if not isinstance(val, (str, int, float, bool)) or ( + isinstance(val, float) and (math.isnan(val) or math.isinf(val)) + ): + raise ConfigError(f"Invalid Doris session variable value for '{k}': {val!r}") + return v + + @property + def _static_connection_kwargs(self) -> t.Dict[str, t.Any]: + kwargs: t.Dict[str, t.Any] = { + "autocommit": True, + "charset": self.charset or "utf8mb4", + "local_infile": False, + } + if not self.ssl_disabled: + ssl_config: t.Dict[str, t.Any] = dict(self.ssl) if isinstance(self.ssl, dict) else {} + for key in ("ca", "cert", "key"): + val = getattr(self, f"ssl_{key}", None) + if val is not None: + ssl_config[key] = val + if self.ssl_verify_cert is not None: + ssl_config["verify_mode"] = self.ssl_verify_cert + if self.ssl_verify_identity is not None: + ssl_config["check_hostname"] = bool(self.ssl_verify_identity) + if not ssl_config and (self.ssl is not None and self.ssl is not False): + ssl_config["check_hostname"] = bool(self.ssl_verify_identity) + if ssl_config: + kwargs["ssl"] = ssl_config + for key in ( + "ssl_ca", + "ssl_cert", + "ssl_key", + "ssl_verify_cert", + "ssl_verify_identity", + ): + val = getattr(self, key, None) + if val is not None: + kwargs[key] = val + return kwargs + + @property + def _connection_kwargs_keys(self) -> t.Set[str]: + is_external_catalog = bool(self.catalog and self.catalog.lower() != "internal") + return { + k + for k in ( + "host", + "user", + "password", + "port", + "database", + "charset", + "collation", + "connect_timeout", + "read_timeout", + "write_timeout", + "ssl_disabled", + ) + if getattr(self, k) is not None and not (k == "database" and is_external_catalog) + } + + @property + def _engine_adapter(self) -> t.Type[EngineAdapter]: + return engine_adapter.DorisEngineAdapter + + @property + def _connection_factory(self) -> t.Callable: + from pymysql import connect + + return connect + + @property + def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]: + if not self.catalog and not self.session_variables: + return None + + def init(cursor: t.Any) -> None: + if self.session_variables: + self._validate_session_variables(self.session_variables) + for k, v in self.session_variables.items(): + val_sql = exp.convert(v).sql("doris", comments=False) + cursor.execute(f"SET `{k}` = {val_sql}") + if self.catalog: + catalog_sql = exp.to_identifier(self.catalog, quoted=True).sql("doris") + cursor.execute(f"SWITCH {catalog_sql}") + if self.database: + db_sql = exp.to_identifier(self.database, quoted=True).sql("doris") + cursor.execute(f"USE {db_sql}") + + return init + + 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/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..3376e7ffc8 --- /dev/null +++ b/sqlmesh/core/engine_adapter/doris.py @@ -0,0 +1,4425 @@ +from __future__ import annotations + +from collections import defaultdict +import datetime +from decimal import Decimal +import itertools +import logging +import re +import threading +import typing as t + +import pandas as pd # noqa: TID253 +from sqlglot import exp, parse, parse_one +from sqlglot.optimizer.qualify_columns import quote_identifiers + +from sqlmesh.core.dialect import to_schema +from sqlmesh.core.engine_adapter.base import ( + 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.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", + "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 catalog_support(self) -> CatalogSupport: + return CatalogSupport.FULL_SUPPORT + + @property + def default_catalog(self) -> t.Optional[str]: + return self._default_catalog or self._extra_config.get("catalog") + + def get_current_catalog(self) -> t.Optional[str]: + try: + catalog = super().get_current_catalog() + if catalog: + return _decoded_str(catalog) + except Exception: + logger.debug("current_catalog() failed on Doris, falling back to default catalog") + return self.default_catalog or "internal" + + def set_current_catalog(self, catalog: str) -> None: + catalog_sql = self._to_sql(exp.to_identifier(catalog, quoted=True)) + self.execute(exp.Command(this="SWITCH", expression=catalog_sql)) + self._active_catalog = catalog + self._needs_catalog_restore = False + + def _restore_active_catalog_on_reconnect(self, sql: t.Optional[str] = None) -> None: + active_cat = self._active_catalog + self._needs_catalog_restore = False + if not active_cat: + return + if sql and re.match( + r"\A(?:\s|/\*.*?\*/|--[^\r\n]*(?:[\r\n]|\Z))*SWITCH\b", + sql, + flags=re.IGNORECASE | re.DOTALL, + ): + return + catalog_sql = self._to_sql(exp.to_identifier(active_cat, quoted=True)) + self.cursor.execute(f"SWITCH {catalog_sql}") + database = self._extra_config.get("database") + default_cat = self.default_catalog or "internal" + if database and active_cat.lower() == str(default_cat).lower(): + db_sql = self._to_sql(exp.to_identifier(str(database), quoted=True)) + self.cursor.execute(f"USE {db_sql}") + + def _is_iceberg_table( + self, + 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 _fetch_native_df( + self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False + ) -> DF: + if self._pre_ping: + try: + self.ping() + except Exception: + logger.info("Connection to the database was lost. Reconnecting...") + self._connection_pool.close() + self._needs_catalog_restore = True + self._connection_pool.get_cursor() + if getattr(self, "_needs_catalog_restore", False): + self._restore_active_catalog_on_reconnect() + try: + return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) + except Exception as e: + 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) + + 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 + if getattr(self, "_needs_catalog_restore", False): + self._restore_active_catalog_on_reconnect(sql) + try: + super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) + except Exception as e: + 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: + return super().get_alter_operations( + current_table_name, + target_table_name, + ignore_destructive=ignore_destructive, + ignore_additive=ignore_additive, + ) + + self._last_described_key_cols = set() + current_cols = self.columns(current_table_name) + current_key_cols = set(getattr(self, "_last_described_key_cols", set())) + if current_cols and not current_key_cols: + current_key_cols = {next(iter(current_cols)).lower()} + target_cols = self.columns(target_table_name) + 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.""" + 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") + ) + ) + ) + if object_names: + lowered_names = [name.lower() for name in object_names] + query = query.where(exp.func("LOWER", exp.column("table_name")).isin(*lowered_names)) + + result: t.List[DataObject] = [] + for row_schema, table_name, table_type in self.fetchall(query): + try: + schema_str = _decoded_str(row_schema) if row_schema is not None else schema_exp.db + name_str = _decoded_str(table_name) if table_name is not None else "unknown" + obj_type = _decoded_str(table_type) if table_type is not None else "table" + + upper_type = obj_type.upper() + if upper_type in ("BASE TABLE", "EXTERNAL TABLE"): + obj_type = "table" + elif upper_type in ("VIEW", "SYSTEM VIEW"): + obj_type = "view" + elif upper_type in ("MATERIALIZED VIEW", "MATERIALIZED_VIEW"): + obj_type = "materialized_view" + + result.append( + DataObject( + catalog=catalog, + schema=schema_str, + name=name_str, + type=DataObjectType.from_str(obj_type), + ) + ) + except (ValueError, AttributeError) as e: + logger.error( + f"Error processing row: {e}, row: {(row_schema, table_name, table_type)}" + ) + continue + + # 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: + if ( + not target_columns_to_types + and column_descriptions + and isinstance(query_or_df, exp.Query) + and query_or_df.named_selects + and "*" not in query_or_df.named_selects + ): + target_columns_to_types = { + col: exp.DataType.build("UNKNOWN") for col in query_or_df.named_selects + } + return super().create_view( + view_name, + query_or_df, + 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 or target_columns_to_types)] + values: t.List[t.Tuple[t.Any, ...]] = list( + ordered_df.itertuples(index=False, name=None) + ) + query_or_df = self._values_to_sql( + values, + source_columns_to_types, + batch_start=0, + batch_end=len(values), + ) + + source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( + query_or_df, + target_columns_to_types, + batch_size=0, + target_table=view_name, + source_columns=source_columns, + ) + if len(source_queries) != 1: + raise SQLMeshError("Only one source query is supported for creating materialized views") + + with source_queries[0] as query: + target_table = exp.to_table(view_name) + + is_sync_mv = bool( + view_properties + and ( + ( + self._extract_str_value(view_properties.get("materialized_type")) or "" + ).upper() + == "SYNC" + ) + ) + + schema_cols_to_types = target_columns_to_types + if ( + not schema_cols_to_types + 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 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) + + # 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") + ): + first_tbl = query.find(exp.Table) + if first_tbl: + drop_view_props["source_table"] = first_tbl + self.drop_view( + view_name, + ignore_if_not_exists=True, + materialized=True, + view_properties=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. + if is_sync_mv and not target_table.db: + source_db: t.Optional[str] = None + if view_properties and view_properties.get("source_table"): + source_db = self._to_source_table(view_properties["source_table"]).db or None + if not source_db: + first_tbl = query.find(exp.Table) + if first_tbl and first_tbl.db: + source_db = first_tbl.db + if source_db: + db_sql = self._to_sql(exp.to_identifier(source_db, quoted=True)) + self.execute(f"USE {db_sql}") + + 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, + ) + build_mode = ( + self._to_upper_str(view_properties["build"]) + if view_properties and view_properties.get("build") is not None + else None + ) + if ( + not is_sync_mv + and build_mode == "DEFERRED" + and bool((materialized_properties or {}).get("has_audits")) + ): + self.execute(f"REFRESH MATERIALIZED VIEW {self._to_sql(target_table)} AUTO") + self._clear_data_object_cache(view_name) + + def _to_source_table(self, raw_src: t.Any) -> exp.Table: + 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) + if materialized and view_properties: + mat_type = self._extract_str_value(view_properties.get("materialized_type")) + source_table = view_properties.get("source_table") + if mat_type and mat_type.upper() == "SYNC" and source_table: + view_table_sql = self._to_sql(target_table) + source_tbl_exp = self._to_source_table(source_table) + # Unqualified synchronous materialized views require the source table's database + # to be active on the session before executing `DROP MATERIALIZED VIEW ON `. + if not target_table.db and source_tbl_exp.db: + db_sql = self._to_sql(exp.to_identifier(source_tbl_exp.db, quoted=True)) + self.execute(f"USE {db_sql}") + source_table_sql = self._to_sql(source_tbl_exp) + exists_clause = "IF EXISTS " if ignore_if_not_exists else "" + self.execute( + f"DROP MATERIALIZED VIEW {exists_clause}{view_table_sql} ON {source_table_sql}" + ) + self._clear_data_object_cache(view_name) + return + 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 = exp.to_table(table_name) + caught_err: t.Optional[Exception] = None + if not self._is_iceberg_table(catalog_name=table.catalog): + try: + super()._truncate_table(table_name) + return + except Exception as e: + err_msg = str(e).lower() + if "truncate" not in err_msg and "not support" not in err_msg: + raise e + caught_err = e + + # When TRUNCATE TABLE is unsupported (e.g., on Iceberg tables), clear all rows via DELETE. + # Doris's DeleteHandler rejects DELETE without WHERE, constant WHERE TRUE, OR predicates, + # and floating-point/complex column predicates, so we pick a key-compatible column and + # execute two complementary `IS NOT NULL` and `IS NULL` deletes. + cols_to_types = self.columns(table_name) + if cols_to_types: + valid_col_name = next( + ( + col + for col, dtype in cols_to_types.items() + if self._is_valid_olap_key_type(dtype) + ), + next(iter(cols_to_types)), + ) + target_col = exp.column(valid_col_name, quoted=True) + self.execute(exp.delete(table, where=target_col.is_(exp.null()).not_())) + self.execute(exp.delete(table, where=target_col.is_(exp.null()))) + return + if caught_err is not None: + raise caught_err + + def _insert_overwrite_by_condition( + self, + table_name: TableName, + source_queries: t.List[SourceQuery], + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + where: t.Optional[exp.Condition] = None, + insert_overwrite_strategy_override: t.Optional[InsertOverwriteStrategy] = None, + **kwargs: t.Any, + ) -> None: + table = exp.to_table(table_name) + if ( + insert_overwrite_strategy_override is None + and where is None + and self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=table.catalog, + table_properties=kwargs.get("table_properties"), + ) + ): + insert_overwrite_strategy_override = InsertOverwriteStrategy.INSERT_OVERWRITE + return super()._insert_overwrite_by_condition( + table_name, + source_queries, + target_columns_to_types=target_columns_to_types, + where=where, + insert_overwrite_strategy_override=insert_overwrite_strategy_override, + **kwargs, + ) + + _TIME_UNIT_PRECEDENCE: t.ClassVar[t.Dict[str, int]] = { + "second": 0, + "minute": 1, + "hour": 2, + "day": 3, + "month": 4, + "year": 5, + } + + @classmethod + def _extract_time_partition_transform( + cls, expr: exp.Expr + ) -> t.Optional[t.Tuple[exp.Column, str]]: + """Extract `(column, unit)` from time-truncation partition expressions (`DATE_TRUNC`, + `TO_DATE`, `DATE`, `YEAR`/`YEARS`, `MONTH`/`MONTHS`, `DAY`/`DAYS`, `HOUR`/`HOURS`). + """ + curr = expr.unnest() + col_node: t.Optional[exp.Column] = None + raw_unit: t.Optional[str] = None + + if isinstance(curr, (exp.TimestampTrunc, exp.DateTrunc)): + unit_arg = curr.args.get("unit") + if isinstance(curr.this, exp.Column): + col_node = curr.this + raw_unit = cls._extract_str_value(unit_arg) + elif isinstance(unit_arg, exp.Column): + col_node = unit_arg + raw_unit = cls._extract_str_value(curr.this) + elif isinstance(curr, (exp.TsOrDsToDate, exp.Date)): + if isinstance(curr.this, exp.Column): + col_node = curr.this + raw_unit = "day" + elif isinstance(curr, (exp.Year, exp.Month, exp.Day, exp.Hour)): + inner_arg = curr.this.unnest() if isinstance(curr.this, exp.Expr) else None + if isinstance(inner_arg, (exp.TsOrDsToDate, exp.Date)) and isinstance( + inner_arg.this, exp.Expr + ): + inner_arg = inner_arg.this.unnest() + if isinstance(inner_arg, exp.Column): + col_node = inner_arg + if isinstance(curr, exp.Year): + raw_unit = "year" + elif isinstance(curr, exp.Month): + raw_unit = "month" + elif isinstance(curr, exp.Day): + raw_unit = "day" + else: + raw_unit = "hour" + elif isinstance(curr, exp.Anonymous): + fn_name = curr.name.upper() + if fn_name == "DATE_TRUNC" and len(curr.expressions) == 2: + arg0, arg1 = curr.expressions + if isinstance(arg0, exp.Column): + col_node = arg0 + raw_unit = cls._extract_str_value(arg1) + elif isinstance(arg1, exp.Column): + col_node = arg1 + raw_unit = cls._extract_str_value(arg0) + elif ( + fn_name in ("TO_DATE", "DATE") + and len(curr.expressions) == 1 + and isinstance(curr.expressions[0], exp.Column) + ): + col_node = curr.expressions[0] + raw_unit = "day" + elif ( + fn_name + in ( + "YEAR", + "YEARS", + "MONTH", + "MONTHS", + "DAY", + "DAYS", + "HOUR", + "HOURS", + "MINUTE", + "MINUTES", + "SECOND", + "SECONDS", + ) + and len(curr.expressions) == 1 + and isinstance(curr.expressions[0], exp.Column) + ): + col_node = curr.expressions[0] + raw_unit = fn_name.lower() + + if col_node is not None and raw_unit: + unit_clean = str(raw_unit).strip().strip("'\"").lower() + if unit_clean.endswith("s") and unit_clean[:-1] in cls._TIME_UNIT_PRECEDENCE: + unit_clean = unit_clean[:-1] + if unit_clean in cls._TIME_UNIT_PRECEDENCE: + clean_col = exp.column( + col_node.name, quoted=bool(getattr(col_node.this, "quoted", False)) + ) + return clean_col, unit_clean + return None + + def insert_overwrite_by_partition( + self, + table_name: TableName, + query_or_df: QueryOrDF, + partitioned_by: t.List[exp.Expr], + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + source_columns: t.Optional[t.List[str]] = None, + ) -> None: + if not partitioned_by: + self.replace_query( + table_name, + query_or_df, + target_columns_to_types=target_columns_to_types, + source_columns=source_columns, + ) + return + parsed_exprs, _ = self._parse_partition_expressions(partitioned_by) + col_order: t.List[str] = [] + col_to_expr: t.Dict[str, exp.Column] = {} + 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) + if ( + target_columns_to_types + and columns_to_types_all_known(target_columns_to_types) + and not self._is_iceberg_table(catalog_name=target_table.catalog) + ): + target_columns_to_types = self._reorder_olap_columns( + target_columns_to_types, {"partitioned_by": partitioned_by} + ) + if source_columns: + source_col_set = set(source_columns) + source_columns = [c for c in target_columns_to_types if c in source_col_set] + [ + c for c in source_columns if c not in target_columns_to_types + ] + + super().insert_overwrite_by_partition( + table_name, + query_or_df, + key_exprs or partitioned_by, + target_columns_to_types=target_columns_to_types, + source_columns=source_columns, + ) + + def delete_from( + self, table_name: TableName, where: t.Optional[t.Union[str, exp.Expr]] = None + ) -> None: + """Delete from a table in Doris.""" + if not where: + 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._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 + + 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 _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.""" + where = where.unnest() + conjuncts: t.List[exp.Expr] = ( + list(where.flatten()) if isinstance(where, exp.And) else [where] + ) + target_in: t.Optional[exp.In] = None + 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, extra_conditions + + def _extract_literal_tuple_in_from_where( + self, where: exp.Expr + ) -> t.Optional[t.Tuple[t.List[exp.Expr], t.List[t.Tuple[exp.Expr, ...]], t.List[exp.Expr]]]: + """Extract positive literal tuple IN predicate `(a, b) IN ((1, 2), (3, 4))` and sibling AND conditions.""" + where = where.unnest() + conjuncts: t.List[exp.Expr] = ( + list(where.flatten()) if isinstance(where, exp.And) else [where] + ) + target_in: t.Optional[exp.In] = None + extra_conditions: t.List[exp.Expr] = [] + + for cond in conjuncts: + unwrapped_cond = cond.unnest() + if ( + target_in is None + and isinstance(unwrapped_cond, exp.In) + and self._get_in_subquery(unwrapped_cond) is None + and isinstance(unwrapped_cond.this, exp.Tuple) + and len(unwrapped_cond.this.expressions) >= 2 + and unwrapped_cond.expressions + and all( + isinstance(item, exp.Tuple) + and len(item.expressions) == len(unwrapped_cond.this.expressions) + for item in unwrapped_cond.expressions + ) + ): + target_in = unwrapped_cond + continue + extra_conditions.append(cond) + + if target_in is None: + return None + + normalized_keys: t.List[exp.Expr] = [] + for k in target_in.this.expressions: + if isinstance(k, exp.Column): + normalized_keys.append(k.copy()) + elif isinstance(k, exp.Identifier): + normalized_keys.append(exp.Column(this=k.copy())) + else: + return None + + tuple_rows: t.List[t.Tuple[exp.Expr, ...]] = [ + tuple(e.copy() for e in item.expressions) + for item in target_in.expressions + if isinstance(item, exp.Tuple) + ] + if not tuple_rows: + return None + + return normalized_keys, tuple_rows, extra_conditions + + @staticmethod + def _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 = list(unwrapped.flatten()) + changed = False + 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. + """ + where = where.unnest() + conjuncts: t.List[exp.Expr] = ( + list(where.flatten()) if isinstance(where, exp.And) else [where] + ) + target_or: t.Optional[exp.Or] = None + extra_conditions: t.List[exp.Expr] = [] + + for cond in conjuncts: + unwrapped_cond = cond.unnest() + if target_or is None and isinstance(unwrapped_cond, exp.Or): + target_or = unwrapped_cond + continue + extra_conditions.append(cond) + + if target_or is None: + return None + + disjuncts = list(target_or.flatten()) + if len(disjuncts) < 2: + return None + + canonical_col_order: t.Optional[t.List[str]] = None + canonical_cols: t.List[exp.Expr] = [] + tuple_rows: t.List[t.Tuple[exp.Expr, ...]] = [] + + for branch in disjuncts: + branch_unwrapped = branch.unnest() + branch_conjuncts = ( + list(branch_unwrapped.flatten()) + if isinstance(branch_unwrapped, exp.And) + else [branch_unwrapped] + ) + col_to_val: t.Dict[str, exp.Expr] = {} + col_to_expr: t.Dict[str, exp.Expr] = {} + for item in branch_conjuncts: + 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, 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.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): + self._delete_by_literal_keys(table_name, key_exprs, subquery, extra_conditions) + return + + join_conditions: t.List[exp.Expr] = [] + orig_subquery = subquery.copy() + subquery = subquery.copy() + first_select = subquery.this + while isinstance(first_select, exp.SetOperation): + first_select = first_select.this + used_t2_aliases: t.Set[str] = set() + for idx, col in enumerate(key_exprs): + col_name = col.name + t1_col = exp.column(col_name, table="_t1") + t2_col_name = col_name + if isinstance(first_select, exp.Select) and idx < len(first_select.expressions): + sel_item = first_select.expressions[idx] + 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)) + + for extra in extra_conditions: + qualified_extra = extra.copy() + for col in qualified_extra.find_all(exp.Column): + if not col.find_ancestor(exp.Subquery, exp.Select) and ( + not col.table or col.table.lower() == target_table.name.lower() + ): + col.set("table", exp.to_identifier("_t1")) + col.set("db", None) + col.set("catalog", None) + join_conditions.append(qualified_extra) + + where_expr = exp.and_(*join_conditions) + + target_sql = f"{self._to_sql(target_table)} AS `_t1`" + subquery_with_alias = exp.alias_(subquery.copy(), "_t2", table=True) + subquery_sql = self._to_sql(subquery_with_alias) + where_sql = self._to_sql(where_expr) + + delete_sql = f"DELETE FROM {target_sql} USING {subquery_sql} WHERE {where_sql}" + try: + self.execute(delete_sql) + except Exception as e: + err_msg = str(e).lower() + if any( + marker in err_msg + for marker in ( + "unique key", + "olaptable", + "olap table", + "not support", + "only support", + "unsupported", + ) + ): + logger.info( + "DELETE ... USING not supported on table '%s' (%s); falling back to batched literal DELETE.", + table_name, + e, + ) + self._delete_by_literal_keys(table_name, key_exprs, 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 < `). + """ + select_query = ( + subquery.this.copy() if isinstance(subquery, exp.Subquery) else subquery.copy() + ) + first_select = select_query + while isinstance(first_select, exp.SetOperation): + first_select = first_select.this + if not isinstance(first_select, exp.Select): + return + + projections: t.List[exp.Expr] = [] + 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") + ): + 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 + k_unnested = k.unnest() + is_year_fn = ( + unit == "year" + and not isinstance(k_unnested, (exp.TimestampTrunc, exp.DateTrunc)) + and not ( + isinstance(k_unnested, exp.Anonymous) + and k_unnested.name.upper() == "DATE_TRUNC" + ) + ) + src_transform = self._extract_time_partition_transform(raw_sel) + if src_transform is not None: + src_arg: exp.Expr = src_transform[0].copy() + elif is_year_fn: + raw_year_indices.add(idx) + projections.append(exp.alias_(raw_sel.copy(), f"_val_{idx}", quoted=True)) + continue + else: + src_arg = raw_sel.copy() + trunc_expr = exp.Anonymous( + this="DATE_TRUNC", + expressions=[src_arg, exp.Literal.string(unit)], + ) + upper_expr = exp.DateAdd( + this=trunc_expr.copy(), + expression=exp.Literal.number(1), + unit=exp.var(unit.upper()), + ) + projections.append(exp.alias_(trunc_expr, f"_low_{idx}", quoted=True)) + projections.append(exp.alias_(upper_expr, f"_high_{idx}", quoted=True)) + else: + projections.append(exp.alias_(raw_sel.copy(), f"_val_{idx}", quoted=True)) + select_query.set("expressions", projections) + fetch_query: exp.Select = select_query.distinct() + elif len(first_select.expressions) == len(key_exprs): + all_branches: t.List[exp.Select] = [] + + def _collect_select_branches(node: exp.Expr) -> None: + if isinstance(node, exp.SetOperation): + _collect_select_branches(node.this) + _collect_select_branches(node.expression) + elif isinstance(node, exp.Select): + all_branches.append(node) + + _collect_select_branches(select_query) + unwrapped_indices: t.Set[int] = set(range(len(key_exprs))) + for branch in all_branches: + if len(branch.expressions) == len(key_exprs): + can_unwrap_group = not branch.args.get("having") and not branch.args.get( + "limit" + ) + branch_exprs = list(branch.expressions) + for idx, k in enumerate(key_exprs): + if self._extract_time_partition_transform(k) is not None: + sel_item = branch_exprs[idx] + raw_sel = sel_item.this if isinstance(sel_item, exp.Alias) else sel_item + src_transform = self._extract_time_partition_transform(raw_sel) + if src_transform is not None and can_unwrap_group: + branch_exprs[idx] = src_transform[0].copy() + if branch.args.get("group"): + branch.set("group", None) + else: + unwrapped_indices.discard(idx) + branch.set("expressions", branch_exprs) + + new_first_exprs = list(first_select.expressions) + for idx, sel_item in enumerate(new_first_exprs): + inner_expr = sel_item.this if isinstance(sel_item, exp.Alias) else sel_item + new_first_exprs[idx] = exp.alias_(inner_expr.copy(), f"_col_{idx}", quoted=True) + first_select.set("expressions", new_first_exprs) + + for idx, k in enumerate(key_exprs): + sq_col = exp.column(f"_col_{idx}", quoted=True) + transform = self._extract_time_partition_transform(k) + if transform is not None: + _, unit = transform + k_unnested = k.unnest() + is_year_fn = ( + unit == "year" + and not isinstance(k_unnested, (exp.TimestampTrunc, exp.DateTrunc)) + and not ( + isinstance(k_unnested, exp.Anonymous) + and k_unnested.name.upper() == "DATE_TRUNC" + ) + ) + if idx not in unwrapped_indices and is_year_fn: + raw_year_indices.add(idx) + projections.append(exp.alias_(sq_col, f"_val_{idx}", quoted=True)) + continue + trunc_expr = exp.Anonymous( + this="DATE_TRUNC", + expressions=[sq_col, exp.Literal.string(unit)], + ) + 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) + + 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: + 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))) + predicates.extend(cleaned_extra_conditions) + self.execute(exp.delete(table_name, where=exp.and_(*predicates))) + + def _delete_by_literal_keys( + self, + table_name: TableName, + key_exprs: t.List[exp.Expr], + subquery: t.Union[exp.Subquery, exp.Select], + extra_conditions: t.Optional[t.List[exp.Expr]] = None, + ) -> None: + """Fallback for DELETE on DUPLICATE KEY / AGGREGATE KEY / Iceberg tables that prohibits OR predicates.""" + select_query = ( + subquery.this.copy() if isinstance(subquery, exp.Subquery) else subquery.copy() + ) + if isinstance(select_query, exp.Select): + select_query = select_query.distinct() + rows = self.fetchall(select_query, quote_identifiers=True) + if not rows: + return + sanitized_rows = [tuple(self._sanitize_key_value(v) for v in row) for row in rows] + self._delete_key_rows( + table_name, + key_exprs, + sanitized_rows, + extra_conditions=extra_conditions, + is_ast_values=False, + ) + + def _delete_key_rows( + self, + table_name: TableName, + key_exprs: t.List[exp.Expr], + rows: t.Sequence[t.Tuple[t.Any, ...]], + extra_conditions: t.Optional[t.List[exp.Expr]] = None, + *, + is_ast_values: bool = False, + ) -> None: + if not rows: + return + + target_table = exp.to_table(table_name) + cleaned_extra_conditions: t.List[exp.Expr] = [] + for extra in extra_conditions or []: + cleaned_extra = extra.copy() + for col in cleaned_extra.find_all(exp.Column): + if not col.find_ancestor(exp.Subquery, exp.Select) and ( + col.table and col.table.lower() == target_table.name.lower() + ): + col.set("table", None) + col.set("db", None) + col.set("catalog", None) + cleaned_extra_conditions.append(cleaned_extra) + extra_conditions = cleaned_extra_conditions + cols = [exp.column(k.name, quoted=True) for k in key_exprs] + batch_size = self.DEFAULT_BATCH_SIZE + + 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] + predicates: t.List[exp.Expr] = [col.isin(*batch), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*predicates))) + + if has_null: + predicates = [col.is_(exp.null()), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*predicates))) + return + + # Composite key (N >= 2): dynamically select the column index to batch via IN (...) + # that minimizes the number of distinct prefix groups, breaking ties in favor of the last column. + batch_idx = min( + range(len(cols)), + key=lambda idx: (len({row[:idx] + row[idx + 1 :] for row in 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] + preds = [*prefix_preds, batch_col.isin(*batch), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*preds))) + + if has_null_batch: + preds = [*prefix_preds, batch_col.is_(exp.null()), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*preds))) + + @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(kwargs.get("table_properties") or {}) + unique_key_prop, duplicate_key_prop, aggregate_key_prop = self._resolve_table_type_and_keys( + table_properties + ) + + if ( + primary_key + and 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=kwargs.get("table_properties"), + ) + if ( + not is_iceberg + and target_columns_to_types + and columns_to_types_all_known(target_columns_to_types) + ): + target_columns_to_types = self._reorder_olap_columns(target_columns_to_types, kwargs) + if source_columns: + source_col_set = set(source_columns) + source_columns = [c for c in target_columns_to_types if c in source_col_set] + [ + c for c in source_columns if c not in target_columns_to_types + ] + super().replace_query( + table_name, + query_or_df, + target_columns_to_types=target_columns_to_types, + table_description=table_description, + column_descriptions=column_descriptions, + source_columns=source_columns, + supports_replace_table_override=supports_replace_table_override, + **kwargs, + ) + + def _scd_type_2( + self, + target_table: TableName, + source_table: QueryOrDF, + unique_key: t.Sequence[exp.Expr], + valid_from_col: exp.Column, + valid_to_col: exp.Column, + execution_time: t.Union[TimeLike, exp.Column], + invalidate_hard_deletes: bool = True, + updated_at_col: t.Optional[exp.Column] = None, + check_columns: t.Optional[t.Union[exp.Star, t.Sequence[exp.Expr]]] = None, + updated_at_as_valid_from: bool = False, + execution_time_as_valid_from: bool = False, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + truncate: bool = False, + source_columns: t.Optional[t.List[str]] = None, + **kwargs: t.Any, + ) -> None: + table = exp.to_table(target_table) + is_iceberg = self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=table.catalog, + table_properties=kwargs.get("table_properties"), + ) + if ( + not is_iceberg + and target_columns_to_types + and columns_to_types_all_known(target_columns_to_types) + ): + target_columns_to_types = self._reorder_olap_columns(target_columns_to_types, kwargs) + if source_columns: + source_col_set = set(source_columns) + source_columns = [c for c in target_columns_to_types if c in source_col_set] + [ + c for c in source_columns if c not in target_columns_to_types + ] + super()._scd_type_2( + target_table=target_table, + source_table=source_table, + unique_key=unique_key, + valid_from_col=valid_from_col, + valid_to_col=valid_to_col, + execution_time=execution_time, + invalidate_hard_deletes=invalidate_hard_deletes, + updated_at_col=updated_at_col, + check_columns=check_columns, + updated_at_as_valid_from=updated_at_as_valid_from, + execution_time_as_valid_from=execution_time_as_valid_from, + target_columns_to_types=target_columns_to_types, + table_description=table_description, + column_descriptions=column_descriptions, + truncate=truncate, + source_columns=source_columns, + **kwargs, + ) + + def _create_table_from_source_queries( + self, + table_name: TableName, + source_queries: t.List[SourceQuery], + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + exists: bool = True, + replace: bool = False, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + table_kind: t.Optional[str] = None, + track_rows_processed: bool = True, + **kwargs: t.Any, + ) -> None: + table = exp.to_table(table_name) + effective_catalog = table.catalog or self.default_catalog + is_iceberg = self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=table.catalog, + table_properties=kwargs.get("table_properties"), + ) + is_internal_iceberg = is_iceberg and ( + not effective_catalog or effective_catalog.lower() == "internal" + ) + + if not is_iceberg and target_columns_to_types: + target_columns_to_types = self._reorder_olap_columns(target_columns_to_types, kwargs) + + target_columns_to_types_known = bool( + target_columns_to_types and columns_to_types_all_known(target_columns_to_types) + ) + first_query = source_queries[0].query_factory() if source_queries else None + + 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, + ) + + def _build_create_table_exp( + self, + table_name_or_schema: t.Union[exp.Schema, TableName], + expression: t.Optional[exp.Expr], + exists: bool = True, + replace: bool = False, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + table_description: t.Optional[str] = None, + table_kind: t.Optional[str] = None, + **kwargs: t.Any, + ) -> exp.Create: + kwargs.setdefault("table_properties", None) + kwargs["is_schema_ddl"] = isinstance(table_name_or_schema, exp.Schema) + + catalog_name = None + if isinstance(table_name_or_schema, exp.Schema): + if isinstance(table_name_or_schema.this, exp.Table): + catalog_name = table_name_or_schema.this.catalog + else: + catalog_name = exp.to_table(table_name_or_schema).catalog + + is_iceberg = self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=catalog_name, + table_properties=kwargs.get("table_properties"), + ) + + if not is_iceberg and target_columns_to_types: + original_cols = dict(target_columns_to_types) + reordered_cols = self._reorder_olap_columns(target_columns_to_types, kwargs) + order_changed = list(reordered_cols) != list(original_cols) + types_changed = any(reordered_cols[c] != original_cols[c] for c in reordered_cols) + if order_changed or types_changed: + 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) + 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): + expression = expression.copy() + updated_projs: t.List[exp.Expr] = [] + for c, new_dtype in target_columns_to_types.items(): + proj = proj_by_name[c.lower()].copy() + if types_changed and original_cols[c] != new_dtype: + if isinstance(proj, exp.Alias) and isinstance( + proj.this, exp.Cast + ): + proj.this.set("to", new_dtype.copy()) + elif isinstance(proj, exp.Alias): + proj.set( + "this", + exp.cast(proj.this.copy(), to=new_dtype.copy()), + ) + elif isinstance(proj, exp.Column): + proj = exp.alias_( + exp.cast(proj.copy(), to=new_dtype.copy()), + proj.name, + quoted=bool(getattr(proj.this, "quoted", True)), + ) + updated_projs.append(proj) + expression.set("expressions", updated_projs) + reordered_in_place = True + if not reordered_in_place: + expression = self._order_projections_and_filter( + expression, target_columns_to_types, coerce_types=types_changed + ) + elif not is_iceberg and isinstance(expression, (exp.Select, exp.SetOperation)): + expression = self._reorder_unannotated_ctas_query(expression, kwargs) + + # Doris AUTO PARTITION BY RANGE requires the partition column(s) to be defined as NOT NULL. + if not is_iceberg: + table_props = kwargs.get("table_properties") or {} + auto_partition = table_props.get("auto_partition") + is_auto = bool( + auto_partition is not None + and str(self._extract_str_value(auto_partition) or "").strip().lower() + in ("true", "1") + ) + raw_part_by = ( + kwargs.get("partitioned_by") + or table_props.get("partitioned_by") + 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.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) -> exp.DataType: + inner = proj.this if isinstance(proj, exp.Alias) else proj + inner = inner.unnest() + if isinstance(inner, (exp.Cast, exp.TryCast)) and isinstance(inner.to, exp.DataType): + return inner.to.copy() + if self._is_integer_indicator_projection(proj): + return exp.DataType.build("INT", dialect=self.dialect) + if isinstance(inner, exp.Literal): + if inner.is_int: + return exp.DataType.build("INT", dialect=self.dialect) + if inner.is_number: + return exp.DataType.build("DOUBLE", dialect=self.dialect) + if inner.is_string: + return exp.DataType.build("VARCHAR(255)", dialect=self.dialect) + if isinstance(inner, exp.Boolean): + return exp.DataType.build("BOOLEAN", dialect=self.dialect) + if isinstance( + inner, + ( + exp.Avg, + exp.Stddev, + exp.StddevPop, + exp.StddevSamp, + exp.Variance, + exp.VariancePop, + ), + ): + return exp.DataType.build("DOUBLE", dialect=self.dialect) + if isinstance(inner, (exp.Array, exp.ArrayAgg)): + return exp.DataType.build("ARRAY", dialect=self.dialect) + if isinstance(inner, (exp.Map, exp.VarMap)): + return exp.DataType.build("MAP", dialect=self.dialect) + if isinstance(inner, exp.Struct): + return exp.DataType.build("STRUCT", dialect=self.dialect, udt=True) + if isinstance(inner, (exp.JSONObject, exp.JSONArray)): + return exp.DataType.build("JSON", dialect=self.dialect) + if isinstance(inner, exp.GroupConcat): + return exp.DataType.build("TEXT", dialect=self.dialect) + if isinstance(inner, (exp.Count, exp.RowNumber, exp.Rank, exp.DenseRank)): + return exp.DataType.build("BIGINT", dialect=self.dialect) + if isinstance( + inner, + ( + exp.Year, + exp.Month, + exp.Day, + exp.Hour, + exp.Minute, + exp.Second, + exp.DayOfWeek, + exp.DayOfYear, + exp.WeekOfYear, + ), + ): + return exp.DataType.build("INT", dialect=self.dialect) + if isinstance(inner, (exp.CurrentDate, exp.Date, exp.TsOrDsToDate)): + return exp.DataType.build("DATE", dialect=self.dialect) + if isinstance(inner, (exp.CurrentTimestamp, exp.DateTrunc, exp.TimestampTrunc)): + return exp.DataType.build("DATETIME", dialect=self.dialect) + if isinstance(inner, exp.Anonymous): + fn_upper = inner.name.upper() + if fn_upper in ( + "AVG", + "STDDEV", + "STDDEV_POP", + "STDDEV_SAMP", + "VARIANCE", + "VAR_POP", + "VAR_SAMP", + ): + return exp.DataType.build("DOUBLE", dialect=self.dialect) + if fn_upper in ("ARRAY", "ARRAY_AGG", "COLLECT_LIST", "COLLECT_SET"): + return exp.DataType.build("ARRAY", dialect=self.dialect) + if fn_upper in ("MAP", "MAP_AGG"): + return exp.DataType.build("MAP", dialect=self.dialect) + if fn_upper in ("STRUCT", "NAMED_STRUCT"): + return exp.DataType.build("STRUCT", dialect=self.dialect, udt=True) + if fn_upper in ("JSON_OBJECT", "JSON_ARRAY", "TO_JSON"): + return exp.DataType.build("JSON", dialect=self.dialect) + if fn_upper == "GROUP_CONCAT": + return exp.DataType.build("TEXT", dialect=self.dialect) + if fn_upper == "DATE_TRUNC": + return exp.DataType.build("DATETIME", dialect=self.dialect) + if fn_upper == "TO_DATE": + return exp.DataType.build("DATE", dialect=self.dialect) + return exp.DataType.build("UNKNOWN") + + def _reorder_unannotated_ctas_query( + self, expression: exp.Expr, kwargs: t.Dict[str, t.Any] + ) -> exp.Expr: + if isinstance(expression, exp.SetOperation): + expression = expression.copy() + branches: t.List[exp.Select] = [] + + def _collect_branches(node: exp.Expr) -> None: + if isinstance(node, exp.SetOperation): + _collect_branches(node.this) + _collect_branches(node.expression) + elif isinstance(node, exp.Select): + branches.append(node) + + _collect_branches(expression) + if not branches: + return expression + first_exprs = branches[0].expressions + num_cols = len(first_exprs) + if ( + num_cols > 0 + and all( + len(b.expressions) == num_cols + and not any(isinstance(p, exp.Star) for p in b.expressions) + for b in branches + ) + and len( + { + p.alias_or_name.lower() + for p in first_exprs + if p.alias_or_name and not isinstance(p, exp.Star) + } + ) + == num_cols + ): + orig_names = [p.alias_or_name for p in first_exprs] + merged_cols_to_types: t.Dict[str, exp.DataType] = {} + for col_idx, col_name in enumerate(orig_names): + inferred = exp.DataType.build("UNKNOWN") + for b in branches: + cand = self._infer_projection_type(b.expressions[col_idx]) + if cand.this != exp.DataType.Type.UNKNOWN: + inferred = cand + break + merged_cols_to_types[col_name] = inferred + reordered_select_cols = self._reorder_olap_columns(merged_cols_to_types, kwargs) + perm = [orig_names.index(c) for c in reordered_select_cols] + coercions = { + orig_names.index(c): reordered_select_cols[c] + for c in reordered_select_cols + if reordered_select_cols[c] != merged_cols_to_types[c] + } + if perm != list(range(num_cols)) or coercions: + for b_idx, branch in enumerate(branches): + old_projs = list(branch.expressions) + updated_branch_projs: t.List[exp.Expr] = [] + for orig_idx in perm: + proj = old_projs[orig_idx].copy() + new_dtype = coercions.get(orig_idx) + if new_dtype is not None: + col_name = orig_names[orig_idx] + if isinstance(proj, exp.Alias) and isinstance( + proj.this, (exp.Cast, exp.TryCast) + ): + proj.this.set("to", new_dtype.copy()) + elif isinstance(proj, (exp.Cast, exp.TryCast)): + cast_expr = exp.cast(proj.this.copy(), to=new_dtype.copy()) + proj = ( + exp.alias_(cast_expr, col_name, quoted=True) + if b_idx == 0 + else cast_expr + ) + elif isinstance(proj, exp.Alias): + proj.set( + "this", exp.cast(proj.this.copy(), to=new_dtype.copy()) + ) + elif isinstance(proj, exp.Column): + cast_expr = exp.cast(proj.copy(), to=new_dtype.copy()) + proj = ( + exp.alias_( + cast_expr, + proj.name, + quoted=bool(getattr(proj.this, "quoted", True)), + ) + if b_idx == 0 + else cast_expr + ) + updated_branch_projs.append(proj) + branch.set("expressions", updated_branch_projs) + return expression + expression.set("this", self._reorder_unannotated_ctas_query(expression.this, kwargs)) + expression.set( + "expression", self._reorder_unannotated_ctas_query(expression.expression, kwargs) + ) + return expression + if isinstance(expression, exp.Select): + proj_by_name = { + proj.alias_or_name.lower(): proj + for proj in expression.expressions + if proj.alias_or_name and not isinstance(proj, exp.Star) + } + if proj_by_name and len(proj_by_name) == len(expression.expressions): + select_cols_to_types = { + proj.alias_or_name: self._infer_projection_type(proj) + for proj in expression.expressions + } + 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: + expression = expression.copy() + updated_projs: t.List[exp.Expr] = [] + for c, new_dtype in reordered_select_cols.items(): + proj = proj_by_name[c.lower()].copy() + if types_changed and select_cols_to_types[c] != new_dtype: + if isinstance(proj, exp.Alias) and isinstance( + proj.this, (exp.Cast, exp.TryCast) + ): + proj.this.set("to", new_dtype.copy()) + elif isinstance(proj, (exp.Cast, exp.TryCast)): + proj = exp.alias_( + exp.cast(proj.this.copy(), to=new_dtype.copy()), + c, + quoted=True, + ) + elif isinstance(proj, exp.Alias): + proj.set("this", exp.cast(proj.this.copy(), to=new_dtype.copy())) + elif isinstance(proj, exp.Column): + proj = exp.alias_( + exp.cast(proj.copy(), to=new_dtype.copy()), + proj.name, + quoted=bool(getattr(proj.this, "quoted", True)), + ) + updated_projs.append(proj) + expression.set("expressions", updated_projs) + elif any(isinstance(proj, exp.Star) for proj in expression.expressions): + table_props = kwargs.get("table_properties") or {} + resolved_uk, resolved_dk, resolved_ak = self._resolve_table_type_and_keys( + 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 ( + not has_explicit_key + and isinstance(with_clause, exp.With) + and isinstance(from_clause, exp.From) + and isinstance(from_clause.this, exp.Table) + ): + target_cte_name = from_clause.this.name.lower() + for cte_idx, cte in enumerate(with_clause.expressions): + if ( + isinstance(cte, exp.CTE) + and cte.alias.lower() == target_cte_name + and isinstance(cte.this, exp.Select) + and len(cte.this.expressions) >= 2 + ): + indicator_projs = [ + p + for p in cte.this.expressions + if self._is_integer_indicator_projection(p) + ] + other_projs = [ + p + for p in cte.this.expressions + if not self._is_integer_indicator_projection(p) + ] + if ( + indicator_projs + and list(cte.this.expressions[: len(indicator_projs)]) + != indicator_projs + ): + expression = expression.copy() + copied_cte = expression.args["with_"].expressions[cte_idx] + copied_select = copied_cte.this + copied_select.set("expressions", indicator_projs + other_projs) + 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.Dict[str, t.Any], + ) -> t.Dict[str, exp.DataType]: + table_props = kwargs.get("table_properties") or {} + if table_props.get("primary_key") is not None or kwargs.get("primary_key") is not None: + raise SQLMeshError( + "Doris does not support 'primary_key' in physical_properties; use 'unique_key', 'duplicate_key', or 'aggregate_key' instead." + ) + if table_props.get("partition") is not None: + raise SQLMeshError( + "Invalid property 'partition' in physical_properties; did you mean 'partition_by', 'partitioned_by', or 'partitions'?" + ) + if ( + sum( + x is not None + for x in ( + kwargs.get("partitioned_by") or None, + table_props.get("partitioned_by"), + table_props.get("partition_by"), + ) + ) + > 1 + ): + raise SQLMeshError( + "Multiple partition definitions found: specify only one of model 'partitioned_by' or physical_properties 'partitioned_by' / 'partition_by'." + ) + 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): + valid_col = next( + ( + c + for c in col_names + if self._is_valid_olap_key_type(updated_types[c]) + and not updated_types[c].is_type("unknown", "null") + ), + None, + ) or next( + (c for c in col_names if self._is_valid_olap_key_type(updated_types[c])), + None, + ) + if valid_col: + col_names = [valid_col] + [c for c in col_names if c != valid_col] + else: + coercible_col = next( + (c for c in col_names if self._is_coercible_string_key_type(updated_types[c])), + None, + ) + if coercible_col: + updated_types[coercible_col] = exp.DataType.build( + "VARCHAR(255)", dialect=self.dialect + ) + if col_names[0] != coercible_col: + col_names = [coercible_col] + [c for c in col_names if c != coercible_col] + + # Ensure 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) + + 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"ORDER 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: + 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", + "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..50376667c5 --- /dev/null +++ b/tests/core/engine_adapter/integration/docker/compose.doris.yaml @@ -0,0 +1,57 @@ +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 = 2G" >> /opt/apache-doris/be/conf/be.conf + echo "storage_flood_stage_usage_percent = 99" >> /opt/apache-doris/be/conf/be.conf + echo "storage_flood_stage_left_capacity_bytes = 104857600" >> /opt/apache-doris/be/conf/be.conf + exec bash /opt/apache-doris/be/bin/start_be.sh --console + ports: + - "8040:8040" + - "9050:9050" + 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..aa60c60563 --- /dev/null +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -0,0 +1,532 @@ +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), + ] diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py new file mode 100644 index 0000000000..0879d9918b --- /dev/null +++ b/tests/core/engine_adapter/test_doris.py @@ -0,0 +1,3580 @@ +import typing as t + +import pytest +from pytest_mock.plugin import MockerFixture +from sqlglot import expressions as exp +from sqlglot import parse_one + +from sqlmesh.core.engine_adapter.doris import DorisEngineAdapter +from sqlmesh.core.engine_adapter.shared import CatalogSupport, DataObject, DataObjectType +from 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`) ORDER BY (`updated_at`, `id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10", + ] + + +def test_create_table_aggregate_key_and_clustered_by_buckets( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "metrics_agg", + target_columns_to_types={ + "total": exp.DataType.build("BIGINT"), + "tenant_id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATE"), + }, + clustered_by=[exp.to_column("tenant_id")], + table_properties={ + "aggregate_key": exp.Tuple( + expressions=[exp.to_column("tenant_id"), exp.to_column("ds")] + ), + "buckets": exp.Literal.number(16), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `metrics_agg` (`tenant_id` INT, `ds` DATE, `total` BIGINT) AGGREGATE KEY (`tenant_id`, `ds`) DISTRIBUTED BY HASH (`tenant_id`) BUCKETS 16", + ] + + +def test_create_table_auto_partition_and_unannotated_partitioned_by_guard( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + # Without partitions or auto_partition, OLAP partitioned_by is safely omitted + adapter.create_table( + "olap_no_partitions", + target_columns_to_types={"id": exp.DataType.build("INT"), "ds": exp.DataType.build("DATE")}, + partitioned_by=[exp.to_column("ds")], + ) + # With auto_partition = true, emits AUTO PARTITION BY RANGE (...) () 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) == [ + "DROP MATERIALIZED VIEW IF EXISTS `db`.`sync_mv` ON `db`.`orders`", + "CREATE MATERIALIZED VIEW `db`.`sync_mv` AS SELECT `region` AS `mv_region`, SUM(`amount`) AS `total` FROM `db`.`orders` GROUP BY `region`", + ] + + adapter.cursor.execute.reset_mock() + + adapter.create_view( + "sync_mv", + parse_one("SELECT id FROM db.orders", read="doris"), + replace=True, + materialized=True, + view_properties={ + "materialized_type": exp.Literal.string("SYNC"), + "source_table": exp.to_table("db.orders"), + }, + ) + + assert to_sql_calls(adapter) == [ + "USE `db`", + "DROP MATERIALIZED VIEW IF EXISTS `sync_mv` ON `db`.`orders`", + "USE `db`", + "CREATE MATERIALIZED VIEW `sync_mv` AS SELECT `id` AS `mv_id` FROM `db`.`orders`", + ] + + 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) + + adapter.delete_from( + "target_table", + where=parse_one( + "(ds BETWEEN '2024-01-01' AND '2024-01-07' OR status = 'expired') AND id NOT BETWEEN 10 AND 20", + read="doris", + ), + ) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `target_table` WHERE ((`ds` >= '2024-01-01' AND `ds` <= '2024-01-07') OR `status` = 'expired') AND NOT `id` BETWEEN 10 AND 20", + ] + + +@pytest.mark.parametrize( + "key_cols,rows", + [ + (["id"], [(1,), (None,), (2,), (3,)]), + (["id", "ds"], [(1, "2024-01-01"), (2, "2024-01-01"), (None, "2024-01-01"), (1, None)]), + ( + ["tenant", "region", "id"], + [ + ("t1", "us", 10), + ("t1", "us", 20), + ("t1", "us", None), + ("t1", None, 30), + (None, None, None), + ], + ), + ], +) +def test_delete_by_literal_keys_doris_delete_handler_invariants( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], + mocker: MockerFixture, + key_cols: t.List[str], + rows: t.List[t.Tuple[t.Any, ...]], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch.object(adapter, "fetchall", return_value=rows) + + adapter._delete_by_literal_keys( + exp.to_table("db.target_tbl"), + [exp.column(c) for c in key_cols], + parse_one(f"SELECT {', '.join(key_cols)} FROM db.src_tbl", read="doris", into=exp.Select), + ) + + sql_calls = to_sql_calls(adapter) + assert len(sql_calls) >= 1 + for sql in sql_calls: + parsed = parse_one(sql, read="doris") + assert isinstance(parsed, exp.Delete) + # Doris DeleteHandler forbids OR and BETWEEN in DELETE WHERE predicates + assert parsed.find(exp.Or) is None + assert parsed.find(exp.Between) is None + + +@pytest.mark.parametrize( + "non_key_type", + [ + "FLOAT", + "DOUBLE", + "ARRAY", + "MAP", + "STRUCT", + "JSON", + "VARIANT", + "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] == ( + f"DROP MATERIALIZED VIEW IF EXISTS `doris_db`.`orders_sync_mv` ON `{physical_tbl.db}`.`{physical_tbl.name}`" + ) + + +def test_dataframe_seed_all_string_coercion_and_batch_size( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + import pandas as pd + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + assert adapter.DEFAULT_BATCH_SIZE == 500 + + # 1. All-TEXT schema (e.g. from columns_to_types_from_df on a string CSV seed) coerces first column to VARCHAR(255) + adapter.create_table( + "db.country_seed", + target_columns_to_types={ + "country_code": exp.DataType.build("text"), + "country_name": exp.DataType.build("text"), + }, + ) + + # 2. Schema with DOUBLE + TEXT moves TEXT column to position 0 and coerces it to VARCHAR(255) + adapter.create_table( + "db.metric_seed", + target_columns_to_types={ + "score": exp.DataType.build("double"), + "metric_name": exp.DataType.build("text"), + }, + ) + + # 3. Explicit unique_key on a TEXT column coerces the key column to VARCHAR(255) + adapter.create_table( + "db.explicit_text_uk", + target_columns_to_types={ + "payload": exp.DataType.build("text"), + "code": exp.DataType.build("text"), + }, + table_properties={"unique_key": exp.to_column("code")}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`country_seed` (`country_code` VARCHAR(255), `country_name` STRING)", + "CREATE TABLE IF NOT EXISTS `db`.`metric_seed` (`metric_name` VARCHAR(255), `score` DOUBLE)", + "CREATE TABLE IF NOT EXISTS `db`.`explicit_text_uk` (`code` VARCHAR(255), `payload` STRING) UNIQUE KEY (`code`) DISTRIBUTED BY HASH (`code`) BUCKETS 10", + ] + + # 4. Explicit unique_key on a non-coercible invalid key type (DOUBLE) raises SQLMeshError + with pytest.raises(SQLMeshError, match="not supported as an OLAP key column"): + adapter.create_table( + "db.bad_double_uk", + target_columns_to_types={ + "score": exp.DataType.build("DOUBLE", dialect="doris"), + "name": exp.DataType.build("VARCHAR(32)", dialect="doris"), + }, + table_properties={"unique_key": exp.to_column("score")}, + ) + + # 5. DataFrame CTAS with pretty_sql=True and 500 rows succeeds without RecursionError + pretty_adapter = make_mocked_engine_adapter(DorisEngineAdapter, pretty_sql=True) + df = pd.DataFrame({"code": [f"c_{i}" for i in range(500)], "val": list(range(500))}) + pretty_adapter.ctas("db.df_seed", df) + df_sql_calls = to_sql_calls(pretty_adapter) + assert len(df_sql_calls) == 1 + assert "CREATE TABLE IF NOT EXISTS `db`.`df_seed`" in df_sql_calls[0] + + +def test_iceberg_bucket_and_truncate_partition_transforms( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + from sqlmesh.core.dialect import parse + from sqlmesh.core.model import SqlModel, load_sql_based_model + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + iceberg_model = t.cast( + SqlModel, + load_sql_based_model( + parse( + """ + MODEL ( + name iceberg_catalog.silver.curated_events, + kind FULL, + dialect doris, + table_format iceberg, + partitioned_by (day(event_ts), bucket(16, user_id), truncate(10, region)) + ); + SELECT + user_id::BIGINT AS user_id, + event_ts::DATETIME AS event_ts, + region::STRING AS region + FROM iceberg_catalog.bronze.raw_events; + """, + default_dialect="doris", + ) + ), + ) + + adapter.create_table( + iceberg_model.name, + target_columns_to_types=iceberg_model.columns_to_types or {}, + table_format=iceberg_model.table_format, + partitioned_by=iceberg_model.partitioned_by, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `iceberg_catalog`.`silver`.`curated_events` (`user_id` BIGINT, `event_ts` DATETIME, `region` STRING) ENGINE=iceberg PARTITION BY LIST (DAY(`event_ts`), BUCKET(16, `user_id`), TRUNCATE(10, `region`)) ()" + ] + + +def test_ctas_auto_partition_non_nullable_and_union_key_reordering( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. CTAS with AUTO PARTITION BY RANGE wraps partition column projection in NON_NULLABLE(...) + adapter.ctas( + "db.ctas_auto_part", + parse_one("SELECT payload, id, ds FROM db.src", read="doris"), + target_columns_to_types={ + "payload": exp.DataType.build("STRING", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + "ds": exp.DataType.build("DATE", dialect="doris"), + }, + partitioned_by=[parse_one("RANGE(DATE_TRUNC(ds, 'DAY'))", read="doris")], + table_properties={"auto_partition": exp.Literal.string("true")}, + ) + + # 2. Unannotated CTAS with UNION ALL reorders projections in all branches for unique_key prefix + adapter.ctas( + "db.ctas_union_reordered", + parse_one( + "SELECT name, id FROM db.src1 UNION ALL SELECT name, id FROM db.src2", read="doris" + ), + table_properties={"unique_key": exp.to_column("id")}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`ctas_auto_part` AUTO PARTITION BY RANGE (DATE_TRUNC(`ds`, 'DAY')) () AS SELECT CAST(`id` AS INT) AS `id`, NON_NULLABLE(CAST(`ds` AS DATE)) AS `ds`, CAST(`payload` AS STRING) AS `payload` FROM (SELECT `payload`, `id`, `ds` FROM `db`.`src`) AS `_subquery`", + "CREATE TABLE IF NOT EXISTS `db`.`ctas_union_reordered` UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10 AS SELECT `id`, `name` FROM `db`.`src1` UNION ALL SELECT `id`, `name` FROM `db`.`src2`", + ] + + +def test_delete_tautology_and_union_subquery_and_dedup_partition_overwrite( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. Literal tautology 1=1 routes to TRUNCATE TABLE; contradiction 1=0 is a no-op + adapter.delete_from("db.tbl", where="1 = 1") + adapter.delete_from("db.tbl", where="1 = 0") + + # 2. IN subquery with UNION resolves the projected column name from the first SELECT branch + adapter.delete_from( + "db.tbl", + where="id IN (SELECT src_id FROM db.a UNION SELECT src_id FROM db.b)", + ) + + # 3. IN subquery with unaliased expression aliases the projection for _t2 join + adapter.delete_from( + "db.tbl", + where="id IN (SELECT other_id + 1 FROM db.a)", + ) + + assert to_sql_calls(adapter) == [ + "TRUNCATE TABLE `db`.`tbl`", + "DELETE FROM `db`.`tbl` AS `_t1` USING (SELECT `src_id` FROM `db`.`a` UNION SELECT `src_id` FROM `db`.`b`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`src_id`", + "DELETE FROM `db`.`tbl` AS `_t1` USING (SELECT `other_id` + 1 AS `_col_0` FROM `db`.`a`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`_col_0`", + ] + + # 4. insert_overwrite_by_partition deduplicates 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` (`id` COMMENT 'Primary user ID', `name`) COMMENT 'Users view' AS SELECT `id`, `name` FROM `db`.`users`", + "CREATE MATERIALIZED VIEW `db`.`unannotated_mv` (`id` COMMENT 'Primary user ID', `name`) COMMENT 'Users MV' AS SELECT `id`, `name` FROM `db`.`users`", + "CREATE TABLE IF NOT EXISTS `db`.`paren_keys_tbl` (`id` INT, `dt` DATE, `name` VARCHAR(50)) UNIQUE KEY (`id`, `dt`) ORDER BY (`dt`, `id`) DISTRIBUTED BY HASH (`id`, `dt`) BUCKETS 12 PROPERTIES ('excluded_trigger_tables'='db.t1,db.t2')", + ] + + # 3. Fail-fast validation for unsupported primary_key, typo `partition`, conflicting partition definitions, + # and invalid RANDOM(col) / column-less HASH distribution. + cols = {"id": exp.DataType.build("INT"), "dt": exp.DataType.build("DATE")} + + with pytest.raises(SQLMeshError, match="does not support 'primary_key'"): + adapter.create_table( + "db.bad_pk", + target_columns_to_types=cols, + table_properties={"primary_key": exp.to_column("id")}, + ) + + with pytest.raises(SQLMeshError, match="Invalid property 'partition'"): + adapter.create_table( + "db.bad_partition_typo", + target_columns_to_types=cols, + table_properties={"partition": exp.to_column("dt")}, + ) + + with pytest.raises(SQLMeshError, match="Multiple partition definitions"): + adapter.create_table( + "db.conflict_part", + target_columns_to_types=cols, + partitioned_by=[exp.to_column("dt")], + table_properties={"partition_by": exp.to_column("dt")}, + ) + + with pytest.raises( + SQLMeshError, match="DISTRIBUTED BY RANDOM does not accept column expressions" + ): + adapter.create_table( + "db.bad_random_cols", + target_columns_to_types=cols, + table_properties={"distributed_by": exp.Literal.string("RANDOM(id)")}, + ) + + with pytest.raises( + SQLMeshError, match="DISTRIBUTED BY HASH requires at least one column expression" + ): + adapter.create_table( + "db.bad_hash_no_cols", + target_columns_to_types=cols, + table_properties={"distributed_by": exp.Literal.string("HASH BUCKETS 10")}, + ) + + +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) diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py index c506d401d9..800fb5fb19 100644 --- a/tests/core/test_connection_config.py +++ b/tests/core/test_connection_config.py @@ -15,6 +15,7 @@ ClickhouseConnectionConfig, ConnectionConfig, DatabricksConnectionConfig, + DorisConnectionConfig, DuckDBAttachOptions, DuckDBConnectionConfig, FabricConnectionConfig, @@ -2479,3 +2480,116 @@ def test_starrocks(make_config): assert advanced_config.charset == "utf8mb4" assert advanced_config.ssl_disabled is True assert advanced_config.concurrent_tasks == 10 + + +def test_doris(make_config): + """Test DorisConnectionConfig basic functionality and multi-catalog/session options""" + # Basic configuration + config = make_config( + type="doris", + host="localhost", + user="root", + password="password", + port=9030, + database="demo", + check_import=False, + ) + assert isinstance(config, DorisConnectionConfig) + assert config.type_ == "doris" + assert config.host == "localhost" + assert config.user == "root" + assert config.password == "password" + assert config.port == 9030 + assert config.database == "demo" + assert config.catalog is None + assert config.get_catalog() is None + assert config.DIALECT == "doris" + assert config.DISPLAY_NAME == "Apache Doris" + assert config.DISPLAY_ORDER == 20 + assert config.is_recommended_for_state_sync is False + assert config.is_forbidden_for_state_sync is True + assert config._cursor_init is None + + # Test with minimal configuration (using default port and optional password) + minimal_config = make_config( + type="doris", + host="fe.doris.cluster", + user="doris_user", + check_import=False, + ) + assert isinstance(minimal_config, DorisConnectionConfig) + assert minimal_config.port == 9030 # Default Doris FE port + assert minimal_config.host == "fe.doris.cluster" + assert minimal_config.user == "doris_user" + assert minimal_config.password is None + + # Test with Iceberg catalog, session variables, and TLS options + advanced_config = make_config( + type="doris", + host="doris-fe", + user="admin", + password="admin123", + port=9030, + database="analytics", + catalog="iceberg_cat", + charset="utf8mb4", + ssl_ca="/path/to/ca.pem", + ssl_verify_cert=True, + session_variables={"query_timeout": 300, "workload_group": "etl_group"}, + concurrent_tasks=10, + check_import=False, + ) + assert isinstance(advanced_config, DorisConnectionConfig) + assert advanced_config.charset == "utf8mb4" + assert advanced_config.catalog == "iceberg_cat" + assert advanced_config.get_catalog() == "iceberg_cat" + assert advanced_config.concurrent_tasks == 10 + assert "database" not in advanced_config._static_connection_kwargs + assert advanced_config._static_connection_kwargs["ssl"] == {"ca": "/path/to/ca.pem"} + assert advanced_config._static_connection_kwargs["ssl_verify_cert"] is True + + mock_cursor = MagicMock() + assert advanced_config._cursor_init is not None + advanced_config._cursor_init(mock_cursor) + assert [call.args[0] for call in mock_cursor.execute.call_args_list] == [ + "SET `query_timeout` = 300", + "SET `workload_group` = 'etl_group'", + "SWITCH `iceberg_cat`", + "USE `analytics`", + ] + + ssl_bool_config = make_config( + type="doris", + host="doris-fe", + user="admin", + ssl=True, + check_import=False, + ) + assert ssl_bool_config._static_connection_kwargs["ssl"] == {} + + ssl_false_config = make_config( + type="doris", + host="doris-fe", + user="admin", + ssl=False, + check_import=False, + ) + assert "ssl" not in ssl_false_config._static_connection_kwargs + + with pytest.raises(ConfigError, match="Invalid Doris session variable name"): + make_config( + type="doris", + host="doris-fe", + user="admin", + session_variables={"bad key; DROP": 1}, + check_import=False, + ) + + with pytest.raises(ConfigError, match="Invalid Doris session variable value"): + make_config( + type="doris", + host="doris-fe", + user="admin", + session_variables={"timeout": float("nan")}, + check_import=False, + )