Skip to content

Commit 41f3f21

Browse files
goldmedalclaude
andauthored
feat(wren): add standalone wren Python SDK package (#1471)
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 312eb20 commit 41f3f21

39 files changed

Lines changed: 5717 additions & 1 deletion

.claude/settings.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,8 @@
1818
"Bash(gh run:*)",
1919
"Bash(gh api:*)",
2020
"Bash(ls:*)",
21-
"Bash(wc:*)"
21+
"Bash(wc:*)",
22+
"Bash(uv sync:*)"
2223
]
2324
}
2425
}

.github/workflows/wren-ci.yml

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
name: Wren SDK CI
2+
permissions:
3+
contents: read
4+
pull-requests: write
5+
6+
on:
7+
pull_request:
8+
paths:
9+
- 'wren/**'
10+
- 'wren-core/**'
11+
- 'wren-core-py/**'
12+
- 'wren-core-base/**'
13+
14+
concurrency:
15+
group: ${{ github.workflow }}-${{ github.event.number }}
16+
cancel-in-progress: true
17+
18+
jobs:
19+
lint:
20+
name: lint
21+
runs-on: ubuntu-latest
22+
defaults:
23+
run:
24+
working-directory: wren
25+
steps:
26+
- uses: actions/checkout@v4
27+
- name: Install uv
28+
uses: astral-sh/setup-uv@v4
29+
- name: Lint
30+
run: |
31+
uvx ruff format --check src/
32+
uvx ruff check src/
33+
34+
test-unit:
35+
name: unit tests
36+
runs-on: ubuntu-latest
37+
defaults:
38+
run:
39+
working-directory: wren
40+
steps:
41+
- uses: actions/checkout@v4
42+
- uses: actions/setup-python@v5
43+
with:
44+
python-version: "3.11"
45+
- name: Install uv
46+
uses: astral-sh/setup-uv@v4
47+
- name: Cache Cargo
48+
uses: actions/cache@v4
49+
with:
50+
path: |
51+
~/.cargo/bin/
52+
~/.cargo/registry/index/
53+
~/.cargo/registry/cache/
54+
~/.cargo/git/db/
55+
./wren-core-py/target/
56+
key: ${{ runner.os }}-cargo-wren-${{ hashFiles('**/Cargo.lock') }}
57+
- name: Install Rust toolchain
58+
uses: dtolnay/rust-toolchain@stable
59+
- name: Build wren-core-py wheel
60+
working-directory: wren-core-py
61+
run: |
62+
pipx install poetry
63+
poetry install --no-root
64+
poetry run maturin build
65+
- name: Install dependencies
66+
run: |
67+
uv lock --find-links ../wren-core-py/target/wheels/ --upgrade-package wren-core-py
68+
uv sync --extra dev --find-links ../wren-core-py/target/wheels/
69+
- name: Run unit tests
70+
run: uv run pytest tests/unit/ -v -m unit
71+
72+
test-postgres:
73+
name: postgres tests
74+
runs-on: ubuntu-latest
75+
defaults:
76+
run:
77+
working-directory: wren
78+
steps:
79+
- uses: actions/checkout@v4
80+
- uses: actions/setup-python@v5
81+
with:
82+
python-version: "3.11"
83+
- name: Install uv
84+
uses: astral-sh/setup-uv@v4
85+
- name: Cache Cargo
86+
uses: actions/cache@v4
87+
with:
88+
path: |
89+
~/.cargo/bin/
90+
~/.cargo/registry/index/
91+
~/.cargo/registry/cache/
92+
~/.cargo/git/db/
93+
./wren-core-py/target/
94+
key: ${{ runner.os }}-cargo-wren-${{ hashFiles('**/Cargo.lock') }}
95+
- name: Install Rust toolchain
96+
uses: dtolnay/rust-toolchain@stable
97+
- name: Build wren-core-py wheel
98+
working-directory: wren-core-py
99+
run: |
100+
pipx install poetry
101+
poetry install --no-root
102+
poetry run maturin build
103+
- name: Install dependencies
104+
run: |
105+
uv lock --find-links ../wren-core-py/target/wheels/ --upgrade-package wren-core-py
106+
uv sync --extra postgres --extra dev --find-links ../wren-core-py/target/wheels/
107+
- name: Run postgres tests
108+
run: uv run pytest tests/connectors/test_postgres.py -v -m postgres

