wren-engine 0.1.0__tar.gz

wren_engine-0.1.0/.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.

wren_engine-0.1.0/.gitignore

@@ -0,0 +1,26 @@
+
+*.iml
+*.ipr
+*.iws
+target/
+**/.idea
+.run
+.DS_Store
+.classpath
+.settings
+.project
+*~
+*.class
+.checkstyle
+.mvn/timing.properties
+.mvn/maven.config
+wren-server/etc
+mcp-server/workspace/*
+!mcp-server/workspace/.gitkeep
+/**/var/
+__pycache__/
+venv/
+**/.env*
+**/*.so
+wren-core-py/dist/
+.claude/worktrees/

wren_engine-0.1.0/PKG-INFO

@@ -0,0 +1,222 @@
+Metadata-Version: 2.4
+Name: wren-engine
+Version: 0.1.0
+Summary: Wren Engine CLI and Python SDK — semantic SQL layer for 20+ data sources
+Project-URL: Homepage, https://getwren.ai
+Project-URL: Repository, https://github.com/Canner/wren-engine
+Project-URL: Issues, https://github.com/Canner/wren-engine/issues
+Author-email: Wren AI <contact@getwren.ai>
+License: Apache-2.0
+Keywords: analytics,cli,data-modeling,database,datafusion,mdl,python,sdk,semantic,semantic-layer,sql,wren,wrenai
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Python: >=3.11
+Requires-Dist: boto3>=1.26
+Requires-Dist: duckdb>=1.0
+Requires-Dist: ibis-framework>=10
+Requires-Dist: loguru>=0.7
+Requires-Dist: opendal>=0.45
+Requires-Dist: pandas>=2
+Requires-Dist: pyarrow-hotfix>=0.6
+Requires-Dist: pyarrow>=14
+Requires-Dist: pyasn1>=0.6.3
+Requires-Dist: pydantic>=2
+Requires-Dist: pyopenssl>=26.0.0
+Requires-Dist: requests>=2.33.0
+Requires-Dist: sqlglot>=27
+Requires-Dist: typer>=0.12
+Requires-Dist: wren-core-py>=0.1
+Provides-Extra: all
+Requires-Dist: databricks-sdk; extra == 'all'
+Requires-Dist: databricks-sql-connector; extra == 'all'
+Requires-Dist: google-auth; extra == 'all'
+Requires-Dist: ibis-framework[athena]; extra == 'all'
+Requires-Dist: ibis-framework[bigquery]; extra == 'all'
+Requires-Dist: ibis-framework[clickhouse]; extra == 'all'
+Requires-Dist: ibis-framework[mssql]; extra == 'all'
+Requires-Dist: ibis-framework[mysql]; extra == 'all'
+Requires-Dist: ibis-framework[postgres]; extra == 'all'
+Requires-Dist: ibis-framework[snowflake]; extra == 'all'
+Requires-Dist: ibis-framework[trino]; extra == 'all'
+Requires-Dist: lancedb>=0.6; extra == 'all'
+Requires-Dist: mysqlclient>=2.2; extra == 'all'
+Requires-Dist: oracledb>=2; extra == 'all'
+Requires-Dist: psycopg>=3; extra == 'all'
+Requires-Dist: pyspark>=3.5; extra == 'all'
+Requires-Dist: redshift-connector; extra == 'all'
+Requires-Dist: sentence-transformers>=2.2; extra == 'all'
+Requires-Dist: trino>=0.321; extra == 'all'
+Provides-Extra: athena
+Requires-Dist: ibis-framework[athena]; extra == 'athena'
+Provides-Extra: bigquery
+Requires-Dist: google-auth; extra == 'bigquery'
+Requires-Dist: ibis-framework[bigquery]; extra == 'bigquery'
+Provides-Extra: clickhouse
+Requires-Dist: ibis-framework[clickhouse]; extra == 'clickhouse'
+Provides-Extra: databricks
+Requires-Dist: databricks-sdk; extra == 'databricks'
+Requires-Dist: databricks-sql-connector; extra == 'databricks'
+Provides-Extra: dev
+Requires-Dist: orjson>=3; extra == 'dev'
+Requires-Dist: pytest>=8; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
+Requires-Dist: testcontainers[mysql,postgres]>=4; extra == 'dev'
+Provides-Extra: memory
+Requires-Dist: lancedb>=0.6; extra == 'memory'
+Requires-Dist: sentence-transformers>=2.2; extra == 'memory'
+Provides-Extra: mssql
+Requires-Dist: ibis-framework[mssql]; extra == 'mssql'
+Provides-Extra: mysql
+Requires-Dist: ibis-framework[mysql]; extra == 'mysql'
+Requires-Dist: mysqlclient>=2.2; extra == 'mysql'
+Provides-Extra: oracle
+Requires-Dist: oracledb>=2; extra == 'oracle'
+Provides-Extra: postgres
+Requires-Dist: ibis-framework[postgres]; extra == 'postgres'
+Requires-Dist: psycopg>=3; extra == 'postgres'
+Provides-Extra: redshift
+Requires-Dist: redshift-connector; extra == 'redshift'
+Provides-Extra: snowflake
+Requires-Dist: ibis-framework[snowflake]; extra == 'snowflake'
+Provides-Extra: spark
+Requires-Dist: pyspark>=3.5; extra == 'spark'
+Provides-Extra: trino
+Requires-Dist: ibis-framework[trino]; extra == 'trino'
+Requires-Dist: trino>=0.321; extra == 'trino'
+Description-Content-Type: text/markdown
+
+# wren-engine
+
+[![PyPI version](https://img.shields.io/pypi/v/wren-engine.svg)](https://pypi.org/project/wren-engine/)
+[![Python](https://img.shields.io/pypi/pyversions/wren-engine.svg)](https://pypi.org/project/wren-engine/)
+[![License](https://img.shields.io/pypi/l/wren-engine.svg)](https://github.com/Canner/wren-engine/blob/main/LICENSE)
+
+Wren Engine CLI and Python SDK — semantic SQL layer for 20+ data sources.
+
+Translate natural SQL queries through an [MDL (Modeling Definition Language)](https://docs.getwren.ai/) semantic layer and execute them against your database. Powered by [Apache DataFusion](https://datafusion.apache.org/) and [Ibis](https://ibis-project.org/).
+
+## Installation
+
+```bash
+pip install wren-engine              # Core (DuckDB included)
+pip install wren-engine[postgres]    # PostgreSQL
+pip install wren-engine[mysql]       # MySQL
+pip install wren-engine[bigquery]    # BigQuery
+pip install wren-engine[snowflake]   # Snowflake
+pip install wren-engine[clickhouse]  # ClickHouse
+pip install wren-engine[trino]       # Trino
+pip install wren-engine[mssql]       # SQL Server
+pip install wren-engine[databricks]  # Databricks
+pip install wren-engine[redshift]    # Redshift
+pip install wren-engine[spark]       # Spark
+pip install wren-engine[athena]      # Athena
+pip install wren-engine[oracle]      # Oracle
+pip install 'wren-engine[memory]'    # Schema & query memory (LanceDB)
+pip install 'wren-engine[all]'       # All connectors + memory
+```
+
+Requires Python 3.11+.
+
+## Quick start
+
+**1. Create `~/.wren/mdl.json`** — your semantic model:
+
+```json
+{
+  "catalog": "wren",
+  "schema": "public",
+  "models": [
+    {
+      "name": "orders",
+      "tableReference": { "schema": "mydb", "table": "orders" },
+      "columns": [
+        { "name": "order_id",    "type": "integer" },
+        { "name": "customer_id", "type": "integer" },
+        { "name": "total",       "type": "double" },
+        { "name": "status",      "type": "varchar" }
+      ],
+      "primaryKey": "order_id"
+    }
+  ]
+}
+```
+
+**2. Create `~/.wren/connection_info.json`** — your connection:
+
+```json
+{
+  "datasource": "mysql",
+  "host": "localhost",
+  "port": 3306,
+  "database": "mydb",
+  "user": "root",
+  "password": "secret"
+}
+```
+
+**3. Run queries** — `wren` auto-discovers both files from `~/.wren`:
+
+```bash
+wren --sql 'SELECT order_id FROM "orders" LIMIT 10'
+```
+
+For the full CLI reference and per-datasource `connection_info.json` formats, see [`docs/cli.md`](docs/cli.md) and [`docs/connections.md`](docs/connections.md).
+
+**4. Index schema for semantic search** (optional, requires `wren-engine[memory]`):
+
+```bash
+wren memory index                              # index MDL schema
+wren memory fetch -q "customer order price"    # fetch relevant schema context
+wren memory store --nl "top customers" --sql "SELECT ..."  # store NL→SQL pair
+wren memory recall -q "best customers"         # retrieve similar past queries
+```
+
+---
+
+## Python SDK
+
+```python
+import base64, orjson
+from wren import WrenEngine, DataSource
+
+manifest = { ... }  # your MDL dict
+manifest_str = base64.b64encode(orjson.dumps(manifest)).decode()
+
+with WrenEngine(manifest_str, DataSource.mysql, {"host": "...", ...}) as engine:
+    result = engine.query('SELECT * FROM "orders" LIMIT 10')
+    print(result.to_pandas())
+```
+
+---
+
+## Development
+
+```bash
+just install-dev    # Install with dev dependencies
+just lint           # Ruff format check + lint
+just format         # Auto-fix
+```
+
+| Command | What it runs | Docker needed |
+|---------|-------------|---------------|
+| `just test-unit` | Unit tests | No |
+| `just test-duckdb` | DuckDB connector tests | No |
+| `just test-postgres` | PostgreSQL connector tests | Yes |
+| `just test-mysql` | MySQL connector tests | Yes |
+| `just test` | All tests | Yes |
+
+## Publishing
+
+```bash
+./scripts/publish.sh            # Build + publish to PyPI
+./scripts/publish.sh --test     # Build + publish to TestPyPI
+./scripts/publish.sh --build    # Build only
+```
+
+## License
+
+Apache-2.0

wren_engine-0.1.0/README.md

@@ -0,0 +1,131 @@
+# wren-engine
+
+[![PyPI version](https://img.shields.io/pypi/v/wren-engine.svg)](https://pypi.org/project/wren-engine/)
+[![Python](https://img.shields.io/pypi/pyversions/wren-engine.svg)](https://pypi.org/project/wren-engine/)
+[![License](https://img.shields.io/pypi/l/wren-engine.svg)](https://github.com/Canner/wren-engine/blob/main/LICENSE)
+
+Wren Engine CLI and Python SDK — semantic SQL layer for 20+ data sources.
+
+Translate natural SQL queries through an [MDL (Modeling Definition Language)](https://docs.getwren.ai/) semantic layer and execute them against your database. Powered by [Apache DataFusion](https://datafusion.apache.org/) and [Ibis](https://ibis-project.org/).
+
+## Installation
+
+```bash
+pip install wren-engine              # Core (DuckDB included)
+pip install wren-engine[postgres]    # PostgreSQL
+pip install wren-engine[mysql]       # MySQL
+pip install wren-engine[bigquery]    # BigQuery
+pip install wren-engine[snowflake]   # Snowflake
+pip install wren-engine[clickhouse]  # ClickHouse
+pip install wren-engine[trino]       # Trino
+pip install wren-engine[mssql]       # SQL Server
+pip install wren-engine[databricks]  # Databricks
+pip install wren-engine[redshift]    # Redshift
+pip install wren-engine[spark]       # Spark
+pip install wren-engine[athena]      # Athena
+pip install wren-engine[oracle]      # Oracle
+pip install 'wren-engine[memory]'    # Schema & query memory (LanceDB)
+pip install 'wren-engine[all]'       # All connectors + memory
+```
+
+Requires Python 3.11+.
+
+## Quick start
+
+**1. Create `~/.wren/mdl.json`** — your semantic model:
+
+```json
+{
+  "catalog": "wren",
+  "schema": "public",
+  "models": [
+    {
+      "name": "orders",
+      "tableReference": { "schema": "mydb", "table": "orders" },
+      "columns": [
+        { "name": "order_id",    "type": "integer" },
+        { "name": "customer_id", "type": "integer" },
+        { "name": "total",       "type": "double" },
+        { "name": "status",      "type": "varchar" }
+      ],
+      "primaryKey": "order_id"
+    }
+  ]
+}
+```
+
+**2. Create `~/.wren/connection_info.json`** — your connection:
+
+```json
+{
+  "datasource": "mysql",
+  "host": "localhost",
+  "port": 3306,
+  "database": "mydb",
+  "user": "root",
+  "password": "secret"
+}
+```
+
+**3. Run queries** — `wren` auto-discovers both files from `~/.wren`:
+
+```bash
+wren --sql 'SELECT order_id FROM "orders" LIMIT 10'
+```
+
+For the full CLI reference and per-datasource `connection_info.json` formats, see [`docs/cli.md`](docs/cli.md) and [`docs/connections.md`](docs/connections.md).
+
+**4. Index schema for semantic search** (optional, requires `wren-engine[memory]`):
+
+```bash
+wren memory index                              # index MDL schema
+wren memory fetch -q "customer order price"    # fetch relevant schema context
+wren memory store --nl "top customers" --sql "SELECT ..."  # store NL→SQL pair
+wren memory recall -q "best customers"         # retrieve similar past queries
+```
+
+---
+
+## Python SDK
+
+```python
+import base64, orjson
+from wren import WrenEngine, DataSource
+
+manifest = { ... }  # your MDL dict
+manifest_str = base64.b64encode(orjson.dumps(manifest)).decode()
+
+with WrenEngine(manifest_str, DataSource.mysql, {"host": "...", ...}) as engine:
+    result = engine.query('SELECT * FROM "orders" LIMIT 10')
+    print(result.to_pandas())
+```
+
+---
+
+## Development
+
+```bash
+just install-dev    # Install with dev dependencies
+just lint           # Ruff format check + lint
+just format         # Auto-fix
+```
+
+| Command | What it runs | Docker needed |
+|---------|-------------|---------------|
+| `just test-unit` | Unit tests | No |
+| `just test-duckdb` | DuckDB connector tests | No |
+| `just test-postgres` | PostgreSQL connector tests | Yes |
+| `just test-mysql` | MySQL connector tests | Yes |
+| `just test` | All tests | Yes |
+
+## Publishing
+
+```bash
+./scripts/publish.sh            # Build + publish to PyPI
+./scripts/publish.sh --test     # Build + publish to TestPyPI
+./scripts/publish.sh --build    # Build only
+```
+
+## License
+
+Apache-2.0

wren_engine-0.1.0/docs/cli.md

@@ -0,0 +1,181 @@
+# CLI reference
+
+## Default command — query
+
+Running `wren --sql '...'` executes a query and prints the result. This is the same as `wren query --sql '...'`.
+
+```bash
+wren --sql 'SELECT COUNT(*) FROM "orders"'
+wren --sql 'SELECT * FROM "orders" LIMIT 5' --output csv
+wren --sql 'SELECT * FROM "orders"' --limit 100 --output json
+```
+
+Output formats: `table` (default), `csv`, `json`.
+
+## `wren query`
+
+Execute SQL and return results.
+
+```bash
+wren query --sql 'SELECT order_id, total FROM "orders" ORDER BY total DESC LIMIT 5'
+```
+
+## `wren dry-plan`
+
+Translate MDL SQL to the native dialect SQL for your data source. No database connection required.
+
+```bash
+wren dry-plan --sql 'SELECT order_id FROM "orders"'
+```
+
+## `wren dry-run`
+
+Validate SQL against the live database without returning rows. Prints `OK` on success.
+
+```bash
+wren dry-run --sql 'SELECT * FROM "orders" LIMIT 1'
+```
+
+## `wren validate`
+
+Same as `dry-run` but prints `Valid` / `Invalid: <reason>`.
+
+```bash
+wren validate --sql 'SELECT * FROM "NonExistent"'
+# Invalid: table not found ...
+```
+
+## Overriding defaults
+
+All flags are optional when `~/.wren/mdl.json` and `~/.wren/connection_info.json` exist:
+
+```bash
+wren --sql '...' \
+  --mdl /path/to/other-mdl.json \
+  --connection-file /path/to/prod-connection_info.json \
+  --datasource postgres
+```
+
+Or pass connection info inline:
+
+```bash
+wren --sql 'SELECT COUNT(*) FROM "orders"' \
+  --connection-info '{"datasource":"mysql","host":"localhost","port":3306,"database":"mydb","user":"root","password":"secret"}'
+```
+
+---
+
+## `wren memory` — Schema & Query Memory
+
+LanceDB-backed semantic memory for MDL schema search and NL→SQL retrieval. Requires the `memory` extra:
+
+```bash
+pip install 'wren-engine[memory]'
+```
+
+All `memory` subcommands accept `--path DIR` to override the default storage location (`~/.wren/memory/`).
+
+### Hybrid strategy: full text vs. embedding search
+
+When providing schema context to an LLM, there is a trade-off:
+
+- **Small schemas** — the full plain-text description fits easily in the LLM context window and gives better results because the LLM sees the complete structure (model→column relationships, join paths, primary keys) rather than isolated fragments from a vector search.
+- **Large schemas** — the full text exceeds what is practical to send in a single prompt, so embedding search is needed to retrieve only the relevant fragments.
+
+`wren memory fetch` automatically picks the right strategy based on the **character length** of the generated plain-text description:
+
+| Schema size | Threshold | Strategy |
+|---|---|---|
+| Below 30,000 chars (~8K tokens) | Default | Returns full plain text |
+| Above 30,000 chars | Default | Returns embedding search results |
+
+The threshold is measured in characters (not tokens) because character length is free to compute, while accurate token counting requires a tokeniser. The 4:1 chars-to-tokens ratio holds for English; CJK text compresses less (~1.5:1), so a CJK-heavy schema switches to embedding search sooner — which is the conservative direction.
+
+The default threshold (30,000 chars) can be overridden with `--threshold`.
+
+### `wren memory index`
+
+Parse the MDL manifest and index all schema items (models, columns, relationships, views) into LanceDB with local embeddings.
+
+```bash
+wren memory index                          # uses ~/.wren/mdl.json
+wren memory index --mdl /path/to/mdl.json  # explicit MDL file
+```
+
+### `wren memory describe`
+
+Print the full schema as structured plain text. No embedding or LanceDB required — this is a pure transformation of the MDL manifest into a human/LLM-readable format.
+
+```bash
+wren memory describe                          # uses ~/.wren/mdl.json
+wren memory describe --mdl /path/to/mdl.json
+```
+
+### `wren memory fetch`
+
+Get schema context for an LLM. Automatically chooses the best strategy based on schema size: full plain text for small schemas, embedding search for large schemas.
+
+When using the search strategy, optional `--type` and `--model` filters narrow the results.
+
+```bash
+wren memory fetch -q "customer order price"
+wren memory fetch -q "revenue" --type column --model orders
+wren memory fetch -q "日期" --threshold 50000 --output json
+```
+
+| Flag | Description |
+|------|-------------|
+| `-q, --query` | Search query (required) |
+| `--mdl` | Path to MDL JSON file |
+| `-l, --limit` | Max results for search strategy (default: 5) |
+| `-t, --type` | Filter: `model`, `column`, `relationship`, `view` (search strategy only) |
+| `--model` | Filter by model name (search strategy only) |
+| `--threshold` | Character threshold for full vs search (default: 30,000) |
+| `-o, --output` | Output format: `table` (default), `json` |
+
+### `wren memory store`
+
+Store a natural-language-to-SQL pair for future few-shot retrieval.
+
+```bash
+wren memory store \
+  --nl "show top customers by revenue" \
+  --sql "SELECT c_name, sum(o_totalprice) FROM orders JOIN customer GROUP BY 1 ORDER BY 2 DESC" \
+  --datasource postgres
+```
+
+### `wren memory recall`
+
+Search stored NL→SQL pairs by semantic similarity to a query.
+
+```bash
+wren memory recall -q "best customers"
+wren memory recall -q "月度營收" --datasource mysql --limit 5 --output json
+```
+
+| Flag | Description |
+|------|-------------|
+| `-q, --query` | Search query (required) |
+| `-l, --limit` | Max results (default: 3) |
+| `-d, --datasource` | Filter by data source |
+| `-o, --output` | Output format: `table` (default), `json` |
+
+### `wren memory status`
+
+Show index statistics: storage path, table names, and row counts.
+
+```bash
+wren memory status
+# Path: /Users/you/.wren/memory
+#   schema_items: 47 rows
+#   query_history: 12 rows
+```
+
+### `wren memory reset`
+
+Drop all memory tables and start fresh.
+
+```bash
+wren memory reset          # prompts for confirmation
+wren memory reset --force  # skip confirmation
+```