wren-engine 0.1.0__py3-none-any.whl

wren/__init__.py

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

wren/cli.py

@@ -0,0 +1,382 @@
+"""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=False)
+
+_WREN_HOME = Path.home() / ".wren"
+_DEFAULT_MDL = _WREN_HOME / "mdl.json"
+_DEFAULT_CONN = _WREN_HOME / "connection_info.json"
+
+
+# ── File discovery helpers ─────────────────────────────────────────────────
+
+
+def _require_mdl(mdl: str | None) -> str:
+    """Return mdl arg if given, else auto-discover mdl.json from ~/.wren."""
+    if mdl is not None:
+        return mdl
+    if _DEFAULT_MDL.exists():
+        return str(_DEFAULT_MDL)
+    typer.echo(
+        f"Error: --mdl not specified and '{_DEFAULT_MDL}' not found.",
+        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).expanduser()
+    if path.exists():
+        import base64  # noqa: PLC0415
+
+        content = path.read_bytes()
+        if path.suffix.lower() == ".json":
+            # Raw JSON file — base64-encode it for WrenEngine
+            return base64.b64encode(content).decode()
+        # Non-.json file — assume it already contains a base64-encoded MDL string
+        return content.decode().strip()
+    # Not a file path — treat as a raw base64 string passed directly
+    return mdl
+
+
+def _load_conn(
+    connection_info: str | None,
+    connection_file: str | None,
+    *,
+    required: bool = True,
+) -> dict:
+    """Load connection dict from inline JSON or file, with ~/.wren auto-discovery.
+
+    If neither --connection-info nor --connection-file is given, looks for
+    connection_info.json in ~/.wren.  Raises typer.Exit(1) if required=True and nothing
+    is found.
+    """
+    if connection_info:
+        try:
+            conn = 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 not isinstance(conn, dict):
+            typer.echo(
+                "Error: --connection-info must decode to a JSON object.", err=True
+            )
+            raise typer.Exit(1)
+        return conn
+
+    path_str = connection_file or (
+        str(_DEFAULT_CONN) if _DEFAULT_CONN.exists() else None
+    )
+    if path_str:
+        path = Path(path_str).expanduser()
+        if not path.exists():
+            typer.echo(f"Error: connection file not found: {path_str}", err=True)
+            raise typer.Exit(1)
+        try:
+            conn = json.loads(path.read_text())
+        except json.JSONDecodeError as e:
+            typer.echo(f"Error: invalid JSON in {path_str}: {e}", err=True)
+            raise typer.Exit(1)
+        if not isinstance(conn, dict):
+            typer.echo(f"Error: {path_str} must contain a JSON object.", err=True)
+            raise typer.Exit(1)
+        return conn
+
+    if required:
+        typer.echo(
+            f"Error: --connection-file not specified and '{_DEFAULT_CONN}' not found.",
+            err=True,
+        )
+        raise typer.Exit(1)
+    return {}
+
+
+def _resolve_datasource(explicit: str | None, conn_dict: dict) -> str:
+    """Return datasource: use explicit --datasource arg first, then pop from conn dict.
+
+    Note: mutates conn_dict by removing the 'datasource' key so it is not
+    forwarded as an unknown field to WrenEngine / the connector.
+    """
+    if explicit:
+        conn_dict.pop("datasource", None)
+        return explicit
+    ds = conn_dict.pop("datasource", None)
+    if ds:
+        return ds
+    typer.echo(
+        "Error: --datasource not specified and 'datasource' key not found in connection info.",
+        err=True,
+    )
+    raise typer.Exit(1)
+
+
+def _build_engine(
+    datasource: str | None,
+    mdl: str | None,
+    connection_info: str | None,
+    connection_file: str | None,
+    *,
+    conn_required: bool = True,
+):
+    from wren.engine import WrenEngine  # noqa: PLC0415
+    from wren.model.data_source import DataSource  # noqa: PLC0415
+
+    manifest_str = _load_manifest(_require_mdl(mdl))
+    conn_dict = _load_conn(connection_info, connection_file, required=conn_required)
+    ds_str = _resolve_datasource(datasource, conn_dict)
+
+    try:
+        ds = DataSource(ds_str.lower())
+    except ValueError:
+        typer.echo(f"Error: unknown datasource '{ds_str}'", err=True)
+        raise typer.Exit(1)
+
+    return WrenEngine(
+        manifest_str=manifest_str, data_source=ds, connection_info=conn_dict
+    )
+
+
+# ── Shared option types ────────────────────────────────────────────────────
+
+DatasourceOpt = Annotated[
+    Optional[str],
+    typer.Option(
+        "--datasource",
+        "-d",
+        help="Data source (e.g. mysql, postgres). Defaults to 'datasource' field in connection_info.json.",
+    ),
+]
+MdlOpt = Annotated[
+    Optional[str],
+    typer.Option(
+        "--mdl",
+        "-m",
+        help=f"Path to MDL JSON file or base64 string. Defaults to {_DEFAULT_MDL}.",
+    ),
+]
+ConnInfoOpt = Annotated[
+    Optional[str],
+    typer.Option("--connection-info", help="Inline JSON connection string"),
+]
+ConnFileOpt = Annotated[
+    Optional[str],
+    typer.Option(
+        "--connection-file",
+        help=f"Path to JSON connection file. Defaults to {_DEFAULT_CONN}.",
+    ),
+]
+LimitOpt = Annotated[
+    Optional[int], typer.Option("--limit", "-l", help="Max rows to return")
+]
+OutputOpt = Annotated[
+    str, typer.Option("--output", "-o", help="Output format: json|csv|table")
+]
+
+
+# ── Default command (no subcommand = query) ────────────────────────────────
+
+
+@app.callback(invoke_without_command=True)
+def main(
+    ctx: typer.Context,
+    sql: Annotated[
+        Optional[str],
+        typer.Option(
+            "--sql", "-s", help="SQL query to execute (runs query by default)"
+        ),
+    ] = None,
+    datasource: DatasourceOpt = None,
+    mdl: MdlOpt = None,
+    connection_info: ConnInfoOpt = None,
+    connection_file: ConnFileOpt = None,
+    limit: LimitOpt = None,
+    output: OutputOpt = "table",
+) -> None:
+    """Wren Engine CLI.
+
+    Run with --sql to execute a query using mdl.json and connection_info.json from
+    ~/.wren.  Use a subcommand (query / dry-run / dry-plan / validate)
+    for explicit control.
+
+    connection_info.json format:
+
+    \b
+      {
+        "datasource": "mysql",
+        "host": "localhost",
+        "port": 3306,
+        "database": "mydb",
+        "user": "root",
+        "password": "secret"
+      }
+    """
+    if ctx.invoked_subcommand is not None:
+        return
+    if sql is None:
+        typer.echo(ctx.get_help())
+        return
+    with _build_engine(datasource, mdl, connection_info, connection_file) as engine:
+        try:
+            result = engine.query(sql, limit=limit)
+        except Exception as e:
+            typer.echo(f"Error: {e}", err=True)
+            raise typer.Exit(1)
+    _print_result(result, output)
+
+
+# ── Subcommands ────────────────────────────────────────────────────────────
+
+
+@app.command()
+def query(
+    sql: Annotated[str, typer.Option("--sql", "-s", help="SQL query to execute")],
+    datasource: DatasourceOpt = None,
+    mdl: MdlOpt = None,
+    connection_info: ConnInfoOpt = None,
+    connection_file: ConnFileOpt = None,
+    limit: LimitOpt = None,
+    output: OutputOpt = "table",
+):
+    """Execute a SQL query through the Wren semantic layer."""
+    with _build_engine(datasource, mdl, connection_info, connection_file) as engine:
+        try:
+            result = engine.query(sql, limit=limit)
+        except Exception as e:
+            typer.echo(f"Error: {e}", err=True)
+            raise typer.Exit(1)
+    _print_result(result, output)
+
+
+@app.command(name="dry-run")
+def dry_run(
+    sql: Annotated[str, typer.Option("--sql", "-s", help="SQL query to validate")],
+    datasource: DatasourceOpt = None,
+    mdl: MdlOpt = None,
+    connection_info: ConnInfoOpt = None,
+    connection_file: ConnFileOpt = None,
+):
+    """Dry-run a SQL query (parse + validate, no results returned)."""
+    with _build_engine(datasource, mdl, connection_info, connection_file) as engine:
+        try:
+            engine.dry_run(sql)
+            typer.echo("OK")
+        except Exception as e:
+            typer.echo(f"Error: {e}", err=True)
+            raise typer.Exit(1)
+
+
+@app.command(name="dry-plan")
+def dry_plan(
+    sql: Annotated[str, typer.Option("--sql", "-s", help="SQL query to plan")],
+    datasource: DatasourceOpt = None,
+    mdl: MdlOpt = None,
+    connection_file: ConnFileOpt = None,
+):
+    """Plan SQL through MDL and print the expanded SQL (no DB required)."""
+    from wren.engine import WrenEngine  # noqa: PLC0415
+    from wren.model.data_source import DataSource  # noqa: PLC0415
+
+    manifest_str = _load_manifest(_require_mdl(mdl))
+    # Read datasource from connection_info.json only when --datasource is not given
+    conn_dict = (
+        _load_conn(None, connection_file, required=False)
+        if connection_file is not None or datasource is None
+        else {}
+    )
+    ds_str = _resolve_datasource(datasource, conn_dict)
+
+    try:
+        ds = DataSource(ds_str.lower())
+    except ValueError:
+        typer.echo(f"Error: unknown datasource '{ds_str}'", err=True)
+        raise typer.Exit(1)
+
+    with WrenEngine(
+        manifest_str=manifest_str, data_source=ds, connection_info={}
+    ) as engine:
+        try:
+            result = engine.dry_plan(sql)
+            typer.echo(result)
+        except Exception as e:
+            typer.echo(f"Error: {e}", err=True)
+            raise typer.Exit(1)
+
+
+@app.command()
+def validate(
+    sql: Annotated[str, typer.Option("--sql", "-s", help="SQL query to validate")],
+    datasource: DatasourceOpt = None,
+    mdl: MdlOpt = None,
+    connection_info: ConnInfoOpt = None,
+    connection_file: ConnFileOpt = None,
+):
+    """Validate SQL can be planned and dry-run against the data source."""
+    with _build_engine(datasource, mdl, connection_info, connection_file) as engine:
+        try:
+            engine.dry_run(sql)
+            typer.echo("Valid")
+        except Exception as e:
+            typer.echo(f"Invalid: {e}", err=True)
+            raise typer.Exit(1)
+
+
+# ── Output formatting ──────────────────────────────────────────────────────
+
+
+def _print_result(table, output: str) -> None:
+    output = output.lower()
+    if output not in {"json", "csv", "table"}:
+        typer.echo(
+            f"Error: unsupported output format '{output}'. Use json, csv, or table.",
+            err=True,
+        )
+        raise typer.Exit(1)
+    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:
+        try:
+            df = table.to_pandas()
+            typer.echo(df.to_string(index=False))
+        except Exception:
+            typer.echo(str(table))
+
+
+@app.command()
+def version():
+    """Print the wren-engine version."""
+    from wren import __version__  # noqa: PLC0415
+
+    typer.echo(f"wren-engine {__version__}")
+
+
+try:
+    import lancedb  # noqa: PLC0415, F401
+    import sentence_transformers  # noqa: PLC0415, F401
+
+    from wren.memory.cli import memory_app  # noqa: PLC0415
+
+    app.add_typer(memory_app)
+except ImportError:
+    pass  # wren[memory] not installed
+
+
+if __name__ == "__main__":
+    app()

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"]

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

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)

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)

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)