feat(wren-engine): add Apache Doris connector support

ibis-server/app/custom_sqlglot/dialects/__init__.py

@@ -1,3 +1,4 @@
 # ruff: noqa: F401
 
+from app.custom_sqlglot.dialects.doris import Doris
 from app.custom_sqlglot.dialects.mysql import MySQL

ibis-server/app/custom_sqlglot/dialects/doris.py

@@ -0,0 +1,10 @@
+from sqlglot import exp
+from sqlglot.dialects import Doris as OriginalDoris
+
+
+class Doris(OriginalDoris):
+    class Generator(OriginalDoris.Generator):
+        TYPE_MAPPING = {
+            **OriginalDoris.Generator.TYPE_MAPPING,
+            exp.DataType.Type.VARBINARY: "BINARY",
+        }

ibis-server/app/model/__init__.py

@@ -53,6 +53,10 @@ class QueryMySqlDTO(QueryDTO):
     connection_info: ConnectionUrl | MySqlConnectionInfo = connection_info_field
 
 
+class QueryDorisDTO(QueryDTO):
+    connection_info: DorisConnectionInfo = connection_info_field
+
+
 class QueryOracleDTO(QueryDTO):
     connection_info: ConnectionUrl | OracleConnectionInfo = connection_info_field
 
@@ -322,6 +326,28 @@ class MySqlConnectionInfo(BaseConnectionInfo):
     )
 
 
+class DorisConnectionInfo(BaseConnectionInfo):
+    host: SecretStr = Field(
+        description="the hostname of your Doris FE", examples=["localhost"]
+    )
+    port: SecretStr = Field(
+        description="the query port of your Doris FE", examples=["9030"]
+    )
+    database: SecretStr = Field(
+        description="the database name of your Doris database", examples=["default"]
+    )
+    user: SecretStr = Field(
+        description="the username of your Doris database", examples=["root"]
+    )
+    password: SecretStr | None = Field(
+        description="the password of your Doris database", examples=["password"],
+        default=None,
+    )
+    kwargs: dict[str, str] | None = Field(
+        description="Additional keyword arguments to pass to PyMySQL", default=None
+    )
+
+
 class PostgresConnectionInfo(BaseConnectionInfo):
     host: SecretStr = Field(
         examples=["localhost"], description="the hostname of your database"
@@ -636,6 +662,7 @@ class GcsFileConnectionInfo(BaseConnectionInfo):
     | ConnectionUrl
     | MSSqlConnectionInfo
     | MySqlConnectionInfo
+    | DorisConnectionInfo
     | OracleConnectionInfo
     | PostgresConnectionInfo
     | RedshiftConnectionInfo

ibis-server/app/model/connector.py

@@ -104,6 +104,8 @@ def __init__(self, data_source: DataSource, connection_info: ConnectionInfo):
             self._connector = DatabricksConnector(connection_info)
         elif data_source == DataSource.mysql:
             self._connector = MySqlConnector(connection_info)
+        elif data_source == DataSource.doris:
+            self._connector = DorisConnector(connection_info)
         else:
             self._connector = IbisConnector(data_source, connection_info)
 
@@ -357,6 +359,43 @@ def _cast_json_columns(self, result_table: Table, col_name: str) -> Table:
         return result_table.mutate(**{col_name: casted_col})
 
 
+class DorisConnector(IbisConnector):
+    """Doris connector - reuses MySQL protocol via ibis.mysql backend.
+
+    Doris is an analytical database that is MySQL-protocol compatible.
+    Autocommit is forced on in get_doris_connection() because Doris may not
+    properly reflect the SERVER_STATUS_AUTOCOMMIT flag, which would cause
+    ibis's raw_sql() to wrap every query in BEGIN/ROLLBACK unnecessarily.
+    """
+
+    def __init__(self, connection_info: ConnectionInfo):
+        super().__init__(DataSource.doris, connection_info)
+
+    def _handle_pyarrow_unsupported_type(self, ibis_table: Table, **kwargs) -> Table:
+        result_table = ibis_table
+        for name, dtype in ibis_table.schema().items():
+            if isinstance(dtype, Decimal):
+                result_table = self._round_decimal_columns(
+                    result_table=result_table, col_name=name, **kwargs
+                )
+            elif isinstance(dtype, UUID):
+                result_table = self._cast_uuid_columns(
+                    result_table=result_table, col_name=name
+                )
+            elif isinstance(dtype, dt.JSON):
+                # Doris JSON columns need the same handling as MySQL
+                result_table = self._cast_json_columns(
+                    result_table=result_table, col_name=name
+                )
+
+        return result_table
+
+    def _cast_json_columns(self, result_table: Table, col_name: str) -> Table:
+        col = result_table[col_name]
+        casted_col = col.cast("string")
+        return result_table.mutate(**{col_name: casted_col})
+
+
 class MSSqlConnector(IbisConnector):
     def __init__(self, connection_info: ConnectionInfo):
         super().__init__(DataSource.mssql, connection_info)

ibis-server/app/model/data_source.py

@@ -24,6 +24,7 @@
     ConnectionUrl,
     DatabricksServicePrincipalConnectionInfo,
     DatabricksTokenConnectionInfo,
+    DorisConnectionInfo,
     GcsFileConnectionInfo,
     LocalFileConnectionInfo,
     MinioFileConnectionInfo,
@@ -36,6 +37,7 @@
     QueryCannerDTO,
     QueryClickHouseDTO,
     QueryDatabricksDTO,
+    QueryDorisDTO,
     QueryDTO,
     QueryGcsFileDTO,
     QueryLocalFileDTO,
@@ -69,6 +71,7 @@ class DataSource(StrEnum):
     clickhouse = auto()
     mssql = auto()
     mysql = auto()
+    doris = auto()
     oracle = auto()
     postgres = auto()
     redshift = auto()
@@ -167,6 +170,8 @@ def _build_connection_info(self, data: dict) -> ConnectionInfo:
                 return MSSqlConnectionInfo.model_validate(data)
             case DataSource.mysql:
                 return MySqlConnectionInfo.model_validate(data)
+            case DataSource.doris:
+                return DorisConnectionInfo.model_validate(data)
             case DataSource.oracle:
                 return OracleConnectionInfo.model_validate(data)
             case DataSource.postgres:
@@ -236,6 +241,7 @@ class DataSourceExtension(Enum):
     clickhouse = QueryClickHouseDTO
     mssql = QueryMSSqlDTO
     mysql = QueryMySqlDTO
+    doris = QueryDorisDTO
     oracle = QueryOracleDTO
     postgres = QueryPostgresDTO
     redshift = QueryRedshiftDTO
@@ -397,6 +403,42 @@ def get_mysql_connection(cls, info: MySqlConnectionInfo) -> BaseBackend:
             **kwargs,
         )
 
