feat(wren): add standalone wren Python SDK package (#1471)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>

core/wren/.claude/CLAUDE.md

@@ -0,0 +1,59 @@
+# CLAUDE.md — wren package
+
+Standalone Python SDK and CLI for Wren Engine. Wraps `wren-core-py` (PyO3 bindings) + Ibis connectors into a single installable package.
+
+## Package Structure
+
+```
+wren/
+  src/wren/
+    engine.py       — WrenEngine facade (transpile / query / dry_run / dry_plan)
+    cli.py          — Typer CLI: wren query|dry-run|transpile|validate
+    mdl/            — wren-core-py session context + manifest extraction helpers
+    connector/      — Per-datasource Ibis connectors (factory.py + one file per source)
+    model/
+      data_source.py — DataSource enum + per-source ConnectionInfo factories
+      error.py       — WrenError, ErrorCode, ErrorPhase
+  tests/
+```
+
+## Build & Development
+
+```bash
+cd wren
+just install          # build wren-core-py wheel + uv sync
+just install-all      # with all optional extras
+just install-extra <extra>   # e.g. just install-extra postgres
+just test             # pytest tests/
+just lint             # ruff format --check + ruff check
+just format           # ruff auto-fix
+just build            # uv build (produces wheel)
+```
+
+Uses `uv` (not Poetry). `pyproject.toml` uses `hatchling` as build backend.
+
+## Key Design Points
+
+- **WrenEngine** is the main entry point. It accepts a base64-encoded MDL JSON string, a `DataSource`, and a connection dict.
+- **Query flow**: `_plan()` → wren-core `SessionContext.transform_sql()` → `_transpile()` via sqlglot → connector `.query()`.
+- **Manifest extraction**: `_plan()` tries to extract a minimal sub-manifest scoped to the query's referenced tables before calling wren-core — this reduces planning overhead. Falls back to the full manifest on error.
+- **`get_session_context` is `@cache`-decorated** — same `(manifest_str, function_path, properties, data_source)` tuple reuses the same SessionContext. Avoid mutating session state.
+- **Write dialect mapping**: `canner` → `trino`; file sources (`local_file`, `s3_file`, `minio_file`, `gcs_file`) → `duckdb`. All others use `data_source.name` directly.
+- **WrenEngine is a context manager** (`__enter__` / `__exit__` call `close()`).
+
+## Connectors
+
+`connector/factory.py` dispatches on `DataSource` to return the right connector. Each connector wraps an Ibis backend and exposes `.query(sql, limit)` and `.dry_run(sql)`. Base class in `connector/base.py`; Ibis-backed connectors share `connector/ibis.py`.
+
+## Optional Extras
+
+Install per data-source extras: `postgres`, `mysql`, `bigquery`, `snowflake`, `clickhouse`, `trino`, `mssql`, `databricks`, `redshift`, `spark`, `athena`, `oracle`, `all`, `dev`.
+
+On macOS, `mysql` extra needs:
+```bash
+PKG_CONFIG_PATH="$(brew --prefix mysql-client)/lib/pkgconfig" just install-extra mysql
+```
+
+## Dependency on wren-core-py
+
+`wren-core-py` wheel is built locally from `../wren-core-py/` and installed via `--find-links`. Run `just build-core` (or `just install`) to rebuild after Rust changes.

core/wren/README.md

@@ -0,0 +1,42 @@
+# wren
+
+Wren Engine CLI and Python SDK — semantic SQL layer for 20+ data sources.
+
+## Installation
+
+```bash
+pip install wren
+```
+
+## Usage
+
+```python
+from wren import WrenEngine, DataSource
+```
+
+See the [Wren Engine documentation](https://getwren.ai) for details.
+
+## Running tests
+
+Install dev dependencies first:
+
+```bash
+just install-dev
+```
+
+| Command | What it runs | Docker needed |
+|---------|-------------|---------------|
+| `just test-unit` | Unit tests (transpile, dry-plan, context manager) | No |
+| `just test-duckdb` | DuckDB connector tests — generates TPCH data via `dbgen` | No |
+| `just test-postgres` | PostgreSQL connector tests — spins up a container | Yes |
+| `just test` | All tests | Yes |
+
+Run a specific connector via marker:
+
+```bash
+just test-connector postgres
+```
+
+To add tests for a new connector, subclass `WrenQueryTestSuite` in
+`tests/connectors/test_<name>.py` and provide a class-scoped `engine` fixture.
+All base tests are inherited automatically.

core/wren/justfile

@@ -0,0 +1,70 @@
+default:
+    @just --list
+
+build-core:
+    cd ../wren-core-py && just install && just build
+
+core-wheel-dir := "../wren-core-py/target/wheels/"
+
+install-core:
+    just build-core
+
+install:
+    just build-core
+    uv sync --find-links {{ core-wheel-dir }}
+
+install-dev:
+    just build-core
+    uv sync --extra dev --find-links {{ core-wheel-dir }}
+
+install-all:
+    just build-core
+    uv sync --all-extras --find-links {{ core-wheel-dir }}
+
+install-extra extra *flags:
+    #!/usr/bin/env bash
+    set -euo pipefail
+    if [ "{{ extra }}" = "mysql" ]; then
+        mysql_prefix=$(brew --prefix mysql-client 2>/dev/null || true)
+        if [ -n "$mysql_prefix" ]; then
+            export PKG_CONFIG_PATH="$mysql_prefix/lib/pkgconfig:${PKG_CONFIG_PATH:-}"
+        fi
+    fi
+    dev_flag=""
+    for flag in {{ flags }}; do
+        if [ "$flag" = "--dev" ]; then
+            dev_flag="--extra dev"
+        fi
+    done
+    uv sync --extra {{ extra }} $dev_flag --find-links {{ core-wheel-dir }}
+
+dev:
+    uv run wren
+
+test:
+    uv run pytest tests/ -v
+
+test-unit:
+    uv run pytest tests/unit/ -v -m unit
+
+test-duckdb:
+    uv run pytest tests/connectors/test_duckdb.py -v -m duckdb
+
+test-postgres:
+    uv run pytest tests/connectors/test_postgres.py -v -m postgres
+
+test-connector marker:
+    uv run pytest tests/connectors/ -v -m {{ marker }}
+
+lint:
+    uv run ruff format --check src/
+    uv run ruff check src/
+
+format:
+    uv run ruff format src/
+    uv run ruff check --fix src/
+
+build:
+    uv build
+
+alias fmt := format

core/wren/pyproject.toml

@@ -0,0 +1,74 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "wren"
+dynamic = ["version"]
+description = "Wren Engine CLI and Python SDK — semantic SQL layer for 20+ data sources"
+readme = "README.md"
+requires-python = ">=3.11"
+license = { text = "Apache-2.0" }
+keywords = ["wren", "sql", "semantic", "mdl", "datafusion"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "License :: OSI Approved :: Apache Software License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+]
+dependencies = [
+    "wren-core-py>=0.1",
+    "ibis-framework>=10",
+    "duckdb>=1.0",
+    "sqlglot>=27",
+    "typer>=0.12",
+    "pydantic>=2",
+    "pyarrow>=14",
+    "pyarrow-hotfix>=0.6",
+    "loguru>=0.7",
+    "opendal>=0.45",
+    "pandas>=2",
+    "boto3>=1.26",
+]
+
+[project.optional-dependencies]
+postgres    = ["psycopg>=3", "ibis-framework[postgres]"]
+mysql       = ["mysqlclient>=2.2", "ibis-framework[mysql]"]
+bigquery    = ["ibis-framework[bigquery]", "google-auth"]
+snowflake   = ["ibis-framework[snowflake]"]
+clickhouse  = ["ibis-framework[clickhouse]"]
+trino       = ["ibis-framework[trino]", "trino>=0.321"]
+mssql       = ["ibis-framework[mssql]"]
+databricks  = ["databricks-sql-connector", "databricks-sdk"]
+redshift    = ["redshift_connector"]
+spark       = ["pyspark>=3.5"]
+athena      = ["ibis-framework[athena]"]
+oracle      = ["ibis-framework[oracle]", "oracledb"]
+all         = ["wren[postgres,mysql,bigquery,snowflake,clickhouse,trino,mssql,databricks,redshift,athena,oracle]"]
+dev         = ["pytest>=8", "ruff>=0.4", "orjson>=3", "testcontainers[postgres]>=4"]
+
+[project.scripts]
+wren = "wren.cli:app"
+
+[project.urls]
+Homepage = "https://getwren.ai"
+Repository = "https://github.com/Canner/wren-engine"
+Issues = "https://github.com/Canner/wren-engine/issues"
+
+[tool.hatch.version]
+path = "src/wren/__init__.py"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/wren"]
+
+[tool.ruff]
+line-length = 88
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "PLC"]
+ignore = [
+  "E501", # line-too-long, enforced by ruff format
+]

core/wren/src/wren/__init__.py

@@ -0,0 +1,9 @@
+"""Wren — semantic SQL layer for 20+ data sources."""
+
+__version__ = "0.0.1"
+
+from wren.engine import WrenEngine
+from wren.model.data_source import DataSource
+from wren.model.error import WrenError
+
+__all__ = ["WrenEngine", "DataSource", "WrenError", "__version__"]

core/wren/src/wren/cli.py

@@ -0,0 +1,225 @@
+"""Wren CLI — SQL transform and execution via the Wren semantic layer."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Annotated, Optional
+
+import typer
+
+app = typer.Typer(name="wren", help="Wren Engine CLI", no_args_is_help=True)
+
+
+def _load_connection_info(
+    connection_info: str | None,
+    connection_file_path: str | None,
+) -> dict:
+    if connection_info:
+        try:
+            return json.loads(connection_info)
+        except json.JSONDecodeError as e:
+            typer.echo(f"Error: invalid JSON in --connection-info: {e}", err=True)
+            raise typer.Exit(1)
+    if connection_file_path:
+        path = Path(connection_file_path)
+        if not path.exists():
+            typer.echo(
+                f"Error: connection file not found: {connection_file_path}", err=True
+            )
+            raise typer.Exit(1)
+        try:
+            return json.loads(path.read_text())
+        except json.JSONDecodeError as e:
+            typer.echo(f"Error: invalid JSON in {connection_file_path}: {e}", err=True)
+            raise typer.Exit(1)
+    typer.echo(
+        "Error: either --connection-info or --connection-file must be provided",
+        err=True,
+    )
+    raise typer.Exit(1)
+
+
+def _load_manifest(mdl: str) -> str:
+    """Load MDL from a file path or treat as base64 string directly."""
+    path = Path(mdl)
+    if path.exists():
+        import base64  # noqa: PLC0415
+
+        content = path.read_bytes()
+        # If it's a JSON file, base64-encode it
+        if mdl.endswith(".json"):
+            return base64.b64encode(content).decode()
+        # Otherwise assume it's already base64
+        return content.decode().strip()
+    # Treat as inline base64 string
+    return mdl
+
+
+def _make_engine(
+    sql: str,
+    datasource: str,
+    mdl: str,
+    connection_info: str | None,
+    connection_file: str | None,
+):
+    from wren.engine import WrenEngine  # noqa: PLC0415
+    from wren.model.data_source import DataSource  # noqa: PLC0415
+
+    manifest_str = _load_manifest(mdl)
+    conn_dict = _load_connection_info(connection_info, connection_file)
+
+    try:
+        ds = DataSource(datasource.lower())
+    except ValueError:
+        typer.echo(f"Error: unknown datasource '{datasource}'", err=True)
+        raise typer.Exit(1)
+
+    return WrenEngine(
+        manifest_str=manifest_str, data_source=ds, connection_info=conn_dict
+    )
+
+
+# ── Common options ─────────────────────────────────────────────────────────
+
+SqlArg = Annotated[str, typer.Option("--sql", "-s", help="SQL query to execute")]
+DatasourceOpt = Annotated[
+    str,
+    typer.Option("--datasource", "-d", help="Data source name (e.g. postgres, duckdb)"),
+]
+MdlOpt = Annotated[
+    str, typer.Option("--mdl", "-m", help="Path to MDL JSON file or base64 MDL string")
+]
+ConnInfoOpt = Annotated[
+    Optional[str], typer.Option("--connection-info", help="JSON connection info string")
+]
+ConnFileOpt = Annotated[
+    Optional[str],
+    typer.Option("--connection-file", help="Path to JSON connection info file"),
+]
+LimitOpt = Annotated[
+    Optional[int], typer.Option("--limit", "-l", help="Max rows to return")
+]
+
+
+@app.command()
+def query(
+    sql: SqlArg,
+    datasource: DatasourceOpt,
+    mdl: MdlOpt,
+    connection_info: ConnInfoOpt = None,
+    connection_file: ConnFileOpt = None,
+    limit: LimitOpt = None,
+    output: Annotated[
+        str, typer.Option("--output", "-o", help="Output format: json|csv|table")
+    ] = "table",
+):
+    """Execute a SQL query through the Wren semantic layer."""
+    engine = _make_engine(sql, datasource, mdl, connection_info, connection_file)
+    try:
+        result = engine.query(sql, limit=limit)
+    except Exception as e:
+        typer.echo(f"Error: {e}", err=True)
+        raise typer.Exit(1)
+    finally:
+        engine.close()
+
+    _print_result(result, output)
+
+
+@app.command(name="dry-run")
+def dry_run(
+    sql: SqlArg,
+    datasource: DatasourceOpt,
+    mdl: MdlOpt,
+    connection_info: ConnInfoOpt = None,
+    connection_file: ConnFileOpt = None,
+):
+    """Dry-run a SQL query (parse + validate, no results returned)."""
+    engine = _make_engine(sql, datasource, mdl, connection_info, connection_file)
+    try:
+        engine.dry_run(sql)
+        typer.echo("OK")
+    except Exception as e:
+        typer.echo(f"Error: {e}", err=True)
+        raise typer.Exit(1)
+    finally:
+        engine.close()
+
+
+@app.command()
+def transpile(
+    sql: SqlArg,
+    datasource: DatasourceOpt,
+    mdl: MdlOpt,
+):
+    """Transform SQL through MDL and emit the target dialect SQL (no DB required)."""
+    from wren.engine import WrenEngine  # noqa: PLC0415
+    from wren.model.data_source import DataSource  # noqa: PLC0415
+
+    manifest_str = _load_manifest(mdl)
+    try:
+        ds = DataSource(datasource.lower())
+    except ValueError:
+        typer.echo(f"Error: unknown datasource '{datasource}'", err=True)
+        raise typer.Exit(1)
+
+    # For transpile we don't need real connection_info — pass a dummy dict
+    engine = WrenEngine(manifest_str=manifest_str, data_source=ds, connection_info={})
+    try:
+        result = engine.transpile(sql)
+        typer.echo(result)
+    except Exception as e:
+        typer.echo(f"Error: {e}", err=True)
+        raise typer.Exit(1)
+    finally:
+        engine.close()
+
+
+@app.command()
+def validate(
+    sql: SqlArg,
+    datasource: DatasourceOpt,
+    mdl: MdlOpt,
+    connection_info: ConnInfoOpt = None,
+    connection_file: ConnFileOpt = None,
+):
+    """Validate SQL can be planned and dry-run against the data source."""
+    engine = _make_engine(sql, datasource, mdl, connection_info, connection_file)
+    try:
+        engine.dry_run(sql)
+        typer.echo("Valid")
+    except Exception as e:
+        typer.echo(f"Invalid: {e}", err=True)
+        raise typer.Exit(1)
+    finally:
+        engine.close()
+
+
+# ── Output formatting ──────────────────────────────────────────────────────
+
+
+def _print_result(table, output: str) -> None:
+    if output == "json":
+        try:
+            df = table.to_pandas()
+            typer.echo(df.to_json(orient="records", lines=True))
+        except Exception:
+            typer.echo(json.dumps(table.to_pydict()))
+    elif output == "csv":
+        try:
+            df = table.to_pandas()
+            typer.echo(df.to_csv(index=False))
+        except Exception:
+            typer.echo(str(table))
+    else:
+        # Default: table format via pandas
+        try:
+            df = table.to_pandas()
+            typer.echo(df.to_string(index=False))
+        except Exception:
+            typer.echo(str(table))
+
+
+if __name__ == "__main__":
+    app()

core/wren/src/wren/connector/__init__.py

@@ -0,0 +1,4 @@
+from wren.connector.base import ConnectorABC, IbisConnector
+from wren.connector.factory import get_connector
+
+__all__ = ["ConnectorABC", "IbisConnector", "get_connector"]

core/wren/src/wren/connector/base.py

@@ -0,0 +1,87 @@
+from abc import ABC, abstractmethod
+
+import pyarrow as pa
+from ibis.expr.datatypes import Decimal
+from ibis.expr.datatypes.core import UUID
+from ibis.expr.types import Table
+from loguru import logger
+
+from wren.model.data_source import DataSource
+
+
+class ConnectorABC(ABC):
+    @abstractmethod
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        pass
+
+    @abstractmethod
+    def dry_run(self, sql: str) -> None:
+        pass
+
+    @abstractmethod
+    def close(self) -> None:
+        pass
+
+
+class IbisConnector(ConnectorABC):
+    def __init__(self, data_source: DataSource, connection_info):
+        self.data_source = data_source
+        self.connection = self.data_source.get_connection(connection_info)
+        self._closed = False
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        ibis_table = self.connection.sql(sql)
+        if limit is not None:
+            ibis_table = ibis_table.limit(limit)
+        ibis_table = self._handle_pyarrow_unsupported_type(ibis_table)
+        return ibis_table.to_pyarrow()
+
+    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
+                )
+        return result_table
+
+    def _cast_uuid_columns(self, result_table: Table, col_name: str) -> Table:
+        return result_table.mutate(**{col_name: result_table[col_name].cast("string")})
+
+    def _round_decimal_columns(
+        self, result_table: Table, col_name: str, scale: int = 9
+    ) -> Table:
+        col = result_table[col_name]
+        decimal_type = Decimal(precision=38, scale=scale)
+        rounded_col = col.cast(decimal_type).round(scale)
+        return result_table.mutate(**{col_name: rounded_col})
+
+    def dry_run(self, sql: str) -> None:
+        self.connection.sql(sql)
+
+    def close(self) -> None:
+        if self._closed or not hasattr(self, "connection") or self.connection is None:
+            return
+        try:
+            if hasattr(self.connection, "con"):
+                if hasattr(self.connection.con, "close"):
+                    self.connection.con.close()
+            elif hasattr(self.connection, "close"):
+                self.connection.close()
+            elif hasattr(self.connection, "disconnect"):
+                self.connection.disconnect()
+            else:
+                logger.warning(
+                    f"Closing connection for {self.data_source.value} is not implemented."
+                )
+        except Exception as e:
+            logger.warning(
+                f"Error closing connection for {self.data_source.value}: {e}"
+            )
+        finally:
+            self._closed = True
+            self.connection = None

core/wren/src/wren/connector/bigquery.py

@@ -0,0 +1,57 @@
+import base64
+from json import loads
+
+import pyarrow as pa
+from loguru import logger
+
+from wren.connector.base import ConnectorABC
+
+
+class BigQueryConnector(ConnectorABC):
+    def __init__(self, connection_info):
+        from google.cloud import bigquery  # noqa: PLC0415
+        from google.oauth2 import service_account  # noqa: PLC0415
+
+        self.connection_info = connection_info
+        credits_json = loads(
+            base64.b64decode(connection_info.credentials.get_secret_value()).decode(
+                "utf-8"
+            )
+        )
+        credentials = service_account.Credentials.from_service_account_info(
+            credits_json
+        )
+        credentials = credentials.with_scopes(
+            [
+                "https://www.googleapis.com/auth/drive",
+                "https://www.googleapis.com/auth/cloud-platform",
+            ]
+        )
+        client = bigquery.Client(
+            credentials=credentials,
+            project=connection_info.get_billing_project_id(),
+        )
+        job_config = bigquery.QueryJobConfig()
+        job_config.job_timeout_ms = connection_info.job_timeout_ms
+        client.default_query_job_config = job_config
+        self.connection = client
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        return self.connection.query(sql).result(max_results=limit).to_arrow()
+
+    def dry_run(self, sql: str) -> None:
+        from google.cloud import bigquery  # noqa: PLC0415
+
+        self.connection.query(
+            sql, job_config=bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
+        )
+
+    def close(self) -> None:
+        try:
+            self.connection.close()
+        except Exception as e:
+            logger.warning(f"Error closing BigQuery connection: {e}")
+
+
+def create_connector(connection_info) -> BigQueryConnector:
+    return BigQueryConnector(connection_info)

core/wren/src/wren/connector/canner.py

@@ -0,0 +1,79 @@
+from contextlib import closing
+from functools import cache
+from typing import Any
+
+import ibis
+import ibis.expr.schema as sch
+import pyarrow as pa
+from ibis import BaseBackend
+from ibis.backends.sql.compilers.postgres import compiler as postgres_compiler
+from ibis.expr.datatypes import Decimal
+from ibis.expr.datatypes.core import UUID
+from ibis.expr.types import Table
+from loguru import logger
+
+from wren.connector.base import ConnectorABC
+from wren.model.data_source import DataSource
+
+
+@cache
+def _get_pg_type_names(connection: BaseBackend) -> dict[int, str]:
+    with closing(connection.raw_sql("SELECT oid, typname FROM pg_type")) as cur:
+        return dict(cur.fetchall())
+
+
+class CannerConnector(ConnectorABC):
+    def __init__(self, connection_info):
+        self.connection = DataSource.canner.get_connection(connection_info)
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        schema = self._get_schema(sql)
+        ibis_table = self.connection.sql(sql, schema=schema)
+        if limit is not None:
+            ibis_table = ibis_table.limit(limit)
+        ibis_table = self._handle_pyarrow_unsupported_type(ibis_table)
+        return ibis_table.to_pyarrow()
+
+    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):
+                col = result_table[name]
+                decimal_type = Decimal(precision=38, scale=9)
+                rounded_col = col.cast(decimal_type).round(9)
+                result_table = result_table.mutate(**{name: rounded_col})
+            elif isinstance(dtype, UUID):
+                result_table = result_table.mutate(
+                    **{name: result_table[name].cast("string")}
+                )
+        return result_table
+
+    def dry_run(self, sql: str) -> Any:
+        return self.connection.raw_sql(f"SELECT * FROM ({sql}) LIMIT 0")
+
+    def close(self) -> None:
+        try:
+            if hasattr(self.connection, "con") and hasattr(
+                self.connection.con, "close"
+            ):
+                self.connection.con.close()
+            elif hasattr(self.connection, "close"):
+                self.connection.close()
+        except Exception as e:
+            logger.warning(f"Error closing Canner connection: {e}")
+
+    def _get_schema(self, sql: str) -> sch.Schema:
+        cur = self.dry_run(sql)
+        type_names = _get_pg_type_names(self.connection)
+        return ibis.schema(
+            {
+                desc.name: postgres_compiler.type_mapper.from_string(
+                    type_names[desc.type_code]
+                )
+                for desc in cur.description
+            }
+        )
+
+
+def create_connector(connection_info) -> CannerConnector:
+    return CannerConnector(connection_info)

core/wren/src/wren/connector/databricks.py

@@ -0,0 +1,65 @@
+from contextlib import closing
+
+import pyarrow as pa
+from loguru import logger
+
+from wren.connector.base import ConnectorABC
+from wren.model import (
+    DatabricksConnectionUnion,
+    DatabricksServicePrincipalConnectionInfo,
+    DatabricksTokenConnectionInfo,
+)
+
+
+class DatabricksConnector(ConnectorABC):
+    def __init__(self, connection_info: DatabricksConnectionUnion):
+        from databricks import sql as dbsql  # noqa: PLC0415
+        from databricks.sdk.core import Config as DbConfig  # noqa: PLC0415
+        from databricks.sdk.core import oauth_service_principal  # noqa: PLC0415
+
+        if isinstance(connection_info, DatabricksTokenConnectionInfo):
+            self.connection = dbsql.connect(
+                server_hostname=connection_info.server_hostname.get_secret_value(),
+                http_path=connection_info.http_path.get_secret_value(),
+                access_token=connection_info.access_token.get_secret_value(),
+            )
+        elif isinstance(connection_info, DatabricksServicePrincipalConnectionInfo):
+            kwargs = {
+                "host": connection_info.server_hostname.get_secret_value(),
+                "client_id": connection_info.client_id.get_secret_value(),
+                "client_secret": connection_info.client_secret.get_secret_value(),
+            }
+            if connection_info.azure_tenant_id is not None:
+                kwargs["azure_tenant_id"] = (
+                    connection_info.azure_tenant_id.get_secret_value()
+                )
+
+            def credential_provider():
+                return oauth_service_principal(DbConfig(**kwargs))
+
+            self.connection = dbsql.connect(
+                server_hostname=connection_info.server_hostname.get_secret_value(),
+                http_path=connection_info.http_path.get_secret_value(),
+                credentials_provider=credential_provider,
+            )
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        with closing(self.connection.cursor()) as cursor:
+            cursor.execute(sql)
+            if limit is not None:
+                return cursor.fetchmany_arrow(limit)
+            return cursor.fetchall_arrow()
+
+    def dry_run(self, sql: str) -> None:
+        with closing(self.connection.cursor()) as cursor:
+            cursor.execute(f"SELECT * FROM ({sql}) AS sub LIMIT 0")
+
+    def close(self) -> None:
+        try:
+            self.connection.close()
+        except Exception as e:
+            logger.warning(f"Error closing Databricks connection: {e}")
+
+
+def create_connector(connection_info) -> DatabricksConnector:
+    return DatabricksConnector(connection_info)

core/wren/src/wren/connector/duckdb.py

@@ -0,0 +1,125 @@
+import os
+
+import opendal
+import pyarrow as pa
+from loguru import logger
+
+from wren.connector.base import ConnectorABC
+from wren.model import (
+    GcsFileConnectionInfo,
+    MinioFileConnectionInfo,
+    S3FileConnectionInfo,
+)
+from wren.model.error import ErrorCode, WrenError
+
+
+def _escape_sql(value: str) -> str:
+    return value.replace("'", "''")
+
+
+def _init_duckdb_s3(connection, info: S3FileConnectionInfo):
+    connection.execute(f"""
+    CREATE SECRET wren_s3 (
+        TYPE S3,
+        KEY_ID '{_escape_sql(info.access_key.get_secret_value())}',
+        SECRET '{_escape_sql(info.secret_key.get_secret_value())}',
+        REGION '{_escape_sql(info.region.get_secret_value())}'
+    )""")
+
+
+def _init_duckdb_minio(connection, info: MinioFileConnectionInfo):
+    connection.execute(f"""
+    CREATE SECRET wren_minio (
+        TYPE S3,
+        KEY_ID '{_escape_sql(info.access_key.get_secret_value())}',
+        SECRET '{_escape_sql(info.secret_key.get_secret_value())}',
+        REGION 'ap-northeast-1'
+    )""")
+    connection.execute("SET s3_endpoint=?", [info.endpoint.get_secret_value()])
+    connection.execute("SET s3_url_style='path'")
+    connection.execute("SET s3_use_ssl=?", [info.ssl_enabled])
+
+
+def _init_duckdb_gcs(connection, info: GcsFileConnectionInfo):
+    connection.execute(f"""
+    CREATE SECRET wren_gcs (
+        TYPE GCS,
+        KEY_ID '{_escape_sql(info.key_id.get_secret_value())}',
+        SECRET '{_escape_sql(info.secret_key.get_secret_value())}'
+    )""")
+
+
+class DuckDBConnector(ConnectorABC):
+    def __init__(self, connection_info):
+        import duckdb  # noqa: PLC0415
+        from duckdb import HTTPException, IOException  # noqa: PLC0415
+
+        self._HTTPException = HTTPException
+        self._IOException = IOException
+        self.connection = duckdb.connect()
+
+        try:
+            if isinstance(connection_info, S3FileConnectionInfo):
+                _init_duckdb_s3(self.connection, connection_info)
+            if isinstance(connection_info, MinioFileConnectionInfo):
+                _init_duckdb_minio(self.connection, connection_info)
+            if isinstance(connection_info, GcsFileConnectionInfo):
+                _init_duckdb_gcs(self.connection, connection_info)
+
+            if connection_info.format == "duckdb":
+                self._attach_database(connection_info)
+        except Exception:
+            self.connection.close()
+            raise
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        if limit is not None:
+            sql = f"SELECT * FROM ({sql}) AS _q LIMIT {int(limit)}"
+        return self.connection.execute(sql).fetch_arrow_table()
+
+    def dry_run(self, sql: str) -> None:
+        self.connection.execute(f"EXPLAIN {sql}")
+
+    def _attach_database(self, connection_info) -> None:
+        db_files = self._list_duckdb_files(connection_info)
+        if not db_files:
+            raise WrenError(ErrorCode.DUCKDB_FILE_NOT_FOUND, "No DuckDB files found.")
+
+        for file in db_files:
+            try:
+                escaped_file = file.replace("'", "''")
+                alias = os.path.splitext(os.path.basename(file))[0].replace('"', '""')
+                self.connection.execute(
+                    f"ATTACH DATABASE '{escaped_file}' AS \"{alias}\" (READ_ONLY);"
+                )
+            except (self._IOException, self._HTTPException) as e:
+                raise WrenError(
+                    ErrorCode.ATTACH_DUCKDB_ERROR, f"Failed to attach: {e!s}"
+                )
+
+    def _list_duckdb_files(self, connection_info) -> list[str]:
+        op = opendal.Operator("fs", root=connection_info.url.get_secret_value())
+        files = []
+        try:
+            for file in op.list("/"):
+                if file.path != "/":
+                    stat = op.stat(file.path)
+                    if not stat.mode.is_dir() and file.path.endswith(".duckdb"):
+                        files.append(
+                            f"{connection_info.url.get_secret_value()}/{file.path}"
+                        )
+        except Exception as e:
+            raise WrenError(
+                ErrorCode.GENERIC_USER_ERROR, f"Failed to list files: {e!s}"
+            )
+        return files
+
+    def close(self) -> None:
+        try:
+            self.connection.close()
+        except Exception as e:
+            logger.warning(f"Error closing DuckDB connection: {e}")
+
+
+def create_connector(connection_info) -> DuckDBConnector:
+    return DuckDBConnector(connection_info)

core/wren/src/wren/connector/factory.py

@@ -0,0 +1,69 @@
+import importlib
+
+from wren.model.data_source import DataSource
+from wren.model.error import ErrorCode, WrenError
+
+_REGISTRY: dict[DataSource, str] = {
+    DataSource.postgres: "wren.connector.postgres",
+    DataSource.mysql: "wren.connector.mysql",
+    DataSource.doris: "wren.connector.mysql",
+    DataSource.mssql: "wren.connector.mssql",
+    DataSource.canner: "wren.connector.canner",
+    DataSource.bigquery: "wren.connector.bigquery",
+    DataSource.local_file: "wren.connector.duckdb",
+    DataSource.s3_file: "wren.connector.duckdb",
+    DataSource.minio_file: "wren.connector.duckdb",
+    DataSource.gcs_file: "wren.connector.duckdb",
+    DataSource.duckdb: "wren.connector.duckdb",
+    DataSource.redshift: "wren.connector.redshift",
+    DataSource.spark: "wren.connector.spark",
+    DataSource.databricks: "wren.connector.databricks",
+    DataSource.trino: "wren.connector.ibis",
+    DataSource.clickhouse: "wren.connector.ibis",
+    DataSource.oracle: "wren.connector.ibis",
+    DataSource.snowflake: "wren.connector.ibis",
+    DataSource.athena: "wren.connector.ibis",
+}
+
+# Map data sources to the correct pip extra when they share a connector module
+_INSTALL_EXTRA: dict[DataSource, str] = {
+    DataSource.doris: "mysql",
+    DataSource.canner: "postgres",
+    DataSource.local_file: "duckdb",
+    DataSource.s3_file: "duckdb",
+    DataSource.minio_file: "duckdb",
+    DataSource.gcs_file: "duckdb",
+}
+
+_NEEDS_DATA_SOURCE = {
+    DataSource.mysql,
+    DataSource.doris,
+    DataSource.trino,
+    DataSource.clickhouse,
+    DataSource.oracle,
+    DataSource.snowflake,
+    DataSource.athena,
+}
+
+
+def get_connector(data_source: DataSource, connection_info):
+    module_path = _REGISTRY.get(data_source)
+    if module_path is None:
+        raise WrenError(
+            ErrorCode.NOT_IMPLEMENTED,
+            f"Unsupported data source: {data_source}",
+        )
+
+    try:
+        module = importlib.import_module(module_path)
+    except ImportError as e:
+        extra = _INSTALL_EXTRA.get(data_source, data_source.value)
+        raise WrenError(
+            ErrorCode.NOT_IMPLEMENTED,
+            f"Connector '{data_source.value}' requires additional dependencies: {e}. "
+            f"Install with: pip install wren[{extra}]",
+        ) from e
+
+    if data_source in _NEEDS_DATA_SOURCE:
+        return module.create_connector(data_source, connection_info)
+    return module.create_connector(connection_info)

core/wren/src/wren/connector/ibis.py

@@ -0,0 +1,103 @@
+"""Generic ibis-backed connectors with data-source-specific error handling."""
+
+import pyarrow as pa
+
+from wren.connector.base import IbisConnector
+from wren.model.data_source import DataSource
+from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError
+
+try:
+    import clickhouse_connect
+
+    _ClickHouseDbError = clickhouse_connect.driver.exceptions.DatabaseError
+except ImportError:
+
+    class _ClickHouseDbError(Exception):
+        pass
+
+
+class TrinoConnector(IbisConnector):
+    def __init__(self, connection_info):
+        super().__init__(DataSource.trino, connection_info)
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        import trino  # noqa: PLC0415
+
+        try:
+            return super().query(sql, limit)
+        except trino.exceptions.TrinoQueryError as e:
+            if not e.error_name == "EXCEEDED_TIME_LIMIT":
+                raise WrenError(
+                    ErrorCode.INVALID_SQL,
+                    str(e),
+                    phase=ErrorPhase.SQL_EXECUTION,
+                    metadata={DIALECT_SQL: sql},
+                ) from e
+            raise
+        except (WrenError, TimeoutError):
+            raise
+
+    def dry_run(self, sql: str) -> None:
+        import trino  # noqa: PLC0415
+
+        try:
+            super().dry_run(sql)
+        except trino.exceptions.TrinoQueryError as e:
+            if not e.error_name == "EXCEEDED_TIME_LIMIT":
+                raise WrenError(
+                    ErrorCode.INVALID_SQL,
+                    str(e),
+                    phase=ErrorPhase.SQL_DRY_RUN,
+                    metadata={DIALECT_SQL: sql},
+                ) from e
+            raise
+        except (WrenError, TimeoutError):
+            raise
+
+
+class ClickHouseConnector(IbisConnector):
+    def __init__(self, connection_info):
+        super().__init__(DataSource.clickhouse, connection_info)
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        try:
+            return super().query(sql, limit)
+        except _ClickHouseDbError as e:
+            if "TIMEOUT_EXCEEDED" not in str(e):
+                raise WrenError(
+                    ErrorCode.INVALID_SQL,
+                    str(e),
+                    phase=ErrorPhase.SQL_EXECUTION,
+                    metadata={DIALECT_SQL: sql},
+                ) from e
+            raise
+        except (WrenError, TimeoutError):
+            raise
+
+    def dry_run(self, sql: str) -> None:
+        try:
+            super().dry_run(sql)
+        except _ClickHouseDbError as e:
+            if "TIMEOUT_EXCEEDED" not in str(e):
+                raise WrenError(
+                    ErrorCode.INVALID_SQL,
+                    str(e),
+                    phase=ErrorPhase.SQL_DRY_RUN,
+                    metadata={DIALECT_SQL: sql},
+                ) from e
+            raise
+        except (WrenError, TimeoutError):
+            raise
+
+
+_DATA_SOURCE_TO_CLASS = {
+    DataSource.trino: TrinoConnector,
+    DataSource.clickhouse: ClickHouseConnector,
+}
+
+
+def create_connector(data_source: DataSource, connection_info) -> IbisConnector:
+    cls = _DATA_SOURCE_TO_CLASS.get(data_source, IbisConnector)
+    if cls is IbisConnector:
+        return IbisConnector(data_source, connection_info)
+    return cls(connection_info)

core/wren/src/wren/connector/mssql.py

@@ -0,0 +1,113 @@
+from contextlib import closing
+from decimal import Decimal as PyDecimal
+
+import pyarrow as pa
+import sqlglot.expressions as sge
+from ibis.expr.datatypes import Decimal
+from ibis.expr.types import Table
+from sqlglot import exp, parse_one
+
+from wren.connector.base import IbisConnector
+from wren.model.data_source import DataSource
+from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError
+
+
+class MSSqlConnector(IbisConnector):
+    def __init__(self, connection_info):
+        super().__init__(DataSource.mssql, connection_info)
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        sql = self._flatten_pagination_limit(sql)
+        ibis_table = self.connection.sql(sql)
+        if limit is not None:
+            ibis_table = ibis_table.limit(limit)
+        ibis_table = self._handle_pyarrow_unsupported_type(ibis_table)
+        return self._round_decimal_columns(ibis_table)
+
+    def _round_decimal_columns(self, ibis_table: Table, scale: int = 9) -> pa.Table:
+        def round_decimal(val):
+            if val is None:
+                return None
+            d = PyDecimal(str(val))
+            return d.quantize(PyDecimal("1." + "0" * scale))
+
+        decimal_columns = [
+            name
+            for name, dtype in ibis_table.schema().items()
+            if isinstance(dtype, Decimal)
+        ]
+        if not decimal_columns:
+            return ibis_table.to_pyarrow()
+
+        pandas_df = ibis_table.to_pandas()
+        for col_name in decimal_columns:
+            pandas_df[col_name] = pandas_df[col_name].apply(round_decimal)
+        return pa.Table.from_pandas(pandas_df)
+
+    def _flatten_pagination_limit(
+        self, sql_query: str, input_dialect: str = "tsql"
+    ) -> str:
+        try:
+            parsed = parse_one(sql_query, dialect=input_dialect)
+            if not isinstance(parsed, exp.Select) or not parsed.args.get("limit"):
+                return sql_query
+
+            from_clause = parsed.find(exp.From)
+            if not from_clause:
+                return sql_query
+
+            subqueries = []
+            if isinstance(from_clause.this, exp.Subquery):
+                subqueries.append(from_clause.this)
+            for join in parsed.args.get("joins") or []:
+                if isinstance(join, exp.Join):
+                    if isinstance(join.this, exp.Subquery):
+                        subqueries.append(join.this)
+                    if join.expression and isinstance(join.expression, exp.Subquery):
+                        subqueries.append(join.expression)
+
+            if len(subqueries) != 1:
+                return sql_query
+
+            inner = subqueries[0].this
+            if not isinstance(inner, exp.Select):
+                return sql_query
+
+            inner.set("limit", exp.Limit(expression=parsed.args["limit"].expression))
+            return inner.sql(dialect="tsql")
+        except Exception:
+            return sql_query
+
+    def dry_run(self, sql: str) -> None:
+        try:
+            super().dry_run(sql)
+        except AttributeError as e:
+            if "NoneType" in str(e) and "lower" in str(e):
+                error_message = self._describe_sql_for_error_message(sql)
+                raise WrenError(
+                    error_code=ErrorCode.INVALID_SQL,
+                    message=f"The sql dry run failed. {error_message}.",
+                    phase=ErrorPhase.SQL_DRY_RUN,
+                    metadata={DIALECT_SQL: sql},
+                ) from e
+            raise WrenError(
+                error_code=ErrorCode.IBIS_PROJECT_ERROR,
+                message=str(e),
+                phase=ErrorPhase.SQL_DRY_RUN,
+            ) from e
+
+    def _describe_sql_for_error_message(self, sql: str) -> str:
+        try:
+            tsql = sge.convert(sql).sql("mssql")
+            describe_sql = f"SELECT error_message FROM sys.dm_exec_describe_first_result_set({tsql}, NULL, 0)"
+            with closing(self.connection.raw_sql(describe_sql)) as cur:
+                rows = cur.fetchall()
+                if not rows:
+                    return "Unknown reason"
+                return rows[0][0]
+        except Exception:
+            return "Unknown reason"
+
+
+def create_connector(connection_info) -> MSSqlConnector:
+    return MSSqlConnector(connection_info)

core/wren/src/wren/connector/mysql.py

@@ -0,0 +1,57 @@
+import ibis.expr.datatypes as dt
+from ibis.expr.datatypes import Decimal
+from ibis.expr.datatypes.core import UUID
+from ibis.expr.types import Table
+
+from wren.connector.base import IbisConnector
+from wren.model.data_source import DataSource
+
+
+class MySqlConnector(IbisConnector):
+    def __init__(self, connection_info):
+        super().__init__(DataSource.mysql, 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):
+                result_table = result_table.mutate(
+                    **{name: result_table[name].cast("string")}
+                )
+        return result_table
+
+
+class DorisConnector(IbisConnector):
+    def __init__(self, connection_info):
+        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):
+                result_table = result_table.mutate(
+                    **{name: result_table[name].cast("string")}
+                )
+        return result_table
+
+
+def create_connector(data_source: DataSource, connection_info):
+    if data_source == DataSource.doris:
+        return DorisConnector(connection_info)
+    return MySqlConnector(connection_info)

core/wren/src/wren/connector/postgres.py

@@ -0,0 +1,76 @@
+from contextlib import suppress
+
+import pyarrow as pa
+from loguru import logger
+
+from wren.connector.base import IbisConnector
+from wren.model.data_source import DataSource
+from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError
+
+
+class PostgresConnector(IbisConnector):
+    def __init__(self, connection_info):
+        super().__init__(DataSource.postgres, connection_info)
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        import psycopg  # noqa: PLC0415
+
+        try:
+            return super().query(sql, limit)
+        except psycopg.errors.QueryCanceled:
+            raise
+        except (WrenError, TimeoutError):
+            raise
+        except Exception as e:
+            raise WrenError(
+                ErrorCode.GENERIC_USER_ERROR,
+                str(e),
+                phase=ErrorPhase.SQL_EXECUTION,
+                metadata={DIALECT_SQL: sql},
+            ) from e
+
+    def dry_run(self, sql: str) -> None:
+        import psycopg  # noqa: PLC0415
+
+        try:
+            super().dry_run(sql)
+        except psycopg.errors.QueryCanceled:
+            raise
+        except (WrenError, TimeoutError):
+            raise
+        except Exception as e:
+            raise WrenError(
+                ErrorCode.GENERIC_USER_ERROR,
+                str(e),
+                phase=ErrorPhase.SQL_DRY_RUN,
+                metadata={DIALECT_SQL: sql},
+            ) from e
+
+    def close(self) -> None:
+        if self._closed or not hasattr(self, "connection") or self.connection is None:
+            return
+        try:
+            if hasattr(self.connection, "con") and self.connection.con is not None:
+                if (
+                    hasattr(self.connection.con, "closed")
+                    and not self.connection.con.closed
+                ):
+                    with suppress(Exception):
+                        self.connection.con.cancel()
+                    import time  # noqa: PLC0415
+
+                    time.sleep(0.1)
+                    self.connection.con.close()
+            elif hasattr(self.connection, "close"):
+                self.connection.close()
+        except Exception as e:
+            logger.warning(f"Error closing postgres connection: {e}")
+            if hasattr(self.connection, "con"):
+                self.connection.con = None
+        finally:
+            self._closed = True
+            self.connection = None
+
+
+def create_connector(connection_info) -> PostgresConnector:
+    return PostgresConnector(connection_info)

core/wren/src/wren/connector/redshift.py

@@ -0,0 +1,68 @@
+from contextlib import closing
+
+import pandas as pd
+import pyarrow as pa
+from loguru import logger
+
+from wren.connector.base import ConnectorABC
+from wren.model import (
+    RedshiftConnectionInfo,
+    RedshiftConnectionUnion,
+    RedshiftIAMConnectionInfo,
+)
+from wren.model.error import ErrorCode, WrenError
+
+
+class RedshiftConnector(ConnectorABC):
+    def __init__(self, connection_info: RedshiftConnectionUnion):
+        import redshift_connector  # noqa: PLC0415
+
+        if isinstance(connection_info, RedshiftIAMConnectionInfo):
+            self.connection = redshift_connector.connect(
+                iam=True,
+                cluster_identifier=connection_info.cluster_identifier.get_secret_value(),
+                database=connection_info.database.get_secret_value(),
+                db_user=connection_info.user.get_secret_value(),
+                access_key_id=connection_info.access_key_id.get_secret_value(),
+                secret_access_key=connection_info.access_key_secret.get_secret_value(),
+                region=connection_info.region.get_secret_value(),
+            )
+        elif isinstance(connection_info, RedshiftConnectionInfo):
+            self.connection = redshift_connector.connect(
+                host=connection_info.host.get_secret_value(),
+                port=int(connection_info.port.get_secret_value()),
+                database=connection_info.database.get_secret_value(),
+                user=connection_info.user.get_secret_value(),
+                password=connection_info.password.get_secret_value(),
+            )
+        else:
+            raise WrenError(
+                ErrorCode.GENERIC_INTERNAL_ERROR,
+                "Invalid Redshift connection_info type",
+            )
+
+        self.connection.autocommit = True
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        if limit is not None:
+            sql = f"SELECT * FROM ({sql}) AS _q LIMIT {int(limit)}"
+        with closing(self.connection.cursor()) as cursor:
+            cursor.execute(sql)
+            cols = [desc[0] for desc in cursor.description]
+            rows = cursor.fetchall()
+            df = pd.DataFrame(rows, columns=cols)
+            return pa.Table.from_pandas(df)
+
+    def dry_run(self, sql: str) -> None:
+        with closing(self.connection.cursor()) as cursor:
+            cursor.execute(f"SELECT * FROM ({sql}) AS sub LIMIT 0")
+
+    def close(self) -> None:
+        try:
+            self.connection.close()
+        except Exception as e:
+            logger.warning(f"Error closing Redshift connection: {e}")
+
+
+def create_connector(connection_info) -> RedshiftConnector:
+    return RedshiftConnector(connection_info)

core/wren/src/wren/connector/spark.py

@@ -0,0 +1,52 @@
+import pyarrow as pa
+
+from wren.connector.base import ConnectorABC
+from wren.model import SparkConnectionInfo
+
+
+class SparkConnector(ConnectorABC):
+    def __init__(self, connection_info: SparkConnectionInfo):
+        self.connection_info = connection_info
+        self.connection = self._create_session()
+        self._closed = False
+
+    def _create_session(self):
+        from pyspark.sql import SparkSession  # noqa: PLC0415
+
+        host = self.connection_info.host.get_secret_value()
+        port = self.connection_info.port.get_secret_value()
+        return (
+            SparkSession.builder.remote(f"sc://{host}:{port}")
+            .appName("wren")
+            .getOrCreate()
+        )
+
+    def query(self, sql: str, limit: int | None = None) -> pa.Table:
+        df = self.connection.sql(sql).toPandas()
+        if hasattr(df, "attrs") and df.attrs:
+            df.attrs = {
+                k: v
+                for k, v in df.attrs.items()
+                if k not in ("metrics", "observed_metrics")
+            }
+        arrow_table = pa.Table.from_pandas(df)
+        if limit is not None:
+            arrow_table = arrow_table.slice(0, limit)
+        return arrow_table
+
+    def dry_run(self, sql: str) -> None:
+        self.connection.sql(sql).limit(0).count()
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        try:
+            self.connection.stop()
+        except Exception:
+            pass
+        finally:
+            self._closed = True
+
+
+def create_connector(connection_info) -> SparkConnector:
+    return SparkConnector(connection_info)

core/wren/src/wren/engine.py

@@ -0,0 +1,207 @@
+"""WrenEngine — SQL transform + execute against a data source.
+
+Example usage:
+
+    from wren.engine import WrenEngine
+    from wren.model.data_source import DataSource
+
+    engine = WrenEngine(
+        manifest_str="<base64-encoded MDL JSON>",
+        data_source=DataSource.postgres,
+        connection_info={"host": "localhost", "port": 5432, ...},
+    )
+
+    # Transform only (no DB required)
+    planned_sql = engine.transpile("SELECT * FROM orders")
+
+    # Execute against the data source
+    arrow_table = engine.query("SELECT * FROM orders", limit=100)
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pyarrow as pa
+import sqlglot
+
+from wren.connector.factory import get_connector
+from wren.mdl import get_manifest_extractor, get_session_context, to_json_base64
+from wren.model.data_source import DataSource
+from wren.model.error import DIALECT_SQL, PLANNED_SQL, ErrorCode, ErrorPhase, WrenError
+
+
+def _get_write_dialect(data_source: DataSource) -> str:
+    if data_source == DataSource.canner:
+        return "trino"
+    if data_source in {
+        DataSource.local_file,
+        DataSource.s3_file,
+        DataSource.minio_file,
+        DataSource.gcs_file,
+    }:
+        return "duckdb"
+    return data_source.name
+
+
+class WrenEngine:
+    """Thin facade over wren-core MDL processing and connector execution.
+
+    Parameters
+    ----------
+    manifest_str:
+        Base64-encoded MDL JSON string (as produced by ``wren_core.to_json_base64``).
+    data_source:
+        Target data source enum value.
+    connection_info:
+        Dict of connection parameters OR a typed ConnectionInfo object.
+    function_path:
+        Optional path to a CSV file of custom function definitions.
+        Passed through to wren-core SessionContext.
+    """
+
+    def __init__(
+        self,
+        manifest_str: str,
+        data_source: DataSource | str,
+        connection_info: dict[str, Any] | object,
+        function_path: str | None = None,
+    ):
+        if isinstance(data_source, str):
+            data_source = DataSource(data_source)
+
+        self.manifest_str = manifest_str
+        self.data_source = data_source
+        self.function_path = function_path
+
+        # Build typed ConnectionInfo if a raw dict was given.
+        # An empty dict is allowed for transpile-only usage (no DB connection).
+        if isinstance(connection_info, dict) and connection_info:
+            self.connection_info = data_source.get_connection_info(connection_info)
+        else:
+            self.connection_info = connection_info
+
+        self._connector = None
+
+    # ------------------------------------------------------------------
+    # SQL transformation (no DB access)
+    # ------------------------------------------------------------------
+
+    def transpile(self, sql: str, properties: dict | None = None) -> str:
+        """Transform SQL through MDL and transpile to the target dialect.
+
+        Returns the dialect SQL string without executing it.
+        """
+        planned = self._plan(sql, properties)
+        return self._transpile(planned)
+
+    def dry_plan(self, sql: str, properties: dict | None = None) -> str:
+        """Return the wren-core planned SQL (DataFusion dialect, before transpile)."""
+        return self._plan(sql, properties)
+
+    # ------------------------------------------------------------------
+    # SQL execution
+    # ------------------------------------------------------------------
+
+    def query(
+        self,
+        sql: str,
+        limit: int | None = None,
+        properties: dict | None = None,
+    ) -> pa.Table:
+        """Transpile and execute SQL, return results as an Arrow table."""
+        dialect_sql = self.transpile(sql, properties)
+        connector = self._get_connector()
+        try:
+            return connector.query(dialect_sql, limit)
+        except WrenError:
+            raise
+        except Exception as e:
+            raise WrenError(
+                ErrorCode.GENERIC_USER_ERROR,
+                str(e),
+                phase=ErrorPhase.SQL_EXECUTION,
+                metadata={DIALECT_SQL: dialect_sql},
+            ) from e
+
+    def dry_run(self, sql: str, properties: dict | None = None) -> None:
+        """Transpile and dry-run SQL without returning results."""
+        dialect_sql = self.transpile(sql, properties)
+        connector = self._get_connector()
+        try:
+            connector.dry_run(dialect_sql)
+        except WrenError:
+            raise
+        except Exception as e:
+            raise WrenError(
+                ErrorCode.GENERIC_USER_ERROR,
+                str(e),
+                phase=ErrorPhase.SQL_DRY_RUN,
+                metadata={DIALECT_SQL: dialect_sql},
+            ) from e
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def close(self) -> None:
+        if self._connector is not None:
+            self._connector.close()
+            self._connector = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_):
+        self.close()
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    # ------------------------------------------------------------------
+
+    def _plan(self, sql: str, properties: dict | None) -> str:
+        processed = None
+        if properties:
+            processed = frozenset(properties.items())
+
+        try:
+            # Extract minimal manifest for the query
+            extractor = get_manifest_extractor(self.manifest_str)
+            tables = extractor.resolve_used_table_names(sql)
+            manifest = extractor.extract_by(tables)
+            effective_manifest = to_json_base64(manifest)
+        except Exception:
+            effective_manifest = self.manifest_str
+
+        try:
+            session = get_session_context(
+                effective_manifest,
+                self.function_path,
+                processed,
+                self.data_source.name,
+            )
+            return session.transform_sql(sql)
+        except Exception as e:
+            raise WrenError(
+                ErrorCode.INVALID_SQL,
+                str(e),
+                phase=ErrorPhase.SQL_PLANNING,
+                metadata={DIALECT_SQL: sql},
+            ) from e
+
+    def _transpile(self, planned_sql: str) -> str:
+        try:
+            write = _get_write_dialect(self.data_source)
+            return sqlglot.transpile(planned_sql, read="duckdb", write=write)[0]
+        except Exception as e:
+            raise WrenError(
+                ErrorCode.SQLGLOT_ERROR,
+                str(e),
+                phase=ErrorPhase.SQL_TRANSPILE,
+                metadata={PLANNED_SQL: planned_sql},
+            ) from e
+
+    def _get_connector(self):
+        if self._connector is None:
+            self._connector = get_connector(self.data_source, self.connection_info)
+        return self._connector

core/wren/src/wren/mdl/__init__.py

@@ -0,0 +1,44 @@
+"""MDL processing utilities backed by wren-core-py."""
+
+from functools import cache
+
+import wren_core
+
+
+@cache
+def get_session_context(
+    manifest_str: str | None,
+    function_path: str | None,
+    properties: frozenset | None = None,
+    data_source: str | None = None,
+) -> wren_core.SessionContext:
+    return wren_core.SessionContext(
+        manifest_str, function_path, properties, data_source
+    )
+
+
+def get_manifest_extractor(manifest_str: str) -> wren_core.ManifestExtractor:
+    return wren_core.ManifestExtractor(manifest_str)
+
+
+def to_json_base64(manifest) -> str:
+    return wren_core.to_json_base64(manifest)
+
+
+def transform_sql(
+    manifest_str: str,
+    sql: str,
+    data_source: str | None = None,
+    function_path: str | None = None,
+    properties: dict | None = None,
+) -> str:
+    """Transform SQL through wren-core MDL processing.
+
+    Returns the planned SQL string (dialect-neutral DataFusion SQL).
+    """
+    processed = None
+    if properties:
+        processed = frozenset(properties.items())
+
+    session = get_session_context(manifest_str, function_path, processed, data_source)
+    return session.transform_sql(sql)

core/wren/src/wren/model/__init__.py

@@ -0,0 +1,296 @@
+"""Connection info models and DTOs for the wren package."""
+
+from __future__ import annotations
+
+from enum import Enum
+from typing import Annotated, Literal, Union
+
+from pydantic import BaseModel, BeforeValidator, Field, SecretStr
+
+from wren.model.error import ErrorCode, WrenError
+
+SecretPort = Annotated[
+    SecretStr, BeforeValidator(lambda v: str(v) if isinstance(v, int) else v)
+]
+
+
+class BaseConnectionInfo(BaseModel):
+    model_config = {"populate_by_name": True}
+
+    def to_key_string(self) -> str:
+        key_parts = []
+        for _, field_value in self:
+            if isinstance(field_value, SecretStr):
+                key_parts.append(field_value.get_secret_value())
+        return "|".join(key_parts)
+
+
+class BigQueryConnectionInfo(BaseConnectionInfo):
+    credentials: SecretStr = Field(
+        description="Base64 encode `credentials.json`", examples=["eyJ..."]
+    )
+    job_timeout_ms: int | None = Field(default=None)
+
+    def get_billing_project_id(self) -> str | None:
+        raise WrenError(
+            ErrorCode.NOT_IMPLEMENTED,
+            "get_billing_project_id not implemented by base class",
+        )
+
+
+class BigQueryDatasetConnectionInfo(BigQueryConnectionInfo):
+    bigquery_type: Literal["dataset"] = "dataset"
+    project_id: SecretStr = Field(examples=["my-project"])
+    dataset_id: SecretStr = Field(examples=["my_dataset"])
+
+    def get_billing_project_id(self):
+        return self.project_id.get_secret_value()
+
+    def __hash__(self):
+        return hash((self.project_id, self.dataset_id, self.credentials))
+
+
+class BigQueryProjectConnectionInfo(BigQueryConnectionInfo):
+    bigquery_type: Literal["project"] = "project"
+    region: SecretStr = Field(examples=["US"])
+    billing_project_id: SecretStr = Field(examples=["billing-project-1"])
+
+    def get_billing_project_id(self):
+        return self.billing_project_id.get_secret_value()
+
+    def __hash__(self):
+        return hash((self.region, self.billing_project_id, self.credentials))
+
+
+BigQueryConnectionUnion = Annotated[
+    Union[BigQueryDatasetConnectionInfo, BigQueryProjectConnectionInfo],
+    Field(discriminator="bigquery_type", default="dataset"),
+]
+
+
+class AthenaConnectionInfo(BaseConnectionInfo):
+    s3_staging_dir: SecretStr = Field(examples=["s3://my-bucket/athena-staging/"])
+    aws_access_key_id: SecretStr | None = Field(default=None)
+    aws_secret_access_key: SecretStr | None = Field(default=None)
+    aws_session_token: SecretStr | None = Field(default=None)
+    web_identity_token: SecretStr | None = Field(default=None)
+    role_arn: SecretStr | None = Field(default=None)
+    role_session_name: SecretStr | None = Field(default=None)
+    region_name: SecretStr = Field(examples=["us-west-2"], default=None)
+    schema_name: SecretStr | None = Field(
+        alias="schema_name", default=SecretStr("default")
+    )
+
+
+class CannerConnectionInfo(BaseConnectionInfo):
+    host: SecretStr = Field(examples=["localhost"])
+    port: SecretPort = Field(examples=["8080"])
+    user: SecretStr = Field(examples=["admin"])
+    pat: SecretStr = Field(examples=["eyJ..."])
+    workspace: SecretStr = Field(examples=["default"])
+    enable_ssl: bool = Field(default=False, alias="enableSSL")
+
+
+class ClickHouseConnectionInfo(BaseConnectionInfo):
+    host: SecretStr = Field(examples=["localhost"])
+    port: SecretPort = Field(examples=["8123"])
+    database: SecretStr = Field(examples=["default"])
+    user: SecretStr = Field(examples=["default"])
+    password: SecretStr | None = Field(default=None)
+    secure: bool = Field(default=False)
+    settings: dict[str, str] | None = Field(default=None)
+    kwargs: dict[str, str] | None = Field(default=None)
+
+
+class MSSqlConnectionInfo(BaseConnectionInfo):
+    host: SecretStr = Field(examples=["localhost"])
+    port: SecretPort = Field(examples=["1433"])
+    database: SecretStr = Field(examples=["master"])
+    user: SecretStr = Field(examples=["sa"])
+    password: SecretStr | None = Field(default=None)
+    driver: str = Field(default="ODBC Driver 18 for SQL Server")
+    tds_version: str = Field(default="8.0", alias="TDS_Version")
+    kwargs: dict[str, str] | None = Field(default=None)
+
+
+class MySqlConnectionInfo(BaseConnectionInfo):
+    host: SecretStr = Field(examples=["localhost"])
+    port: SecretPort = Field(examples=["3306"])
+    database: SecretStr = Field(examples=["default"])
+    user: SecretStr = Field(examples=["root"])
+    password: SecretStr | None = Field(default=None)
+    ssl_mode: SecretStr | None = Field(alias="sslMode", default=SecretStr("ENABLED"))
+    ssl_ca: SecretStr | None = Field(alias="sslCA", default=None)
+    kwargs: dict[str, str] | None = Field(default=None)
+
+
+class DorisConnectionInfo(BaseConnectionInfo):
+    host: SecretStr = Field(examples=["localhost"])
+    port: SecretPort = Field(examples=["9030"])
+    database: SecretStr = Field(examples=["default"])
+    user: SecretStr = Field(examples=["root"])
+    password: SecretStr | None = Field(default=None)
+    kwargs: dict[str, str] | None = Field(default=None)
+
+
+class PostgresConnectionInfo(BaseConnectionInfo):
+    host: SecretStr = Field(examples=["localhost"])
+    port: SecretPort = Field(examples=["5432"])
+    database: SecretStr = Field(examples=["postgres"])
+    user: SecretStr = Field(examples=["postgres"])
+    password: SecretStr | None = Field(default=None)
+    kwargs: dict[str, str] | None = Field(default=None)
+
+
+class OracleConnectionInfo(BaseConnectionInfo):
+    host: SecretStr = Field(default="localhost", examples=["localhost"])
+    port: SecretPort = Field(default="1521", examples=[1521])
+    database: SecretStr = Field(default="orcl", examples=["orcl"])
+    user: SecretStr = Field(examples=["admin"])
+    password: SecretStr | None = Field(default=None)
+    dsn: SecretStr | None = Field(default=None)
+
+
+class RedshiftConnectionInfo(BaseConnectionInfo):
+    redshift_type: Literal["redshift"] = "redshift"
+    host: SecretStr = Field(examples=["localhost"])
+    port: SecretPort = Field(examples=["5439"])
+    database: SecretStr = Field(examples=["dev"])
+    user: SecretStr = Field(examples=["awsuser"])
+    password: SecretStr = Field(examples=["password"])
+
+
+class RedshiftIAMConnectionInfo(BaseConnectionInfo):
+    redshift_type: Literal["redshift_iam"] = "redshift_iam"
+    cluster_identifier: SecretStr = Field(examples=["my-redshift-cluster"])
+    database: SecretStr = Field(examples=["dev"])
+    user: SecretStr = Field(examples=["awsuser"])
+    region: SecretStr = Field(examples=["us-west-2"])
+    access_key_id: SecretStr = Field(examples=["AKIA..."])
+    access_key_secret: SecretStr = Field(examples=["my-secret-key"])
+
+
+RedshiftConnectionUnion = Annotated[
+    Union[RedshiftConnectionInfo, RedshiftIAMConnectionInfo],
+    Field(discriminator="redshift_type"),
+]
+
+
+class SnowflakeConnectionInfo(BaseConnectionInfo):
+    user: SecretStr = Field(examples=["admin"])
+    password: SecretStr | None = Field(default=None)
+    account: SecretStr = Field(examples=["myaccount"])
+    database: SecretStr = Field(examples=["mydb"])
+    sf_schema: SecretStr = Field(alias="schema", examples=["myschema"])
+    warehouse: SecretStr | None = Field(default=None)
+    private_key: SecretStr | None = Field(default=None)
+    kwargs: dict[str, str] | None = Field(default=None)
+
+
+class SparkConnectionInfo(BaseConnectionInfo):
+    host: SecretStr = Field(examples=["localhost"])
+    port: SecretPort = Field(examples=["15002"])
+
+
+class DatabricksTokenConnectionInfo(BaseConnectionInfo):
+    databricks_type: Literal["token"] = "token"
+    server_hostname: SecretStr = Field(
+        alias="serverHostname", examples=["dbc-xxx.cloud.databricks.com"]
+    )
+    http_path: SecretStr = Field(alias="httpPath", examples=["/sql/1.0/warehouses/xxx"])
+    access_token: SecretStr = Field(alias="accessToken", examples=["dapi..."])
+
+
+class DatabricksServicePrincipalConnectionInfo(BaseConnectionInfo):
+    databricks_type: Literal["service_principal"] = "service_principal"
+    server_hostname: SecretStr = Field(alias="serverHostname")
+    http_path: SecretStr = Field(alias="httpPath")
+    client_id: SecretStr = Field(alias="clientId")
+    client_secret: SecretStr = Field(alias="clientSecret")
+    azure_tenant_id: SecretStr | None = Field(alias="azureTenantId", default=None)
+
+
+DatabricksConnectionUnion = Annotated[
+    Union[DatabricksTokenConnectionInfo, DatabricksServicePrincipalConnectionInfo],
+    Field(discriminator="databricks_type"),
+]
+
+
+class TrinoConnectionInfo(BaseConnectionInfo):
+    host: SecretStr = Field(examples=["localhost"])
+    port: SecretPort = Field(default="8080")
+    catalog: SecretStr = Field(examples=["hive"])
+    trino_schema: SecretStr = Field(alias="schema", examples=["default"])
+    user: SecretStr | None = Field(default=None)
+    password: SecretStr | None = Field(default=None)
+    kwargs: dict[str, str] | None = Field(default=None)
+
+
+class LocalFileConnectionInfo(BaseConnectionInfo):
+    url: SecretStr = Field(default="/", examples=["/data"])
+    format: str = Field(default="csv", examples=["csv", "parquet", "json", "duckdb"])
+
+
+class S3FileConnectionInfo(BaseConnectionInfo):
+    url: SecretStr = Field(default="/", examples=["/data"])
+    format: str = Field(default="csv")
+    bucket: SecretStr = Field(examples=["my-bucket"])
+    region: SecretStr = Field(examples=["us-west-2"])
+    access_key: SecretStr = Field(examples=["my-access-key"])
+    secret_key: SecretStr = Field(examples=["my-secret-key"])
+
+
+class MinioFileConnectionInfo(BaseConnectionInfo):
+    url: SecretStr = Field(default="/", examples=["/data"])
+    format: str = Field(default="csv")
+    ssl_enabled: bool = Field(default=False)
+    endpoint: SecretStr = Field(examples=["localhost:9000"])
+    bucket: SecretStr = Field(examples=["my-bucket"])
+    access_key: SecretStr = Field(examples=["my-account"])
+    secret_key: SecretStr = Field(examples=["my-password"])
+
+
+class GcsFileConnectionInfo(BaseConnectionInfo):
+    url: SecretStr = Field(default="/", examples=["/data"])
+    format: str = Field(default="csv")
+    bucket: SecretStr = Field(examples=["my-bucket"])
+    key_id: SecretStr = Field(examples=["my-key-id"])
+    secret_key: SecretStr = Field(examples=["my-secret-key"])
+    credentials: SecretStr | None = Field(default=None, examples=["eyJ..."])
+
+
+class ConnectionUrl(BaseConnectionInfo):
+    connection_url: SecretStr = Field(alias="connectionUrl")
+    kwargs: dict[str, str] | None = Field(default=None)
+
+
+ConnectionInfo = (
+    AthenaConnectionInfo
+    | BigQueryDatasetConnectionInfo
+    | BigQueryProjectConnectionInfo
+    | CannerConnectionInfo
+    | ClickHouseConnectionInfo
+    | ConnectionUrl
+    | MSSqlConnectionInfo
+    | MySqlConnectionInfo
+    | DorisConnectionInfo
+    | OracleConnectionInfo
+    | PostgresConnectionInfo
+    | RedshiftConnectionInfo
+    | RedshiftIAMConnectionInfo
+    | SnowflakeConnectionInfo
+    | SparkConnectionInfo
+    | DatabricksTokenConnectionInfo
+    | DatabricksServicePrincipalConnectionInfo
+    | TrinoConnectionInfo
+    | LocalFileConnectionInfo
+    | S3FileConnectionInfo
+    | MinioFileConnectionInfo
+    | GcsFileConnectionInfo
+)
+
+
+class SSLMode(str, Enum):
+    DISABLED = "disabled"
+    ENABLED = "enabled"
+    VERIFY_CA = "verify_ca"

core/wren/src/wren/model/data_source.py

@@ -0,0 +1,477 @@
+from __future__ import annotations
+
+import base64
+import ssl
+import urllib
+from enum import Enum, StrEnum, auto
+from json import loads
+from typing import Any
+from urllib.parse import unquote_plus
+
+import boto3
+import ibis
+from ibis import BaseBackend
+
+from wren.model import (
+    AthenaConnectionInfo,
+    BaseConnectionInfo,
+    BigQueryDatasetConnectionInfo,
+    BigQueryProjectConnectionInfo,
+    CannerConnectionInfo,
+    ClickHouseConnectionInfo,
+    ConnectionInfo,
+    ConnectionUrl,
+    DatabricksServicePrincipalConnectionInfo,
+    DatabricksTokenConnectionInfo,
+    DorisConnectionInfo,
+    GcsFileConnectionInfo,
+    LocalFileConnectionInfo,
+    MinioFileConnectionInfo,
+    MSSqlConnectionInfo,
+    MySqlConnectionInfo,
+    OracleConnectionInfo,
+    PostgresConnectionInfo,
+    RedshiftConnectionInfo,
+    RedshiftIAMConnectionInfo,
+    S3FileConnectionInfo,
+    SnowflakeConnectionInfo,
+    SparkConnectionInfo,
+    SSLMode,
+    TrinoConnectionInfo,
+)
+from wren.model.error import ErrorCode, WrenError
+
+X_WREN_DB_STATEMENT_TIMEOUT = "x-wren-db-statement_timeout"
+
+
+class DataSource(StrEnum):
+    athena = auto()
+    bigquery = auto()
+    canner = auto()
+    clickhouse = auto()
+    mssql = auto()
+    mysql = auto()
+    doris = auto()
+    oracle = auto()
+    postgres = auto()
+    redshift = auto()
+    snowflake = auto()
+    trino = auto()
+    local_file = auto()
+    s3_file = auto()
+    minio_file = auto()
+    gcs_file = auto()
+    duckdb = auto()
+    spark = auto()
+    databricks = auto()
+
+    def get_connection(self, info: ConnectionInfo) -> BaseBackend:
+        try:
+            return DataSourceExtension[self].get_connection(info)
+        except KeyError:
+            raise NotImplementedError(f"Unsupported data source: {self}")
+
+    def get_connection_info(
+        self,
+        data: dict[str, Any] | ConnectionInfo,
+        headers: dict[str, str] | None = None,
+    ) -> ConnectionInfo:
+        headers = headers or {}
+        if isinstance(data, BaseConnectionInfo):
+            info = data
+        else:
+            info = self._build_connection_info(data)
+        match self:
+            case DataSource.postgres:
+                kwargs = info.kwargs if info.kwargs else {}
+                if "connect_timeout" not in kwargs:
+                    kwargs["connect_timeout"] = 120
+                options = kwargs.get("options", "")
+                if "statement_timeout" not in options:
+                    if options:
+                        options += " "
+                    options += f"-c statement_timeout={headers.get(X_WREN_DB_STATEMENT_TIMEOUT, 180)}s"
+                    kwargs["options"] = options
+                info.kwargs = kwargs
+            case DataSource.clickhouse:
+                session_timeout = headers.get(X_WREN_DB_STATEMENT_TIMEOUT, 180)
+                if info.settings is None:
+                    info.settings = {}
+                if "max_execution_time" not in info.settings:
+                    info.settings["max_execution_time"] = int(session_timeout)
+            case DataSource.trino:
+                session_timeout = headers.get(X_WREN_DB_STATEMENT_TIMEOUT, 180)
+                if info.kwargs is None:
+                    info.kwargs = {}
+                session_properties = info.kwargs.get("session_properties", {})
+                if "query_max_execution_time" not in session_properties:
+                    session_properties["query_max_execution_time"] = (
+                        f"{session_timeout}s"
+                    )
+                info.kwargs["session_properties"] = session_properties
+            case DataSource.bigquery:
+                session_timeout = headers.get(X_WREN_DB_STATEMENT_TIMEOUT, 180)
+                if not hasattr(info, "job_timeout_ms") or info.job_timeout_ms is None:
+                    info.job_timeout_ms = int(session_timeout) * 1000
+        return info
+
+    def _build_connection_info(self, data: dict) -> ConnectionInfo:
+        if "connectionUrl" in data or "connection_url" in data:
+            if self == DataSource.clickhouse:
+                return self._handle_clickhouse_url(
+                    urllib.parse.urlparse(
+                        data.get("connectionUrl", data.get("connection_url"))
+                    )
+                )
+            return ConnectionUrl.model_validate(data)
+
+        match self:
+            case DataSource.athena:
+                return AthenaConnectionInfo.model_validate(data)
+            case DataSource.bigquery:
+                if "bigquery_type" in data and data["bigquery_type"] == "project":
+                    return BigQueryProjectConnectionInfo.model_validate(data)
+                return BigQueryDatasetConnectionInfo.model_validate(data)
+            case DataSource.canner:
+                return CannerConnectionInfo.model_validate(data)
+            case DataSource.clickhouse:
+                return ClickHouseConnectionInfo.model_validate(data)
+            case DataSource.mssql:
+                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:
+                return PostgresConnectionInfo.model_validate(data)
+            case DataSource.redshift:
+                if "redshift_type" in data and data["redshift_type"] == "redshift_iam":
+                    return RedshiftIAMConnectionInfo.model_validate(data)
+                return RedshiftConnectionInfo.model_validate(data)
+            case DataSource.snowflake:
+                return SnowflakeConnectionInfo.model_validate(data)
+            case DataSource.trino:
+                return TrinoConnectionInfo.model_validate(data)
+            case DataSource.duckdb | DataSource.local_file:
+                return LocalFileConnectionInfo.model_validate(data)
+            case DataSource.s3_file:
+                return S3FileConnectionInfo.model_validate(data)
+            case DataSource.minio_file:
+                return MinioFileConnectionInfo.model_validate(data)
+            case DataSource.gcs_file:
+                return GcsFileConnectionInfo.model_validate(data)
+            case DataSource.spark:
+                return SparkConnectionInfo.model_validate(data)
+            case DataSource.databricks:
+                if (
+                    "databricks_type" in data
+                    and data["databricks_type"] == "service_principal"
+                ):
+                    return DatabricksServicePrincipalConnectionInfo.model_validate(data)
+                return DatabricksTokenConnectionInfo.model_validate(data)
+            case _:
+                raise NotImplementedError(f"Unsupported data source: {self}")
+
+    def _handle_clickhouse_url(
+        self, parsed: urllib.parse.ParseResult
+    ) -> ClickHouseConnectionInfo:
+        if not parsed.scheme or parsed.scheme != "clickhouse":
+            raise WrenError(
+                ErrorCode.INVALID_CONNECTION_INFO,
+                "Invalid connection URL for ClickHouse",
+            )
+        kwargs = {}
+        if parsed.username:
+            kwargs["user"] = parsed.username
+        if parsed.password:
+            kwargs["password"] = unquote_plus(parsed.password)
+        if parsed.hostname:
+            kwargs["host"] = parsed.hostname
+        if parsed.port:
+            kwargs["port"] = str(parsed.port)
+        if database := parsed.path[1:]:
+            kwargs["database"] = database
+        parsed_kwargs = dict(urllib.parse.parse_qsl(parsed.query))
+        if "secure" in parsed_kwargs:
+            kwargs["secure"] = self._safe_strtobool(parsed_kwargs["secure"])
+            parsed_kwargs.pop("secure")
+        kwargs["kwargs"] = parsed_kwargs
+        return ClickHouseConnectionInfo(**kwargs)
+
+    def _safe_strtobool(self, val: str) -> bool:
+        return val.lower() in {"1", "true", "yes", "y"}
+
+
+class DataSourceExtension(Enum):
+    athena = "athena"
+    bigquery = "bigquery"
+    canner = "canner"
+    clickhouse = "clickhouse"
+    mssql = "mssql"
+    mysql = "mysql"
+    doris = "doris"
+    oracle = "oracle"
+    postgres = "postgres"
+    redshift = "redshift"
+    snowflake = "snowflake"
+    trino = "trino"
+    local_file = "local_file"
+    duckdb = "duckdb"
+    s3_file = "s3_file"
+    minio_file = "minio_file"
+    gcs_file = "gcs_file"
+    databricks = "databricks"
+    spark = "spark"
+
+    def get_connection(self, info: ConnectionInfo) -> BaseBackend:
+        try:
+            if hasattr(info, "connection_url"):
+                kwargs = info.kwargs if info.kwargs else {}
+                return ibis.connect(info.connection_url.get_secret_value(), **kwargs)
+            if self.name in {"local_file", "redshift", "spark", "duckdb"}:
+                raise NotImplementedError(
+                    f"{self.name} connection is not implemented to get ibis backend"
+                )
+            return getattr(self, f"get_{self.name}_connection")(info)
+        except KeyError:
+            raise NotImplementedError(f"Unsupported data source: {self}")
+        except WrenError:
+            raise
+        except Exception as e:
+            raise WrenError(ErrorCode.GET_CONNECTION_ERROR, f"{e!s}") from e
+
+    @staticmethod
+    def get_athena_connection(info: AthenaConnectionInfo) -> BaseBackend:
+        kwargs: dict[str, Any] = {
+            "s3_staging_dir": info.s3_staging_dir.get_secret_value(),
+            "schema_name": info.schema_name.get_secret_value(),
+        }
+        if info.region_name:
+            kwargs["region_name"] = info.region_name.get_secret_value()
+
+        if info.web_identity_token and info.role_arn:
+            oidc_token = info.web_identity_token.get_secret_value()
+            role_arn = info.role_arn.get_secret_value()
+            session_name = (
+                info.role_session_name.get_secret_value()
+                if info.role_session_name
+                else "wren-oidc-session"
+            )
+            region = info.region_name.get_secret_value() if info.region_name else None
+            sts = boto3.client("sts", region_name=region)
+            resp = sts.assume_role_with_web_identity(
+                RoleArn=role_arn,
+                RoleSessionName=session_name,
+                WebIdentityToken=oidc_token,
+            )
+            creds = resp["Credentials"]
+            kwargs["aws_access_key_id"] = creds["AccessKeyId"]
+            kwargs["aws_secret_access_key"] = creds["SecretAccessKey"]
+            kwargs["aws_session_token"] = creds["SessionToken"]
+        elif info.aws_access_key_id and info.aws_secret_access_key:
+            kwargs["aws_access_key_id"] = info.aws_access_key_id.get_secret_value()
+            kwargs["aws_secret_access_key"] = (
+                info.aws_secret_access_key.get_secret_value()
+            )
+            if info.aws_session_token:
+                kwargs["aws_session_token"] = info.aws_session_token.get_secret_value()
+
+        return ibis.athena.connect(**kwargs)
+
+    @staticmethod
+    def get_bigquery_connection(info: BigQueryDatasetConnectionInfo) -> BaseBackend:
+        from google.cloud import bigquery  # noqa: PLC0415
+        from google.oauth2 import service_account  # noqa: PLC0415
+
+        credits_json = loads(
+            base64.b64decode(info.credentials.get_secret_value()).decode("utf-8")
+        )
+        credentials = service_account.Credentials.from_service_account_info(
+            credits_json
+        )
+        credentials = credentials.with_scopes(
+            [
+                "https://www.googleapis.com/auth/drive",
+                "https://www.googleapis.com/auth/cloud-platform",
+            ]
+        )
+        bq_client = bigquery.Client(
+            project=info.project_id.get_secret_value(), credentials=credentials
+        )
+        job_config = bigquery.QueryJobConfig()
+        job_config.job_timeout_ms = info.job_timeout_ms
+        bq_client.default_query_job_config = job_config
+        return ibis.bigquery.connect(client=bq_client, credentials=credentials)
+
+    @staticmethod
+    def get_canner_connection(info: CannerConnectionInfo) -> BaseBackend:
+        return ibis.postgres.connect(
+            host=info.host.get_secret_value(),
+            port=int(info.port.get_secret_value()),
+            database=info.workspace.get_secret_value(),
+            user=info.user.get_secret_value(),
+            password=info.pat.get_secret_value(),
+        )
+
+    @staticmethod
+    def get_clickhouse_connection(info: ClickHouseConnectionInfo) -> BaseBackend:
+        return ibis.clickhouse.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 and info.password.get_secret_value()),
+            settings=info.settings if info.settings else {},
+            **info.kwargs if info.kwargs else {},
+        )
+
+    @classmethod
+    def get_mssql_connection(cls, info: MSSqlConnectionInfo) -> BaseBackend:
+        return ibis.mssql.connect(
+            host=info.host.get_secret_value(),
+            port=info.port.get_secret_value(),
+            database=info.database.get_secret_value(),
+            user=info.user.get_secret_value(),
+            password=info.password.get_secret_value(),
+            driver=info.driver,
+            TDS_Version=info.tds_version,
+            **info.kwargs if info.kwargs else {},
+        )
+
+    @classmethod
+    def get_mysql_connection(cls, info: MySqlConnectionInfo) -> BaseBackend:
+        ssl_context = cls._create_ssl_context(info)
+        kwargs = {"ssl": ssl_context} if ssl_context else {}
+        kwargs.setdefault("charset", "utf8mb4")
+        if info.kwargs:
+            kwargs.update(info.kwargs)
+        return 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,
+        )
+
+    @classmethod
+    def get_doris_connection(cls, info: DorisConnectionInfo) -> BaseBackend:
+        kwargs: dict = {}
+        kwargs.setdefault("charset", "utf8mb4")
+        if info.kwargs:
+            kwargs.update(info.kwargs)
+        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,
+        )
+        connection.con.get_autocommit = lambda: True
+        return connection
+
+    @staticmethod
+    def get_postgres_connection(info: PostgresConnectionInfo) -> BaseBackend:
+        return ibis.postgres.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 and info.password.get_secret_value()),
+            **info.kwargs if info.kwargs else {},
+        )
+
+    @staticmethod
+    def get_oracle_connection(info: OracleConnectionInfo) -> BaseBackend:
+        if hasattr(info, "dsn") and info.dsn:
+            return ibis.oracle.connect(
+                dsn=info.dsn.get_secret_value(),
+                user=info.user.get_secret_value(),
+                password=(info.password and info.password.get_secret_value()),
+            )
+        return ibis.oracle.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 and info.password.get_secret_value()),
+        )
+
+    @staticmethod
+    def get_snowflake_connection(info: SnowflakeConnectionInfo) -> BaseBackend:
+        if hasattr(info, "private_key") and info.private_key:
+            params = {
+                "user": info.user.get_secret_value(),
+                "private_key": info.private_key.get_secret_value(),
+                "account": info.account.get_secret_value(),
+                "database": info.database.get_secret_value(),
+                "schema": info.sf_schema.get_secret_value(),
+            }
+        else:
+            params = {
+                "user": info.user.get_secret_value(),
+                "password": info.password.get_secret_value(),
+                "account": info.account.get_secret_value(),
+                "database": info.database.get_secret_value(),
+                "schema": info.sf_schema.get_secret_value(),
+            }
+        if hasattr(info, "warehouse") and info.warehouse:
+            params["warehouse"] = info.warehouse.get_secret_value()
+        if info.kwargs:
+            params.update(info.kwargs)
+        return ibis.snowflake.connect(**params)
+
+    @staticmethod
+    def get_trino_connection(info: TrinoConnectionInfo) -> BaseBackend:
+        return ibis.trino.connect(
+            host=info.host.get_secret_value(),
+            port=int(info.port.get_secret_value()),
+            database=info.catalog.get_secret_value(),
+            schema=info.trino_schema.get_secret_value(),
+            user=(info.user and info.user.get_secret_value()),
+            password=(info.password and info.password.get_secret_value()),
+            **info.kwargs if info.kwargs else {},
+        )
+
+    @staticmethod
+    def get_databricks_connection(info: DatabricksTokenConnectionInfo) -> BaseBackend:
+        return ibis.databricks.connect(
+            server_hostname=info.server_hostname.get_secret_value(),
+            http_path=info.http_path.get_secret_value(),
+            access_token=info.access_token.get_secret_value(),
+        )
+
+    @staticmethod
+    def _create_ssl_context(info: ConnectionInfo) -> ssl.SSLContext | None:
+        ssl_mode = (
+            info.ssl_mode.get_secret_value()
+            if hasattr(info, "ssl_mode") and info.ssl_mode
+            else None
+        )
+
+        if ssl_mode == SSLMode.VERIFY_CA and not info.ssl_ca:
+            raise WrenError(
+                ErrorCode.INVALID_CONNECTION_INFO,
+                "SSL CA must be provided when SSL mode is VERIFY CA",
+            )
+
+        if not ssl_mode or ssl_mode == SSLMode.DISABLED:
+            return None
+
+        ctx = ssl.create_default_context()
+        ctx.check_hostname = False
+
+        if ssl_mode == SSLMode.ENABLED:
+            ctx.verify_mode = ssl.CERT_NONE
+        elif ssl_mode == SSLMode.VERIFY_CA:
+            ctx.verify_mode = ssl.CERT_REQUIRED
+            ctx.load_verify_locations(
+                cadata=base64.b64decode(info.ssl_ca.get_secret_value()).decode("utf-8")
+                if info.ssl_ca
+                else None
+            )
+
+        return ctx

