diff --git a/.github/workflows/wren-ci.yml b/.github/workflows/wren-ci.yml index 42ecc5b1c..49ac80734 100644 --- a/.github/workflows/wren-ci.yml +++ b/.github/workflows/wren-ci.yml @@ -93,6 +93,10 @@ jobs: tests/connectors/test_mysql.py tests/connectors/test_mysql_connector.py marker: mysql + - datasource: starrocks + extra: mysql + test_file: tests/connectors/test_starrocks.py + marker: starrocks defaults: run: working-directory: core/wren @@ -104,7 +108,7 @@ jobs: - name: Install uv uses: astral-sh/setup-uv@v4 - name: Install system dependencies for mysqlclient - if: matrix.datasource == 'mysql' + if: matrix.extra == 'mysql' run: | sudo apt-get update sudo apt-get install -y default-libmysqlclient-dev pkg-config diff --git a/core/wren-core-base/manifest-macro/src/lib.rs b/core/wren-core-base/manifest-macro/src/lib.rs index 95b007031..16e70b3ea 100644 --- a/core/wren-core-base/manifest-macro/src/lib.rs +++ b/core/wren-core-base/manifest-macro/src/lib.rs @@ -93,6 +93,8 @@ pub fn data_source(python_binding: proc_macro::TokenStream) -> proc_macro::Token MySQL, #[serde(alias = "doris")] Doris, + #[serde(alias = "starrocks")] + StarRocks, #[serde(alias = "postgres")] Postgres, #[serde(alias = "snowflake")] diff --git a/core/wren-core-base/src/mdl/manifest.rs b/core/wren-core-base/src/mdl/manifest.rs index 362e3e044..897d902e8 100644 --- a/core/wren-core-base/src/mdl/manifest.rs +++ b/core/wren-core-base/src/mdl/manifest.rs @@ -185,6 +185,7 @@ impl Display for DataSource { DataSource::MSSQL => write!(f, "MSSQL"), DataSource::MySQL => write!(f, "MYSQL"), DataSource::Doris => write!(f, "DORIS"), + DataSource::StarRocks => write!(f, "STARROCKS"), DataSource::Postgres => write!(f, "POSTGRES"), DataSource::Snowflake => write!(f, "SNOWFLAKE"), DataSource::Datafusion => write!(f, "DATAFUSION"), @@ -214,6 +215,7 @@ impl FromStr for DataSource { "MSSQL" => Ok(DataSource::MSSQL), "MYSQL" => Ok(DataSource::MySQL), "DORIS" => Ok(DataSource::Doris), + "STARROCKS" => Ok(DataSource::StarRocks), "POSTGRES" => Ok(DataSource::Postgres), "SNOWFLAKE" => Ok(DataSource::Snowflake), "DATAFUSION" => Ok(DataSource::Datafusion), diff --git a/core/wren-core/core/src/mdl/dialect/inner_dialect.rs b/core/wren-core/core/src/mdl/dialect/inner_dialect.rs index 035916ca1..1c13b0cef 100644 --- a/core/wren-core/core/src/mdl/dialect/inner_dialect.rs +++ b/core/wren-core/core/src/mdl/dialect/inner_dialect.rs @@ -98,6 +98,7 @@ pub fn get_inner_dialect(data_source: &DataSource) -> Box { match data_source { DataSource::MySQL => Box::new(MySQLDialect {}), DataSource::Doris => Box::new(MySQLDialect {}), + DataSource::StarRocks => Box::new(MySQLDialect {}), DataSource::BigQuery => Box::new(BigQueryDialect {}), DataSource::Oracle => Box::new(OracleDialect {}), DataSource::MSSQL => Box::new(MsSqlDialect {}), diff --git a/core/wren-mdl/mdl.schema.json b/core/wren-mdl/mdl.schema.json index b9c7ea9bd..3f3b54f17 100644 --- a/core/wren-mdl/mdl.schema.json +++ b/core/wren-mdl/mdl.schema.json @@ -273,6 +273,7 @@ "MSSQL", "MYSQL", "DORIS", + "STARROCKS", "POSTGRES", "SNOWFLAKE", "DATAFUSION", @@ -293,6 +294,7 @@ "mssql", "mysql", "doris", + "starrocks", "postgres", "snowflake", "datafusion", diff --git a/core/wren/.claude/CLAUDE.md b/core/wren/.claude/CLAUDE.md index fab489f5e..b826a8e08 100644 --- a/core/wren/.claude/CLAUDE.md +++ b/core/wren/.claude/CLAUDE.md @@ -57,7 +57,7 @@ Uses `uv` (not Poetry). `pyproject.toml` uses `hatchling` as build backend. ## Connectors -`connector/factory.py` dispatches on `DataSource` to return the right connector. Each connector wraps a native driver/client and returns Arrow tables, exposing `.query(sql, limit)` and `.dry_run(sql)`. Base class in `connector/base.py` (`ConnectorABC`, plus the shared `coerce_limit` helper). There is **no `connector/ibis.py`**, and no module under `connector/` imports `ibis` — `ibis-framework` is not a dependency of this package. Most data sources have their own module; the exceptions are `doris`, which shares `mysql.py` as `DorisConnector(MySqlConnector)` (MySQL-compatible protocol), and the file sources `local_file`, `s3_file`, `minio_file` and `gcs_file`, which share `duckdb.py`. +`connector/factory.py` dispatches on `DataSource` to return the right connector. Each connector wraps a native driver/client and returns Arrow tables, exposing `.query(sql, limit)` and `.dry_run(sql)`. Base class in `connector/base.py` (`ConnectorABC`, plus the shared `coerce_limit` helper). There is **no `connector/ibis.py`**, and no module under `connector/` imports `ibis` — `ibis-framework` is not a dependency of this package. Most data sources have their own module; the exceptions are `doris` and `starrocks`, which share `mysql.py` as `DorisConnector(MySqlConnector)` (MySQL-compatible protocol), and the file sources `local_file`, `s3_file`, `minio_file` and `gcs_file`, which share `duckdb.py`. - **Dedicated modules**: `postgres.py`, `mysql.py`, `mssql.py`, `bigquery.py`, `duckdb.py`, `oracle.py`, `redshift.py`, `spark.py`, `databricks.py`, `trino.py`, `clickhouse.py`, `snowflake.py`, `athena.py`, `datafusion.py`, `canner.py` - **canner** dispatches to its own `canner.py` (Postgres wire protocol over `psycopg`). The `canner: postgres` entry in `factory._INSTALL_EXTRA` is only the pip extra to suggest when an import fails — it is not a dispatch target diff --git a/core/wren/README.md b/core/wren/README.md index 2f388747b..f7f47009d 100644 --- a/core/wren/README.md +++ b/core/wren/README.md @@ -270,6 +270,7 @@ dependencies change, re-run an install recipe first. | `just test-duckdb` | DuckDB connector tests | No | | `just test-postgres` | PostgreSQL connector tests | Yes | | `just test-mysql` | MySQL connector tests | Yes | +| `just test-starrocks` | StarRocks connector tests (needs the `mysql` extra) | Yes | | `just test` | All tests | Yes | Profile web tests (`test_profile_web.py`) require `wrenai[ui]`: diff --git a/core/wren/justfile b/core/wren/justfile index 6858c1466..d3c27138b 100644 --- a/core/wren/justfile +++ b/core/wren/justfile @@ -77,6 +77,9 @@ test-postgres: test-mysql: uv run --no-sync pytest tests/connectors/test_mysql.py tests/connectors/test_mysql_connector.py -v -m mysql +test-starrocks: + uv run --no-sync pytest tests/connectors/test_starrocks.py -v -m starrocks + test-snowflake: uv run --no-sync pytest tests/connectors/test_snowflake.py -v -m snowflake diff --git a/core/wren/src/wren/connector/factory.py b/core/wren/src/wren/connector/factory.py index 7528c9069..97722066c 100644 --- a/core/wren/src/wren/connector/factory.py +++ b/core/wren/src/wren/connector/factory.py @@ -8,6 +8,7 @@ _REGISTRY: dict[DataSource, str] = { DataSource.postgres: "wren.connector.postgres", DataSource.mysql: "wren.connector.mysql", DataSource.doris: "wren.connector.mysql", + DataSource.starrocks: "wren.connector.mysql", DataSource.mssql: "wren.connector.mssql", DataSource.canner: "wren.connector.canner", DataSource.bigquery: "wren.connector.bigquery", @@ -30,6 +31,7 @@ _REGISTRY: dict[DataSource, str] = { # Map data sources to the correct pip extra when they share a connector module _INSTALL_EXTRA: dict[DataSource, str] = { DataSource.doris: "mysql", + DataSource.starrocks: "mysql", DataSource.canner: "postgres", DataSource.local_file: "duckdb", DataSource.s3_file: "duckdb", @@ -40,6 +42,7 @@ _INSTALL_EXTRA: dict[DataSource, str] = { _NEEDS_DATA_SOURCE = { DataSource.mysql, DataSource.doris, + DataSource.starrocks, DataSource.trino, } diff --git a/core/wren/src/wren/connector/mysql.py b/core/wren/src/wren/connector/mysql.py index c7b00c7a9..2e88ee90d 100644 --- a/core/wren/src/wren/connector/mysql.py +++ b/core/wren/src/wren/connector/mysql.py @@ -1,11 +1,11 @@ -"""Native MySQLdb connector for MySQL and Doris. +"""Native MySQLdb connector for MySQL, Doris and StarRocks. Replaces the previous ibis-based implementation. Uses the ``mysqlclient`` (``MySQLdb``) driver directly and builds PyArrow tables from cursor descriptions so no ibis backend is required. -Doris speaks the MySQL wire protocol and reuses the same query path; it -only differs in how the connection is opened. +Doris and StarRocks speak the MySQL wire protocol and reuse the same query +path; they only differ in how the connection is opened. """ from __future__ import annotations @@ -111,7 +111,9 @@ class MySqlConnector(ConnectorABC): class DorisConnector(MySqlConnector): - """Doris connector. Speaks MySQL protocol; routes through Doris connection.""" + """Doris / StarRocks connector. Speaks MySQL protocol; routes through the + Doris connection (also used for StarRocks, which shares the same wire + protocol and ``ANSI_QUOTES`` limitation).""" def __init__(self, connection_info): import MySQLdb # noqa: PLC0415 @@ -125,7 +127,7 @@ class DorisConnector(MySqlConnector): def create_connector(data_source: DataSource, connection_info) -> MySqlConnector: - if data_source == DataSource.doris: + if data_source in (DataSource.doris, DataSource.starrocks): return DorisConnector(connection_info) return MySqlConnector(connection_info) @@ -687,7 +689,8 @@ def _build_mysql_connect_kwargs(connection_info) -> dict: def _build_doris_connect_kwargs(connection_info) -> dict: - """Translate ``DorisConnectionInfo`` / ``ConnectionUrl`` into MySQLdb kwargs.""" + """Translate ``DorisConnectionInfo`` / ``StarRocksConnectionInfo`` / + ``ConnectionUrl`` into MySQLdb kwargs.""" if hasattr(connection_info, "connection_url"): return _build_mysql_connect_kwargs(connection_info) diff --git a/core/wren/src/wren/context.py b/core/wren/src/wren/context.py index a5f74a3ca..e301af2eb 100644 --- a/core/wren/src/wren/context.py +++ b/core/wren/src/wren/context.py @@ -529,6 +529,7 @@ _VALID_DIALECTS = { "s3_file", "snowflake", "spark", + "starrocks", "trino", } diff --git a/core/wren/src/wren/dbt.py b/core/wren/src/wren/dbt.py index 04e46ab32..f29930d59 100644 --- a/core/wren/src/wren/dbt.py +++ b/core/wren/src/wren/dbt.py @@ -54,6 +54,7 @@ DBT_ADAPTER_TO_WREN_DATASOURCE = { "snowflake": "snowflake", "spark": "spark", "sqlserver": "mssql", + "starrocks": "starrocks", "trino": "trino", } @@ -475,6 +476,30 @@ def convert_dbt_target_to_wren_profile(target: DbtTarget) -> dict[str, Any]: }, ) + if datasource == "starrocks": + # Wren's StarRocks connection has no catalog field, so a profile that + # targets a catalog other than the default would silently connect to + # ``default_catalog``. Refuse it instead of querying the wrong place. + catalog = output.get("catalog") + if catalog not in (None, "", "default_catalog"): + raise DbtLoadError( + f"dbt target uses StarRocks catalog {catalog!r}, which Wren does " + "not support yet; only 'default_catalog' can be imported." + ) + # dbt-starrocks names the login ``username`` and the database ``schema``. + return _build_wren_profile( + "starrocks", + { + "host": str(_require_output_field(output, "host")), + "port": str(output.get("port", "9030")), + "database": str( + _require_output_field(output, "schema", "database", "dbname") + ), + "user": str(_require_output_field(output, "username", "user")), + "password": str(output["password"]) if output.get("password") else None, + }, + ) + if datasource in {"mysql", "redshift", "mssql", "clickhouse"}: return _build_wren_profile( datasource, diff --git a/core/wren/src/wren/model/__init__.py b/core/wren/src/wren/model/__init__.py index b0ba2b4da..4e7b49368 100644 --- a/core/wren/src/wren/model/__init__.py +++ b/core/wren/src/wren/model/__init__.py @@ -152,6 +152,10 @@ class DorisConnectionInfo(BaseConnectionInfo): kwargs: dict[str, str] | None = Field(default=None) +class StarRocksConnectionInfo(DorisConnectionInfo): + """StarRocks speaks the MySQL wire protocol (FE query port, 9030 by default).""" + + class PostgresConnectionInfo(BaseConnectionInfo): host: str = Field(examples=["localhost"]) port: StrPort = Field(examples=["5432"]) @@ -307,6 +311,7 @@ ConnectionInfo = ( | MSSqlConnectionInfo | MySqlConnectionInfo | DorisConnectionInfo + | StarRocksConnectionInfo | OracleConnectionInfo | PostgresConnectionInfo | RedshiftConnectionInfo diff --git a/core/wren/src/wren/model/data_source.py b/core/wren/src/wren/model/data_source.py index 8630de69b..59dd666e1 100644 --- a/core/wren/src/wren/model/data_source.py +++ b/core/wren/src/wren/model/data_source.py @@ -30,6 +30,7 @@ from wren.model import ( S3FileConnectionInfo, SnowflakeConnectionInfo, SparkConnectionInfo, + StarRocksConnectionInfo, TrinoConnectionInfo, ) from wren.model.error import ErrorCode, WrenError @@ -46,6 +47,7 @@ class DataSource(StrEnum): mssql = auto() mysql = auto() doris = auto() + starrocks = auto() oracle = auto() postgres = auto() redshift = auto() @@ -130,6 +132,8 @@ class DataSource(StrEnum): return MySqlConnectionInfo.model_validate(data) case DataSource.doris: return DorisConnectionInfo.model_validate(data) + case DataSource.starrocks: + return StarRocksConnectionInfo.model_validate(data) case DataSource.oracle: return OracleConnectionInfo.model_validate(data) case DataSource.postgres: diff --git a/core/wren/src/wren/model/field_registry.py b/core/wren/src/wren/model/field_registry.py index d227beec2..dde970bcd 100644 --- a/core/wren/src/wren/model/field_registry.py +++ b/core/wren/src/wren/model/field_registry.py @@ -37,6 +37,7 @@ from wren.model import ( S3FileConnectionInfo, SnowflakeConnectionInfo, SparkConnectionInfo, + StarRocksConnectionInfo, TrinoConnectionInfo, ) @@ -65,6 +66,7 @@ DATASOURCE_MODELS: dict[str, list[type[BaseConnectionInfo]]] = { "s3_file": [S3FileConnectionInfo], "snowflake": [SnowflakeConnectionInfo], "spark": [SparkConnectionInfo], + "starrocks": [StarRocksConnectionInfo], "trino": [TrinoConnectionInfo], "connection_url": [ConnectionUrl], } diff --git a/core/wren/tests/conftest.py b/core/wren/tests/conftest.py index d170cd102..bde9fc111 100644 --- a/core/wren/tests/conftest.py +++ b/core/wren/tests/conftest.py @@ -18,6 +18,9 @@ def pytest_configure(config: pytest.Config) -> None: "markers", "postgres: PostgreSQL connector tests — requires Docker" ) config.addinivalue_line("markers", "mysql: MySQL connector tests — requires Docker") + config.addinivalue_line( + "markers", "starrocks: StarRocks connector tests — requires Docker" + ) config.addinivalue_line( "markers", "snowflake: Snowflake connector tests — mocked, no Docker required" ) diff --git a/core/wren/tests/connectors/test_starrocks.py b/core/wren/tests/connectors/test_starrocks.py new file mode 100644 index 000000000..7ac6b8166 --- /dev/null +++ b/core/wren/tests/connectors/test_starrocks.py @@ -0,0 +1,237 @@ +"""StarRocks connector tests. + +Uses ``testcontainers`` to spin up a real single-node StarRocks (FE + BE in one +``allin1-ubuntu`` image). TPCH-shaped fixture data is fabricated inline in +Python (no network downloads) and loaded over the MySQL protocol (FE query port +9030), the same way the connector talks to the server. + +The StarRocks-specific tests at the bottom guard the reason ``starrocks`` is its +own data source: SQL must be transpiled with the sqlglot ``starrocks`` dialect, +not ``doris`` (``DATE_TRUNC`` argument order, ``ARRAY_AGG``, ``GROUP_CONCAT``). +""" + +from __future__ import annotations + +import base64 +import datetime as _dt +import json +import time + +import orjson +import pytest +from testcontainers.core.container import DockerContainer + +from tests.suite.manifests import make_tpch_manifest +from tests.suite.query import WrenQueryTestSuite +from wren import WrenEngine +from wren.model.data_source import DataSource + +pytestmark = pytest.mark.starrocks + +_IMAGE = "starrocks/allin1-ubuntu:3.5.21" +_DATABASE = "wren_test" +_QUERY_PORT = 9030 +_ORDER_COUNT = 15000 +_CUSTOMER_COUNT = 1500 +_ORDER_STATUSES = ("O", "F", "P") +_BASE_DATE = _dt.date(1992, 1, 1) +_INSERT_BATCH = 1000 + + +def _make_fixture_rows() -> tuple[list[tuple], list[tuple]]: + """Fabricate TPCH-shaped orders + customer rows without network access. + + Row counts match TPCH sf=0.01 so the shared ``WrenQueryTestSuite`` + assertions (15000 orders, 1500 customers, first orderkey == 1) hold. + """ + customers = [(i, f"Customer#{i:09d}") for i in range(1, _CUSTOMER_COUNT + 1)] + orders = [ + ( + i, + ((i - 1) % _CUSTOMER_COUNT) + 1, + _ORDER_STATUSES[i % len(_ORDER_STATUSES)], + float(100 + i), + _BASE_DATE + _dt.timedelta(days=i % 3650), + ) + for i in range(1, _ORDER_COUNT + 1) + ] + return orders, customers + + +class _StarRocksContainer(DockerContainer): + """Minimal single-node StarRocks container — exposes the FE query port.""" + + def __init__(self, image: str = _IMAGE): + super().__init__(image) + self.with_exposed_ports(_QUERY_PORT) + + def get_host_ip(self) -> str: + return self.get_container_host_ip() + + def get_query_port(self) -> int: + return int(self.get_exposed_port(_QUERY_PORT)) + + +def _connect(host: str, port: int, database: str | None = None): + import MySQLdb # noqa: PLC0415 + + # MySQLdb treats ``host="localhost"`` as the unix socket /tmp/mysql.sock; + # testcontainers returns "localhost" when Docker is reached over a socket. + # The connector itself rewrites it the same way (_build_mysql_connect_kwargs). + kwargs = { + "host": "127.0.0.1" if host == "localhost" else host, + "port": port, + "user": "root", + "passwd": "", + } + if database: + kwargs["db"] = database + conn = MySQLdb.connect(**kwargs) + conn.autocommit(True) + return conn + + +def _wait_until_ready(host: str, port: int, timeout: float = 300.0) -> None: + """Wait for the FE to accept connections and for the BE to be usable. + + The FE answers on 9030 well before a backend is available; creating a table + before then fails with "no alive backend". ``Alive=true`` is also reported + about two seconds before the BE has reported its disk capacity, and a + ``CREATE TABLE`` inside that window fails with "backends without enough + disk space", so also wait until ``TotalCapacity`` is non-zero. + """ + deadline = time.time() + timeout + last_err: Exception | None = None + while time.time() < deadline: + conn = None + try: + conn = _connect(host, port) + cur = conn.cursor() + cur.execute("SHOW BACKENDS") + columns = [d[0] for d in cur.description] + alive = columns.index("Alive") + capacity = columns.index("TotalCapacity") + rows = cur.fetchall() + if rows and all( + str(r[alive]).lower() == "true" + and not str(r[capacity]).startswith("0.000") + for r in rows + ): + return + last_err = RuntimeError(f"backends not ready yet: {rows!r}") + except Exception as e: # noqa: BLE001 + last_err = e + finally: + if conn is not None: + conn.close() + time.sleep(2) + raise RuntimeError(f"StarRocks did not become ready: {last_err}") + + +def _load_tpch(host: str, port: int) -> None: + """Create the schema and bulk-load fabricated TPCH-shaped data.""" + orders_rows, customer_rows = _make_fixture_rows() + + conn = _connect(host, port) + try: + cur = conn.cursor() + cur.execute(f"CREATE DATABASE IF NOT EXISTS {_DATABASE}") + cur.execute(f"USE {_DATABASE}") + cur.execute( + "CREATE TABLE orders (" + " o_orderkey INT NOT NULL," + " o_custkey INT NOT NULL," + " o_orderstatus VARCHAR(1) NOT NULL," + " o_totalprice DOUBLE NOT NULL," + " o_orderdate DATE NOT NULL" + ") DUPLICATE KEY (o_orderkey) " + "DISTRIBUTED BY HASH (o_orderkey) BUCKETS 1 " + 'PROPERTIES ("replication_num" = "1")' + ) + cur.execute( + "CREATE TABLE customer (" + " c_custkey INT NOT NULL," + " c_name VARCHAR(25) NOT NULL" + ") DUPLICATE KEY (c_custkey) " + "DISTRIBUTED BY HASH (c_custkey) BUCKETS 1 " + 'PROPERTIES ("replication_num" = "1")' + ) + for i in range(0, len(orders_rows), _INSERT_BATCH): + cur.executemany( + "INSERT INTO orders VALUES (%s, %s, %s, %s, %s)", + orders_rows[i : i + _INSERT_BATCH], + ) + cur.executemany("INSERT INTO customer VALUES (%s, %s)", customer_rows) + finally: + conn.close() + + +@pytest.fixture(scope="module") +def _starrocks_endpoint(): + """Start StarRocks once per module and load the fixture data.""" + with _StarRocksContainer() as sr: + host = sr.get_host_ip() + port = sr.get_query_port() + _wait_until_ready(host, port) + _load_tpch(host, port) + yield host, port + + +class TestStarRocks(WrenQueryTestSuite): + manifest = make_tpch_manifest(table_catalog=None, table_schema=_DATABASE) + + @pytest.fixture(scope="class") + def engine(self, _starrocks_endpoint) -> WrenEngine: # type: ignore[override] + host, port = _starrocks_endpoint + conn_info = { + "host": host, + "port": port, + "database": _DATABASE, + "user": "root", + "password": "", + } + manifest_str = base64.b64encode(orjson.dumps(self.manifest)).decode() + with WrenEngine( + manifest_str, DataSource.starrocks, conn_info, fallback=False + ) as e: + yield e + + # ------------------------------------------------------------------ + # StarRocks-specific dialect tests (these fail with the ``doris`` dialect) + # ------------------------------------------------------------------ + + def test_date_trunc_uses_starrocks_argument_order(self, engine: WrenEngine) -> None: + result = engine.query( + "SELECT date_trunc('month', o_orderdate) AS m FROM \"orders\" " + "WHERE o_orderkey = 1" + ) + assert result.num_rows == 1 + value = result["m"][0].as_py() + expected = _make_fixture_rows()[0][0][4].replace(day=1) + # StarRocks returns DATE for a DATE input; tolerate a datetime too. + assert getattr(value, "date", lambda: value)() == expected + + def test_array_agg_is_not_rewritten_to_collect_list( + self, engine: WrenEngine + ) -> None: + planned = engine.dry_plan( + 'SELECT array_agg(o_orderkey) AS a FROM "orders" WHERE o_orderkey <= 3' + ) + assert "collect_list" not in planned.lower() + + result = engine.query( + 'SELECT array_agg(o_orderkey) AS a FROM "orders" WHERE o_orderkey <= 3' + ) + value = result["a"][0].as_py() + if isinstance(value, str): + value = json.loads(value) + assert sorted(int(v) for v in value) == [1, 2, 3] + + def test_group_concat_has_a_single_separator(self, engine: WrenEngine) -> None: + result = engine.query( + 'SELECT group_concat(o_orderstatus) AS s FROM "orders" ' + "WHERE o_orderkey <= 3" + ) + value = result["s"][0].as_py() + assert ",," not in value + assert len(value.split(",")) == 3 diff --git a/core/wren/tests/unit/test_connector_factory.py b/core/wren/tests/unit/test_connector_factory.py index 6e2808210..b9fcde9b2 100644 --- a/core/wren/tests/unit/test_connector_factory.py +++ b/core/wren/tests/unit/test_connector_factory.py @@ -9,6 +9,29 @@ from wren.model.error import ErrorCode, WrenError pytestmark = pytest.mark.unit +def test_starrocks_uses_mysql_connector_module_and_extra() -> None: + assert factory._REGISTRY[DataSource.starrocks] == "wren.connector.mysql" + assert factory._INSTALL_EXTRA[DataSource.starrocks] == "mysql" + assert DataSource.starrocks in factory._NEEDS_DATA_SOURCE + + +def test_get_connector_passes_starrocks_to_create_connector(monkeypatch) -> None: + calls: list[tuple] = [] + + class _FakeModule: + @staticmethod + def create_connector(data_source, connection_info): + calls.append((data_source, connection_info)) + return type("_Connector", (), {})() + + monkeypatch.setattr(factory.importlib, "import_module", lambda name: _FakeModule) + + connector = factory.get_connector(DataSource.starrocks, {"host": "sr"}) + + assert calls == [(DataSource.starrocks, {"host": "sr"})] + assert connector.dialect == "starrocks" + + def test_connector_import_error_has_quoted_wrenai_extra_hint(monkeypatch) -> None: def _fake_import_module(name: str): if name == "wren.connector.mysql": diff --git a/core/wren/tests/unit/test_cte_rewriter.py b/core/wren/tests/unit/test_cte_rewriter.py index fa9169e71..41e354a67 100644 --- a/core/wren/tests/unit/test_cte_rewriter.py +++ b/core/wren/tests/unit/test_cte_rewriter.py @@ -520,6 +520,22 @@ class TestDialectMapping: def test_mssql_maps_to_tsql(self): assert get_sqlglot_dialect(DataSource.mssql) == "tsql" + def test_starrocks_uses_its_own_sqlglot_dialect(self): + # Doris and StarRocks differ in DATE_TRUNC argument order and in + # ARRAY_AGG / GROUP_CONCAT spelling, so StarRocks must not reuse "doris". + assert get_sqlglot_dialect(DataSource.starrocks) == "starrocks" + assert get_sqlglot_dialect(DataSource.doris) == "doris" + + def test_starrocks_date_trunc_argument_order(self): + sql = "SELECT date_trunc('month', d) FROM t" + starrocks = get_sqlglot_dialect(DataSource.starrocks) + doris = get_sqlglot_dialect(DataSource.doris) + out = sqlglot.transpile(sql, read=starrocks, write=starrocks)[0] + assert out == "SELECT DATE_TRUNC('MONTH', d) FROM t" + # The doris dialect swaps the arguments, which StarRocks rejects. + out_doris = sqlglot.transpile(sql, read=doris, write=doris)[0] + assert out_doris == "SELECT DATE_TRUNC(d, 'MONTH') FROM t" + def test_canner_maps_to_trino(self): assert get_sqlglot_dialect(DataSource.canner) == "trino" diff --git a/core/wren/tests/unit/test_dbt.py b/core/wren/tests/unit/test_dbt.py index a5852e7fd..a9bc65451 100644 --- a/core/wren/tests/unit/test_dbt.py +++ b/core/wren/tests/unit/test_dbt.py @@ -279,6 +279,52 @@ class TestConvertDbtTargetToWrenProfile: assert profile["dataset_id"] == "analytics" assert profile["credentials"] + def test_map_starrocks_adapter(self): + assert map_dbt_adapter_to_wren("starrocks") == "starrocks" + + def test_convert_starrocks_profile(self, tmp_path): + project_dir, profiles_path = _write_basic_dbt_project(tmp_path) + profiles_path.write_text( + "jaffle_shop:\n" + " target: dev\n" + " outputs:\n" + " dev:\n" + " type: starrocks\n" + " host: sr.example.com\n" + " port: 9030\n" + " schema: analytics\n" + " username: etl\n" + " password: secret\n" + ) + target = resolve_dbt_target(project_dir, profiles_path=profiles_path) + + profile = convert_dbt_target_to_wren_profile(target) + + assert profile["datasource"] == "starrocks" + assert profile["host"] == "sr.example.com" + assert profile["port"] == "9030" + assert profile["database"] == "analytics" + assert profile["user"] == "etl" + assert profile["password"] == "secret" + + def test_convert_starrocks_profile_rejects_custom_catalog(self, tmp_path): + project_dir, profiles_path = _write_basic_dbt_project(tmp_path) + profiles_path.write_text( + "jaffle_shop:\n" + " target: dev\n" + " outputs:\n" + " dev:\n" + " type: starrocks\n" + " host: sr.example.com\n" + " schema: analytics\n" + " username: etl\n" + " catalog: hive_catalog\n" + ) + target = resolve_dbt_target(project_dir, profiles_path=profiles_path) + + with pytest.raises(DbtLoadError, match="catalog"): + convert_dbt_target_to_wren_profile(target) + def test_convert_databricks_profile_preserves_catalog(self, tmp_path): project_dir, profiles_path = _write_basic_dbt_project(tmp_path) profiles_path.write_text( diff --git a/core/wren/tests/unit/test_mysql_helpers.py b/core/wren/tests/unit/test_mysql_helpers.py index 7d30941f2..5add9bdde 100644 --- a/core/wren/tests/unit/test_mysql_helpers.py +++ b/core/wren/tests/unit/test_mysql_helpers.py @@ -554,3 +554,37 @@ def test_lazy_init_thread_safe() -> None: # Sanity: the field-type map is non-empty (MySQLdb constants exist). if fn is _mysql_field_type_map: assert len(got) > 0 + + +def test_create_connector_routes_starrocks_to_doris_connector(monkeypatch) -> None: + pytest.importorskip("MySQLdb") + from unittest.mock import MagicMock # noqa: PLC0415 + + import MySQLdb # noqa: PLC0415 + + from wren.model.data_source import DataSource # noqa: PLC0415 + + connection = MagicMock() + monkeypatch.setattr(MySQLdb, "connect", lambda **kwargs: connection) + info = DataSource.starrocks.get_connection_info( + {"host": "sr", "port": "9030", "database": "olymp", "user": "u"} + ) + + connector = mysql_connector.create_connector(DataSource.starrocks, info) + + assert isinstance(connector, mysql_connector.DorisConnector) + # MySqlConnector would run ``SET sql_mode=...ANSI_QUOTES`` on connect. + connection.cursor.assert_not_called() + + +def test_starrocks_connection_info_shares_doris_fields() -> None: + from wren.model.data_source import DataSource # noqa: PLC0415 + + info = DataSource.starrocks.get_connection_info( + {"host": "sr", "port": "9030", "database": "olymp", "user": "u"} + ) + kwargs = mysql_connector._build_doris_connect_kwargs(info) + assert kwargs["host"] == "sr" + assert kwargs["port"] == 9030 + assert kwargs["db"] == "olymp" + assert kwargs["autocommit"] is True diff --git a/docs/core/guides/connect.md b/docs/core/guides/connect.md index fa13ea5d0..e6032f8ad 100644 --- a/docs/core/guides/connect.md +++ b/docs/core/guides/connect.md @@ -23,6 +23,8 @@ If you have an AI coding agent installed, the `/wren` skill drives the whole thi | DuckDB | _(included)_ | | PostgreSQL | `postgres` | | MySQL | `mysql` | +| Doris | `mysql` | +| StarRocks | `mysql` | | BigQuery | `bigquery` | | Snowflake | `snowflake` | | ClickHouse | `clickhouse` | diff --git a/docs/core/guides/dbt-integration.md b/docs/core/guides/dbt-integration.md index 1cca71def..fc63a2b53 100644 --- a/docs/core/guides/dbt-integration.md +++ b/docs/core/guides/dbt-integration.md @@ -71,7 +71,7 @@ Useful options: | `--name` | Destination Wren profile name | | `--no-activate` | Save the profile without making it active | -Supported adapters include `postgres`, `bigquery`, `snowflake`, `databricks`, `trino`, `clickhouse`, `duckdb`, `mysql`, `redshift`, `spark`, `athena`, `mssql`, and `doris`. +Supported adapters include `postgres`, `bigquery`, `snowflake`, `databricks`, `trino`, `clickhouse`, `duckdb`, `mysql`, `redshift`, `spark`, `athena`, `mssql`, `doris`, and `starrocks`. ## Import the dbt Project diff --git a/docs/core/reference/mdl.md b/docs/core/reference/mdl.md index d2196528b..667f25f78 100644 --- a/docs/core/reference/mdl.md +++ b/docs/core/reference/mdl.md @@ -382,7 +382,7 @@ Models and views support an optional `dialect` field declaring which SQL dialect | `dialect` omitted | Falls back to project-level `data_source`. Default. | | `dialect` set | Engine uses the named dialect parser for this object's SQL. | -Valid dialect values: `athena`, `bigquery`, `canner`, `clickhouse`, `databricks`, `datafusion`, `doris`, `duckdb`, `gcs_file`, `local_file`, `minio_file`, `mssql`, `mysql`, `oracle`, `postgres`, `redshift`, `s3_file`, `snowflake`, `spark`, `trino`. +Valid dialect values: `athena`, `bigquery`, `canner`, `clickhouse`, `databricks`, `datafusion`, `doris`, `duckdb`, `gcs_file`, `local_file`, `minio_file`, `mssql`, `mysql`, `oracle`, `postgres`, `redshift`, `s3_file`, `snowflake`, `spark`, `starrocks`, `trino`. ### Schema version vs layout version