Feat(doris): Add support for Doris

Author: xinge-ji

.circleci/continue_config.yml

@@ -285,6 +285,7 @@ workflows:
           matrix:
             parameters:
               engine:
+                - doris
                 - duckdb
                 - postgres
                 - mysql

.circleci/wait-for-db.sh

@@ -34,6 +34,34 @@ clickhouse_ready() {
     probe_port 8123
 }
 
+doris_ready() {
+    probe_port 9030
+
+    echo "Checking for 3 alive Doris backends..."
+    sleep 15
+
+    while true; do
+        echo "Checking Doris backends..."
+        ALIVE_BACKENDS=$(docker exec -i doris-fe-01 mysql -h127.0.0.1 -P9030 -uroot -e "show backends \G" | grep -c "^ *Alive: true$")
+
+        # fallback value if failed to get number
+        if ! [[ "$ALIVE_BACKENDS" =~ ^[0-9]+$ ]]; then
+            echo "WARN: Unable to parse number of alive backends, got: '$ALIVE_BACKENDS'"
+            ALIVE_BACKENDS=0
+        fi
+
+        echo "Found $ALIVE_BACKENDS alive backends"
+
+        if [ "$ALIVE_BACKENDS" -ge 3 ]; then
+            echo "Doris has 3 or more alive backends"
+            break
+        fi
+
+        echo "Waiting for more backends to become alive..."
+        sleep 5
+    done
+}
+
 postgres_ready() {
     probe_port 5432
 }

.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,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,trino,web,risingwave]"
       - make api-docs
 
 mkdocs:

Makefile

@@ -148,7 +148,9 @@ trino-test: engine-trino-up
 
 risingwave-test: engine-risingwave-up
 	pytest -n auto -m "risingwave" --retries 3 --junitxml=test-results/junit-risingwave.xml
-
+	
+doris-test: engine-doris-up
+	pytest -n auto -m "doris" --retries 3 --junitxml=test-results/junit-doris.xml
 #################
 # Cloud Engines #
 #################

pyproject.toml

@@ -102,6 +102,7 @@ dev = [
 ]
 dbt = ["dbt-core<2"]
 dlt = ["dlt"]
+doris = ["pymysql"]
 duckdb = []
 fabric = ["pyodbc>=5.0.0"]
 gcppostgres = ["cloud-sql-python-connector[pg8000]>=1.8.0"]
@@ -263,7 +264,8 @@ markers = [
     "snowflake: test for Snowflake",
     "spark: test for Spark",
     "trino: test for Trino (all connectors)",
-    "risingwave: test for Risingwave"
+    "risingwave: test for Risingwave",
+    "doris: test for Doris",
 ]
 addopts = "-n 0 --dist=loadgroup"
 asyncio_default_fixture_loop_scope = "session"

sqlmesh/cli/project_init.py

@@ -167,7 +167,7 @@ class ExampleObjects:
     python_macros: t.Dict[str, str]
 
 
-def _gen_example_objects(schema_name: str) -> ExampleObjects:
+def _gen_example_objects(schema_name: str, dialect: str) -> ExampleObjects:
     sql_models: t.Dict[str, str] = {}
     python_models: t.Dict[str, str] = {}
     seeds: t.Dict[str, str] = {}
@@ -200,6 +200,7 @@ def _gen_example_objects(schema_name: str) -> ExampleObjects:
   name {incremental_model_name},
   kind INCREMENTAL_BY_TIME_RANGE (
     time_column event_date
+    {"partition_by_time_column false" if dialect == "doris" else ""}
   ),
   start '2020-01-01',
   cron '@daily',
@@ -352,7 +353,10 @@ def init_example_project(
         )
         return config_path
 
-    example_objects = _gen_example_objects(schema_name=schema_name)
+    example_objects = _gen_example_objects(
+        schema_name=schema_name,
+        dialect=dialect if dialect else DIALECT_TO_TYPE.get(engine_type, "duckdb"),
+    )
 
     if template != ProjectTemplate.EMPTY:
         _create_object_files(models_path, example_objects.sql_models, "sql")

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,