core/wren/src/wren/model/error.py

@@ -0,0 +1,86 @@
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+DIALECT_SQL = "dialectSql"
+PLANNED_SQL = "plannedSql"
+
+
+class ErrorCode(int, Enum):
+    GENERIC_USER_ERROR = 1
+    NOT_FOUND = 2
+    MDL_NOT_FOUND = 3
+    INVALID_SQL = 4
+    INVALID_MDL = 5
+    DUCKDB_FILE_NOT_FOUND = 6
+    ATTACH_DUCKDB_ERROR = 7
+    VALIDATION_RULE_NOT_FOUND = 8
+    VALIDATION_ERROR = 9
+    VALIDATION_PARAMETER_ERROR = 10
+    GET_CONNECTION_ERROR = 11
+    INVALID_CONNECTION_INFO = 12
+    GENERIC_INTERNAL_ERROR = 100
+    LEGACY_ENGINE_ERROR = 101
+    NOT_IMPLEMENTED = 102
+    IBIS_PROJECT_ERROR = 103
+    SQLGLOT_ERROR = 104
+    GENERIC_EXTERNAL_ERROR = 200
+    DATABASE_TIMEOUT = 201
+
+
+class ErrorPhase(int, Enum):
+    REQUEST_RECEIVED = 1
+    MDL_EXTRACTION = 2
+    SQL_PARSING = 3
+    SQL_PLANNING = 4
+    SQL_TRANSPILE = 5
+    SQL_EXECUTION = 6
+    SQL_DRY_RUN = 7
+    RESPONSE_GENERATION = 8
+    METADATA_FETCHING = 9
+    VALIDATION = 10
+    SQL_SUBSTITUTE = 11
+
+
+class WrenError(Exception):
+    error_code: ErrorCode
+    message: str
+    phase: ErrorPhase | None = None
+    metadata: dict[str, Any] | None = None
+    timestamp: str | None = None
+
+    def __init__(
+        self,
+        error_code: ErrorCode,
+        message: str,
+        phase: ErrorPhase | None = None,
+        metadata: dict[str, Any] | None = None,
+        cause: Exception | None = None,
+    ):
+        self.error_code = error_code
+        self.message = message
+        self.phase = phase
+        self.metadata = metadata
+        self.timestamp = datetime.now().isoformat()
+        super().__init__(message)
+        if cause is not None:
+            self.__cause__ = cause
+
+    def __str__(self) -> str:
+        parts = [f"[{self.error_code.name}] {self.message}"]
+        if self.phase:
+            parts.append(f"phase={self.phase.name}")
+        return " ".join(parts)
+
+
+class DatabaseTimeoutError(WrenError):
+    def __init__(self, message: str):
+        enhanced_message = (
+            f"{message!s}.\n"
+            "It seems your database is not responding or the query is taking too long. "
+            "Please check your database status and query performance."
+        )
+        super().__init__(
+            error_code=ErrorCode.DATABASE_TIMEOUT,
+            message=enhanced_message,
+        )