ibis-server/docs/development.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@ To start the server:
5353
```
5454
WREN_ENGINE_ENDPOINT=http://localhost:8080
5555
```
56+
Install the dependencies for testing by `just install --with dev`. Then you can run the tests by `just test`. \
5657
More information about the environment variables can be found in the [Environment Variables](#Environment-Variables) section.
5758
- Run specific data source test using [pytest marker](https://docs.pytest.org/en/stable/example/markers.html). There are some markers for different data sources. See the list
5859
in [pyproject.toml](../pyproject.toml).

ibis-server/poetry.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
[virtualenvs]
2+
in-project = true

wren/.claude/CLAUDE.md

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
# CLAUDE.md — wren package
2+
3+
Standalone Python SDK and CLI for Wren Engine. Wraps `wren-core-py` (PyO3 bindings) + Ibis connectors into a single installable package.
4+
5+
## Package Structure
6+
7+
```
8+
wren/
9+
src/wren/
10+
engine.py — WrenEngine facade (transpile / query / dry_run / dry_plan)
11+
cli.py — Typer CLI: wren query|dry-run|transpile|validate
12+
mdl/ — wren-core-py session context + manifest extraction helpers
13+
connector/ — Per-datasource Ibis connectors (factory.py + one file per source)
14+
model/
15+
data_source.py — DataSource enum + per-source ConnectionInfo factories
16+
error.py — WrenError, ErrorCode, ErrorPhase
17+
tests/
18+
```
19+
20+
## Build & Development
21+
22+
```bash
23+
cd wren
24+
just install # build wren-core-py wheel + uv sync
25+
just install-all # with all optional extras
26+
just install-extra <extra> # e.g. just install-extra postgres
27+
just test # pytest tests/
28+
just lint # ruff format --check + ruff check
29+
just format # ruff auto-fix
30+
just build # uv build (produces wheel)
31+
```
32+
33+
Uses `uv` (not Poetry). `pyproject.toml` uses `hatchling` as build backend.
34+
35+
## Key Design Points
36+
37+
- **WrenEngine** is the main entry point. It accepts a base64-encoded MDL JSON string, a `DataSource`, and a connection dict.
38+
- **Query flow**: `_plan()` → wren-core `SessionContext.transform_sql()``_transpile()` via sqlglot → connector `.query()`.
39+
- **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.
40+
- **`get_session_context` is `@cache`-decorated** — same `(manifest_str, function_path, properties, data_source)` tuple reuses the same SessionContext. Avoid mutating session state.
41+
- **Write dialect mapping**: `canner``trino`; file sources (`local_file`, `s3_file`, `minio_file`, `gcs_file`) → `duckdb`. All others use `data_source.name` directly.
42+
- **WrenEngine is a context manager** (`__enter__` / `__exit__` call `close()`).
43+
44+
## Connectors
45+
46+
`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`.
47+
48+
## Optional Extras
49+
50+
Install per data-source extras: `postgres`, `mysql`, `bigquery`, `snowflake`, `clickhouse`, `trino`, `mssql`, `databricks`, `redshift`, `spark`, `athena`, `oracle`, `all`, `dev`.
51+
52+
On macOS, `mysql` extra needs:
53+
```bash
54+
PKG_CONFIG_PATH="$(brew --prefix mysql-client)/lib/pkgconfig" just install-extra mysql
55+
```
56+
57+
## Dependency on wren-core-py
58+
59+
`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.

wren/README.md

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# wren
2+
3+
Wren Engine CLI and Python SDK — semantic SQL layer for 20+ data sources.
4+
5+
## Installation
6+
7+
```bash
8+
pip install wren
9+
```
10+
11+
## Usage
12+
13+
```python
14+
from wren import WrenEngine, DataSource
15+
```
16+
17+
See the [Wren Engine documentation](https://getwren.ai) for details.
18+
19+
## Running tests
20+
21+
Install dev dependencies first:
22+
23+
```bash
24+
just install-dev
25+
```
26+
27+
| Command | What it runs | Docker needed |
28+
|---------|-------------|---------------|
29+
| `just test-unit` | Unit tests (transpile, dry-plan, context manager) | No |
30+
| `just test-duckdb` | DuckDB connector tests — generates TPCH data via `dbgen` | No |
31+
| `just test-postgres` | PostgreSQL connector tests — spins up a container | Yes |
32+
| `just test` | All tests | Yes |
33+
34+
Run a specific connector via marker:
35+
36+
```bash
37+
just test-connector postgres
38+
```
39+
40+
To add tests for a new connector, subclass `WrenQueryTestSuite` in
41+
`tests/connectors/test_<name>.py` and provide a class-scoped `engine` fixture.
42+
All base tests are inherited automatically.

wren/justfile

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
default:
2+
@just --list
3+
4+
build-core:
5+
cd ../wren-core-py && just install && just build
6+
7+
core-wheel-dir := "../wren-core-py/target/wheels/"
8+
9+
install-core:
10+
just build-core
11+
12+
install:
13+
just build-core
14+
uv sync --find-links {{ core-wheel-dir }}
15+
16+
install-dev:
17+
just build-core
18+
uv sync --extra dev --find-links {{ core-wheel-dir }}
19+
20+
install-all:
21+
just build-core
22+
uv sync --all-extras --find-links {{ core-wheel-dir }}
23+
24+
install-extra extra *flags:
25+
#!/usr/bin/env bash
26+
set -euo pipefail
27+
if [ "{{ extra }}" = "mysql" ]; then
28+
mysql_prefix=$(brew --prefix mysql-client 2>/dev/null || true)
29+
if [ -n "$mysql_prefix" ]; then
30+
export PKG_CONFIG_PATH="$mysql_prefix/lib/pkgconfig:${PKG_CONFIG_PATH:-}"
31+
fi
32+
fi
33+
dev_flag=""
34+
for flag in {{ flags }}; do
35+
if [ "$flag" = "--dev" ]; then
36+
dev_flag="--extra dev"
37+
fi
38+
done
39+
uv sync --extra {{ extra }} $dev_flag --find-links {{ core-wheel-dir }}
40+
41+
dev:
42+
uv run wren
43+
44+
test:
45+
uv run pytest tests/ -v
46+
47+
test-unit:
48+
uv run pytest tests/unit/ -v -m unit
49+
50+
test-duckdb:
51+
uv run pytest tests/connectors/test_duckdb.py -v -m duckdb
52+
53+
test-postgres:
54+
uv run pytest tests/connectors/test_postgres.py -v -m postgres
55+
56+
test-connector marker:
57+
uv run pytest tests/connectors/ -v -m {{ marker }}
58+
59+
lint:
60+
uv run ruff format --check src/
61+
uv run ruff check src/
62+
63+
format:
64+
uv run ruff format src/
65+
uv run ruff check --fix src/
66+
67+
build:
68+
uv build
69+
70+
alias fmt := format

wren/pyproject.toml

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
[build-system]
2+
requires = ["hatchling"]
3+
build-backend = "hatchling.build"
4+
5+
[project]
6+
name = "wren"
7+
dynamic = ["version"]
8+
description = "Wren Engine CLI and Python SDK — semantic SQL layer for 20+ data sources"
9+
readme = "README.md"
10+
requires-python = ">=3.11"
11+
license = { text = "Apache-2.0" }
12+
keywords = ["wren", "sql", "semantic", "mdl", "datafusion"]
13+
classifiers = [
14+
"Development Status :: 4 - Beta",
15+
"Intended Audience :: Developers",
16+
"License :: OSI Approved :: Apache Software License",
17+
"Programming Language :: Python :: 3",
18+
"Programming Language :: Python :: 3.11",
19+
"Programming Language :: Python :: 3.12",
20+
]
21+
dependencies = [
22+
"wren-core-py>=0.1",
23+
"ibis-framework>=10",
24+
"duckdb>=1.0",
25+
"sqlglot>=27",
26+
"typer>=0.12",
27+
"pydantic>=2",
28+
"pyarrow>=14",
29+
"pyarrow-hotfix>=0.6",
30+
"loguru>=0.7",
31+
"opendal>=0.45",
32+
"pandas>=2",
33+
"boto3>=1.26",
34+
]
35+
36+
[project.optional-dependencies]
37+
postgres = ["psycopg>=3", "ibis-framework[postgres]"]
38+
mysql = ["mysqlclient>=2.2", "ibis-framework[mysql]"]
39+
bigquery = ["ibis-framework[bigquery]", "google-auth"]
40+
snowflake = ["ibis-framework[snowflake]"]
41+
clickhouse = ["ibis-framework[clickhouse]"]
42+
trino = ["ibis-framework[trino]", "trino>=0.321"]
43+
mssql = ["ibis-framework[mssql]"]
44+
databricks = ["databricks-sql-connector", "databricks-sdk"]
45+
redshift = ["redshift_connector"]
46+
spark = ["pyspark>=3.5"]
47+
athena = ["ibis-framework[athena]"]
48+
oracle = ["ibis-framework[oracle]", "oracledb"]
49+
all = ["wren[postgres,mysql,bigquery,snowflake,clickhouse,trino,mssql,databricks,redshift,athena,oracle]"]
50+
dev = ["pytest>=8", "ruff>=0.4", "orjson>=3", "testcontainers[postgres]>=4"]
51+
52+
[project.scripts]
53+
wren = "wren.cli:app"
54+
55+
[project.urls]
56+
Homepage = "https://getwren.ai"
57+
Repository = "https://github.com/Canner/wren-engine"
58+
Issues = "https://github.com/Canner/wren-engine/issues"
59+
60+
[tool.hatch.version]
61+
path = "src/wren/__init__.py"
62+
63+
[tool.hatch.build.targets.wheel]
64+
packages = ["src/wren"]
65+
66+
[tool.ruff]
67+
line-length = 88
68+
target-version = "py311"
69+
70+
[tool.ruff.lint]
71+
select = ["E", "F", "I", "PLC"]
72+
ignore = [
73+
"E501", # line-too-long, enforced by ruff format
74+
]

wren/src/wren/__init__.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
"""Wren — semantic SQL layer for 20+ data sources."""
2+
3+
__version__ = "0.0.1"
4+
5+
from wren.engine import WrenEngine
6+
from wren.model.data_source import DataSource
7+
from wren.model.error import WrenError
8+
9+
__all__ = ["WrenEngine", "DataSource", "WrenError", "__version__"]

0 commit comments

Comments
 (0)