+    @classmethod
+    def get_doris_connection(cls, info: DorisConnectionInfo) -> BaseBackend:
+        kwargs = {}
+
+        # utf8mb4 is the actual charset used by Doris (MySQL-compatible)
+        kwargs.setdefault("charset", "utf8mb4")
+
+        if info.kwargs:
+            kwargs.update(info.kwargs)
+        # Doris is MySQL-protocol compatible, reuse ibis.mysql.connect()
+        connection = ibis.mysql.connect(
+            host=info.host.get_secret_value(),
+            port=int(info.port.get_secret_value()),
+            database=info.database.get_secret_value(),
+            user=info.user.get_secret_value(),
+            password=info.password.get_secret_value() if info.password else "",
+            **kwargs,
+        )
+        # Doris does not properly reflect the SERVER_STATUS_AUTOCOMMIT flag
+        # in its MySQL-protocol handshake/OK packets. As a result, the
+        # underlying mysqlclient driver's get_autocommit() always returns
+        # False — even after explicitly calling autocommit(True).
+        #
+        # ibis's raw_sql() checks get_autocommit() and, when it returns
+        # False, wraps every query in BEGIN/ROLLBACK. Doris (an OLAP engine)
+        # does not support transactional SELECT inside BEGIN and will reject
+        # with: "This is in a transaction, only insert, update, delete,
+        # commit, rollback is acceptable."
+        #
+        # Fix: override get_autocommit on THIS connection instance only so
+        # that ibis skips the BEGIN/ROLLBACK wrapping. This is a per-object
+        # attribute override — it does NOT affect the MySQLdb class, other
+        # MySQL connections, or any other data-source driver.
+        connection.con.get_autocommit = lambda: True
+        return connection
+
     @staticmethod
     def get_postgres_connection(info: PostgresConnectionInfo) -> BaseBackend:
         return ibis.postgres.connect(