core/wren/tests/conftest.py

@@ -0,0 +1,9 @@
+"""Root pytest configuration for the wren package test suite."""
+
+import pytest
+
+
+def pytest_configure(config: pytest.Config) -> None:
+    config.addinivalue_line("markers", "unit: unit tests — no database required")
+    config.addinivalue_line("markers", "duckdb: DuckDB connector tests — no Docker required")
+    config.addinivalue_line("markers", "postgres: PostgreSQL connector tests — requires Docker")

core/wren/tests/connectors/test_duckdb.py

@@ -0,0 +1,48 @@
+"""DuckDB connector tests.
+
+Uses DuckDB's built-in TPCH extension to generate test data — no Docker needed.
+The data is written to a temp file so ``DuckDBConnector`` can attach it.
+"""
+
+from __future__ import annotations
+
+import base64
+
+import duckdb
+import orjson
+import pytest
+
+from wren import WrenEngine
+from wren.model.data_source import DataSource
+
+from tests.suite.manifests import make_tpch_manifest
+from tests.suite.query import WrenQueryTestSuite
+
+pytestmark = pytest.mark.duckdb
+
+# DuckDB TPCH tables live in the "main" schema of the attached file.
+# The DuckDBConnector attaches the file as catalog = stem of the filename,
+# so "tpch.duckdb" → catalog "tpch".
+_CATALOG = "tpch"
+_SCHEMA = "main"
+
+
+class TestDuckDB(WrenQueryTestSuite):
+    manifest = make_tpch_manifest(table_catalog=_CATALOG, table_schema=_SCHEMA)
+    # DuckDB TPCH dbgen produces INTEGER as int64 in Arrow
+    order_id_dtype = "int64"
+
+    @pytest.fixture(scope="class")
+    def engine(self, tmp_path_factory) -> WrenEngine:  # type: ignore[override]
+        db_dir = tmp_path_factory.mktemp("duckdb")
+        db_path = db_dir / "tpch.duckdb"
+
+        # Generate TPCH sf=0.01 (1500 orders, 150 customers) into the file.
+        con = duckdb.connect(str(db_path))
+        con.execute("INSTALL tpch; LOAD tpch; CALL dbgen(sf=0.01)")
+        con.close()
+
+        manifest_str = base64.b64encode(orjson.dumps(self.manifest)).decode()
+        conn_info = {"url": str(db_dir), "format": "duckdb"}
+        with WrenEngine(manifest_str, DataSource.duckdb, conn_info) as e:
+            yield e