sqlmesh/core/config/connection.py

@@ -56,6 +56,8 @@
     "trino",
     # Nullable types are problematic
     "clickhouse",
+    # Do not support table name starts with "_"
+    "doris",
 }
 MOTHERDUCK_TOKEN_REGEX = re.compile(r"(\?|\&)(motherduck_token=)(\S*)")
 
@@ -2274,6 +2276,80 @@ def init(cursor: t.Any) -> None:
         return init
 
 
+class DorisConnectionConfig(ConnectionConfig):
+    """Configuration for the Apache Doris connection.
+
+    Apache Doris uses MySQL network protocol and is compatible with MySQL ecosystem tools,
+    JDBC/ODBC drivers, and various visualization tools.
+
+    Args:
+        host: The hostname of the Doris FE (Frontend) node.
+        user: The Doris username.
+        password: The Doris password.
+        port: The port number of the Doris FE node. Default is 9030.
+        database: The optional database name.
+        charset: The optional character set.
+        collation: The optional collation.
+        ssl_disabled: Whether to disable SSL 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.
+        local_infile: Whether or not to allow local file access.
+        pre_ping: Whether or not to pre-ping the connection before starting a new transaction to ensure it is still alive.
+    """
+
+    host: str
+    user: str
+    password: str
+    port: t.Optional[int] = 9030  # Default Doris FE port
+    database: t.Optional[str] = None
+    charset: t.Optional[str] = None
+    collation: t.Optional[str] = None
+    ssl_disabled: t.Optional[bool] = None
+
+    concurrent_tasks: int = 4
+    register_comments: bool = True
+    local_infile: bool = False
+    pre_ping: bool = True
+
+    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[17]] = 17
+
+    _engine_import_validator = _get_engine_import_validator("pymysql", "doris")
+
+    @property
+    def _connection_kwargs_keys(self) -> t.Set[str]:
+        connection_keys = {
+            "host",
+            "user",
+            "password",
+        }
+        if self.port is not None:
+            connection_keys.add("port")
+        if self.database is not None:
+            connection_keys.add("database")
+        if self.charset is not None:
+            connection_keys.add("charset")
+        if self.collation is not None:
+            connection_keys.add("collation")
+        if self.ssl_disabled is not None:
+            connection_keys.add("ssl_disabled")
+        if self.local_infile is not None:
+            connection_keys.add("local_infile")
+        return connection_keys
+
+    @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
+
+
 CONNECTION_CONFIG_TO_TYPE = {
     # Map all subclasses of ConnectionConfig to the value of their `type_` field.
     tpe.all_field_infos()["type_"].default: tpe

sqlmesh/core/engine_adapter/__init__.py

@@ -9,6 +9,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
@@ -37,6 +38,7 @@
     "athena": AthenaEngineAdapter,
     "risingwave": RisingwaveEngineAdapter,
     "fabric": FabricEngineAdapter,
+    "doris": DorisEngineAdapter,
 }
 
 DIALECT_ALIASES = {

sqlmesh/core/engine_adapter/base.py

@@ -2586,6 +2586,13 @@ def _to_sql(self, expression: exp.Expression, quote: bool = True, **kwargs: t.An
 
         return expression.sql(**sql_gen_kwargs, copy=False)  # type: ignore
 
+    def _get_temp_table_name(self, table: TableName) -> str:
+        """
+        Get the name of the temp table.
+        """
+        table_obj = exp.to_table(table)
+        return f"__temp_{table_obj.name}_{random_id(short=True)}"
+    
     def _get_data_objects(
         self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None
     ) -> t.List[DataObject]:
@@ -2602,7 +2609,8 @@ def _get_temp_table(
         """
         table = t.cast(exp.Table, exp.to_table(table).copy())
         table.set(
-            "this", exp.to_identifier(f"__temp_{table.name}_{random_id(short=True)}", quoted=quoted)
+            "this",
+            exp.to_identifier(self._get_temp_table_name(table), quoted=quoted),
         )
 
         if table_only: