diff --git a/.github/scripts/install-prerequisites.sh b/.github/scripts/install-prerequisites.sh index 6997633a31..a955f0c47c 100755 --- a/.github/scripts/install-prerequisites.sh +++ b/.github/scripts/install-prerequisites.sh @@ -28,6 +28,10 @@ elif [ "$ENGINE" == "fabric" ]; then rm packages-microsoft-prod.deb ENGINE_DEPENDENCIES="msodbcsql18" +elif [ "$ENGINE" == "doris" ]; then + ENGINE_DEPENDENCIES="mysql-client" + sudo sysctl -w vm.max_map_count=2000000 + sudo swapoff -a || true fi ALL_DEPENDENCIES="$COMMON_DEPENDENCIES $ENGINE_DEPENDENCIES" diff --git a/.github/scripts/wait-for-db.sh b/.github/scripts/wait-for-db.sh index 4a076f31f8..33105ffd5b 100755 --- a/.github/scripts/wait-for-db.sh +++ b/.github/scripts/wait-for-db.sh @@ -104,6 +104,35 @@ db2_ready() { done } +doris_ready() { + probe_port 9030 || return 1 + HOSTNAME=${DOCKER_HOSTNAME:-127.0.0.1} + if [ "$HOSTNAME" = "localhost" ]; then + HOSTNAME="127.0.0.1" + fi + INIT_SQL="ADMIN SET FRONTEND CONFIG ('force_olap_table_replication_num' = '1');" + if docker ps --format '{{.Names}}' 2>/dev/null | grep -qx 'doris-fe-01'; then + docker exec -i doris-fe-01 mysql -h127.0.0.1 -P9030 -uroot -e "SHOW BACKENDS\G" 2>/dev/null | grep -q "Alive: true" && \ + docker exec -i doris-fe-01 mysql -h127.0.0.1 -P9030 -uroot -e "$INIT_SQL" 2>/dev/null + elif command -v mysql > /dev/null 2>&1; then + mysql -h "$HOSTNAME" -P 9030 -u root -e "SHOW BACKENDS\G" 2>/dev/null | grep -q "Alive: true" && \ + mysql -h "$HOSTNAME" -P 9030 -u root -e "$INIT_SQL" 2>/dev/null + else + python3 -c " +import pymysql, sys +conn = pymysql.connect(host='$HOSTNAME', port=9030, user='root', autocommit=True) +with conn.cursor() as cur: + cur.execute('SHOW BACKENDS') + rows = cur.fetchall() + cols = [d[0] for d in cur.description] + alive_idx = cols.index('Alive') + if not rows or not any(str(r[alive_idx]).lower() == 'true' for r in rows): + sys.exit(1) + cur.execute(\"ADMIN SET FRONTEND CONFIG ('force_olap_table_replication_num' = '1')\") +" 2>/dev/null + fi +} + echo "Waiting for $ENGINE to be ready..." READINESS_FUNC="${ENGINE}_ready" diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 2c18fe1e04..75a15f2b88 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -269,6 +269,7 @@ jobs: clickhouse, risingwave, starrocks, + doris, ] env: PYTEST_XDIST_AUTO_NUM_WORKERS: 2 diff --git a/.readthedocs.yaml b/.readthedocs.yaml index ee4794538f..09262c1577 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -6,7 +6,7 @@ build: python: "3.10" jobs: pre_build: - - pip install -e ".[athena,azuresql,bigframes,bigquery,clickhouse,databricks,dbt,dlt,gcppostgres,github,llm,mssql,mysql,mwaa,postgres,redshift,slack,snowflake,starrocks,trino,web,risingwave]" + - pip install -e ".[athena,azuresql,bigframes,bigquery,clickhouse,databricks,dbt,dlt,doris,gcppostgres,github,llm,mssql,mysql,mwaa,postgres,redshift,slack,snowflake,starrocks,trino,web,risingwave]" - make api-docs mkdocs: diff --git a/Makefile b/Makefile index 0d8571b33f..b1ce245486 100644 --- a/Makefile +++ b/Makefile @@ -226,6 +226,9 @@ starrocks-test: engine-starrocks-up db2-test: engine-db2-up pytest -n auto -m "db2" --reruns 3 --junitxml=test-results/junit-db2.xml +doris-test: engine-doris-up + pytest -n auto -m "doris" --reruns 3 --junitxml=test-results/junit-doris.xml + ################# # Cloud Engines # ################# diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md index 38f6814c8b..656478d8d6 100644 --- a/docs/guides/configuration.md +++ b/docs/guides/configuration.md @@ -933,6 +933,7 @@ These pages describe the connection configuration options for each execution eng * [BigQuery](../integrations/engines/bigquery.md) * [ClickHouse](../integrations/engines/clickhouse.md) * [Databricks](../integrations/engines/databricks.md) +* [Doris](../integrations/engines/doris.md) * [DuckDB](../integrations/engines/duckdb.md) * [Fabric](../integrations/engines/fabric.md) * [MotherDuck](../integrations/engines/motherduck.md) @@ -978,6 +979,7 @@ Unsupported state engines, even for development: * [Spark](../integrations/engines/spark.md) * [StarRocks](../integrations/engines/starrocks.md) * [Trino](../integrations/engines/trino.md) +* [Doris](../integrations/engines/doris.md) This example gateway configuration uses Snowflake for the data warehouse connection and Postgres for the state backend connection: diff --git a/docs/guides/connections.md b/docs/guides/connections.md index 1345905bf7..bf6d4811d6 100644 --- a/docs/guides/connections.md +++ b/docs/guides/connections.md @@ -85,6 +85,7 @@ default_gateway: local_db * [ClickHouse](../integrations/engines/clickhouse.md) * [Databricks](../integrations/engines/databricks.md) * [Db2](../integrations/engines/db2.md) +* [Doris](../integrations/engines/doris.md) * [DuckDB](../integrations/engines/duckdb.md) * [Fabric](../integrations/engines/fabric.md) * [MotherDuck](../integrations/engines/motherduck.md) diff --git a/docs/installation.md b/docs/installation.md index 6bcac928e5..221a40fcce 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -54,6 +54,7 @@ Other extras are required to use specific SQL engines, like Bigquery or Postgres | Bigquery | `pip install "sqlmesh[bigquery]"` | | ClickHouse | `pip install "sqlmesh[clickhouse]"` | | Databricks | `pip install "sqlmesh[databricks]"` | + | Doris | `pip install "sqlmesh[doris]"` | | Fabric | `pip install "sqlmesh[fabric]"` | | GCP Postgres | `pip install "sqlmesh[gcppostgres]"` | | MS SQL Server | `pip install "sqlmesh[mssql]"` | diff --git a/docs/integrations/engines/doris.md b/docs/integrations/engines/doris.md new file mode 100644 index 0000000000..86def97b47 --- /dev/null +++ b/docs/integrations/engines/doris.md @@ -0,0 +1,489 @@ +# Apache Doris + +This page describes SQLMesh support for the [Apache Doris](https://doris.apache.org/) engine (versions 2.1, 3.x, 4.x, and 5.0), including native Doris OLAP tables and Apache Iceberg multi-catalog lakehouse tables. + +!!! note + Doris may not be used for the SQLMesh [state connection](../../reference/configuration.md#connections). Configure a transactional database such as Postgres, MySQL, or DuckDB (for local development) as your `state_connection`. + +## Local/Built-in Scheduler + +**Engine Adapter Type**: `doris` + +### Installation + +```bash +pip install "sqlmesh[doris]" +``` + +### Connection options + +| Option | Description | Type | Required | +| --- | --- | :---: | :---: | +| `type` | Engine type name - must be `doris` | string | Y | +| `host` | The hostname of the Doris FE (Frontend) server | string | Y | +| `user` | The username to use for authentication with Doris FE | string | Y | +| `password` | The password to use for authentication with Doris FE (Default: `None`) | string | N | +| `port` | The query port number of the Doris FE server (Default: `9030`) | int | N | +| `database` | The optional default database (schema) name | string | N | +| `catalog` | The optional default catalog name (e.g., `internal` or an external `iceberg` catalog) | string | N | +| `charset` | The character set used for the connection (Default: `utf8mb4`) | string | N | +| `collation` | The collation used for the connection | string | N | +| `connect_timeout` | Connection timeout in seconds (Default: `10`) | int | N | +| `read_timeout` | Read timeout in seconds | int | N | +| `write_timeout` | Write timeout in seconds | int | N | +| `session_variables` | Session variables to set on each new connection via `SET = ` (e.g., `workload_group`, `query_timeout`) | dict | N | +| `ssl` | Explicit SSL configuration dictionary passed to `PyMySQL` | dict | N | +| `ssl_disabled` | Whether SSL is disabled | bool | N | +| `ssl_ca` | Path to the file that contains a PEM-formatted CA certificate | string | N | +| `ssl_cert` | Path to the file that contains a PEM-formatted client certificate | string | N | +| `ssl_key` | Path to the file that contains a PEM-formatted private key for the client certificate | string | N | +| `ssl_verify_cert` | Set to `true` to verify the server's certificate against the CA certificate | bool | N | +| `ssl_verify_identity` | Set to `true` to verify the server's identity in addition to the certificate | bool | N | + +### Connection configuration example + +```yaml +gateways: + doris: + connection: + type: doris + host: fe.doris.cluster + port: 9030 + user: doris_user + password: {{ env_var('DORIS_PASSWORD') }} + database: analytics + catalog: internal + connect_timeout: 30 + session_variables: + workload_group: etl_group + query_timeout: 3600 + state_connection: + type: postgres + host: state-db.internal + port: 5432 + user: sqlmesh + password: {{ env_var('POSTGRES_PASSWORD') }} + database: sqlmesh_state +``` + +--- + +## OLAP Table Models + +Apache Doris supports three native OLAP table data models: **DUPLICATE KEY**, **UNIQUE KEY**, and **AGGREGATE KEY**. + +!!! tip "Automatic Key Column Prefix Reordering" + Doris requires key columns (`UNIQUE KEY`, `DUPLICATE KEY`, or `AGGREGATE KEY`) to appear as the **first columns** in a `CREATE TABLE` definition in the exact order of the key. SQLMesh's `DorisEngineAdapter` automatically reorders your schema definition during `CREATE TABLE` so key columns appear first, regardless of the column order in your `SELECT` query. + +### DUPLICATE KEY model (default) + +By default, OLAP tables use the `DUPLICATE KEY` model. If no key is specified, Doris uses the first column as the duplicate key and hash distribution column. + +```sql +MODEL ( + name analytics.user_events, + kind FULL, + clustered_by (user_id), + physical_properties ( + duplicate_key = ('user_id', 'event_time'), + buckets = 10 + ) +); + +SELECT + user_id::BIGINT AS user_id, + event_time::DATETIME AS event_time, + event_type::VARCHAR(64) AS event_type +FROM raw.events; +``` + +### UNIQUE KEY model (Merge-on-Write & custom sort keys) + +For `INCREMENTAL_BY_UNIQUE_KEY` models, SQLMesh automatically configures the model's `unique_key` as the Doris table's `UNIQUE KEY` and `DISTRIBUTED BY HASH(...)` key—no extra `physical_properties` required: + +```sql +MODEL ( + name analytics.orders_upsert, + kind INCREMENTAL_BY_UNIQUE_KEY ( + unique_key order_id + ), + physical_properties ( + order_by = ('updated_at', 'order_id'), + enable_unique_key_merge_on_write = 'true' + ) +); + +SELECT + order_id::BIGINT AS order_id, + customer_id::BIGINT AS customer_id, + status::VARCHAR(32) AS status, + updated_at::DATETIME AS updated_at +FROM raw.orders +WHERE updated_at BETWEEN @start_ds AND @end_ds; +``` + +You can also specify `unique_key` explicitly in `physical_properties` for `FULL` or `INCREMENTAL_BY_TIME_RANGE` models, and optionally decouple the sort key from the unique key using `order_by` (creating a `SEQUENCE MAP COL(col1, col2)` sort key on Doris 2.1+ Merge-on-Write tables). + +### AGGREGATE KEY model + +For pre-aggregated rollup tables, specify the key columns using `aggregate_key` in `physical_properties`: + +```sql +MODEL ( + name analytics.daily_site_metrics, + kind FULL, + clustered_by (site_id), + physical_properties ( + aggregate_key = ('site_id', 'event_date'), + buckets = 16 + ) +); + +SELECT + site_id::INT AS site_id, + event_date::DATE AS event_date, + SUM(pv)::BIGINT AS total_pv, + MAX(dwell_seconds)::INT AS max_dwell +FROM analytics.user_events +GROUP BY site_id, event_date; +``` + +--- + +## Distribution & Bucketing + +Doris requires a `DISTRIBUTED BY` clause on every OLAP table. SQLMesh supports three ways to configure distribution: + +1. **SQLMesh `clustered_by` + optional `buckets`** (recommended for cross-engine portability): + ```sql + MODEL ( + name analytics.events, + kind FULL, + clustered_by (user_id), + physical_properties ( + buckets = 16 + ) + ); + ``` +2. **Shorthand string `distributed_by`**: + ```sql + physical_properties ( + distributed_by = 'HASH(user_id) BUCKETS 16' + -- or 'RANDOM BUCKETS AUTO' + ) + ``` +3. **Structured dictionary `distributed_by`**: + ```sql + physical_properties ( + distributed_by = ( + kind = 'HASH', + expressions = 'user_id', + buckets = 'AUTO' + ) + ) + ``` + +If none of `clustered_by` or `distributed_by` is provided: +- When `unique_key` is present, SQLMesh defaults to `DISTRIBUTED BY HASH() BUCKETS `. +- When `buckets` is specified with `duplicate_key` or `aggregate_key`, SQLMesh emits `DISTRIBUTED BY HASH() BUCKETS ` (or `DISTRIBUTED BY RANDOM BUCKETS ` when no key is specified). + +--- + +## Partitioning + +Doris OLAP tables support **Multi-Range Batch Partitioning**, **Explicit Partition Definitions**, **Dynamic Partitioning**, and **Auto Partitioning** (Doris 2.1+). + +### Multi-range & explicit partitions + +```sql +MODEL ( + name analytics.partitioned_events, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column event_date + ), + partitioned_by RANGE(event_date), + physical_properties ( + partitions = 'FROM ("2023-01-01") TO ("2027-01-01") INTERVAL 1 MONTH' + ) +); +``` + +### Auto partitioning (`auto_partition = true`) + +In Doris 2.1+, `AUTO PARTITION` automatically creates partitions on demand as new data arrives: + +```sql +MODEL ( + name analytics.auto_partitioned_events, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column event_time + ), + partitioned_by RANGE(DATE_TRUNC(event_time, 'day')), + physical_properties ( + auto_partition = 'true' + ) +); +``` + +--- + +## Apache Iceberg & Multi-Catalog Lakehouse Integration + +Apache Doris 2.1–5.0 acts as a federated lakehouse engine that can both query and write to external **Apache Iceberg** catalogs (REST/Polaris, AWS Glue, Hive Metastore, Hadoop/S3, DLF). + +### Pattern 1: Reading Iceberg Bronze tables into Doris OLAP Silver/Gold models + +Because `DorisEngineAdapter` supports multi-catalog 3-part naming (`catalog.db.table`), your SQLMesh models can query external Iceberg tables directly while materializing sub-second serving tables in Doris's `internal` OLAP catalog: + +```sql +MODEL ( + name silver.user_sessions, + kind INCREMENTAL_BY_TIME_RANGE ( + time_column session_date + ) +); + +SELECT + user_id, + session_id, + CAST(event_ts AS DATE) AS session_date +FROM iceberg_catalog.bronze.clickstream +WHERE event_ts BETWEEN @start_ts AND @end_ts; +``` + +### Pattern 2: Materializing Apache Iceberg tables via Doris + +SQLMesh can also create, partition, and overwrite Apache Iceberg tables directly through Doris (`ENGINE=iceberg` and `INSERT OVERWRITE TABLE ... VALUES (*)`). + +!!! important "Virtual Environments & Catalog Bootstrap on External Iceberg Catalogs" + Doris only supports `CREATE VIEW` inside the `internal` catalog; external Iceberg catalogs do not support Doris views, and Doris does not support `environment_suffix_target: catalog`. + When materializing models directly inside an external Iceberg catalog: + + - Set `virtual_environment_mode: dev_only` (and keep `environment_suffix_target: schema` or `table`) in `config.yaml` so production tables are materialized directly in the Iceberg catalog without virtual views. + - Configure `environment_catalog_mapping` in `config.yaml` to route non-`prod` virtual environment views to the `internal` catalog (otherwise `DorisEngineAdapter` will log a warning and skip view creation in external catalogs during `dev` promotion): + ```yaml + virtual_environment_mode: dev_only + environment_suffix_target: schema + environment_catalog_mapping: + "^(?!prod$).*": internal + before_all: + - > + CREATE CATALOG IF NOT EXISTS iceberg_catalog PROPERTIES ( + 'type' = 'iceberg', + 'iceberg.catalog.type' = 'rest', + 'uri' = 'http://rest-catalog:8181', + 'warehouse' = 's3://lakehouse/warehouse' + ) + ``` + +```sql +MODEL ( + name iceberg_catalog.silver.curated_events, + kind FULL, + table_format iceberg, + partitioned_by (day(event_ts), bucket(16, user_id)), + physical_properties ( + 'write-format' = 'parquet', + 'compression-codec' = 'zstd' + ) +); + +SELECT + user_id::BIGINT AS user_id, + event_ts::DATETIME AS event_ts, + payload::STRING AS payload +FROM iceberg_catalog.bronze.raw_events; +``` + +--- + +## DDL-Only / Schema Contract Management for Streaming Tables (`doris_ddl_only`) + +In real-time architectures where **Apache Flink (`flink-doris-connector`)**, **Kafka Routine Load**, **Stream Load**, or **Spark** streams data continuously into a Doris OLAP or Iceberg table, you may want SQLMesh to manage **only the table DDL (`CREATE TABLE`) and schema evolution (`ALTER TABLE`)** without executing `INSERT` or `REPLACE` statements during `sqlmesh run`. + +You can achieve this by adding a lightweight `CustomMaterialization` plugin (`materializations/doris_ddl_only.py`) to your SQLMesh project: + +```python +from __future__ import annotations + +import typing as t +from sqlmesh.core.model import Model +from sqlmesh.core.snapshot.evaluator import CustomMaterialization +from sqlmesh.utils.errors import SQLMeshError + +if t.TYPE_CHECKING: + from sqlmesh.core.engine_adapter._typing import QueryOrDF + + +class DorisDDLOnlyMaterialization(CustomMaterialization): + """Manages CREATE TABLE and ALTER TABLE schema migrations without inserting data.""" + + NAME = "doris_ddl_only" + + def create( + self, + table_name: str, + model: Model, + is_table_deployable: bool, + render_kwargs: t.Dict[str, t.Any], + skip_grants: bool = False, + **kwargs: t.Any, + ) -> None: + if not model.annotated: + raise SQLMeshError( + f"Model '{model.name}' uses 'doris_ddl_only' and must explicitly declare `columns (...)` in MODEL (...)." + ) + super().create( + table_name=table_name, + model=model, + is_table_deployable=is_table_deployable, + render_kwargs=render_kwargs, + skip_grants=skip_grants, + **kwargs, + ) + + def insert( + self, + table_name: str, + query_or_df: QueryOrDF, + model: Model, + is_first_insert: bool, + render_kwargs: t.Dict[str, t.Any], + **kwargs: t.Any, + ) -> None: + # No-op: external Flink / Routine Load / Spark pipelines write data to this table. + return + + def append( + self, + table_name: str, + query_or_df: QueryOrDF, + model: Model, + render_kwargs: t.Dict[str, t.Any], + **kwargs: t.Any, + ) -> None: + # No-op: external Flink / Routine Load / Spark pipelines write data to this table. + return +``` + +Then define your streaming target table in `models/` with `virtual_environment_mode: dev_only`: + +```sql +MODEL ( + name realtime.flink_order_events, + kind CUSTOM ( + materialization 'doris_ddl_only', + forward_only true, + disable_restatement true + ), + columns ( + order_id BIGINT, + user_id BIGINT, + amount DECIMAL(18, 2), + updated_at DATETIME + ), + physical_properties ( + unique_key = ('order_id'), + order_by = ('updated_at', 'order_id'), + buckets = 16, + enable_unique_key_merge_on_write = 'true', + enable_light_schema_change = 'true' + ) +); + +SELECT + CAST(NULL AS BIGINT) AS order_id, + CAST(NULL AS BIGINT) AS user_id, + CAST(NULL AS DECIMAL(18, 2)) AS amount, + CAST(NULL AS DATETIME) AS updated_at; +``` + +!!! tip "Zero-Downtime Schema Evolution with Flink / Routine Load" + Setting `forward_only true` and `disable_restatement true` on a `doris_ddl_only` model ensures that schema additions (`ALTER TABLE ... ADD COLUMN`) are applied in-place on the existing table (`enable_light_schema_change = 'true'`) rather than creating a new empty physical table, and prevents accidental restatement wipes. + +--- + +## Compute-Storage Decoupled Mode (Doris 3.0 / 4.x / 5.0) + +For Doris 3.0+ Compute-Storage Decoupled deployments on object storage (S3, GCS, Azure Blob, OSS) or HDFS: + +- Pass `storage_vault_name` in `physical_properties` (or `model_defaults.physical_properties` in `config.yaml`) to bind tables to a specific storage vault. +- Use `session_variables` in your `doris` connection config to bind SQLMesh workloads to a specific `workload_group` or cloud cluster. + +```sql +MODEL ( + name analytics.cloud_events, + kind FULL, + clustered_by (event_id), + physical_properties ( + storage_vault_name = 'prod_s3_vault', + file_cache_ttl_seconds = '86400' + ) +); +``` + +--- + +## Materialized Views & Indexing + +### Asynchronous Materialized Views (MTMV) + +```sql +MODEL ( + name analytics.user_summary_mv, + kind VIEW ( + materialized true + ), + partitioned_by ds, + physical_properties ( + build = 'IMMEDIATE', + refresh = 'AUTO', + refresh_trigger = 'ON SCHEDULE EVERY 1 HOUR', + distributed_by = 'HASH(user_id) BUCKETS 10' + ) +); + +SELECT + user_id, + CAST(event_time AS DATE) AS ds, + COUNT(*) AS event_count +FROM analytics.user_events +GROUP BY user_id, CAST(event_time AS DATE); +``` + +### Inverted Indexes (Doris 2.1–5.0) + +You can attach Doris inverted indexes (`USING INVERTED`) for full-text search and fast point filtering using post-statements gated on `@runtime_stage = 'creating'`: + +```sql +MODEL ( + name analytics.searchable_logs, + kind FULL +); + +SELECT + log_id::BIGINT AS log_id, + message::STRING AS message +FROM raw.logs; + +@IF( + @runtime_stage = 'creating', + CREATE INDEX idx_message ON @this_model (message) USING INVERTED PROPERTIES("parser" = "english") COMMENT 'Full-text inverted index' +); +``` + +--- + +## Migrating from `dbt-doris` + +If you are migrating an existing `dbt-doris` project to native SQLMesh models, map your `dbt-doris` config keys as follows: + +| `dbt-doris` Config | SQLMesh Model Equivalent | +| --- | --- | +| `materialized='table'` | `kind FULL` | +| `materialized='incremental'` (with `unique_key`) | `kind INCREMENTAL_BY_UNIQUE_KEY (unique_key ...)` | +| `materialized='incremental'` (partition overwrite) | `kind INCREMENTAL_BY_TIME_RANGE (time_column ...)` or `kind INCREMENTAL_BY_PARTITION` | +| `duplicate_key=['c1', 'c2']` | `physical_properties (duplicate_key = ('c1', 'c2'))` | +| `unique_key=['c1']` | `kind INCREMENTAL_BY_UNIQUE_KEY (unique_key c1)` or `physical_properties (unique_key = ('c1'))` | +| `distributed_by=['c1']`, `buckets=10` | `clustered_by (c1), physical_properties (buckets = 10)` | +| `properties={...}` | Key-value entries inside `physical_properties (...)` | diff --git a/docs/integrations/overview.md b/docs/integrations/overview.md index 1c9d56b7e2..6e0398eb5e 100644 --- a/docs/integrations/overview.md +++ b/docs/integrations/overview.md @@ -17,6 +17,7 @@ SQLMesh supports the following execution engines for running SQLMesh projects (e * [ClickHouse](./engines/clickhouse.md) (clickhouse) * [Databricks](./engines/databricks.md) (databricks) * [Db2](./engines/db2.md) (db2) +* [Doris](./engines/doris.md) (doris) * [DuckDB](./engines/duckdb.md) (duckdb) * [Fabric](./engines/fabric.md) (fabric) * [MotherDuck](./engines/motherduck.md) (motherduck) diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index ae036a7774..a31eea0ccf 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -222,6 +222,7 @@ These pages describe the connection configuration options for each execution eng * [BigQuery](../integrations/engines/bigquery.md) * [ClickHouse](../integrations/engines/clickhouse.md) * [Databricks](../integrations/engines/databricks.md) +* [Doris](../integrations/engines/doris.md) * [DuckDB](../integrations/engines/duckdb.md) * [Fabric](../integrations/engines/fabric.md) * [MotherDuck](../integrations/engines/motherduck.md) diff --git a/mkdocs.yml b/mkdocs.yml index 49c4b9163b..334feb3019 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -83,6 +83,7 @@ nav: - integrations/engines/clickhouse.md - integrations/engines/databricks.md - integrations/engines/db2.md + - integrations/engines/doris.md - integrations/engines/duckdb.md - integrations/engines/fabric.md - integrations/engines/motherduck.md diff --git a/pyproject.toml b/pyproject.toml index 378d3ffe43..41b33339e0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -116,6 +116,7 @@ db2 = [ "ibm_db", "db2-sqlglot-dialect;python_version>=\"3.10\"" ] +doris = ["pymysql"] duckdb = [] fabric = ["pyodbc>=5.0.0"] fabric-mssql-python = ["mssql-python>=1.1.0;python_version>=\"3.10\""] @@ -290,6 +291,7 @@ markers = [ "trino: test for Trino (all connectors)", "risingwave: test for Risingwave", "starrocks: test for StarRocks", + "doris: test for Doris", # Other "set_default_connection", diff --git a/sqlmesh/core/config/__init__.py b/sqlmesh/core/config/__init__.py index 50d2d9a5a2..3113c61435 100644 --- a/sqlmesh/core/config/__init__.py +++ b/sqlmesh/core/config/__init__.py @@ -12,6 +12,7 @@ BigQueryConnectionConfig as BigQueryConnectionConfig, ConnectionConfig as ConnectionConfig, DatabricksConnectionConfig as DatabricksConnectionConfig, + DorisConnectionConfig as DorisConnectionConfig, DuckDBConnectionConfig as DuckDBConnectionConfig, FabricConnectionConfig as FabricConnectionConfig, GCPPostgresConnectionConfig as GCPPostgresConnectionConfig, diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 29cb1ef467..1334810a1c 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,167 @@ def connect_db2(**kwargs: t.Any) -> t.Any: return connect_db2 +class DorisConnectionConfig(ConnectionConfig): + """Configuration for the Apache Doris connection. + + Apache Doris uses the MySQL network protocol and supports both internal OLAP tables + and external Multi-Catalog lakehouse catalogs (e.g., Apache Iceberg). + + Args: + host: The hostname of the Doris FE (Frontend) node. + user: The Doris username. + password: The optional Doris password. + port: The query port number of the Doris FE node. Default is 9030. + database: The optional database name. + catalog: The optional default catalog name (e.g., 'internal' or an external Iceberg catalog). + charset: The optional character set (defaults to 'utf8mb4'). + collation: The optional collation. + connect_timeout: Timeout in seconds for establishing a connection. Default is 10. + read_timeout: Timeout in seconds for reading from the connection. + write_timeout: Timeout in seconds for writing to the connection. + ssl: Optional dictionary of SSL arguments passed to PyMySQL. + ssl_ca: Path to the file that contains a PEM-formatted CA certificate. + ssl_cert: Path to the file that contains a PEM-formatted client certificate. + ssl_key: Path to the file that contains a PEM-formatted private key for the client certificate. + ssl_verify_cert: Whether to verify the server certificate against the CA certificate. + ssl_verify_identity: Whether to verify the server hostname against the certificate. + ssl_disabled: Whether to disable SSL connection. + session_variables: Optional dictionary of Doris session variables to set on each connection. + concurrent_tasks: The maximum number of tasks that can use this connection concurrently. + register_comments: Whether or not to register model comments with the SQL engine. + pre_ping: Whether or not to pre-ping the connection before starting a new statement. + """ + + host: str + user: str + password: t.Optional[str] = None + port: t.Optional[int] = 9030 + database: t.Optional[str] = None + catalog: t.Optional[str] = None + charset: t.Optional[str] = None + collation: t.Optional[str] = None + connect_timeout: t.Optional[int] = 10 + read_timeout: t.Optional[int] = None + write_timeout: t.Optional[int] = None + ssl: t.Optional[t.Union[bool, t.Dict[str, t.Any]]] = None + ssl_ca: t.Optional[str] = None + ssl_cert: t.Optional[str] = None + ssl_key: t.Optional[str] = None + ssl_verify_cert: t.Optional[bool] = None + ssl_verify_identity: t.Optional[bool] = None + ssl_disabled: t.Optional[bool] = None + session_variables: t.Optional[t.Dict[str, t.Any]] = None + + concurrent_tasks: int = 4 + register_comments: bool = True + pre_ping: bool = False + + type_: t.Literal["doris"] = Field(alias="type", default="doris") + DIALECT: t.ClassVar[t.Literal["doris"]] = "doris" + DISPLAY_NAME: t.ClassVar[t.Literal["Apache Doris"]] = "Apache Doris" + DISPLAY_ORDER: t.ClassVar[t.Literal[20]] = 20 + + _engine_import_validator = _get_engine_import_validator("pymysql", "doris") + + @field_validator("session_variables", mode="after") + @classmethod + def _validate_session_variables( + cls, v: t.Optional[t.Dict[str, t.Any]] + ) -> t.Optional[t.Dict[str, t.Any]]: + if v: + import math + + for k, val in v.items(): + if not isinstance(k, str) or not re.fullmatch(r"[a-zA-Z_][a-zA-Z0-9_]*\Z", k): + raise ConfigError(f"Invalid Doris session variable name: '{k}'") + if not isinstance(val, (str, int, float, bool)) or ( + isinstance(val, float) and (math.isnan(val) or math.isinf(val)) + ): + raise ConfigError(f"Invalid Doris session variable value for '{k}': {val!r}") + return v + + @property + def _static_connection_kwargs(self) -> t.Dict[str, t.Any]: + kwargs: t.Dict[str, t.Any] = { + "autocommit": True, + "charset": self.charset or "utf8mb4", + "local_infile": False, + } + if not self.ssl_disabled: + ssl_config: t.Dict[str, t.Any] = dict(self.ssl) if isinstance(self.ssl, dict) else {} + for key in ("ca", "cert", "key"): + val = getattr(self, f"ssl_{key}", None) + if val is not None: + ssl_config[key] = val + if ssl_config or (self.ssl is not None and self.ssl is not False): + kwargs["ssl"] = ssl_config + for key in ( + "ssl_ca", + "ssl_cert", + "ssl_key", + "ssl_verify_cert", + "ssl_verify_identity", + ): + val = getattr(self, key, None) + if val is not None: + kwargs[key] = val + return kwargs + + @property + def _connection_kwargs_keys(self) -> t.Set[str]: + is_external_catalog = bool(self.catalog and self.catalog.lower() != "internal") + return { + k + for k in ( + "host", + "user", + "password", + "port", + "database", + "charset", + "collation", + "connect_timeout", + "read_timeout", + "write_timeout", + "ssl_disabled", + ) + if getattr(self, k) is not None and not (k == "database" and is_external_catalog) + } + + @property + def _engine_adapter(self) -> t.Type[EngineAdapter]: + return engine_adapter.DorisEngineAdapter + + @property + def _connection_factory(self) -> t.Callable: + from pymysql import connect + + return connect + + @property + def _cursor_init(self) -> t.Optional[t.Callable[[t.Any], None]]: + if not self.catalog and not self.session_variables: + return None + + def init(cursor: t.Any) -> None: + if self.session_variables: + self._validate_session_variables(self.session_variables) + for k, v in self.session_variables.items(): + val_sql = exp.convert(v).sql("doris", comments=False) + cursor.execute(f"SET `{k}` = {val_sql}") + if self.catalog: + catalog_sql = exp.to_identifier(self.catalog, quoted=True).sql("doris") + cursor.execute(f"SWITCH {catalog_sql}") + if self.database: + db_sql = exp.to_identifier(self.database, quoted=True).sql("doris") + cursor.execute(f"USE {db_sql}") + + return init + + 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/doris.py b/sqlmesh/core/engine_adapter/doris.py new file mode 100644 index 0000000000..64c59e12c5 --- /dev/null +++ b/sqlmesh/core/engine_adapter/doris.py @@ -0,0 +1,1962 @@ +from __future__ import annotations + +from collections import defaultdict +import itertools +import logging +import re +import typing as t + +import pandas as pd # noqa: TID253 +from sqlglot import exp, parse, parse_one + +from sqlmesh.core.dialect import to_schema +from sqlmesh.core.engine_adapter.base import ( + EngineAdapterWithIndexSupport, + InsertOverwriteStrategy, + _decoded_str, +) +from sqlmesh.core.engine_adapter.mixins import ( + GetCurrentCatalogFromFunctionMixin, + LogicalMergeMixin, + PandasNativeFetchDFSupportMixin, + RowDiffMixin, +) +from sqlmesh.core.engine_adapter.shared import ( + CatalogSupport, + CommentCreationTable, + CommentCreationView, + DataObject, + DataObjectType, + SourceQuery, + set_catalog, +) +from sqlmesh.utils import get_source_columns_to_types, random_id +from sqlmesh.utils.errors import SQLMeshError + +if t.TYPE_CHECKING: + from sqlmesh.core._typing import SchemaName, TableName + from sqlmesh.core.engine_adapter._typing import DF, QueryOrDF + from sqlmesh.core.node import IntervalUnit + +logger = logging.getLogger(__name__) + + +@set_catalog() +class DorisEngineAdapter( + EngineAdapterWithIndexSupport, + LogicalMergeMixin, + GetCurrentCatalogFromFunctionMixin, + PandasNativeFetchDFSupportMixin, + RowDiffMixin, +): + DIALECT = "doris" + DEFAULT_BATCH_SIZE = 5000 + SUPPORTS_TRANSACTIONS = False + COMMENT_CREATION_TABLE = CommentCreationTable.IN_SCHEMA_DEF_NO_CTAS + COMMENT_CREATION_VIEW = CommentCreationView.IN_SCHEMA_DEF_NO_COMMANDS + MAX_TABLE_COMMENT_LENGTH = 2048 + MAX_COLUMN_COMMENT_LENGTH = 255 + SUPPORTS_INDEXES = True + SUPPORTS_REPLACE_TABLE = False + MAX_IDENTIFIER_LENGTH = 64 + SUPPORTS_MATERIALIZED_VIEWS = True + SUPPORTS_MATERIALIZED_VIEW_SCHEMA = True + SUPPORTS_CREATE_DROP_CATALOG = False + INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.DELETE_INSERT + CURRENT_CATALOG_EXPRESSION = exp.func("current_catalog") + SCHEMA_DIFFER_KWARGS = { + "parameterized_type_defaults": { + exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(9, 0), (0,)], + exp.DataType.build("CHAR", dialect=DIALECT).this: [(1,), (255,)], + exp.DataType.build("VARCHAR", dialect=DIALECT).this: [(65533,)], + exp.DataType.build("DATETIME", dialect=DIALECT).this: [(0,)], + exp.DataType.build("INT", dialect=DIALECT).this: [(11,)], + exp.DataType.build("BIGINT", dialect=DIALECT).this: [(20,)], + exp.DataType.build("SMALLINT", dialect=DIALECT).this: [(6,)], + exp.DataType.build("TINYINT", dialect=DIALECT).this: [(4,), (1,)], + }, + "types_with_unlimited_length": { + exp.DataType.build("TEXT", dialect=DIALECT).this: { + exp.DataType.build("VARCHAR", dialect=DIALECT).this, + exp.DataType.build("CHAR", dialect=DIALECT).this, + }, + exp.DataType.build("VARCHAR", dialect=DIALECT).this: { + exp.DataType.build("VARCHAR", dialect=DIALECT).this, + exp.DataType.build("CHAR", dialect=DIALECT).this, + exp.DataType.build("TEXT", dialect=DIALECT).this, + }, + }, + "compatible_types": { + exp.DataType.build("BOOLEAN", dialect=DIALECT): { + exp.DataType.build("TINYINT", dialect=DIALECT), + exp.DataType.build("SMALLINT", dialect=DIALECT), + exp.DataType.build("INT", dialect=DIALECT), + exp.DataType.build("BIGINT", dialect=DIALECT), + exp.DataType.build("INT128", dialect=DIALECT), + }, + exp.DataType.build("TINYINT", dialect=DIALECT): { + exp.DataType.build("SMALLINT", dialect=DIALECT), + exp.DataType.build("INT", dialect=DIALECT), + exp.DataType.build("BIGINT", dialect=DIALECT), + exp.DataType.build("INT128", dialect=DIALECT), + exp.DataType.build("DOUBLE", dialect=DIALECT), + }, + exp.DataType.build("SMALLINT", dialect=DIALECT): { + exp.DataType.build("INT", dialect=DIALECT), + exp.DataType.build("BIGINT", dialect=DIALECT), + exp.DataType.build("INT128", dialect=DIALECT), + exp.DataType.build("DOUBLE", dialect=DIALECT), + }, + exp.DataType.build("INT", dialect=DIALECT): { + exp.DataType.build("BIGINT", dialect=DIALECT), + exp.DataType.build("INT128", dialect=DIALECT), + exp.DataType.build("DOUBLE", dialect=DIALECT), + }, + exp.DataType.build("BIGINT", dialect=DIALECT): { + exp.DataType.build("INT128", dialect=DIALECT), + }, + exp.DataType.build("FLOAT", dialect=DIALECT): { + exp.DataType.build("DOUBLE", dialect=DIALECT), + }, + exp.DataType.build("DATE", dialect=DIALECT): { + exp.DataType.build("DATETIME", dialect=DIALECT), + }, + }, + } + + INVALID_KEY_TYPES = { + "FLOAT", + "DOUBLE", + "STRING", + "TEXT", + "JSON", + "JSONB", + "VARIANT", + "OBJECT", + "ARRAY", + "MAP", + "STRUCT", + "BITMAP", + "HLL", + "QUANTILE_STATE", + "AGG_STATE", + } + + @property + def catalog_support(self) -> CatalogSupport: + return CatalogSupport.FULL_SUPPORT + + @property + def default_catalog(self) -> t.Optional[str]: + return self._default_catalog or self._extra_config.get("catalog") + + def get_current_catalog(self) -> t.Optional[str]: + try: + catalog = super().get_current_catalog() + if catalog: + return _decoded_str(catalog) + except Exception: + logger.debug("current_catalog() failed on Doris, falling back to default catalog") + return self.default_catalog or "internal" + + def set_current_catalog(self, catalog: str) -> None: + catalog_sql = self._to_sql(exp.to_identifier(catalog, quoted=True)) + self.execute(exp.Command(this="SWITCH", expression=catalog_sql)) + + def _is_iceberg_table( + self, + table_format: t.Optional[str] = None, + catalog_name: t.Optional[str] = None, + table_properties: t.Optional[t.Dict[str, t.Any]] = None, + ) -> bool: + if table_format and str(table_format).strip().lower() == "iceberg": + return True + if table_properties: + engine_val = self._extract_str_value(table_properties.get("engine")) + if engine_val and engine_val.strip().lower() == "iceberg": + return True + effective_catalog = catalog_name or self.default_catalog + if effective_catalog and str(effective_catalog).strip().lower() != "internal": + return True + return False + + @staticmethod + def _extract_str_value(val: t.Any) -> t.Optional[str]: + if val is None: + return None + if isinstance(val, exp.Literal): + return str(val.this) + if isinstance(val, exp.Identifier): + return str(val.this) + if isinstance(val, exp.Column): + return str(val.name) + if isinstance(val, exp.Var): + return str(val.this) + if isinstance(val, str): + return val + if isinstance(val, exp.Expr): + return val.name or str(val) + return str(val) + + @classmethod + def _extract_column_names(cls, key_expr: t.Any) -> t.List[str]: + if key_expr is None: + return [] + if isinstance(key_expr, (exp.Tuple, exp.Array)): + names: t.List[str] = [] + for e in key_expr.expressions: + names.extend(cls._extract_column_names(e)) + return names + if isinstance(key_expr, exp.Paren): + return cls._extract_column_names(key_expr.this) + if isinstance(key_expr, (list, tuple)): + names = [] + for e in key_expr: + names.extend(cls._extract_column_names(e)) + return names + if isinstance(key_expr, exp.Column): + return [key_expr.name] + if isinstance(key_expr, (exp.Identifier, exp.Literal, exp.Var)): + raw = str(key_expr.this) + return [c.strip().strip("`") for c in raw.split(",") if c.strip()] + if isinstance(key_expr, str): + return [c.strip().strip("`") for c in key_expr.split(",") if c.strip()] + return [str(key_expr).strip().strip("`")] + + @classmethod + def _is_valid_olap_key_type(cls, dtype: exp.DataType) -> bool: + type_name = ( + dtype.this.value.upper() + if isinstance(dtype.this, exp.DataType.Type) + else str(dtype.this).upper() + ) + if type_name in cls.INVALID_KEY_TYPES: + return False + if dtype.is_type("userdefined"): + kind = str(dtype.args.get("kind") or "").upper() + if kind in cls.INVALID_KEY_TYPES: + return False + return True + + def _fetch_native_df( + self, query: t.Union[exp.Expr, str], quote_identifiers: bool = False + ) -> DF: + self._connection_pool.get_cursor() + return super()._fetch_native_df(query, quote_identifiers=quote_identifiers) + + def _execute(self, sql: str, track_rows_processed: bool = False, **kwargs: t.Any) -> None: + try: + super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) + except Exception as e: + err_code = e.args[0] if getattr(e, "args", None) else None + err_name = e.__class__.__name__ + is_disconnect = (err_code in (2006, 2013, 2055) and err_name == "OperationalError") or ( + err_code == 0 and err_name == "InterfaceError" + ) + if is_disconnect and not self._connection_pool.is_transaction_active: + logger.warning( + "Doris connection lost (%s: %s), reconnecting and retrying once...", + err_name, + err_code, + ) + self._connection_pool.close_cursor() + conn = self._connection_pool.get() + conn.ping(reconnect=True) + super()._execute(sql, track_rows_processed=track_rows_processed, **kwargs) + else: + raise + + def columns( + self, table_name: TableName, include_pseudo_columns: bool = False + ) -> t.Dict[str, exp.DataType]: + """Fetches column names and types for the target table, normalizing Doris type aliases.""" + self.execute(exp.Describe(this=exp.to_table(table_name), kind="TABLE")) + describe_output = self.cursor.fetchall() + cols: t.Dict[str, exp.DataType] = {} + for column_name, column_type, *_ in itertools.takewhile( + lambda row: not str(row[0]).startswith("#"), + describe_output, + ): + col_name_str = _decoded_str(column_name) if column_name else "" + raw_type = _decoded_str(column_type) if column_type else "" + if not col_name_str.strip() or not raw_type.strip(): + continue + normalized_type = re.sub( + r"\bDECIMALV[23]\b(?!\s*`?\s*:)", "DECIMAL", raw_type, flags=re.IGNORECASE + ) + normalized_type = re.sub( + r"\bDATETIMEV2\b(?!\s*`?\s*:)", "DATETIME", normalized_type, flags=re.IGNORECASE + ) + normalized_type = re.sub( + r"\bDATEV2\b(?!\s*`?\s*:)", "DATE", normalized_type, flags=re.IGNORECASE + ) + normalized_type = re.sub( + r"\bLARGEINT(?:\(\d+\)|\b)(?!\s*`?\s*:)", + "INT128", + normalized_type, + flags=re.IGNORECASE, + ) + normalized_type = re.sub( + r"\bTINYINT\(1\)", "BOOLEAN", normalized_type, flags=re.IGNORECASE + ) + normalized_type = re.sub( + r"\b(TINYINT|SMALLINT|INT|BIGINT)\(\d+\)", + r"\1", + normalized_type, + flags=re.IGNORECASE, + ) + cols[col_name_str] = exp.DataType.build(normalized_type, dialect=self.dialect, udt=True) + return cols + + def create_schema( + self, + schema_name: SchemaName, + ignore_if_exists: bool = True, + warn_on_error: bool = True, + properties: t.Optional[t.List[exp.Expr]] = None, + ) -> None: + """Create a schema (database) in Doris.""" + properties = properties or [] + return super()._create_schema( + schema_name=schema_name, + ignore_if_exists=ignore_if_exists, + warn_on_error=warn_on_error, + properties=properties, + kind="DATABASE", + ) + + def drop_schema( + self, + schema_name: SchemaName, + ignore_if_not_exists: bool = True, + cascade: bool = False, + **drop_args: t.Dict[str, exp.Expr], + ) -> None: + """Drop schema in Doris. Note: Doris doesn't support CASCADE clause.""" + return self._drop_object( + name=schema_name, + exists=ignore_if_not_exists, + kind="DATABASE", + cascade=False, + **drop_args, + ) + + def _get_data_objects( + self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None + ) -> t.List[DataObject]: + """Returns all the data objects that exist in the given schema and optional catalog.""" + schema_exp = to_schema(schema_name) + catalog = schema_exp.catalog or None + query_catalog = catalog or self.default_catalog + + query = ( + exp.select( + exp.column("table_schema").as_("schema_name"), + exp.column("table_name").as_("name"), + exp.case() + .when( + exp.column("table_type").eq("BASE TABLE"), + exp.Literal.string("table"), + ) + .when( + exp.column("table_type").eq("SYSTEM VIEW"), + exp.Literal.string("view"), + ) + .when( + exp.column("table_type").eq("VIEW"), + exp.Literal.string("view"), + ) + .when( + exp.column("table_type").eq("MATERIALIZED VIEW"), + exp.Literal.string("materialized_view"), + ) + .else_(exp.column("table_type")) + .as_("type"), + ) + .from_(exp.table_("tables", db="information_schema", catalog=query_catalog)) + .where(exp.column("table_schema").eq(schema_exp.db)) + ) + if object_names: + lowered_names = [name.lower() for name in object_names] + query = query.where(exp.func("LOWER", exp.column("table_name")).isin(*lowered_names)) + + result: t.List[DataObject] = [] + for row_schema, table_name, table_type in self.fetchall(query): + try: + schema_str = _decoded_str(row_schema) if row_schema is not None else schema_exp.db + name_str = _decoded_str(table_name) if table_name is not None else "unknown" + obj_type = _decoded_str(table_type) if table_type is not None else "table" + + upper_type = obj_type.upper() + if upper_type in ("BASE TABLE", "EXTERNAL TABLE"): + obj_type = "table" + elif upper_type in ("VIEW", "SYSTEM VIEW"): + obj_type = "view" + elif upper_type in ("MATERIALIZED VIEW", "MATERIALIZED_VIEW"): + obj_type = "materialized_view" + + result.append( + DataObject( + catalog=catalog, + schema=schema_str, + name=name_str, + type=DataObjectType.from_str(obj_type), + ) + ) + except (ValueError, AttributeError) as e: + logger.error( + f"Error processing row: {e}, row: {(row_schema, table_name, table_type)}" + ) + continue + + return result + + def create_view( + self, + view_name: TableName, + query_or_df: QueryOrDF, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + replace: bool = True, + materialized: bool = False, + materialized_properties: t.Optional[t.Dict[str, t.Any]] = None, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, + source_columns: t.Optional[t.List[str]] = None, + **create_kwargs: t.Any, + ) -> None: + target_table = exp.to_table(view_name) + effective_catalog = target_table.catalog or self.default_catalog + if effective_catalog and effective_catalog.lower() != "internal": + if materialized: + raise SQLMeshError( + f"Doris only supports creating materialized views in the 'internal' catalog, got '{effective_catalog}' " + f"for view '{self._to_sql(target_table)}'." + ) + logger.warning( + "Skipping view creation for '%s' because Doris does not support creating views in external catalog '%s'. " + "Configure `environment_catalog_mapping` in your SQLMesh configuration to route virtual environment views to the 'internal' catalog.", + self._to_sql(target_table), + effective_catalog, + ) + return + + if replace: + self.drop_view( + view_name, + ignore_if_not_exists=True, + materialized=materialized, + view_properties=view_properties, + ) + if not materialized: + return super().create_view( + view_name, + query_or_df, + target_columns_to_types, + replace=False, + materialized=False, + materialized_properties=materialized_properties, + table_description=table_description, + column_descriptions=column_descriptions, + view_properties=view_properties, + source_columns=source_columns, + **create_kwargs, + ) + self._create_materialized_view( + view_name, + query_or_df, + target_columns_to_types=target_columns_to_types, + materialized_properties=materialized_properties, + table_description=table_description, + column_descriptions=column_descriptions, + view_properties=view_properties, + source_columns=source_columns, + **create_kwargs, + ) + + def _create_materialized_view( + self, + view_name: TableName, + query_or_df: QueryOrDF, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + materialized_properties: t.Optional[t.Dict[str, t.Any]] = None, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + view_properties: t.Optional[t.Dict[str, exp.Expr]] = None, + source_columns: t.Optional[t.List[str]] = None, + **create_kwargs: t.Any, + ) -> None: + query_or_df = self._native_df_to_pandas_df(query_or_df) + + # Inline DataFrames as a SELECT ... FROM VALUES query rather than delegating DataFrame + # handling to _get_source_queries_and_columns_to_types, which would otherwise materialize + # a temporary table that gets dropped as soon as the SourceQuery context manager exits. + if isinstance(query_or_df, pd.DataFrame): + values: t.List[t.Tuple[t.Any, ...]] = list( + query_or_df.itertuples(index=False, name=None) + ) + target_columns_to_types, source_columns = self._columns_to_types( + query_or_df, target_columns_to_types, source_columns + ) + if not target_columns_to_types: + raise SQLMeshError("columns_to_types must be provided for dataframes") + source_columns_to_types = get_source_columns_to_types( + target_columns_to_types, source_columns + ) + query_or_df = self._values_to_sql( + values, + source_columns_to_types, + batch_start=0, + batch_end=len(values), + ) + + source_queries, target_columns_to_types = self._get_source_queries_and_columns_to_types( + query_or_df, + target_columns_to_types, + batch_size=0, + target_table=view_name, + source_columns=source_columns, + ) + if len(source_queries) != 1: + raise SQLMeshError("Only one source query is supported for creating materialized views") + + with source_queries[0] as query: + target_table = exp.to_table(view_name) + + schema: t.Union[exp.Table, exp.Schema] = target_table + if target_columns_to_types: + schema = self._build_schema_exp( + target_table, + target_columns_to_types, + column_descriptions, + is_view=True, + ) + + partitioned_by = None + if materialized_properties and materialized_properties.get("partitioned_by"): + partitioned_by = materialized_properties.get("partitioned_by") + elif create_kwargs.get("partitioned_by"): + partitioned_by = create_kwargs.get("partitioned_by") + elif view_properties and view_properties.get("partitioned_by"): + raw_vp = view_properties.get("partitioned_by") + if raw_vp is not None: + partitioned_by = raw_vp if isinstance(raw_vp, list) else [raw_vp] + + clustered_by = None + if materialized_properties and materialized_properties.get("clustered_by"): + clustered_by = materialized_properties.get("clustered_by") + elif create_kwargs.get("clustered_by"): + clustered_by = create_kwargs.get("clustered_by") + + properties_exp = self._build_table_properties_exp( + catalog_name=target_table.catalog, + table_properties=view_properties, + target_columns_to_types=target_columns_to_types, + table_description=( + table_description + if self.COMMENT_CREATION_VIEW.supports_schema_def and self.comments_enabled + else None + ), + partitioned_by=partitioned_by, + clustered_by=clustered_by, + partition_interval_unit=( + materialized_properties.get("partition_interval_unit") + if materialized_properties + else None + ), + table_kind="MATERIALIZED_VIEW", + ) + + is_sync_mv = bool( + view_properties + and ( + ( + self._extract_str_value(view_properties.get("materialized_type")) or "" + ).upper() + == "SYNC" + ) + ) + if is_sync_mv and isinstance(query, exp.Select): + query = query.copy() + aliased_selects: t.List[exp.Expr] = [] + for sel_expr in query.expressions: + if isinstance(sel_expr, exp.Column): + aliased_selects.append( + exp.alias_(sel_expr.copy(), f"mv_{sel_expr.name}", quoted=True) + ) + elif ( + isinstance(sel_expr, exp.Alias) + and isinstance(sel_expr.this, exp.Column) + and sel_expr.alias == sel_expr.this.name + ): + aliased_selects.append( + exp.alias_(sel_expr.this.copy(), f"mv_{sel_expr.alias}", quoted=True) + ) + else: + aliased_selects.append(sel_expr) + query.set("expressions", aliased_selects) + + # Doris synchronous materialized views are single-table rollup indexes whose DDL syntax + # (`CREATE MATERIALIZED VIEW AS SELECT ... FROM .`) forbids a + # `.` qualifier on `` and requires an active database on the session. + if is_sync_mv and not target_table.db: + source_db: t.Optional[str] = None + if view_properties and view_properties.get("source_table"): + raw_src = view_properties["source_table"] + if isinstance(raw_src, exp.Column): + source_db = raw_src.db or None + elif isinstance(raw_src, exp.Table): + source_db = raw_src.db or None + else: + src_str = self._extract_str_value(raw_src) or str(raw_src) + source_db = exp.to_table(src_str, dialect=self.dialect).db or None + if not source_db: + first_tbl = query.find(exp.Table) + if first_tbl and first_tbl.db: + source_db = first_tbl.db + if source_db: + db_sql = 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, + expression=query, + properties=properties_exp, + ), + quote_identifiers=self.QUOTE_IDENTIFIERS_IN_VIEWS, + ) + self._clear_data_object_cache(view_name) + + def drop_view( + self, + view_name: TableName, + ignore_if_not_exists: bool = True, + materialized: bool = False, + **kwargs: t.Any, + ) -> None: + """Drop view in Doris.""" + target_table = exp.to_table(view_name) + effective_catalog = target_table.catalog or self.default_catalog + if effective_catalog and effective_catalog.lower() != "internal": + return + + view_properties = kwargs.pop("view_properties", None) + if materialized and view_properties: + mat_type = self._extract_str_value(view_properties.get("materialized_type")) + source_table = view_properties.get("source_table") + if mat_type and mat_type.upper() == "SYNC" and source_table: + view_table_sql = self._to_sql(target_table) + if isinstance(source_table, exp.Column): + source_tbl_exp = exp.table_( + source_table.name, + db=source_table.db or None, + catalog=source_table.catalog or None, + ) + elif isinstance(source_table, exp.Table): + source_tbl_exp = source_table + else: + source_table_str = self._extract_str_value(source_table) or str(source_table) + source_tbl_exp = exp.to_table(source_table_str, dialect=self.dialect) + # 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_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, + ) + + def insert_overwrite_by_partition( + self, + table_name: TableName, + query_or_df: QueryOrDF, + partitioned_by: t.List[exp.Expr], + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + source_columns: t.Optional[t.List[str]] = None, + ) -> None: + parsed_exprs, _ = self._parse_partition_expressions(partitioned_by) + key_exprs: t.List[exp.Expr] = [] + for expr in parsed_exprs or partitioned_by: + if isinstance(expr, exp.Column): + key_exprs.append(expr) + else: + cols = list(expr.find_all(exp.Column)) + if cols: + key_exprs.extend(cols) + else: + key_exprs.append(expr) + super().insert_overwrite_by_partition( + table_name, + query_or_df, + key_exprs or partitioned_by, + target_columns_to_types=target_columns_to_types, + source_columns=source_columns, + ) + + def delete_from( + self, table_name: TableName, where: t.Optional[t.Union[str, exp.Expr]] = None + ) -> None: + """Delete from a table in Doris.""" + if not where or where == exp.true(): + self._truncate_table(table_name) + return + + if isinstance(where, str): + where = parse_one(where, read=self.dialect) + else: + where = where.copy() + + where = self._rewrite_between_in_where(where) + + subquery_info = self._extract_subquery_from_where(where) + if subquery_info is not None: + self._execute_delete_with_subquery(table_name, *subquery_info) + else: + super().delete_from(table_name, where) + + @staticmethod + def _rewrite_between_in_where(where: exp.Expr) -> exp.Expr: + if isinstance(where, exp.Between): + return exp.and_( + exp.GTE(this=where.this.copy(), expression=where.args["low"].copy()), + exp.LTE(this=where.this.copy(), expression=where.args["high"].copy()), + copy=False, + ) + for between in list(where.find_all(exp.Between)): + if between.find_ancestor(exp.Subquery, exp.Select, exp.Not): + continue + replacement: exp.Expr = exp.and_( + exp.GTE(this=between.this.copy(), expression=between.args["low"].copy()), + exp.LTE(this=between.this.copy(), expression=between.args["high"].copy()), + copy=False, + ) + if not isinstance(between.parent, (exp.Where, exp.Paren)): + replacement = exp.Paren(this=replacement) + between.replace(replacement) + return where + + def _extract_subquery_from_where( + self, where: exp.Expr + ) -> t.Optional[t.Tuple[t.List[exp.Expr], exp.Subquery, t.List[exp.Expr]]]: + """Extract positive IN (SELECT ...) subquery and any sibling AND conditions.""" + conjuncts: t.List[exp.Expr] = ( + list(where.flatten()) if isinstance(where, exp.And) else [where] + ) + target_in: t.Optional[exp.In] = None + subquery: t.Optional[exp.Subquery] = None + extra_conditions: t.List[exp.Expr] = [] + + for cond in conjuncts: + unwrapped_cond = cond.unnest() + if target_in is None and isinstance(unwrapped_cond, exp.In): + sq = self._get_in_subquery(unwrapped_cond) + if sq is not None: + target_in = unwrapped_cond + subquery = sq + continue + extra_conditions.append(cond) + + if target_in is None or subquery is None: + return None + + lhs = target_in.this + select_expr = subquery.this + key_exprs: t.List[exp.Expr] = [] + + lhs_concat = self._unwrap_concat_ws(lhs) + if lhs_concat is not None: + key_exprs = [e.copy() for e in lhs_concat.expressions[1:]] + if isinstance(select_expr, exp.Select) and len(select_expr.expressions) == 1: + first_sel = select_expr.expressions[0] + raw_sel = first_sel.this if isinstance(first_sel, exp.Alias) else first_sel + sel_concat = self._unwrap_concat_ws(raw_sel) + if sel_concat is not None: + select_expr = select_expr.copy() + select_expr.set("expressions", [e.copy() for e in sel_concat.expressions[1:]]) + subquery = exp.Subquery(this=select_expr) + elif isinstance(lhs, exp.Tuple): + key_exprs = [e.copy() for e in lhs.expressions] + else: + key_exprs = [lhs.copy()] + + normalized_keys: t.List[exp.Expr] = [] + for k in key_exprs: + if isinstance(k, exp.Column): + normalized_keys.append(k) + elif isinstance(k, exp.Identifier): + normalized_keys.append(exp.Column(this=k.copy())) + else: + return None + + if not normalized_keys: + return None + + return normalized_keys, subquery, extra_conditions + + @staticmethod + def _unwrap_concat_ws(expr: exp.Expr) -> t.Optional[exp.Expr]: + curr = expr.unnest() + if isinstance(curr, exp.Column) and isinstance(curr.this, exp.Expr): + curr = curr.this.unnest() + if isinstance(curr, exp.ConcatWs) or ( + isinstance(curr, exp.Anonymous) and curr.name.upper() == "CONCAT_WS" + ): + return curr + return None + + @staticmethod + def _get_in_subquery(in_expr: exp.In) -> t.Optional[exp.Subquery]: + query = in_expr.args.get("query") + if isinstance(query, exp.Subquery): + return query.copy() + if isinstance(query, exp.Select): + return exp.Subquery(this=query.copy()) + if len(in_expr.expressions) == 1: + first = in_expr.expressions[0] + if isinstance(first, exp.Subquery): + return first.copy() + if isinstance(first, exp.Select): + return exp.Subquery(this=first.copy()) + return None + + def _execute_delete_with_subquery( + self, + table_name: TableName, + key_exprs: t.List[exp.Expr], + subquery: exp.Subquery, + extra_conditions: t.List[exp.Expr], + ) -> None: + """Execute DELETE FROM ... USING (SELECT ...) in Doris, with literal batch fallback for DUPLICATE KEY/Iceberg.""" + target_table = exp.to_table(table_name) + if self._is_iceberg_table(catalog_name=target_table.catalog): + self._delete_by_literal_keys(table_name, key_exprs, subquery, extra_conditions) + return + + join_conditions: t.List[exp.Expr] = [] + for idx, col in enumerate(key_exprs): + col_name = col.name + t1_col = exp.column(col_name, table="_t1") + t2_col_name = col_name + if isinstance(subquery.this, exp.Select) and idx < len(subquery.this.expressions): + sel_item = subquery.this.expressions[idx] + if sel_item.alias_or_name: + t2_col_name = sel_item.alias_or_name + t2_col = exp.column(t2_col_name, table="_t2") + join_conditions.append(t1_col.eq(t2_col)) + + for extra in extra_conditions: + qualified_extra = extra.copy() + for col in qualified_extra.find_all(exp.Column): + if not col.find_ancestor(exp.Subquery, exp.Select) and not col.table: + col.set("table", exp.to_identifier("_t1")) + join_conditions.append(qualified_extra) + + where_expr = exp.and_(*join_conditions) + + target_sql = f"{self._to_sql(target_table)} AS `_t1`" + subquery_with_alias = exp.alias_(subquery.copy(), "_t2", table=True) + subquery_sql = self._to_sql(subquery_with_alias) + where_sql = self._to_sql(where_expr) + + delete_sql = f"DELETE FROM {target_sql} USING {subquery_sql} WHERE {where_sql}" + try: + self.execute(delete_sql) + except Exception as e: + err_msg = str(e).lower() + if any( + marker in err_msg + for marker in ( + "unique key", + "olaptable", + "olap table", + "not support", + "only support", + "unsupported", + ) + ): + logger.info( + "DELETE ... USING not supported on table '%s' (%s); falling back to batched literal DELETE.", + table_name, + e, + ) + self._delete_by_literal_keys(table_name, key_exprs, subquery, extra_conditions) + else: + raise + + @staticmethod + def _sanitize_key_value(val: t.Any) -> t.Any: + if val is None: + return None + try: + if pd.isna(val): + return None + except (TypeError, ValueError): + pass + return val + + def _delete_by_literal_keys( + self, + table_name: TableName, + key_exprs: t.List[exp.Expr], + subquery: t.Union[exp.Subquery, exp.Select], + extra_conditions: t.Optional[t.List[exp.Expr]] = None, + ) -> None: + """Fallback for DELETE on DUPLICATE KEY / AGGREGATE KEY / Iceberg tables that prohibits OR predicates.""" + select_query = ( + subquery.this.copy() if isinstance(subquery, exp.Subquery) else subquery.copy() + ) + if isinstance(select_query, exp.Select): + select_query = select_query.distinct() + rows = self.fetchall(select_query) + if not rows: + return + + extra_conditions = extra_conditions or [] + cols = [exp.column(k.name, quoted=True) for k in key_exprs] + batch_size = self.DEFAULT_BATCH_SIZE + + sanitized_rows = [tuple(self._sanitize_key_value(v) for v in row) for row in rows] + + if len(cols) == 1: + col = cols[0] + non_null_vals = [exp.convert(r[0]) for r in sanitized_rows if r[0] is not None] + has_null = any(r[0] is None for r in sanitized_rows) + + for i in range(0, len(non_null_vals), batch_size): + batch = non_null_vals[i : i + batch_size] + predicates: t.List[exp.Expr] = [col.isin(*batch), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*predicates))) + + if has_null: + predicates = [col.is_(exp.null()), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*predicates))) + return + + # Composite key (N >= 2): dynamically select the column index to batch via IN (...) + # that minimizes the number of distinct prefix groups, breaking ties in favor of the last column. + batch_idx = min( + range(len(cols)), + key=lambda idx: (len({row[:idx] + row[idx + 1 :] for row in sanitized_rows}), -idx), + ) + grouped: t.Dict[t.Tuple[t.Any, ...], t.List[t.Any]] = defaultdict(list) + for row in sanitized_rows: + prefix_key = row[:batch_idx] + row[batch_idx + 1 :] + grouped[prefix_key].append(row[batch_idx]) + + batch_col = cols[batch_idx] + prefix_cols = cols[:batch_idx] + cols[batch_idx + 1 :] + + for prefix_vals, batch_vals in grouped.items(): + prefix_preds: t.List[exp.Expr] = [ + p_col.eq(exp.convert(p_val)) if p_val is not None else p_col.is_(exp.null()) + for p_col, p_val in zip(prefix_cols, prefix_vals) + ] + non_null_batch = [exp.convert(v) for v in batch_vals if v is not None] + has_null_batch = any(v is None for v in batch_vals) + + for i in range(0, len(non_null_batch), batch_size): + batch = non_null_batch[i : i + batch_size] + preds = [*prefix_preds, batch_col.isin(*batch), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*preds))) + + if has_null_batch: + preds = [*prefix_preds, batch_col.is_(exp.null()), *extra_conditions] + self.execute(exp.delete(table_name, where=exp.and_(*preds))) + + def adjust_physical_properties_for_incremental( + self, + physical_properties: t.Dict[str, t.Any], + *, + requires_delete_capable_table: bool, + unique_key: t.Optional[t.List[exp.Expr]], + model_name: str, + ) -> t.Dict[str, t.Any]: + """Promote ``unique_key`` from ``INCREMENTAL_BY_UNIQUE_KEY`` models into ``physical_properties`` + when no explicit OLAP key model (``unique_key``, ``duplicate_key``, ``aggregate_key``) is set. + """ + if ( + unique_key + and "unique_key" not in physical_properties + and "duplicate_key" not in physical_properties + and "aggregate_key" not in physical_properties + ): + physical_properties["unique_key"] = ( + unique_key[0] if len(unique_key) == 1 else exp.Tuple(expressions=list(unique_key)) + ) + return physical_properties + + def _create_table_from_columns( + self, + table_name: TableName, + target_columns_to_types: t.Dict[str, exp.DataType], + primary_key: t.Optional[t.Tuple[str, ...]] = None, + exists: bool = True, + table_description: t.Optional[str] = None, + column_descriptions: t.Optional[t.Dict[str, str]] = None, + **kwargs: t.Any, + ) -> None: + table_properties = dict(kwargs.get("table_properties") or {}) + + if primary_key and "unique_key" not in table_properties: + table_properties["unique_key"] = exp.Tuple( + expressions=[exp.to_column(col) for col in primary_key] + ) + + kwargs["table_properties"] = table_properties + + super()._create_table_from_columns( + table_name=table_name, + target_columns_to_types=target_columns_to_types, + primary_key=None, + exists=exists, + table_description=table_description, + column_descriptions=column_descriptions, + **kwargs, + ) + + def _build_create_table_exp( + self, + table_name_or_schema: t.Union[exp.Schema, TableName], + expression: t.Optional[exp.Expr], + exists: bool = True, + replace: bool = False, + target_columns_to_types: t.Optional[t.Dict[str, exp.DataType]] = None, + table_description: t.Optional[str] = None, + table_kind: t.Optional[str] = None, + **kwargs: t.Any, + ) -> exp.Create: + kwargs.setdefault("table_properties", None) + kwargs["is_schema_ddl"] = isinstance(table_name_or_schema, exp.Schema) + + catalog_name = None + if isinstance(table_name_or_schema, exp.Schema): + if isinstance(table_name_or_schema.this, exp.Table): + catalog_name = table_name_or_schema.this.catalog + else: + catalog_name = exp.to_table(table_name_or_schema).catalog + + is_iceberg = self._is_iceberg_table( + table_format=kwargs.get("table_format"), + catalog_name=catalog_name, + table_properties=kwargs.get("table_properties"), + ) + + if not is_iceberg and target_columns_to_types: + reordered_cols = self._reorder_olap_columns(target_columns_to_types, kwargs) + if list(reordered_cols) != list(target_columns_to_types): + target_columns_to_types = reordered_cols + if isinstance(table_name_or_schema, exp.Schema): + col_def_map = { + col_def.name: col_def + for col_def in table_name_or_schema.expressions + if isinstance(col_def, exp.ColumnDef) + } + other_exprs = [ + e + for e in table_name_or_schema.expressions + if not isinstance(e, exp.ColumnDef) + ] + if len(col_def_map) == len(target_columns_to_types) and all( + c in col_def_map for c in target_columns_to_types + ): + table_name_or_schema = table_name_or_schema.copy() + table_name_or_schema.set( + "expressions", + [col_def_map[c] for c in target_columns_to_types] + other_exprs, + ) + if expression is not None and isinstance(expression, exp.Query): + expression = self._order_projections_and_filter( + expression, target_columns_to_types + ) + + return super()._build_create_table_exp( + table_name_or_schema, + expression=expression, + exists=exists, + replace=replace, + target_columns_to_types=target_columns_to_types, + table_description=table_description, + table_kind=table_kind, + **kwargs, + ) + + def _reorder_olap_columns( + self, + target_columns_to_types: t.Dict[str, exp.DataType], + kwargs: t.Dict[str, t.Any], + ) -> t.Dict[str, exp.DataType]: + table_props = kwargs.get("table_properties") or {} + explicit_key = ( + table_props.get("unique_key") + or kwargs.get("unique_key") + or table_props.get("duplicate_key") + or table_props.get("aggregate_key") + ) + col_names = list(target_columns_to_types) + if not col_names: + return target_columns_to_types + + if explicit_key: + key_cols = [ + k for k in self._extract_column_names(explicit_key) if k in target_columns_to_types + ] + if key_cols and col_names[: len(key_cols)] != key_cols: + new_order = key_cols + [c for c in col_names if c not in key_cols] + return {c: target_columns_to_types[c] for c in new_order} + return target_columns_to_types + + first_dtype = target_columns_to_types[col_names[0]] + if not self._is_valid_olap_key_type(first_dtype): + valid_col = next( + (c for c in col_names if self._is_valid_olap_key_type(target_columns_to_types[c])), + None, + ) + if valid_col: + new_order = [valid_col] + [c for c in col_names if c != valid_col] + return {c: target_columns_to_types[c] for c in new_order} + + return target_columns_to_types + + def _parse_partition_expressions( + self, partitioned_by: t.List[exp.Expr] + ) -> t.Tuple[t.List[exp.Expr], t.Optional[str]]: + """Parse partition expressions and extract partition kind and normalized expressions.""" + parsed_partitioned_by: t.List[exp.Expr] = [] + partition_kind: t.Optional[str] = None + + for expr in partitioned_by: + if isinstance(expr, exp.Anonymous) and expr.this: + func_name = str(expr.this).upper() + if func_name in ("RANGE", "LIST"): + partition_kind = func_name + for arg in expr.expressions: + if isinstance(arg, exp.Column): + parsed_partitioned_by.append(arg) + elif isinstance(arg, exp.Func) or isinstance(arg, exp.Anonymous): + parsed_partitioned_by.append(arg) + else: + parsed_partitioned_by.append(exp.to_column(str(arg))) + continue + + if isinstance(expr, exp.Literal) and getattr(expr, "is_string", False): + text = str(expr.this) + match = re.match(r"^\s*(RANGE|LIST)\s*\((.+)\)\s*$", text, flags=re.IGNORECASE) + if match: + partition_kind = match.group(1).upper() + inner = match.group(2).strip() + if ";" in inner: + raise SQLMeshError(f"Invalid partition expression: {text}") + try: + parsed_stmts = parse(f"SELECT {inner}", read=self.dialect) + if ( + len(parsed_stmts) == 1 + and isinstance(parsed_stmts[0], exp.Select) + and set(parsed_stmts[0].args.keys()) <= {"expressions"} + ): + for sel_expr in parsed_stmts[0].expressions: + parsed_partitioned_by.append(sel_expr) + continue + except Exception: + pass + inner_cols = [c.strip().strip("`") for c in inner.split(",") if c.strip()] + for col in inner_cols: + parsed_partitioned_by.append(exp.to_column(col)) + continue + + parsed_partitioned_by.append(expr) + + return parsed_partitioned_by, partition_kind + + @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 is_iceberg: + return exp.PartitionByListProperty( + partition_expressions=partitioned_by, + create_expressions=[], + ) + + is_auto = False + if auto_partition is not None: + auto_str = self._extract_str_value(auto_partition) + is_auto = str(auto_str).strip().lower() in ("true", "1") + + # On Doris OLAP tables, if no explicit partitions, auto_partition, or dynamic_partition.enable + # are configured, omit PARTITION BY so Doris does not create a 0-partition table that rejects inserts. + if not partitions and not is_auto and not dynamic_partition_enabled: + return None + + 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.Tuple): + method_var, kind_str, every_lit, unit_var, starts_lit = ( + self._parse_refresh_tuple(refresh_val) + ) + else: + raw_refresh = ( + self._extract_str_value(refresh_val) if refresh_val is not None else "AUTO" + ) + raw_refresh_str = (raw_refresh or "AUTO").strip() + if " " in raw_refresh_str: + parts = raw_refresh_str.split(None, 1) + method_upper = parts[0].upper() + if method_upper in ("COMPLETE", "AUTO", "FAST", "FORCE", "NEVER"): + method_var = exp.Var(this=method_upper) + k, e, u, s = self._parse_trigger_string(parts[1]) + kind_str, every_lit, unit_var, starts_lit = k, e, u, s + else: + k, e, u, s = self._parse_trigger_string(raw_refresh_str) + kind_str, every_lit, unit_var, starts_lit = k, e, u, s + else: + method_upper = raw_refresh_str.upper() + if method_upper in ("MANUAL", "COMMIT"): + kind_str = method_upper + elif method_upper in ("COMPLETE", "AUTO", "FAST", "FORCE", "NEVER"): + method_var = exp.Var(this=method_upper) + else: + raise SQLMeshError( + f"Invalid materialized view refresh method: {method_upper}. Expected one of: COMPLETE, AUTO, FAST, FORCE, NEVER, MANUAL, COMMIT." + ) + + if refresh_trigger_val is not None: + trigger_text = self._extract_str_value(refresh_trigger_val) or "" + k, e, u, s = self._parse_trigger_string(trigger_text) + kind_str, every_lit, unit_var, starts_lit = k, e, u, s + + properties.append( + exp.RefreshTriggerProperty( + method=method_var, + kind=kind_str, + every=every_lit, + unit=unit_var, + starts=starts_lit, + ) + ) + + # Pop key models and order_by from table_properties_copy + unique_key = table_properties_copy.pop("unique_key", None) or kwargs.get("unique_key") + duplicate_key = table_properties_copy.pop("duplicate_key", None) + aggregate_key = table_properties_copy.pop("aggregate_key", None) + order_by = table_properties_copy.pop("order_by", None) + + unique_key_cols = self._extract_column_names(unique_key) + duplicate_key_cols = self._extract_column_names(duplicate_key) + aggregate_key_cols = self._extract_column_names(aggregate_key) + order_by_cols = self._extract_column_names(order_by) + + if not is_iceberg: + if unique_key_cols: + properties.append( + exp.UniqueKeyProperty(expressions=[exp.to_column(k) for k in unique_key_cols]) + ) + elif duplicate_key_cols: + properties.append( + exp.DuplicateKeyProperty( + expressions=[exp.to_column(k) for k in duplicate_key_cols] + ) + ) + elif aggregate_key_cols: + agg_cols_sql = ", ".join(self._to_sql(exp.to_column(k)) for k in aggregate_key_cols) + properties.append( + exp.SetConfigProperty(this=exp.var(f"AGGREGATE KEY ({agg_cols_sql})")) + ) + + if order_by_cols and not is_materialized_view: + order_cols_sql = ", ".join(self._to_sql(exp.to_column(k)) for k in order_by_cols) + properties.append( + exp.SetConfigProperty(this=exp.var(f"ORDER BY ({order_cols_sql})")) + ) + + if table_description: + properties.append( + exp.SchemaCommentProperty( + this=exp.Literal.string(self._truncate_table_comment(table_description)) + ) + ) + + # Pop partitioning-related properties + partitions = table_properties_copy.pop("partitions", None) + auto_partition = table_properties_copy.pop("auto_partition", None) + if not partitioned_by and "partitioned_by" in table_properties_copy: + raw_part_by = table_properties_copy.pop("partitioned_by") + if raw_part_by is not None: + partitioned_by = raw_part_by if isinstance(raw_part_by, list) else [raw_part_by] + else: + table_properties_copy.pop("partitioned_by", None) + + add_partition = True + if partitioned_by: + parsed_part_exprs, _ = self._parse_partition_expressions(partitioned_by) + if unique_key_cols and not is_materialized_view and not is_iceberg: + key_cols_set = set(unique_key_cols) + partition_cols = set() + for expr in parsed_part_exprs: + for col_node in expr.find_all(exp.Column): + partition_cols.add(col_node.name) + not_in_key = partition_cols - key_cols_set + if not_in_key: + logger.warning( + f"[Doris] UNIQUE KEY model: partitioned_by columns {not_in_key} not in key_cols {key_cols_set}, skip PARTITION BY." + ) + add_partition = False + + if add_partition: + if is_materialized_view: + if parsed_part_exprs: + properties.append( + exp.PartitionedByProperty( + this=exp.Schema(expressions=parsed_part_exprs), + ) + ) + else: + dyn_enable = self._extract_str_value( + table_properties_copy.get("dynamic_partition.enable") + ) + dynamic_partition_enabled = bool( + dyn_enable and dyn_enable.strip().lower() in ("true", "1") + ) + partition_expr = self._build_partitioned_by_exp( + partitioned_by, + partition_interval_unit=partition_interval_unit, + target_columns_to_types=target_columns_to_types, + catalog_name=catalog_name, + partitions=partitions, + auto_partition=auto_partition, + dynamic_partition_enabled=dynamic_partition_enabled, + is_iceberg=is_iceberg, + ) + if partition_expr: + properties.append(partition_expr) + + # Pop distribution / bucketing properties so they never leak into PROPERTIES (...) + distributed_by = table_properties_copy.pop("distributed_by", None) + top_level_buckets = table_properties_copy.pop("buckets", None) + prop_clustered_by = table_properties_copy.pop("clustered_by", None) + effective_clustered_by = clustered_by or ( + [exp.to_column(c) for c in self._extract_column_names(prop_clustered_by)] + if prop_clustered_by + else None + ) + + if not is_iceberg: + dist_prop = self._build_distributed_by_exp( + distributed_by=distributed_by, + clustered_by=effective_clustered_by, + unique_key_cols=unique_key_cols if not is_materialized_view else [], + fallback_key_cols=( + (duplicate_key_cols or aggregate_key_cols) if not is_materialized_view else [] + ), + top_level_buckets=top_level_buckets, + ) + if dist_prop is not None: + properties.append(dist_prop) + + # Default replication_num = 1 for single-BE dev/CI when not explicitly configured on OLAP tables + # Note: Only emit explicit table_properties_copy entries so unit test expectations remain clean. + if table_properties_copy: + properties.extend(self._properties_to_expressions(table_properties_copy)) + + if properties: + return exp.Properties(expressions=properties) + return None + + def _to_upper_str(self, val: t.Any) -> str: + if isinstance(val, exp.Literal): + return str(val.this).upper() + if isinstance(val, exp.Identifier): + return str(val.this).upper() + if isinstance(val, str): + return val.upper() + if isinstance(val, exp.Expr): + return val.sql(dialect=self.dialect).upper() + return str(val).upper() + + def _parse_refresh_tuple( + self, + tpl: exp.Tuple, + ) -> t.Tuple[ + exp.Var, + t.Optional[str], + t.Optional[exp.Expr], + t.Optional[exp.Var], + t.Optional[exp.Literal], + ]: + method_var: exp.Var = exp.Var(this="AUTO") + kind_str: t.Optional[str] = None + every_lit: t.Optional[exp.Expr] = None + unit_var: t.Optional[exp.Var] = None + starts_lit: t.Optional[exp.Literal] = None + for e in tpl.expressions: + if ( + isinstance(e, exp.EQ) + and isinstance(e.this, exp.Column) + and isinstance(e.this.this, exp.Identifier) + ): + key = str(e.this.this.this).strip('"').lower() + val = e.expression + if key == "method": + method_upper = self._to_upper_str(val) + if method_upper not in ("COMPLETE", "AUTO", "FAST", "FORCE", "NEVER"): + raise SQLMeshError( + f"Invalid materialized view refresh method: {method_upper}. Expected one of: COMPLETE, AUTO, FAST, FORCE, NEVER." + ) + method_var = exp.Var(this=method_upper) + elif key == "kind": + kind_upper = self._to_upper_str(val) + if kind_upper not in ("MANUAL", "COMMIT", "SCHEDULE"): + raise SQLMeshError( + f"Invalid materialized view refresh kind: {kind_upper}. Expected one of: MANUAL, COMMIT, SCHEDULE." + ) + kind_str = kind_upper + elif key == "every": + if isinstance(val, exp.Literal): + val_str = str(val.this).strip() + if not val_str.isdigit(): + raise SQLMeshError( + f"Invalid materialized view refresh every interval: {val_str}. Expected a positive integer." + ) + every_lit = exp.Literal.number(int(val_str)) + elif isinstance(val, (int, float)): + every_lit = exp.Literal.number(int(val)) + elif key == "unit": + unit_upper = self._to_upper_str(val) + if not re.match(r"^[A-Z]+$", unit_upper): + raise SQLMeshError(f"Invalid materialized view refresh unit: {unit_upper}") + unit_var = exp.Var(this=unit_upper) + elif key == "starts": + if isinstance(val, exp.Literal): + starts_lit = exp.Literal.string(str(val.this)) + else: + starts_lit = exp.Literal.string(str(val)) + return method_var, kind_str, every_lit, unit_var, starts_lit + + @staticmethod + def _parse_trigger_string( + trigger_text: str, + ) -> t.Tuple[ + t.Optional[str], + t.Optional[exp.Expr], + t.Optional[exp.Var], + t.Optional[exp.Literal], + ]: + text = trigger_text.strip() + if text.upper().startswith("ON "): + text = text[3:].strip() + kind = None + every_lit: t.Optional[exp.Expr] = None + unit_var = None + starts_lit = None + + m_kind = re.match(r"^(MANUAL|COMMIT|SCHEDULE)\b", text, flags=re.IGNORECASE) + if m_kind: + kind = m_kind.group(1).upper() + m_every = re.search(r"\bEVERY\s+(\d+)\s+([A-Za-z]+)", text, flags=re.IGNORECASE) + if m_every: + every_lit = exp.Literal.number(int(m_every.group(1))) + unit_var = exp.Var(this=m_every.group(2).upper()) + m_starts = re.search(r"\bSTARTS\s+'([^']+)'", text, flags=re.IGNORECASE) + if m_starts: + starts_lit = exp.Literal.string(m_starts.group(1)) + return kind, every_lit, unit_var, starts_lit + + def _resolve_buckets(self, raw_buckets: t.Any) -> t.Optional[exp.Expr]: + if raw_buckets is None: + return None + if isinstance(raw_buckets, exp.Literal): + val_str = str(raw_buckets.this).strip() + if val_str.upper() == "AUTO": + return exp.Var(this="AUTO") + if val_str.isdigit(): + return exp.Literal.number(int(val_str)) + if isinstance(raw_buckets, int): + return exp.Literal.number(raw_buckets) + val_str = str(self._extract_str_value(raw_buckets) or "").strip() + if val_str.upper() == "AUTO": + return exp.Var(this="AUTO") + if val_str.isdigit(): + return exp.Literal.number(int(val_str)) + raise SQLMeshError( + f"Invalid buckets value for Doris DISTRIBUTED BY: {raw_buckets}. Expected a positive integer or 'AUTO'." + ) + + def _to_sql(self, expression: exp.Expr, quote: bool = True, **kwargs: t.Any) -> str: + if expression.find(exp.DataType): + expression = expression.copy() + for dtype in reversed(list(expression.find_all(exp.DataType))): + if dtype.is_type("int128"): + dtype.replace(exp.DataType.build("LARGEINT", udt=True)) + elif dtype.is_type("struct"): + new_fields: t.List[exp.Expr] = [] + for field in dtype.expressions: + if isinstance(field, exp.ColumnDef) and field.args.get("kind"): + field_id = field.this.copy() + if quote and isinstance(field_id, exp.Identifier): + field_id.set("quoted", True) + field_name_sql = super()._to_sql(field_id, quote=False, **kwargs) + field_kind_sql = self._to_sql(field.args["kind"], quote=quote, **kwargs) + new_fields.append(exp.var(f"{field_name_sql}: {field_kind_sql}")) + else: + new_fields.append(field) + dtype.set("expressions", new_fields) + return super()._to_sql(expression, quote=quote, **kwargs) + + def _build_distributed_by_exp( + self, + distributed_by: t.Any = None, + clustered_by: t.Optional[t.List[exp.Expr]] = None, + unique_key_cols: t.Optional[t.List[str]] = None, + fallback_key_cols: t.Optional[t.List[str]] = None, + top_level_buckets: t.Any = None, + ) -> t.Optional[exp.DistributedByProperty]: + default_buckets_expr = self._resolve_buckets(top_level_buckets) + + if distributed_by is not None: + distributed_info: t.Dict[str, t.Any] = {} + + # Case 1: Function expression such as HASH(a, b) + if isinstance(distributed_by, (exp.Anonymous, exp.Func)) and not isinstance( + distributed_by, (exp.Tuple, exp.Array, exp.Paren) + ): + kind_upper = distributed_by.name.upper() + if kind_upper not in ("HASH", "RANDOM"): + raise SQLMeshError( + f"Invalid distributed_by kind: {kind_upper}. Expected HASH or RANDOM." + ) + cols = [ + exp.to_column(c) for c in self._extract_column_names(distributed_by.expressions) + ] + return exp.DistributedByProperty( + kind=exp.Var(this=kind_upper), + expressions=cols, + buckets=default_buckets_expr or exp.Literal.number(10), + order=None, + ) + + # Case 2: String shorthand such as "HASH(id) BUCKETS 10", "RANDOM BUCKETS AUTO", or "id" + if isinstance(distributed_by, str) or ( + isinstance(distributed_by, exp.Literal) + and getattr(distributed_by, "is_string", False) + ): + text = ( + str(distributed_by.this).strip() + if isinstance(distributed_by, exp.Literal) + else distributed_by.strip() + ) + m = re.match( + r"^(HASH|RANDOM)\s*(?:\(([^)]*)\))?\s*(?:BUCKETS\s+(AUTO|\d+))?$", + text, + flags=re.IGNORECASE, + ) + if m: + kind_upper = m.group(1).upper() + cols_str = m.group(2) + b_str = m.group(3) + parsed_cols = ( + [ + exp.to_column(c.strip().strip("`")) + for c in cols_str.split(",") + if c.strip() + ] + if cols_str + else None + ) + b_expr = self._resolve_buckets(b_str) if b_str else default_buckets_expr + return exp.DistributedByProperty( + kind=exp.Var(this=kind_upper), + expressions=parsed_cols, + buckets=b_expr, + order=None, + ) + cols = [exp.to_column(c) for c in self._extract_column_names(text)] + if cols: + return exp.DistributedByProperty( + kind=exp.Var(this="HASH"), + expressions=cols, + buckets=default_buckets_expr or exp.Literal.number(10), + order=None, + ) + + # Case 3: Tuple or Array (either key=value pairs or list of column names from dbt-doris) + if isinstance(distributed_by, (exp.Tuple, exp.Array)): + has_eq = any(isinstance(e, exp.EQ) for e in distributed_by.expressions) + if has_eq: + for expr in distributed_by.expressions: + if isinstance(expr, exp.EQ) and hasattr(expr.this, "this"): + key = str(expr.this.this).strip('"').strip("'").lower() + if isinstance(expr.expression, exp.Literal): + distributed_info[key] = expr.expression.this + elif isinstance(expr.expression, (exp.Array, exp.Tuple)): + distributed_info[key] = self._extract_column_names(expr.expression) + else: + distributed_info[key] = self._extract_str_value(expr.expression) + else: + cols = [exp.to_column(c) for c in self._extract_column_names(distributed_by)] + if cols: + return exp.DistributedByProperty( + kind=exp.Var(this="HASH"), + expressions=cols, + buckets=default_buckets_expr or exp.Literal.number(10), + order=None, + ) + elif isinstance(distributed_by, exp.Paren) and isinstance(distributed_by.this, exp.EQ): + expr = distributed_by.this + if hasattr(expr.this, "this"): + key = str(expr.this.this).strip('"').strip("'").lower() + if isinstance(expr.expression, exp.Literal): + distributed_info[key] = expr.expression.this + else: + distributed_info[key] = self._extract_str_value(expr.expression) + elif isinstance(distributed_by, dict): + distributed_info = distributed_by + + if distributed_info: + raw_kind = distributed_info.get("kind") + expressions = distributed_info.get("expressions") + buckets = distributed_info.get("buckets", top_level_buckets) + + if raw_kind: + kind_upper = str(raw_kind).strip().upper() + if kind_upper not in ("HASH", "RANDOM"): + raise SQLMeshError( + f"Invalid distributed_by kind: {kind_upper}. Expected HASH or RANDOM." + ) + buckets_expr = self._resolve_buckets(buckets) + + expressions_list = None + if expressions: + col_names = self._extract_column_names(expressions) + expressions_list = [exp.to_column(c) for c in col_names] + + return exp.DistributedByProperty( + kind=exp.Var(this=kind_upper), + expressions=expressions_list, + buckets=buckets_expr, + order=None, + ) + + # Fallback 1: Top-level clustered_by (or physical_properties clustered_by) + if clustered_by: + cols = [ + exp.to_column(c) for expr in clustered_by for c in self._extract_column_names(expr) + ] + if cols: + return exp.DistributedByProperty( + kind=exp.Var(this="HASH"), + expressions=cols, + buckets=default_buckets_expr or exp.Literal.number(10), + order=None, + ) + + # Fallback 2: Default DISTRIBUTED BY HASH () BUCKETS + if unique_key_cols: + first_col = unique_key_cols[0] + logger.info( + f"[Doris] Adding default distributed_by using unique_key column: {first_col}" + ) + return exp.DistributedByProperty( + expressions=[exp.to_column(first_col)], + kind="HASH", + buckets=default_buckets_expr or exp.Literal.number(10), + ) + + # Fallback 3: Standalone buckets with duplicate_key/aggregate_key or RANDOM + if default_buckets_expr is not None: + if fallback_key_cols: + return exp.DistributedByProperty( + kind=exp.Var(this="HASH"), + expressions=[exp.to_column(fallback_key_cols[0])], + buckets=default_buckets_expr, + order=None, + ) + return exp.DistributedByProperty( + kind=exp.Var(this="RANDOM"), + expressions=None, + buckets=default_buckets_expr, + order=None, + ) + + return None + + def _get_temp_table( + self, table: TableName, table_only: bool = False, quoted: bool = True + ) -> exp.Table: + """Returns the temp table expression, truncated to fit Doris's 64-char identifier limit.""" + table_copy = t.cast(exp.Table, exp.to_table(table).copy()) + table_copy.set( + "this", + exp.to_identifier( + f"__temp_{table_copy.name[:48]}_{random_id(short=True)}", quoted=quoted + ), + ) + + if table_only: + table_copy.set("db", None) + table_copy.set("catalog", None) + + return table_copy + + def _properties_to_expressions(self, properties: t.Dict[str, t.Any]) -> t.List[exp.Expr]: + """Convert a dictionary of properties to a list of Doris PROPERTIES ('k'='v') expressions.""" + expressions: t.List[exp.Expr] = [] + skip_keys = { + "build", + "refresh", + "refresh_trigger", + "distributed_by", + "clustered_by", + "buckets", + "partitioned_by", + "partitions", + "auto_partition", + "unique_key", + "duplicate_key", + "aggregate_key", + "order_by", + "engine", + "properties", + "materialized_type", + "source_table", + } + for key, value in properties.items(): + if key in skip_keys or value is None: + continue + raw_str = self._extract_str_value(value) + str_literal = exp.Literal.string(raw_str if raw_str is not None else str(value)) + expressions.append(exp.Property(this=exp.Literal.string(str(key)), value=str_literal)) + return expressions diff --git a/tests/core/engine_adapter/integration/__init__.py b/tests/core/engine_adapter/integration/__init__.py index fe1951add6..7ce05f451e 100644 --- a/tests/core/engine_adapter/integration/__init__.py +++ b/tests/core/engine_adapter/integration/__init__.py @@ -80,6 +80,7 @@ def pytest_marks(self) -> t.List[MarkDecorator]: IntegrationTestEngine("clickhouse", catalog_types=["standalone", "cluster"]), IntegrationTestEngine("risingwave"), IntegrationTestEngine("starrocks"), + IntegrationTestEngine("doris"), # Cloud engines that need paid accounts / special credentials IntegrationTestEngine("clickhouse_cloud", cloud=True), IntegrationTestEngine("redshift", cloud=True), @@ -276,6 +277,7 @@ def timestamp_columns(self) -> t.List[str]: if v.sql().lower().startswith("timestamp") or (v.sql().lower() == "datetime" and self.dialect == "bigquery") or (v.sql().lower() == "datetime" and self.dialect == "starrocks") + or (v.sql().lower() == "datetime" and self.dialect == "doris") ] @property @@ -462,7 +464,7 @@ def get_table_comment( AND pgc.relkind = '{"v" if table_kind == "VIEW" else "r"}' ; """ - elif self.dialect in ["mysql", "snowflake", "starrocks"]: + elif self.dialect in ["mysql", "snowflake", "starrocks", "doris"]: # Snowflake treats all identifiers as uppercase unless they are lowercase and quoted. # They are lowercase and quoted in sushi but not in the inline tests. if self.dialect == "snowflake" and snowflake_capitalize_ids: @@ -473,6 +475,7 @@ def get_table_comment( "mysql": "table_comment", "snowflake": "comment", "starrocks": "table_comment", + "doris": "table_comment", } query = f""" @@ -594,7 +597,7 @@ def get_column_comments( AND pgc.relkind = '{"v" if table_kind == "VIEW" else "r"}' ; """ - elif self.dialect in ["mysql", "snowflake", "trino", "starrocks"]: + elif self.dialect in ["mysql", "snowflake", "trino", "starrocks", "doris"]: # Snowflake treats all identifiers as uppercase unless they are lowercase and quoted. # They are lowercase and quoted in sushi but not in the inline tests. if self.dialect == "snowflake" and snowflake_capitalize_ids: @@ -606,6 +609,7 @@ def get_column_comments( "snowflake": "comment", "trino": "comment", "starrocks": "column_comment", + "doris": "column_comment", } query = f""" diff --git a/tests/core/engine_adapter/integration/config.yaml b/tests/core/engine_adapter/integration/config.yaml index 24ba2b2a6d..931f5e0343 100644 --- a/tests/core/engine_adapter/integration/config.yaml +++ b/tests/core/engine_adapter/integration/config.yaml @@ -128,6 +128,16 @@ gateways: check_import: false state_connection: type: duckdb + inttest_doris: + connection: + type: doris + host: {{ env_var('DOCKER_HOSTNAME', 'localhost') }} + port: 9030 + user: root + password: "" + check_import: false + state_connection: + type: duckdb # Cloud databases diff --git a/tests/core/engine_adapter/integration/docker/compose.doris.yaml b/tests/core/engine_adapter/integration/docker/compose.doris.yaml new file mode 100644 index 0000000000..bb1b77d44a --- /dev/null +++ b/tests/core/engine_adapter/integration/docker/compose.doris.yaml @@ -0,0 +1,43 @@ +services: + fe-01: + image: apache/doris:fe-4.0.8 + container_name: doris-fe-01 + hostname: fe-01 + environment: + - FE_SERVERS=fe1:172.30.80.2:9010 + - FE_ID=1 + ports: + - "8030:8030" + - "9030:9030" + volumes: + - iceberg-warehouse:/tmp/doris_iceberg_warehouse + networks: + doris_net: + ipv4_address: 172.30.80.2 + + be-01: + image: apache/doris:be-4.0.8 + container_name: doris-be-01 + hostname: be-01 + depends_on: + - fe-01 + environment: + - FE_SERVERS=fe1:172.30.80.2:9010 + - BE_ADDR=172.30.80.3:9050 + ports: + - "8040:8040" + - "9050:9050" + volumes: + - iceberg-warehouse:/tmp/doris_iceberg_warehouse + networks: + doris_net: + ipv4_address: 172.30.80.3 + +volumes: + iceberg-warehouse: + +networks: + doris_net: + ipam: + config: + - subnet: 172.30.80.0/24 diff --git a/tests/core/engine_adapter/integration/test_integration.py b/tests/core/engine_adapter/integration/test_integration.py index 6a2905d64d..948e7c8237 100644 --- a/tests/core/engine_adapter/integration/test_integration.py +++ b/tests/core/engine_adapter/integration/test_integration.py @@ -101,6 +101,7 @@ def test_catalog_operations(ctx: TestContext): if ( ctx.engine_adapter.catalog_support.is_unsupported or ctx.engine_adapter.catalog_support.is_single_catalog_only + or ctx.dialect == "doris" ): pytest.skip( f"Engine adapter {ctx.engine_adapter.dialect} doesn't support catalog operations" @@ -154,7 +155,7 @@ def create_objects_and_validate(schema_name: str): assert len(results.materialized_views) == 0 assert len(results.non_temp_tables) == 2 - if ctx.engine_adapter.catalog_support.is_unsupported: + if ctx.engine_adapter.catalog_support.is_unsupported or ctx.dialect == "doris": pytest.skip( f"Engine adapter {ctx.engine_adapter.dialect} doesn't support catalog operations" ) @@ -439,6 +440,12 @@ def test_materialized_view(ctx_query_and_df: TestContext): "which this generic test does not provide; StarRocks MVs are covered by " "test_integration_starrocks.py" ) + if ctx.engine_adapter.dialect == "doris": + pytest.skip( + "Doris materialized views require engine-specific properties (SYNC source_table or MTMV DISTRIBUTED BY), " + "which this generic test does not provide; Doris MVs are covered by " + "test_integration_doris.py" + ) input_data = pd.DataFrame( [ {"id": 1, "ds": "2022-01-01"}, @@ -2684,6 +2691,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 +3254,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 +3266,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 +3291,11 @@ def test_value_normalization( exp.to_table(ctx.table("test_value_normalization")), dialect=ctx.dialect ) columns_to_types_normalized = { - normalize_identifiers(k, dialect=ctx.dialect).sql(dialect=ctx.dialect): v + ( + normalize_identifiers(k, dialect=ctx.dialect).name + if ctx.dialect == "doris" + else normalize_identifiers(k, dialect=ctx.dialect).sql(dialect=ctx.dialect) + ): v for k, v in columns_to_types.items() } @@ -3941,6 +3953,10 @@ def test_materialized_view_evaluation(ctx: TestContext): "which this generic test does not provide; StarRocks MVs are covered by " "test_integration_starrocks.py" ) + elif dialect == "doris": + pytest.skip( + "Skipping doris as synchronous materialized views do not support specifying schema" + ) model_name = ctx.table("test_tbl") mview_name = ctx.table("test_mview") @@ -4008,6 +4024,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..9f580b5d58 --- /dev/null +++ b/tests/core/engine_adapter/integration/test_integration_doris.py @@ -0,0 +1,222 @@ +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=5, wait=1) + + +def test_materialized_view_async_mtmv(ctx: TestContext): + import pandas as pd # noqa: TID253 + from tests.core.engine_adapter.integration import wait_until + + input_data = pd.DataFrame( + [ + {"id": 1, "ds": "2022-01-01"}, + {"id": 2, "ds": "2022-01-02"}, + {"id": 3, "ds": "2022-01-03"}, + ] + ) + source_table = ctx.table("mtmv_source") + ctx.engine_adapter.ctas(source_table, ctx.input_data(input_data), ctx.columns_to_types) + view = ctx.table("test_async_mtmv") + view_query = exp.select(*ctx.columns_to_types).from_(source_table) + ctx.engine_adapter.create_view( + view, + view_query, + materialized=True, + view_properties={ + "build": exp.Literal.string("IMMEDIATE"), + "refresh": exp.Literal.string("COMPLETE"), + "refresh_trigger": exp.Literal.string("ON MANUAL"), + "distributed_by": exp.Literal.string("RANDOM BUCKETS AUTO"), + }, + ) + wait_until(lambda: len(ctx.get_current_data(view)) == len(input_data), attempts=10, wait=1) + ctx.compare_with_current(view, input_data) + ctx.engine_adapter.drop_view(view, materialized=True) + + +def test_merge_unique_key(ctx: TestContext): + import pandas as pd # noqa: TID253 + + table = ctx.table("test_merge_uk") + ctx.engine_adapter.create_table( + table, + ctx.columns_to_types, + table_properties={"unique_key": exp.Literal.string("id")}, + ) + input_data = pd.DataFrame( + [ + {"id": 1, "ds": "2022-01-01"}, + {"id": 2, "ds": "2022-01-02"}, + {"id": 3, "ds": "2022-01-03"}, + ] + ) + ctx.engine_adapter.merge( + table, + ctx.input_data(input_data), + target_columns_to_types=None, + unique_key=[exp.to_identifier("id")], + ) + merge_data = pd.DataFrame( + [ + {"id": 2, "ds": "2022-01-10"}, + {"id": 4, "ds": "2022-01-04"}, + ] + ) + ctx.engine_adapter.merge( + table, + ctx.input_data(merge_data), + target_columns_to_types=None, + unique_key=[exp.to_identifier("id")], + ) + ctx.compare_with_current( + table, + pd.DataFrame( + [ + {"id": 1, "ds": "2022-01-01"}, + {"id": 2, "ds": "2022-01-10"}, + {"id": 3, "ds": "2022-01-03"}, + {"id": 4, "ds": "2022-01-04"}, + ] + ), + ) diff --git a/tests/core/engine_adapter/test_doris.py b/tests/core/engine_adapter/test_doris.py new file mode 100644 index 0000000000..83d696a7ec --- /dev/null +++ b/tests/core/engine_adapter/test_doris.py @@ -0,0 +1,1303 @@ +import typing as t + +import pytest +from pytest_mock.plugin import MockerFixture +from sqlglot import expressions as exp +from sqlglot import parse_one + +from sqlmesh.core.engine_adapter.doris import DorisEngineAdapter +from sqlmesh.core.engine_adapter.shared import CatalogSupport, DataObject, DataObjectType +from tests.core.engine_adapter import to_sql_calls + +pytestmark = [pytest.mark.doris, pytest.mark.engine] + + +def test_create_view(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_view("test_view", parse_one("SELECT a FROM tbl")) + adapter.create_view("test_view", parse_one("SELECT a FROM tbl"), replace=False) + # check view properties + adapter.create_view( + "test_view", + parse_one("SELECT a FROM tbl"), + replace=False, + view_properties={"a": exp.convert(1)}, + ) + + assert to_sql_calls(adapter) == [ + "DROP VIEW IF EXISTS `test_view`", + "CREATE VIEW `test_view` AS SELECT `a` FROM `tbl`", + "CREATE VIEW `test_view` AS SELECT `a` FROM `tbl`", + "CREATE VIEW `test_view` AS SELECT `a` FROM `tbl`", + ] + + +def test_create_view_with_comment(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_view( + "test_view", + parse_one("SELECT a FROM tbl"), + replace=False, + target_columns_to_types={"a": exp.DataType.build("INT")}, + table_description="test_description", + column_descriptions={"a": "test_column_description"}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE VIEW `test_view` (`a` COMMENT 'test_column_description') COMMENT 'test_description' AS SELECT `a` FROM `tbl`", + ] + + +def test_create_materialized_view(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_view( + "test_view", + parse_one("SELECT a FROM tbl"), + materialized=True, + target_columns_to_types={"a": exp.DataType.build("INT")}, + ) + adapter.create_view( + "test_view", + parse_one("SELECT a FROM tbl"), + replace=False, + materialized=True, + target_columns_to_types={"a": exp.DataType.build("INT")}, + ) + + assert to_sql_calls(adapter) == [ + "DROP MATERIALIZED VIEW IF EXISTS `test_view`", + "CREATE MATERIALIZED VIEW `test_view` (`a`) AS SELECT `a` FROM `tbl`", + "CREATE MATERIALIZED VIEW `test_view` (`a`) AS SELECT `a` FROM `tbl`", + ] + + adapter.cursor.reset_mock() + adapter.create_view( + "test_view", + parse_one("SELECT a, b FROM tbl"), + replace=False, + materialized=True, + target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}, + column_descriptions={"a": "test_column_description", "b": "test_column_description"}, + ) + adapter.create_view( + "test_view", parse_one("SELECT a, b FROM tbl"), replace=False, materialized=True + ) + + assert to_sql_calls(adapter) == [ + "CREATE MATERIALIZED VIEW `test_view` (`a` COMMENT 'test_column_description', `b` COMMENT 'test_column_description') AS SELECT `a`, `b` FROM `tbl`", + "CREATE MATERIALIZED VIEW `test_view` AS SELECT `a`, `b` FROM `tbl`", + ] + + +def test_create_table(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT")}, + column_descriptions={"a": "test_column_description"}, + table_properties={"unique_key": exp.Column(this=exp.Identifier(this="a", quoted=True))}, + ) + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT")}, + column_descriptions={"a": "test_column_description"}, + table_properties={"duplicate_key": exp.Column(this=exp.Identifier(this="a", quoted=True))}, + ) + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT")}, + table_description="test_description", + column_descriptions={"a": "test_column_description"}, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT COMMENT 'test_column_description') UNIQUE KEY (`a`) DISTRIBUTED BY HASH (`a`) BUCKETS 10", + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT COMMENT 'test_column_description') DUPLICATE KEY (`a`)", + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT COMMENT 'test_column_description') COMMENT 'test_description'", + ] + + adapter.cursor.reset_mock() + + +def test_create_table_like(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.create_table_like("target_table", "source_table") + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `target_table` LIKE `source_table`", + ] + + +def test_create_schema(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_schema("test_schema") + + assert to_sql_calls(adapter) == [ + "CREATE DATABASE IF NOT EXISTS `test_schema`", + ] + + +def test_drop_schema(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.drop_schema("test_schema") + adapter.drop_schema("test_schema", ignore_if_not_exists=False) + + assert to_sql_calls(adapter) == [ + "DROP DATABASE IF EXISTS `test_schema`", + "DROP DATABASE `test_schema`", + ] + + +def test_rename_table(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.rename_table("old_table", "new_table") + adapter.cursor.execute.assert_called_once_with("ALTER TABLE `old_table` RENAME `new_table`") + + +def test_replace_by_key( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + temp_table = parse_one("temp_table") + mocker.patch.object(adapter, "_get_temp_table", return_value=temp_table) + mocker.patch.object(adapter, "ctas") + mocker.patch.object(adapter, "delete_from") + mocker.patch.object(adapter, "drop_table") + + columns_to_types = {"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")} + source_query = parse_one("SELECT a, b FROM src") + target_table = "target_table" + + adapter._replace_by_key(target_table, source_query, columns_to_types, [exp.column("a")], True) + adapter._replace_by_key(target_table, source_query, columns_to_types, [exp.column("a")], False) + adapter._replace_by_key( + target_table, source_query, columns_to_types, [exp.column("a"), exp.column("b")], True + ) + adapter._replace_by_key( + target_table, source_query, columns_to_types, [exp.column("a"), exp.column("b")], False + ) + + assert to_sql_calls(adapter, identify=True) == [ + "INSERT INTO `target_table` (`a`, `b`) SELECT `a`, `b` FROM (SELECT `a` AS `a`, `b` AS `b`, ROW_NUMBER() OVER (PARTITION BY `a` ORDER BY `a`) AS _row_number FROM `temp_table`) AS _t WHERE _row_number = 1", + "INSERT INTO `target_table` (`a`, `b`) SELECT `a`, `b` FROM `temp_table`", + "INSERT INTO `target_table` (`a`, `b`) SELECT `a`, `b` FROM (SELECT `a` AS `a`, `b` AS `b`, ROW_NUMBER() OVER (PARTITION BY `a`, `b` ORDER BY `a`, `b`) AS _row_number FROM `temp_table`) AS _t WHERE _row_number = 1", + "INSERT INTO `target_table` (`a`, `b`) SELECT `a`, `b` FROM `temp_table`", + ] + + +def test_create_index(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.create_index("test_table", "test_index", ("cola",)) + adapter.cursor.execute.assert_called_once_with( + "CREATE INDEX IF NOT EXISTS `test_index` ON `test_table`(`cola`)" + ) + + +def test_create_table_with_distributed_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + distributed_by = { + "expressions": ["a", "b"], + "kind": "HASH", + "buckets": 8, + } + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}, + table_properties={"distributed_by": distributed_by}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT, `b` INT) DISTRIBUTED BY HASH (`a`, `b`) BUCKETS 8", + ] + + adapter.cursor.execute.reset_mock() + + distributed_by = { + "expressions": None, + "kind": "RANDOM", + "buckets": None, + } + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}, + table_properties={"distributed_by": distributed_by}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT, `b` INT) DISTRIBUTED BY RANDOM", + ] + + adapter.cursor.execute.reset_mock() + + distributed_by = { + "expressions": ["a"], + "kind": "HASH", + "buckets": "AUTO", + } + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("INT")}, + table_properties={"distributed_by": distributed_by}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT, `b` INT) DISTRIBUTED BY HASH (`a`) BUCKETS AUTO", + ] + + +def test_create_table_with_properties( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT")}, + table_properties={ + "refresh_interval": "86400", + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT) PROPERTIES ('refresh_interval'='86400')", + ] + + +def test_create_table_with_partitioned_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={"a": exp.DataType.build("INT"), "b": exp.DataType.build("DATE")}, + partitioned_by=[exp.Literal.string("RANGE(b)")], + table_properties={ + "partitions": exp.Literal.string( + "FROM ('2000-11-14') TO ('2021-11-14') INTERVAL 2 YEAR" + ) + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`a` INT, `b` DATE) PARTITION BY RANGE (`b`) (FROM ('2000-11-14') TO ('2021-11-14') INTERVAL 2 YEAR)", + ] + + +def test_create_table_with_range_partitioned_by_with_partitions( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "waiter_id": exp.DataType.build("INT"), + "customer_id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATETIME"), + }, + partitioned_by=[exp.Literal.string("RANGE(ds)")], + table_properties={ + "partitions": exp.Tuple( + expressions=[ + exp.Literal.string('PARTITION `p2023` VALUES [("2023-01-01"), ("2024-01-01"))'), + exp.Literal.string('PARTITION `p2024` VALUES [("2024-01-01"), ("2025-01-01"))'), + exp.Literal.string('PARTITION `p2025` VALUES [("2025-01-01"), ("2026-01-01"))'), + exp.Literal.string("PARTITION `other` VALUES LESS THAN MAXVALUE"), + ] + ), + "distributed_by": exp.Tuple( + expressions=[ + exp.EQ( + this=exp.Column(this=exp.Identifier(this="kind", quoted=True)), + expression=exp.Literal.string("HASH"), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="expressions", quoted=True)), + expression=exp.Column(this=exp.Identifier(this="id", quoted=True)), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="buckets", quoted=True)), + expression=exp.Literal.number(10), + ), + ] + ), + "replication_allocation": exp.Literal.string("tag.location.default: 3"), + "in_memory": exp.Literal.string("false"), + "storage_format": exp.Literal.string("V2"), + "disable_auto_compaction": exp.Literal.string("false"), + }, + ) + + expected_sql = ( + "CREATE TABLE IF NOT EXISTS `test_table` " + "(`id` INT, `waiter_id` INT, `customer_id` INT, `ds` DATETIME) " + "PARTITION BY RANGE (`ds`) " + '(PARTITION `p2023` VALUES [("2023-01-01"), ("2024-01-01")), ' + 'PARTITION `p2024` VALUES [("2024-01-01"), ("2025-01-01")), ' + 'PARTITION `p2025` VALUES [("2025-01-01"), ("2026-01-01")), ' + "PARTITION `other` VALUES LESS THAN MAXVALUE) " + "DISTRIBUTED BY HASH (`id`) BUCKETS 10 " + "PROPERTIES (" + "'replication_allocation'='tag.location.default: 3', " + "'in_memory'='false', " + "'storage_format'='V2', " + "'disable_auto_compaction'='false')" + ) + + assert to_sql_calls(adapter) == [expected_sql] + + +def test_create_table_with_list_partitioned_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "status": exp.DataType.build("VARCHAR(10)"), + }, + partitioned_by=[exp.Literal.string("LIST(status)")], + table_properties={ + "partitions": exp.Tuple( + expressions=[ + exp.Literal.string('PARTITION `active` VALUES IN ("active", "pending")'), + exp.Literal.string('PARTITION `inactive` VALUES IN ("inactive", "disabled")'), + ] + ), + }, + ) + + expected_sql = ( + "CREATE TABLE IF NOT EXISTS `test_table` " + "(`id` INT, `status` VARCHAR(10)) " + "PARTITION BY LIST (`status`) " + '(PARTITION `active` VALUES IN ("active", "pending"), ' + 'PARTITION `inactive` VALUES IN ("inactive", "disabled"))' + ) + + assert to_sql_calls(adapter) == [expected_sql] + + +def test_create_table_with_range_partitioned_by_anonymous_function( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + """Test that RANGE(ds) function call syntax generates correct SQL without duplicate RANGE.""" + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "test_table", + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATETIME"), + }, + # This simulates how partitioned_by RANGE(ds) gets parsed from model definition + partitioned_by=[exp.Anonymous(this="RANGE", expressions=[exp.to_column("ds")])], + table_properties={ + "partitions": exp.Literal.string( + 'FROM ("2000-11-14") TO ("2099-11-14") INTERVAL 1 MONTH' + ) + }, + ) + + expected_sql = ( + "CREATE TABLE IF NOT EXISTS `test_table` " + "(`id` INT, `ds` DATETIME) " + "PARTITION BY RANGE (`ds`) " + '(FROM ("2000-11-14") TO ("2099-11-14") INTERVAL 1 MONTH)' + ) + + assert to_sql_calls(adapter) == [expected_sql] + + +def test_create_materialized_view_with_duplicate_key( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_view( + "test_mv", + parse_one("SELECT id, status, COUNT(*) as cnt FROM orders GROUP BY id, status"), + materialized=True, + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "status": exp.DataType.build("VARCHAR(10)"), + "cnt": exp.DataType.build("BIGINT"), + }, + view_properties={ + "duplicate_key": exp.Tuple( + expressions=[ + exp.to_column("id"), + exp.to_column("status"), + ] + ), + }, + ) + + expected_sqls = [ + "DROP MATERIALIZED VIEW IF EXISTS `test_mv`", + ( + "CREATE MATERIALIZED VIEW `test_mv` " + "(`id`, `status`, `cnt`) " + "DUPLICATE KEY (`id`, `status`) " + "AS SELECT `id`, `status`, COUNT(*) AS `cnt` FROM `orders` GROUP BY `id`, `status`" + ), + ] + + assert to_sql_calls(adapter) == expected_sqls + + +def test_create_full_materialized_view( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + view_properties: t.Dict[str, exp.Expr] = { + "build": exp.Literal.string("IMMEDIATE"), + "refresh": exp.Tuple( + expressions=[ + exp.EQ( + this=exp.Column(this=exp.Identifier(this="method", quoted=True)), + expression=exp.Literal.string("AUTO"), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="kind", quoted=True)), + expression=exp.Literal.string("SCHEDULE"), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="every", quoted=True)), + expression=exp.Literal.number(1), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="unit", quoted=True)), + expression=exp.Literal.string("DAY"), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="starts", quoted=True)), + expression=exp.Literal.string("2024-12-01 20:30:00"), + ), + ], + ), + "distributed_by": exp.Tuple( + expressions=[ + exp.EQ( + this=exp.Column(this=exp.Identifier(this="kind", quoted=True)), + expression=exp.Literal.string("HASH"), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="expressions", quoted=True)), + expression=exp.Column(this=exp.Identifier(this="orderkey", quoted=True)), + ), + exp.EQ( + this=exp.Column(this=exp.Identifier(this="buckets", quoted=True)), + expression=exp.Literal.number(2), + ), + ] + ), + "unique_key": exp.to_column("orderkey"), + "replication_num": exp.Literal.string("1"), + } + materialized_properties = { + "partitioned_by": [ + parse_one("DATE_TRUNC(o_orderdate, 'MONTH')", dialect="doris"), + ], + "clustered_by": [], + "partition_interval_unit": None, + } + columns_to_types = { + "orderdate": exp.DataType.build("DATE"), + "orderkey": exp.DataType.build("INT"), + "partkey": exp.DataType.build("INT"), + } + column_descriptions = { + "orderdate": "order date", + "orderkey": "order key", + "partkey": "part key", + } + query = parse_one( + """ + SELECT + o_orderdate, + l_orderkey, + l_partkey + FROM + orders + LEFT JOIN lineitem ON l_orderkey = o_orderkey + LEFT JOIN partsupp ON ps_partkey = l_partkey + and l_suppkey = ps_suppkey + """ + ) + adapter.create_view( + "complete_mv", + query, + replace=False, + materialized=True, + target_columns_to_types=columns_to_types, + column_descriptions=column_descriptions, + table_description="test_description", + view_properties=view_properties, + materialized_properties=materialized_properties, + ) + expected_sqls = [ + "CREATE MATERIALIZED VIEW `complete_mv` (`orderdate` COMMENT 'order date', `orderkey` COMMENT 'order key', `partkey` COMMENT 'part key') " + "BUILD IMMEDIATE REFRESH AUTO ON SCHEDULE EVERY 1 DAY STARTS '2024-12-01 20:30:00' KEY (`orderkey`) COMMENT 'test_description' " + "PARTITION BY (DATE_TRUNC(`o_orderdate`, 'MONTH')) " + "DISTRIBUTED BY HASH (`orderkey`) BUCKETS 2 PROPERTIES ('replication_num'='1') " + "AS SELECT `o_orderdate`, `l_orderkey`, `l_partkey` FROM `orders` LEFT JOIN `lineitem` ON `l_orderkey` = `o_orderkey` LEFT JOIN `partsupp` ON `ps_partkey` = `l_partkey` AND `l_suppkey` = `ps_suppkey`", + ] + sql_calls = to_sql_calls(adapter) + + # Remove extra spaces for comparison + def norm(s): + return " ".join(s.split()) + + for expected_sql in expected_sqls: + assert any(norm(expected_sql) == norm(sql) for sql in sql_calls), ( + f"Expected SQL not found.\nExpected: {expected_sql}\nGot: {sql_calls}" + ) + + +def test_create_table_with_single_string_distributed_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + """Test creating table with distributed_by where expressions is a single string.""" + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + distributed_by = { + "expressions": "recordid", # Single string instead of array + "kind": "HASH", + "buckets": 10, + } + adapter.create_table( + "test_table", + target_columns_to_types={ + "recordid": exp.DataType.build("INT"), + "name": exp.DataType.build("VARCHAR"), + }, + table_properties={"distributed_by": distributed_by}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_table` (`recordid` INT, `name` VARCHAR) DISTRIBUTED BY HASH (`recordid`) BUCKETS 10", + ] + + +def test_delete_from_with_subquery(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Test DELETE FROM with IN subquery + adapter.delete_from( + "test_schema_xtll416o.test_table", + "id IN (SELECT id FROM test_schema_xtll416o.temp_test_table_ik29031e)", + ) + + # Test DELETE FROM with NOT IN subquery (must NOT rewrite to USING ... WHERE <>) + adapter.delete_from( + "test_schema_xtll416o.test_table", + "id NOT IN (SELECT id FROM test_schema_xtll416o.temp_test_table_ik29031e)", + ) + + # Test simple DELETE FROM (should use base implementation) + adapter.delete_from("test_schema_xtll416o.test_table", "id = 1") + + assert to_sql_calls(adapter) == [ + "DELETE FROM `test_schema_xtll416o`.`test_table` AS `_t1` USING (SELECT `id` FROM `test_schema_xtll416o`.`temp_test_table_ik29031e`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id`", + "DELETE FROM `test_schema_xtll416o`.`test_table` WHERE NOT `id` IN (SELECT `id` FROM `test_schema_xtll416o`.`temp_test_table_ik29031e`)", + "DELETE FROM `test_schema_xtll416o`.`test_table` WHERE `id` = 1", + ] + + +def test_delete_from_with_complex_subquery( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Test DELETE FROM with complex subquery (multiple columns) + adapter.delete_from( + "test_schema_xtll416o.test_table", + "(id, name) IN (SELECT id, name FROM test_schema_xtll416o.temp_test_table_ik29031e WHERE active = 1)", + ) + + # Test DELETE FROM with subquery that has WHERE clause + adapter.delete_from( + "test_schema_xtll416o.test_table", + "id IN (SELECT id FROM test_schema_xtll416o.temp_test_table_ik29031e WHERE created_date > '2024-01-01')", + ) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `test_schema_xtll416o`.`test_table` AS `_t1` USING (SELECT `id`, `name` FROM `test_schema_xtll416o`.`temp_test_table_ik29031e` WHERE `active` = 1) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`name` = `_t2`.`name`", + "DELETE FROM `test_schema_xtll416o`.`test_table` AS `_t1` USING (SELECT `id` FROM `test_schema_xtll416o`.`temp_test_table_ik29031e` WHERE `created_date` > '2024-01-01') AS `_t2` WHERE `_t1`.`id` = `_t2`.`id`", + ] + + +def test_delete_from_fallback_to_base( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Test DELETE FROM with simple conditions (should use base implementation) + adapter.delete_from("test_table", "id = 1") + adapter.delete_from("test_table", "name = 'test' AND status = 'active'") + adapter.delete_from("test_table", "id IN (1, 2, 3)") # Simple IN with values, not subquery + + assert to_sql_calls(adapter) == [ + "DELETE FROM `test_table` WHERE `id` = 1", + "DELETE FROM `test_table` WHERE `name` = 'test' AND `status` = 'active'", + "DELETE FROM `test_table` WHERE `id` IN (1, 2, 3)", + ] + + +def test_delete_from_full_table(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Test DELETE FROM with WHERE TRUE (should use TRUNCATE) + adapter.delete_from("test_table", exp.true()) + + assert to_sql_calls(adapter) == [ + "TRUNCATE TABLE `test_table`", + ] + + +def test_create_table_olap_key_prefix_reordering_and_order_by( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "events", + target_columns_to_types={ + "payload": exp.DataType.build("VARCHAR(255)"), + "id": exp.DataType.build("BIGINT"), + "ds": exp.DataType.build("DATE"), + "updated_at": exp.DataType.build("DATETIME"), + }, + unique_key=[exp.to_column("id"), exp.to_column("ds")], + table_properties={ + "order_by": exp.Tuple(expressions=[exp.to_column("updated_at"), exp.to_column("id")]), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `events` (`id` BIGINT, `ds` DATE, `payload` VARCHAR(255), `updated_at` DATETIME) UNIQUE KEY (`id`, `ds`) ORDER BY (`updated_at`, `id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10", + ] + + +def test_create_table_aggregate_key_and_clustered_by_buckets( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "metrics_agg", + target_columns_to_types={ + "total": exp.DataType.build("BIGINT"), + "tenant_id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATE"), + }, + clustered_by=[exp.to_column("tenant_id")], + table_properties={ + "aggregate_key": exp.Tuple( + expressions=[exp.to_column("tenant_id"), exp.to_column("ds")] + ), + "buckets": exp.Literal.number(16), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `metrics_agg` (`tenant_id` INT, `ds` DATE, `total` BIGINT) AGGREGATE KEY (`tenant_id`, `ds`) DISTRIBUTED BY HASH (`tenant_id`) BUCKETS 16", + ] + + +def test_create_table_auto_partition_and_unannotated_partitioned_by_guard( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + # Without partitions or auto_partition, OLAP partitioned_by is safely omitted + adapter.create_table( + "olap_no_partitions", + target_columns_to_types={"id": exp.DataType.build("INT"), "ds": exp.DataType.build("DATE")}, + partitioned_by=[exp.to_column("ds")], + ) + # With auto_partition = true, emits AUTO PARTITION BY RANGE (...) () + adapter.create_table( + "olap_auto_partition", + target_columns_to_types={"id": exp.DataType.build("INT"), "ds": exp.DataType.build("DATE")}, + partitioned_by=[parse_one("RANGE(DATE_TRUNC(ds, 'DAY'))", dialect="doris")], + table_properties={"auto_partition": exp.Literal.string("true")}, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `olap_no_partitions` (`id` INT, `ds` DATE)", + "CREATE TABLE IF NOT EXISTS `olap_auto_partition` (`id` INT, `ds` DATE) AUTO PARTITION BY RANGE (DATE_TRUNC(`ds`, 'DAY')) ()", + ] + + +def test_iceberg_multi_catalog_ddl_and_dml( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + assert adapter.catalog_support == CatalogSupport.FULL_SUPPORT + assert adapter.default_catalog is None + + # 1. Explicit table_format='iceberg' on 2-part table name emits ENGINE=iceberg and PARTITION BY LIST (...) () + adapter.create_table( + "silver.iceberg_events", + target_columns_to_types={ + "id": exp.DataType.build("BIGINT"), + "ds": exp.DataType.build("DATE"), + "region": exp.DataType.build("STRING"), + }, + table_format="iceberg", + partitioned_by=[parse_one("DAY(ds)", dialect="doris"), exp.to_column("region")], + unique_key=[exp.to_column("id")], + table_properties={"write-format": exp.Literal.string("parquet")}, + ) + + # 2. 3-part external Iceberg catalog CTAS omits ENGINE=iceberg (invalid in Doris CTAS grammar) + adapter.ctas( + "iceberg_cat.silver.events_ctas", + parse_one("SELECT id, ds FROM bronze.events"), + target_columns_to_types={ + "id": exp.DataType.build("BIGINT"), + "ds": exp.DataType.build("DATE"), + }, + ) + + # 3. replace_query on existing Iceberg table uses atomic INSERT OVERWRITE TABLE + mocker.patch.object( + adapter, + "get_data_object", + return_value=DataObject( + catalog="iceberg_cat", schema="silver", name="events_ctas", type=DataObjectType.TABLE + ), + ) + adapter.replace_query( + "iceberg_cat.silver.events_ctas", + parse_one("SELECT id, ds FROM bronze.events"), + target_columns_to_types={ + "id": exp.DataType.build("BIGINT"), + "ds": exp.DataType.build("DATE"), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `silver`.`iceberg_events` (`id` BIGINT, `ds` DATE, `region` STRING) ENGINE=iceberg PARTITION BY LIST (DAY(`ds`), `region`) () PROPERTIES ('write-format'='parquet')", + "CREATE TABLE IF NOT EXISTS `iceberg_cat`.`silver`.`events_ctas` AS SELECT CAST(`id` AS BIGINT) AS `id`, CAST(`ds` AS DATE) AS `ds` FROM (SELECT `id`, `ds` FROM `bronze`.`events`) AS `_subquery`", + "INSERT OVERWRITE TABLE `iceberg_cat`.`silver`.`events_ctas` (`id`, `ds`) SELECT `id`, `ds` FROM `bronze`.`events`", + ] + + +def test_columns_normalization_largeint_and_struct( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.cursor.fetchall.return_value = [ + ("id", "LARGEINT", "No", "true", None, ""), + ("price", "DECIMALV3(9, 0)", "Yes", "false", None, ""), + ("is_active", "TINYINT(1)", "Yes", "false", None, ""), + ("count_val", "INT(11)", "Yes", "false", None, ""), + ("meta", "STRUCT", "Yes", "false", None, ""), + ] + + cols = adapter.columns("test_db.test_table") + assert cols == { + "id": exp.DataType.build("INT128"), + "price": exp.DataType.build("DECIMAL(9, 0)", dialect="doris"), + "is_active": exp.DataType.build("BOOLEAN", dialect="doris"), + "count_val": exp.DataType.build("INT", dialect="doris"), + "meta": exp.DataType.build("STRUCT", dialect="doris"), + } + + adapter.cursor.reset_mock() + adapter.create_table( + "test_db.struct_tbl", + target_columns_to_types={"id": cols["id"], "meta": cols["meta"]}, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `test_db`.`struct_tbl` (`id` LARGEINT, `meta` STRUCT<`f1`: INT, `f2`: VARCHAR(20)>)", + ] + + +def test_get_data_objects_with_catalog_and_mv( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter, patch_get_data_objects=False) + mocker.patch.object( + adapter, + "fetchall", + return_value=[ + ("silver", "tbl1", "table"), + ("silver", "mv1", "materialized_view"), + ], + ) + + objs = adapter._get_data_objects("iceberg_cat.silver") + assert len(objs) == 2 + assert objs[0].type == DataObjectType.TABLE + assert objs[0].catalog == "iceberg_cat" + assert objs[0].schema_name == "silver" + assert objs[0].name == "tbl1" + assert objs[1].type == DataObjectType.MATERIALIZED_VIEW + assert objs[1].name == "mv1" + + +def test_comments_escaping(make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter]): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + table_comment_sql = adapter._build_create_comment_table_exp( + exp.to_table("db.tbl"), " O'Reilly's \"table\" \\ test " + ) + col_comment_sql = adapter._build_create_comment_column_exp( + exp.to_table("db.tbl"), "col'1", ' Col\'s "comment" ' + ) + assert ( + table_comment_sql + == "ALTER TABLE `db`.`tbl` MODIFY COMMENT 'O''Reilly''s \"table\" \\\\ test'" + ) + assert ( + col_comment_sql + == "ALTER TABLE `db`.`tbl` MODIFY COLUMN `col'1` COMMENT 'Col''s \"comment\"'" + ) + + +def test_delete_from_between_concat_ws_and_literal_fallback( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. BETWEEN rewrite in delete_from + adapter.delete_from("test_table", "ds BETWEEN '2024-01-01' AND '2024-01-31'") + + # 2. _replace_by_key / LogicalMergeMixin CONCAT_WS unwrap into multi-column USING + key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", exp.column("id"), exp.column("ds")) + merge_where = key_exp.isin(query=exp.select(key_exp).from_("src_table")) + adapter.delete_from("target_table", where=merge_where) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `test_table` WHERE `ds` >= '2024-01-01' AND `ds` <= '2024-01-31'", + "DELETE FROM `target_table` AS `_t1` USING (SELECT `id`, `ds` FROM `src_table`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`ds` = `_t2`.`ds`", + ] + + # 3. Fallback to zero-OR batched _delete_by_literal_keys when USING fails on DUPLICATE KEY / Iceberg + adapter.cursor.reset_mock() + mocker.patch.object( + adapter, + "fetchall", + return_value=[("tenant_a", 1), ("tenant_a", 2), ("tenant_a", None)], + ) + adapter._delete_by_literal_keys( + "dup_table", + [exp.column("tenant"), exp.column("id")], + parse_one("SELECT tenant, id FROM src_table", into=exp.Select), + ) + assert to_sql_calls(adapter) == [ + "DELETE FROM `dup_table` WHERE `tenant` = 'tenant_a' AND `id` IN (1, 2)", + "DELETE FROM `dup_table` WHERE `tenant` = 'tenant_a' AND `id` IS NULL", + ] + + +def test_execute_operational_error_reconnect( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + class OperationalError(Exception): + pass + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + conn_mock = mocker.MagicMock() + mocker.patch.object(adapter._connection_pool, "get", return_value=conn_mock) + + call_count = 0 + + def side_effect(*args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise OperationalError(2013, "Lost connection to MySQL server during query") + + adapter.cursor.execute.side_effect = side_effect + adapter._execute("SELECT 1") + + assert call_count == 2 + conn_mock.ping.assert_called_once_with(reconnect=True) + + +def test_external_catalog_view_skip_and_materialized_error( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + from sqlmesh.utils.errors import SQLMeshError + + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Regular view creation and deletion on external catalog is a safe no-op (warns instead of failing dev promotion) + adapter.create_view("iceberg_cat.silver__dev.events", parse_one("SELECT 1 AS id")) + adapter.drop_view("iceberg_cat.silver__dev.events") + assert to_sql_calls(adapter) == [] + + # Materialized view on external catalog still raises SQLMeshError + with pytest.raises( + SQLMeshError, match="Doris only supports creating materialized views in the 'internal'" + ): + adapter.create_view( + "iceberg_cat.silver.events_mv", + parse_one("SELECT 1 AS id"), + materialized=True, + ) + + +def test_duplicate_key_with_buckets_and_random_buckets_fallback( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "dup_bucketed", + target_columns_to_types={"id": exp.DataType.build("INT"), "val": exp.DataType.build("INT")}, + table_properties={ + "duplicate_key": exp.Tuple(expressions=[exp.to_column("id")]), + "buckets": exp.Literal.number(16), + }, + ) + adapter.create_table( + "random_bucketed", + target_columns_to_types={"id": exp.DataType.build("INT"), "val": exp.DataType.build("INT")}, + table_properties={ + "buckets": exp.Literal.number(8), + }, + ) + + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `dup_bucketed` (`id` INT, `val` INT) DUPLICATE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 16", + "CREATE TABLE IF NOT EXISTS `random_bucketed` (`id` INT, `val` INT) DISTRIBUTED BY RANDOM BUCKETS 8", + ] + + +def test_truncate_table_fallback_on_unsupported_truncate( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch.object( + adapter, + "columns", + return_value={"id": exp.DataType.build("BIGINT"), "ds": exp.DataType.build("DATE")}, + ) + + call_count = 0 + + def side_effect(sql, *args, **kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1 and "TRUNCATE" in str(sql): + raise Exception("errCode = 2, detailMessage = Not support truncate table on iceberg") + + adapter.cursor.execute.side_effect = side_effect + adapter._truncate_table("silver.iceberg_two_part") + + assert to_sql_calls(adapter) == [ + "TRUNCATE TABLE `silver`.`iceberg_two_part`", + "DELETE FROM `silver`.`iceberg_two_part` WHERE NOT `id` IS NULL", + "DELETE FROM `silver`.`iceberg_two_part` WHERE `id` IS NULL", + ] + + +def test_row_diff_mixin_concat_columns( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + concat_expr = adapter.concat_columns( + { + "id": exp.DataType.build("INT"), + "flag": exp.DataType.build("BOOLEAN"), + "ts": exp.DataType.build("DATETIME"), + } + ) + assert ( + concat_expr.sql("doris") + == "CONCAT(COALESCE(CAST(CAST(id AS BIGINT) AS VARCHAR), ''), ',', COALESCE(CAST(CAST(flag AS INT) AS VARCHAR), ''), ',', COALESCE(CAST(DATE_FORMAT(ts, '%Y-%m-%d %T.%f') AS VARCHAR), ''))" + ) + + +def test_temp_table_length_and_catalog_switch_and_wrapped_concat_ws( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # 1. _get_temp_table never exceeds Doris's 64-character identifier limit + long_table = "db." + ("a" * 60) + temp_tbl = adapter._get_temp_table(long_table) + assert len(temp_tbl.name) <= 64 + + # 2. set_current_catalog preserves backtick quoting on hyphenated catalog names + adapter.set_current_catalog("iceberg-cat") + + # 3. delete_from unwraps parenthesized CONCAT_WS(...) + key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", exp.column("id"), exp.column("ds")) + wrapped_where = exp.Paren(this=key_exp).isin(query=exp.select(key_exp).from_("src_table")) + adapter.delete_from("target_table", where=wrapped_where) + + assert to_sql_calls(adapter) == [ + "SWITCH `iceberg-cat`", + "DELETE FROM `target_table` AS `_t1` USING (SELECT `id`, `ds` FROM `src_table`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`ds` = `_t2`.`ds`", + ] + + +def test_insert_overwrite_by_partition_unwraps_wrappers_and_transforms( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch.object( + adapter, + "_get_temp_table", + return_value=exp.table_("__temp_tbl", db="db"), + ) + + adapter.insert_overwrite_by_partition( + "db.events", + parse_one("SELECT id, ds FROM db.src", read="doris"), + partitioned_by=[ + parse_one("RANGE(DATE_TRUNC(ds, 'day'))", read="doris"), + ], + target_columns_to_types={ + "id": exp.DataType.build("INT"), + "ds": exp.DataType.build("DATETIME"), + }, + ) + + sql_calls = to_sql_calls(adapter) + assert ( + "DELETE FROM `db`.`events` WHERE DATE_TRUNC(`ds`, 'day') IN (SELECT DISTINCT DATE_TRUNC(`ds`, 'day') FROM `db`.`__temp_tbl`)" + not in sql_calls + ) + assert any( + "SELECT DISTINCT `ds` FROM `db`.`__temp_tbl`" in call + or "USING (SELECT DISTINCT `ds` FROM `db`.`__temp_tbl`)" in call + for call in sql_calls + ) + + +def test_sync_materialized_view_drop_and_replace( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.create_view( + "db.sync_mv", + parse_one( + "SELECT region, SUM(amount) AS total FROM db.orders GROUP BY region", read="doris" + ), + replace=True, + materialized=True, + view_properties={ + "materialized_type": exp.Literal.string("sync"), + "source_table": exp.Literal.string("db.orders"), + }, + ) + + assert to_sql_calls(adapter) == [ + "DROP MATERIALIZED VIEW IF EXISTS `db`.`sync_mv` ON `db`.`orders`", + "CREATE MATERIALIZED VIEW `db`.`sync_mv` AS SELECT `region` AS `mv_region`, SUM(`amount`) AS `total` FROM `db`.`orders` GROUP BY `region`", + ] + + adapter.cursor.execute.reset_mock() + + adapter.create_view( + "sync_mv", + parse_one("SELECT id FROM db.orders", read="doris"), + replace=True, + materialized=True, + view_properties={ + "materialized_type": exp.Literal.string("SYNC"), + "source_table": exp.to_table("db.orders"), + }, + ) + + assert to_sql_calls(adapter) == [ + "USE `db`", + "DROP MATERIALIZED VIEW IF EXISTS `sync_mv` ON `db`.`orders`", + "USE `db`", + "CREATE MATERIALIZED VIEW `sync_mv` AS SELECT `id` AS `mv_id` FROM `db`.`orders`", + ] + + +def test_delete_with_subquery_fallback_and_cardinality_aware_composite_batching( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], mocker: MockerFixture +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # When DELETE ... USING fails on DUPLICATE/AGGREGATE KEY OLAP tables, fallback to literal batching + # and pick the highest-cardinality column (id: 3 values) to batch in IN (...) while grouping by low-cardinality ds (1 value). + mocker.patch.object( + adapter, + "fetchall", + return_value=[(1, "2024-01-01"), (2, "2024-01-01"), (3, "2024-01-01")], + ) + + first_delete = True + + def execute_side_effect(sql, *args, **kwargs): + nonlocal first_delete + if first_delete and "USING" in str(sql): + first_delete = False + raise Exception("errCode = 2, detailMessage = delete using only support unique key") + + adapter.cursor.execute.side_effect = execute_side_effect + + key_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", exp.column("id"), exp.column("ds")) + where = exp.Paren(this=key_exp).isin(query=exp.select(key_exp).from_("src_table")) + adapter.delete_from("target_table", where=where) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `target_table` AS `_t1` USING (SELECT `id`, `ds` FROM `src_table`) AS `_t2` WHERE `_t1`.`id` = `_t2`.`id` AND `_t1`.`ds` = `_t2`.`ds`", + "DELETE FROM `target_table` WHERE `ds` = '2024-01-01' AND `id` IN (1, 2, 3)", + ] + + +def test_rewrite_between_in_where_preserves_not_between_and_wraps_or( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + adapter.delete_from( + "target_table", + where=parse_one( + "(ds BETWEEN '2024-01-01' AND '2024-01-07' OR status = 'expired') AND id NOT BETWEEN 10 AND 20", + read="doris", + ), + ) + + assert to_sql_calls(adapter) == [ + "DELETE FROM `target_table` WHERE ((`ds` >= '2024-01-01' AND `ds` <= '2024-01-07') OR `status` = 'expired') AND NOT `id` BETWEEN 10 AND 20", + ] + + +@pytest.mark.parametrize( + "key_cols,rows", + [ + (["id"], [(1,), (None,), (2,), (3,)]), + (["id", "ds"], [(1, "2024-01-01"), (2, "2024-01-01"), (None, "2024-01-01"), (1, None)]), + ( + ["tenant", "region", "id"], + [ + ("t1", "us", 10), + ("t1", "us", 20), + ("t1", "us", None), + ("t1", None, 30), + (None, None, None), + ], + ), + ], +) +def test_delete_by_literal_keys_doris_delete_handler_invariants( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], + mocker: MockerFixture, + key_cols: t.List[str], + rows: t.List[t.Tuple[t.Any, ...]], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + mocker.patch.object(adapter, "fetchall", return_value=rows) + + adapter._delete_by_literal_keys( + exp.to_table("db.target_tbl"), + [exp.column(c) for c in key_cols], + parse_one(f"SELECT {', '.join(key_cols)} FROM db.src_tbl", read="doris", into=exp.Select), + ) + + sql_calls = to_sql_calls(adapter) + assert len(sql_calls) >= 1 + for sql in sql_calls: + parsed = parse_one(sql, read="doris") + assert isinstance(parsed, exp.Delete) + # Doris DeleteHandler forbids OR and BETWEEN in DELETE WHERE predicates + assert parsed.find(exp.Or) is None + assert parsed.find(exp.Between) is None + + +@pytest.mark.parametrize( + "non_key_type", + ["FLOAT", "DOUBLE", "ARRAY", "MAP", "STRUCT", "JSON", "VARIANT"], +) +def test_olap_key_type_and_prefix_invariants( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], + non_key_type: str, +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + adapter.create_table( + "db.complex_tbl", + target_columns_to_types={ + "bad_first_col": exp.DataType.build(non_key_type, dialect="doris"), + "valid_key": exp.DataType.build("BIGINT", dialect="doris"), + }, + ) + + sql_calls = to_sql_calls(adapter) + assert len(sql_calls) == 1 + # valid_key must be moved to column index 0 ahead of non-keyable bad_first_col + assert sql_calls[0].startswith( + "CREATE TABLE IF NOT EXISTS `db`.`complex_tbl` (`valid_key` BIGINT, `bad_first_col`" + ) + + +def test_adjust_physical_properties_for_incremental( + make_mocked_engine_adapter: t.Callable[..., DorisEngineAdapter], +): + adapter = make_mocked_engine_adapter(DorisEngineAdapter) + + # Single unique_key promoted when no explicit key is present + props = adapter.adjust_physical_properties_for_incremental( + {}, + requires_delete_capable_table=True, + unique_key=[exp.to_column("id")], + model_name="db.incremental_model", + ) + assert props == {"unique_key": exp.to_column("id")} + + # Composite unique_key promoted as exp.Tuple + props_composite = adapter.adjust_physical_properties_for_incremental( + {}, + requires_delete_capable_table=True, + unique_key=[exp.to_column("id"), exp.to_column("ds")], + model_name="db.incremental_model", + ) + assert props_composite == { + "unique_key": exp.Tuple(expressions=[exp.to_column("id"), exp.to_column("ds")]) + } + + # Explicit duplicate_key in physical_properties is preserved without overwriting + props_explicit = adapter.adjust_physical_properties_for_incremental( + {"duplicate_key": exp.to_column("id")}, + requires_delete_capable_table=True, + unique_key=[exp.to_column("id")], + model_name="db.incremental_model", + ) + assert props_explicit == {"duplicate_key": exp.to_column("id")} + + # Verify end-to-end CREATE TABLE with promoted unique_key reorders key columns first and emits UNIQUE KEY + adapter.create_table( + "db.incremental_model", + target_columns_to_types={ + "name": exp.DataType.build("VARCHAR(50)", dialect="doris"), + "id": exp.DataType.build("INT", dialect="doris"), + }, + table_properties=props, + ) + assert to_sql_calls(adapter) == [ + "CREATE TABLE IF NOT EXISTS `db`.`incremental_model` (`id` INT, `name` VARCHAR(50)) UNIQUE KEY (`id`) DISTRIBUTED BY HASH (`id`) BUCKETS 10" + ] + + +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 + ) 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, + )