core/wren/tests/connectors/test_postgres.py

@@ -0,0 +1,90 @@
+"""PostgreSQL connector tests.
+
+Uses testcontainers to spin up a real Postgres instance.
+TPCH data is generated via DuckDB's built-in extension and loaded via psycopg.
+"""
+
+from __future__ import annotations
+
+import base64
+from urllib.parse import urlparse
+
+import duckdb
+import orjson
+import psycopg
+import pytest
+from testcontainers.postgres import PostgresContainer
+
+from wren import WrenEngine
+from wren.model.data_source import DataSource
+
+from tests.suite.manifests import make_tpch_manifest
+from tests.suite.query import WrenQueryTestSuite
+
+pytestmark = pytest.mark.postgres
+
+_SCHEMA = "public"
+
+
+def _load_tpch(conn_str: str) -> None:
+    """Generate TPCH sf=0.01 via DuckDB and bulk-load into Postgres."""
+    duck = duckdb.connect()
+    duck.execute("INSTALL tpch; LOAD tpch; CALL dbgen(sf=0.01)")
+
+    orders_rows = duck.execute(
+        "SELECT o_orderkey, o_custkey, o_orderstatus, "
+        "cast(o_totalprice as double), o_orderdate FROM orders"
+    ).fetchall()
+    customer_rows = duck.execute(
+        "SELECT c_custkey, c_name FROM customer"
+    ).fetchall()
+    duck.close()
+
+    with psycopg.connect(conn_str) as pg:
+        with pg.cursor() as cur:
+            cur.execute("""
+                CREATE TABLE orders (
+                    o_orderkey   INTEGER PRIMARY KEY,
+                    o_custkey    INTEGER NOT NULL,
+                    o_orderstatus CHAR(1) NOT NULL,
+                    o_totalprice  DOUBLE PRECISION NOT NULL,
+                    o_orderdate   DATE NOT NULL
+                )
+            """)
+            cur.executemany(
+                "INSERT INTO orders VALUES (%s, %s, %s, %s, %s)", orders_rows
+            )
+
+            cur.execute("""
+                CREATE TABLE customer (
+                    c_custkey INTEGER PRIMARY KEY,
+                    c_name    VARCHAR(25) NOT NULL
+                )
+            """)
+            cur.executemany(
+                "INSERT INTO customer VALUES (%s, %s)", customer_rows
+            )
+
+
+class TestPostgres(WrenQueryTestSuite):
+    manifest = make_tpch_manifest(table_catalog=None, table_schema=_SCHEMA)
+
+    @pytest.fixture(scope="class")
+    def engine(self) -> WrenEngine:  # type: ignore[override]
+        with PostgresContainer("postgres:16") as pg:
+            # testcontainers returns a SQLAlchemy-style URL; psycopg wants
+            # the plain postgresql:// form.
+            url = pg.get_connection_url().replace("+psycopg2", "")
+            _load_tpch(url)
+
+            parsed = urlparse(url)
+            conn_info = {
+                "host": parsed.hostname,
+                "port": parsed.port,
+                "database": parsed.path.lstrip("/"),
+                "user": parsed.username,
+                "password": parsed.password,
+            }
+            manifest_str = base64.b64encode(orjson.dumps(self.manifest)).decode()
+            with WrenEngine(manifest_str, DataSource.postgres, conn_info) as e:
+                yield e

core/wren/tests/suite/manifests.py

@@ -0,0 +1,70 @@
+"""Shared MDL manifest factories for connector tests."""
+
+from __future__ import annotations
+
+
+def make_tpch_manifest(table_catalog: str | None, table_schema: str) -> dict:
+    """Return a minimal TPCH MDL manifest for orders + customer.
+
+    The manifest includes:
+    - orders model with basic TPCH columns + a calculated field
+    - customer model with basic TPCH columns
+    - a MANY_TO_ONE relationship from orders to customer
+    - a relationship column on orders (customer → c_name)
+
+    Args:
+        table_catalog: Database catalog for tableReference (None to omit).
+        table_schema: Database schema for tableReference (e.g. "main", "public").
+    """
+
+    def table_ref(table: str) -> dict:
+        ref: dict = {"schema": table_schema, "table": table}
+        if table_catalog is not None:
+            ref["catalog"] = table_catalog
+        return ref
+
+    return {
+        "catalog": "wren",
+        "schema": "public",
+        "models": [
+            {
+                "name": "orders",
+                "tableReference": table_ref("orders"),
+                "columns": [
+                    {"name": "o_orderkey", "type": "integer"},
+                    {"name": "o_custkey", "type": "integer"},
+                    {"name": "o_orderstatus", "type": "varchar"},
+                    {"name": "o_totalprice", "type": "double"},
+                    {"name": "o_orderdate", "type": "date"},
+                    {
+                        "name": "order_cust_key",
+                        "type": "varchar",
+                        "expression": "concat(cast(o_orderkey as varchar), '_', cast(o_custkey as varchar))",
+                    },
+                    {
+                        "name": "customer",
+                        "type": "customer",
+                        "relationship": "orders_customer",
+                    },
+                ],
+                "primaryKey": "o_orderkey",
+            },
+            {
+                "name": "customer",
+                "tableReference": table_ref("customer"),
+                "columns": [
+                    {"name": "c_custkey", "type": "integer"},
+                    {"name": "c_name", "type": "varchar"},
+                ],
+                "primaryKey": "c_custkey",
+            },
+        ],
+        "relationships": [
+            {
+                "name": "orders_customer",
+                "models": ["orders", "customer"],
+                "joinType": "many_to_one",
+                "condition": '"orders".o_custkey = "customer".c_custkey',
+            }
+        ],
+    }

core/wren/tests/suite/query.py

@@ -0,0 +1,217 @@
+"""Shared query test suite for WrenEngine connector tests.
+
+How to add tests for a new data source
+=======================================
+
+Step 1 — Create the test file
+------------------------------
+Add ``tests/connectors/test_<name>.py``.  Copy the skeleton below and fill in
+the three required parts: pytest mark, manifest, and engine fixture.
+
+    # tests/connectors/test_clickhouse.py
+    import base64
+    import orjson
+    import pytest
+    from testcontainers.clickhouse import ClickHouseContainer
+
+    from wren import WrenEngine
+    from wren.model.data_source import DataSource
+    from tests.suite.manifests import make_tpch_manifest
+    from tests.suite.query import WrenQueryTestSuite
+
+    pytestmark = pytest.mark.clickhouse          # (1) marker for `just test-connector clickhouse`
+
+    class TestClickHouse(WrenQueryTestSuite):
+        manifest = make_tpch_manifest(           # (2) manifest — adjust catalog/schema
+            table_catalog=None,                  #     to match where TPCH data lands
+            table_schema="default",
+        )
+        order_id_dtype = "int64"                 # (3) override any differing expectations
+
+        @pytest.fixture(scope="class")
+        def engine(self):                        # (4) engine fixture — class-scoped
+            with ClickHouseContainer() as ch:
+                _load_tpch(ch.get_connection_url())
+                conn_info = { ... }
+                manifest_str = base64.b64encode(orjson.dumps(self.manifest)).decode()
+                with WrenEngine(manifest_str, DataSource.clickhouse, conn_info) as e:
+                    yield e
+
+Step 2 — Register the pytest marker
+-------------------------------------
+Add one line to ``tests/conftest.py``::
+
+    config.addinivalue_line("markers", "clickhouse: ClickHouse connector tests — requires Docker")
+
+Step 3 — Run the new tests
+----------------------------
+Install the connector extra if needed::
+
+    just install-extra clickhouse
+
+Then run::
+
+    just test-connector clickhouse
+
+All ``WrenQueryTestSuite`` tests run automatically.  Any ``test_*`` methods you
+add directly to ``TestClickHouse`` run alongside them.
+
+Overridable class variables
+----------------------------
+Override these in the subclass to match what the connector actually returns:
+
+    order_count    int   Total rows in TPCH orders table (default 15000, sf=0.01)
+    customer_count int   Total rows in TPCH customer table (default 1500, sf=0.01)
+    order_id_dtype str   Arrow dtype string for o_orderkey (default "int32")
+
+Example::
+
+    class TestClickHouse(WrenQueryTestSuite):
+        order_id_dtype = "int64"   # ClickHouse INT32 → Arrow int64
+
+Adding connector-specific tests
+---------------------------------
+Add ``test_*`` methods directly to the subclass — pytest discovers them
+alongside all inherited tests::
+
+    class TestClickHouse(WrenQueryTestSuite):
+        ...
+
+        def test_array_column(self, engine):
+            result = engine.query('SELECT array_col FROM "orders" LIMIT 1')
+            assert result.num_rows == 1
+
+Sharing tests across a subset of connectors (mix-ins)
+-------------------------------------------------------
+For capabilities shared by *some* connectors (e.g. timezone, window functions),
+define a separate mix-in class and include it only where relevant::
+
+    # tests/suite/timezone.py
+    class TimezoneTestSuite:
+        def test_timestamptz(self, engine): ...
+
+    # tests/connectors/test_postgres.py
+    class TestPostgres(WrenQueryTestSuite, TimezoneTestSuite):
+        ...   # gets core tests + timezone tests
+
+    # tests/connectors/test_duckdb.py
+    class TestDuckDB(WrenQueryTestSuite):
+        ...   # gets only core tests
+"""
+
+from __future__ import annotations
+
+import base64
+from typing import ClassVar
+
+import orjson
+import pyarrow as pa
+import pytest
+
+from wren import WrenEngine
+from wren.model.error import WrenError
+
+
+class WrenQueryTestSuite:
+    """Abstract base class providing shared query tests for all connectors.
+
+    Each subclass must:
+    - Set ``manifest`` to a connector-appropriate MDL dict.
+    - Provide a class-scoped ``engine`` fixture returning a ``WrenEngine``.
+    """
+
+    # Subclass must set this
+    manifest: ClassVar[dict]
+
+    # Overridable expectations — connectors may differ on counts or dtypes
+    order_count: ClassVar[int] = 15000      # TPCH sf=0.01
+    customer_count: ClassVar[int] = 1500    # TPCH sf=0.01
+    order_id_dtype: ClassVar[str] = "int32" # Postgres INTEGER → Arrow int32
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    @pytest.fixture(scope="class")
+    def manifest_str(self) -> str:
+        return base64.b64encode(orjson.dumps(self.manifest)).decode()
+
+    # ------------------------------------------------------------------
+    # Query execution tests
+    # ------------------------------------------------------------------
+
+    def test_basic_select(self, engine: WrenEngine) -> None:
+        result = engine.query(
+            'SELECT o_orderkey, o_custkey, o_orderstatus FROM "orders" ORDER BY o_orderkey LIMIT 1'
+        )
+        assert isinstance(result, pa.Table)
+        assert result.num_rows == 1
+        assert result.column_names == ["o_orderkey", "o_custkey", "o_orderstatus"]
+        # TPCH sf=0.01: first order row (orderkey=1)
+        assert result["o_orderkey"][0].as_py() == 1
+
+    def test_count(self, engine: WrenEngine) -> None:
+        result = engine.query('SELECT COUNT(*) AS cnt FROM "orders"')
+        assert result["cnt"][0].as_py() == self.order_count
+
+    def test_query_with_limit(self, engine: WrenEngine) -> None:
+        # engine.query limit= parameter truncates the result
+        result = engine.query('SELECT o_orderkey FROM "orders" ORDER BY o_orderkey', limit=3)
+        assert result.num_rows == 3
+
+    def test_calculated_field(self, engine: WrenEngine) -> None:
+        result = engine.query(
+            'SELECT o_orderkey, o_custkey, order_cust_key FROM "orders" ORDER BY o_orderkey LIMIT 1'
+        )
+        assert result.num_rows == 1
+        orderkey = result["o_orderkey"][0].as_py()
+        custkey = result["o_custkey"][0].as_py()
+        calc = result["order_cust_key"][0].as_py()
+        assert calc == f"{orderkey}_{custkey}"
+
+    def test_explicit_join(self, engine: WrenEngine) -> None:
+        result = engine.query(
+            """
+            SELECT o.o_orderkey, c.c_name
+            FROM "orders" o
+            JOIN "customer" c ON o.o_custkey = c.c_custkey
+            ORDER BY o.o_orderkey
+            LIMIT 5
+            """
+        )
+        assert result.num_rows == 5
+        assert "o_orderkey" in result.column_names
+        assert "c_name" in result.column_names
+
+    def test_order_id_dtype(self, engine: WrenEngine) -> None:
+        result = engine.query('SELECT o_orderkey FROM "orders" LIMIT 1')
+        field = result.schema.field("o_orderkey")
+        assert str(field.type) == self.order_id_dtype
+
+    # ------------------------------------------------------------------
+    # Dry run tests
+    # ------------------------------------------------------------------
+
+    def test_dry_run_valid(self, engine: WrenEngine) -> None:
+        # Should not raise
+        engine.dry_run('SELECT * FROM "orders" LIMIT 1')
+
+    def test_dry_run_invalid_table(self, engine: WrenEngine) -> None:
+        with pytest.raises(WrenError):
+            engine.dry_run('SELECT * FROM "NotFound"')
+
+    # ------------------------------------------------------------------
+    # Transpile / dry-plan tests (no DB access)
+    # ------------------------------------------------------------------
+
+    def test_transpile_returns_sql(self, engine: WrenEngine) -> None:
+        sql = engine.transpile('SELECT o_orderkey FROM "orders" LIMIT 1')
+        assert isinstance(sql, str)
+        assert len(sql) > 0
+
+    def test_dry_plan_returns_datafusion_sql(self, engine: WrenEngine) -> None:
+        planned = engine.dry_plan('SELECT o_orderkey FROM "orders" LIMIT 1')
+        assert isinstance(planned, str)
+        assert len(planned) > 0
+        # dry_plan returns wren-core / DataFusion SQL, not dialect SQL
+        assert "orders" in planned.lower()

core/wren/tests/unit/test_engine.py

@@ -0,0 +1,127 @@
+"""Unit tests for WrenEngine — no database required.
+
+transpile() and dry_plan() exercise the wren-core MDL planning + sqlglot
+transpile path without connecting to any data source.
+"""
+
+from __future__ import annotations
+
+import base64
+import math
+
+import orjson
+import pytest
+
+from wren import WrenEngine
+from wren.model.data_source import DataSource
+from wren.model.error import WrenError
+
+pytestmark = pytest.mark.unit
+
+# Minimal manifest with a single model.  No real DB needed for planning.
+_MANIFEST = {
+    "catalog": "wren",
+    "schema": "public",
+    "models": [
+        {
+            "name": "orders",
+            "tableReference": {"schema": "main", "table": "orders"},
+            "columns": [
+                {"name": "o_orderkey", "type": "integer"},
+                {"name": "o_custkey", "type": "integer"},
+                {"name": "o_orderstatus", "type": "varchar"},
+                {
+                    "name": "order_cust_key",
+                    "type": "varchar",
+                    "expression": "concat(cast(o_orderkey as varchar), '_', cast(o_custkey as varchar))",
+                },
+            ],
+            "primaryKey": "o_orderkey",
+        }
+    ],
+}
+_MANIFEST_STR = base64.b64encode(orjson.dumps(_MANIFEST)).decode()
+
+
+@pytest.fixture(scope="module")
+def duckdb_engine(tmp_path_factory):
+    """A WrenEngine pointed at a temporary DuckDB file (not queried by unit tests)."""
+    db_dir = tmp_path_factory.mktemp("unit_duckdb")
+    conn_info = {"url": str(db_dir), "format": "duckdb"}
+    with WrenEngine(_MANIFEST_STR, DataSource.duckdb, conn_info) as e:
+        yield e
+
+
+@pytest.fixture(scope="module")
+def pg_engine():
+    """A WrenEngine configured for Postgres (no real connection opened for planning)."""
+    conn_info = {
+        "host": "localhost",
+        "port": 5432,
+        "database": "test",
+        "user": "test",
+        "password": "test",
+    }
+    with WrenEngine(_MANIFEST_STR, DataSource.postgres, conn_info) as e:
+        yield e
+
+
+# ------------------------------------------------------------------
+# transpile
+# ------------------------------------------------------------------
+
+
+def test_transpile_returns_string(duckdb_engine: WrenEngine) -> None:
+    sql = duckdb_engine.transpile('SELECT o_orderkey FROM "orders" LIMIT 1')
+    assert isinstance(sql, str)
+    assert len(sql) > 0
+
+
+def test_transpile_postgres_dialect(pg_engine: WrenEngine) -> None:
+    """Transpile should produce Postgres-flavoured SQL (no backtick quoting, etc.)."""
+    sql = pg_engine.transpile('SELECT o_orderkey FROM "orders" LIMIT 1')
+    assert isinstance(sql, str)
+    # sqlglot Postgres output uses double-quote identifiers, not backticks
+    assert "`" not in sql
+
+
+def test_transpile_calculated_field(duckdb_engine: WrenEngine) -> None:
+    sql = duckdb_engine.transpile('SELECT order_cust_key FROM "orders" LIMIT 1')
+    assert isinstance(sql, str)
+    # The calculated column expression should be expanded in the SQL
+    assert "concat" in sql.lower() or "||" in sql.lower()
+
+
+def test_transpile_invalid_sql_raises(duckdb_engine: WrenEngine) -> None:
+    with pytest.raises(WrenError):
+        duckdb_engine.transpile("SELECT * FROM not_a_model_in_manifest")
+
+
+# ------------------------------------------------------------------
+# dry_plan
+# ------------------------------------------------------------------
+
+
+def test_dry_plan_returns_datafusion_sql(duckdb_engine: WrenEngine) -> None:
+    planned = duckdb_engine.dry_plan('SELECT o_orderkey FROM "orders" LIMIT 1')
+    assert isinstance(planned, str)
+    assert len(planned) > 0
+
+
+def test_dry_plan_invalid_sql_raises(duckdb_engine: WrenEngine) -> None:
+    with pytest.raises(WrenError):
+        duckdb_engine.dry_plan("SELECT * FROM nonexistent_model")
+
+
+# ------------------------------------------------------------------
+# Context manager
+# ------------------------------------------------------------------
+
+
+def test_context_manager_closes_connector() -> None:
+    conn_info = {"url": "/tmp", "format": "duckdb"}
+    with WrenEngine(_MANIFEST_STR, DataSource.duckdb, conn_info) as e:
+        assert e._connector is None  # connector is lazily initialized
+
+    # After __exit__, internal state is cleaned up
+    assert e._connector is None