Skip to content

Commit c76918f

Browse files
authored
Merge pull request #229 from CausalInferenceLab/feat/db-explorer-adapters
feat(db): Cloudflare D1 + 범용 SQLAlchemy explorer + 자동 라우팅 팩토리
2 parents 25c1162 + be29bb5 commit c76918f

9 files changed

Lines changed: 1296 additions & 8 deletions

File tree

.env.example

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,3 +14,25 @@ OPENAI_API_KEY=
1414
# in production so secrets decrypt across restarts and machines. Generate one:
1515
# python -c "from cryptography.fernet import Fernet; print(Fernet.generate_key().decode())"
1616
LANG2SQL_SECRET_KEY=
17+
18+
# ── Database connection ──────────────────────────────────────────────────
19+
# The default explorer is chosen from these (precedence: LANG2SQL_DB_URL, then
20+
# Cloudflare D1, else an offline canned stub). Install the matching driver
21+
# extra, e.g. uv sync --extra postgres (or bigquery/snowflake/mysql/duckdb).
22+
#
23+
# Any SQLAlchemy URL works (one adapter, many engines):
24+
# postgresql+psycopg://user:pass@host:5432/dbname
25+
# bigquery://project/dataset
26+
# snowflake://user:pass@account/db/schema?warehouse=wh
27+
# mysql+pymysql://user:pass@host/dbname
28+
# duckdb:////absolute/path/to/file.duckdb
29+
LANG2SQL_DB_URL=
30+
# Optional default schema for the SQLAlchemy explorer.
31+
LANG2SQL_DB_SCHEMA=
32+
33+
# Cloudflare D1 (used when LANG2SQL_DB_URL is unset). D1 is SQLite over an HTTP
34+
# API — no driver needed. Find IDs in the Cloudflare dashboard → D1.
35+
CLOUDFLARE_D1_ACCOUNT_ID=
36+
CLOUDFLARE_D1_DATABASE_ID=
37+
# API token with D1 read access (Account → API Tokens). Required for D1 queries.
38+
CLOUDFLARE_API_TOKEN=

pyproject.toml

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,24 @@ authors = [
1414
dependencies = [
1515
"discord.py>=2.3,<3.0", # Phase 1 frontend transport
1616
"cryptography>=42.0", # EncryptedSecrets at-rest encryption
17+
"sqlalchemy>=2.0", # generic DB explorer (one adapter, many engines)
1718
]
1819

1920
[project.optional-dependencies]
20-
# Real outbound adapters (V1 ships urllib OpenAI + stub PG; these enable v1.5 swaps)
21-
postgres = ["psycopg[binary]>=3.2,<4.0"]
21+
# DB driver extras. The SQLAlchemyExplorer is dialect-agnostic; install only the
22+
# drivers you connect to. Cloudflare D1 needs no driver (HTTP API via stdlib).
23+
postgres = ["psycopg[binary]>=3.2,<4.0"]
24+
bigquery = ["sqlalchemy-bigquery>=1.11"]
25+
snowflake = ["snowflake-sqlalchemy>=1.6"]
26+
mysql = ["pymysql>=1.1"]
27+
duckdb = ["duckdb-engine>=0.13"]
28+
all-db = [
29+
"psycopg[binary]>=3.2,<4.0",
30+
"sqlalchemy-bigquery>=1.11",
31+
"snowflake-sqlalchemy>=1.6",
32+
"pymysql>=1.1",
33+
"duckdb-engine>=0.13",
34+
]
2235

2336
[project.scripts]
2437
lang2sql = "lang2sql.frontends.cli.app:main"
Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,20 @@
1-
"""DB adapters — :class:`ExplorerPort` impls."""
1+
"""DB adapters — :class:`ExplorerPort` impls + the connection factory.
2+
3+
``build_explorer`` routes a connection string to the right adapter:
4+
Cloudflare D1 over its HTTP API, everything else over generic SQLAlchemy.
5+
"""
26

37
from __future__ import annotations
48

9+
from .d1_explorer import D1Explorer
10+
from .factory import build_explorer, explorer_from_env
511
from .postgres_explorer import PostgresExplorer
12+
from .sqlalchemy_explorer import SqlAlchemyExplorer
613

7-
__all__ = ["PostgresExplorer"]
14+
__all__ = [
15+
"build_explorer",
16+
"explorer_from_env",
17+
"D1Explorer",
18+
"SqlAlchemyExplorer",
19+
"PostgresExplorer",
20+
]
Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
"""Cloudflare D1 explorer — read-only introspection over the D1 HTTP API.
2+
3+
D1 is SQLite that lives on Cloudflare's edge and is only reachable from a Worker
4+
or over the REST query endpoint. A Python process (our Discord bot) uses the
5+
**HTTP Query API**:
6+
7+
POST /client/v4/accounts/{account_id}/d1/database/{database_id}/query
8+
Authorization: Bearer <token>
9+
{"sql": "...", "params": [...]}
10+
11+
Since D1 *is* SQLite, schema introspection uses ``sqlite_master`` / ``PRAGMA``.
12+
The HTTP call is injectable (``transport``) so the adapter is unit-testable with
13+
no network.
14+
"""
15+
16+
from __future__ import annotations
17+
18+
import asyncio
19+
import json
20+
import os
21+
import urllib.request
22+
from typing import Any, Callable
23+
24+
from ...core.ports.explorer import Column, Table
25+
26+
_API_ROOT = "https://api.cloudflare.com/client/v4"
27+
28+
# A transport takes (sql, params) and returns the parsed D1 JSON response.
29+
Transport = Callable[[str, list], dict]
30+
31+
32+
class D1Explorer:
33+
"""ExplorerPort backed by Cloudflare D1's HTTP query API."""
34+
35+
def __init__(
36+
self,
37+
account_id: str,
38+
database_id: str,
39+
token: str | None = None,
40+
*,
41+
transport: Transport | None = None,
42+
timeout: float = 30.0,
43+
) -> None:
44+
self.account_id = account_id
45+
self.database_id = database_id
46+
self._token = token if token is not None else os.environ.get("CLOUDFLARE_API_TOKEN")
47+
self._timeout = timeout
48+
self._transport = transport or self._http_transport
49+
50+
# --- ExplorerPort ----------------------------------------------------
51+
52+
async def list_tables(self) -> list[Table]:
53+
rows = await self._query(
54+
"SELECT name FROM sqlite_master WHERE type='table' "
55+
"AND name NOT LIKE 'sqlite_%' AND name NOT LIKE '_cf_%' ORDER BY name"
56+
)
57+
return [Table(name=r["name"], schema="") for r in rows]
58+
59+
async def describe_table(self, name: str) -> Table:
60+
rows = await self._query(f"PRAGMA table_info({_ident(name)})")
61+
cols = [
62+
Column(name=r["name"], type=r["type"] or "", nullable=not bool(r["notnull"]))
63+
for r in rows
64+
]
65+
return Table(name=name, schema="", columns=cols)
66+
67+
async def sample_rows(self, name: str, limit: int = 5) -> list[dict]:
68+
return await self._query(f"SELECT * FROM {_ident(name)} LIMIT {int(limit)}")
69+
70+
async def execute(self, sql: str, limit: int = 1000) -> list[dict]:
71+
rows = await self._query(sql)
72+
return rows[: int(limit)]
73+
74+
# --- internals -------------------------------------------------------
75+
76+
async def _query(self, sql: str, params: list | None = None) -> list[dict]:
77+
resp = await asyncio.to_thread(self._transport, sql, params or [])
78+
if not resp.get("success", False):
79+
errors = resp.get("errors") or resp.get("messages") or "unknown D1 error"
80+
raise RuntimeError(f"D1 query failed: {errors}")
81+
result = resp.get("result") or []
82+
if not result:
83+
return []
84+
# The query endpoint returns one result object per statement.
85+
return result[0].get("results", []) or []
86+
87+
def _http_transport(self, sql: str, params: list) -> dict:
88+
if not self._token:
89+
raise RuntimeError("CLOUDFLARE_API_TOKEN not set (D1 requires an API token)")
90+
url = (
91+
f"{_API_ROOT}/accounts/{self.account_id}"
92+
f"/d1/database/{self.database_id}/query"
93+
)
94+
body = json.dumps({"sql": sql, "params": params}).encode("utf-8")
95+
req = urllib.request.Request(
96+
url,
97+
data=body,
98+
method="POST",
99+
headers={
100+
"Authorization": f"Bearer {self._token}",
101+
"Content-Type": "application/json",
102+
},
103+
)
104+
with urllib.request.urlopen(req, timeout=self._timeout) as resp:
105+
return json.loads(resp.read().decode("utf-8"))
106+
107+
108+
def _ident(name: str) -> str:
109+
"""Quote a SQLite identifier, rejecting anything that isn't a plain name.
110+
111+
introspection helpers interpolate the table name into PRAGMA/SELECT where
112+
binds aren't allowed, so we hard-validate to avoid injection.
113+
"""
114+
if not name.replace("_", "").isalnum():
115+
raise ValueError(f"unsafe table identifier: {name!r}")
116+
return f'"{name}"'
Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
"""build_explorer — turn a connection string into the right ExplorerPort.
2+
3+
This is what makes ``/connect`` trivial: the user (or env) gives one URL and the
4+
factory routes it. Cloudflare D1 has its own HTTP adapter; everything else with
5+
a normal SQLAlchemy URL goes through the generic SQLAlchemy explorer.
6+
7+
d1://<account_id>/<database_id> → D1Explorer (token from env)
8+
postgresql+psycopg://user:…/db → SqlAlchemyExplorer
9+
bigquery://project/dataset → SqlAlchemyExplorer
10+
snowflake://user:…@account/db → SqlAlchemyExplorer
11+
mysql+pymysql://… / duckdb:///… → SqlAlchemyExplorer
12+
"""
13+
14+
from __future__ import annotations
15+
16+
import os
17+
from urllib.parse import urlsplit
18+
19+
from ...core.ports.explorer import ExplorerPort
20+
from .d1_explorer import D1Explorer
21+
from .sqlalchemy_explorer import SqlAlchemyExplorer
22+
23+
24+
def build_explorer(connection: str, *, schema: str | None = None) -> ExplorerPort:
25+
"""Route a connection string to the matching explorer adapter.
26+
27+
``schema`` is forwarded to the SQLAlchemy explorer (ignored by D1, which is
28+
schema-less SQLite). Raises ``ValueError`` on an empty/unparseable string.
29+
"""
30+
if not connection or not connection.strip():
31+
raise ValueError("empty connection string")
32+
33+
scheme = urlsplit(connection).scheme.lower()
34+
if not scheme:
35+
raise ValueError(f"connection string has no scheme: {connection!r}")
36+
37+
if scheme == "d1":
38+
parts = urlsplit(connection)
39+
account_id = parts.netloc
40+
database_id = parts.path.lstrip("/")
41+
if not account_id or not database_id:
42+
raise ValueError("d1 URL must be d1://<account_id>/<database_id>")
43+
return D1Explorer(account_id=account_id, database_id=database_id)
44+
45+
# Anything else is assumed to be a SQLAlchemy URL (driver loaded lazily).
46+
return SqlAlchemyExplorer(connection, schema=schema)
47+
48+
49+
def explorer_from_env() -> ExplorerPort | None:
50+
"""Build an explorer from environment, or ``None`` if nothing is configured.
51+
52+
Precedence: an explicit ``LANG2SQL_DB_URL`` wins; otherwise a pair of
53+
``CLOUDFLARE_D1_ACCOUNT_ID`` + ``CLOUDFLARE_D1_DATABASE_ID`` selects D1.
54+
"""
55+
url = os.environ.get("LANG2SQL_DB_URL")
56+
if url:
57+
return build_explorer(url, schema=os.environ.get("LANG2SQL_DB_SCHEMA"))
58+
59+
account = os.environ.get("CLOUDFLARE_D1_ACCOUNT_ID")
60+
database = os.environ.get("CLOUDFLARE_D1_DATABASE_ID")
61+
if account and database:
62+
return build_explorer(f"d1://{account}/{database}")
63+
64+
return None
Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
"""Generic SQLAlchemy explorer — one adapter, many engines.
2+
3+
A single :class:`ExplorerPort` implementation that connects to anything
4+
SQLAlchemy speaks (PostgreSQL, MySQL, Snowflake, BigQuery, DuckDB, SQLite, …)
5+
purely from a connection URL. This is the "사용성" win: adding a new warehouse is
6+
``pip install <driver>`` + a DSN, not a new adapter class.
7+
8+
The engine is created lazily on first use so constructing the explorer (and
9+
routing to it in the factory) never imports a driver that isn't installed.
10+
Blocking DB calls run in a worker thread to keep the async event loop free.
11+
"""
12+
13+
from __future__ import annotations
14+
15+
import asyncio
16+
from typing import Any
17+
18+
from ...core.ports.explorer import Column, Table
19+
20+
21+
class SqlAlchemyExplorer:
22+
"""ExplorerPort over a SQLAlchemy Engine, built from a connection URL."""
23+
24+
def __init__(self, url: str, *, schema: str | None = None) -> None:
25+
self.url = url
26+
self._schema = schema
27+
self._engine: Any = None # created lazily
28+
29+
def _get_engine(self) -> Any:
30+
if self._engine is None:
31+
from sqlalchemy import create_engine # imported here = lazy driver load
32+
33+
self._engine = create_engine(self.url)
34+
return self._engine
35+
36+
# --- ExplorerPort ----------------------------------------------------
37+
38+
async def list_tables(self) -> list[Table]:
39+
return await asyncio.to_thread(self._list_tables_sync)
40+
41+
async def describe_table(self, name: str) -> Table:
42+
return await asyncio.to_thread(self._describe_table_sync, name)
43+
44+
async def sample_rows(self, name: str, limit: int = 5) -> list[dict]:
45+
# Bind the limit; quote the identifier via the dialect's preparer.
46+
eng = self._get_engine()
47+
qname = eng.dialect.identifier_preparer.quote(name)
48+
return await self.execute(f"SELECT * FROM {qname}", limit=limit)
49+
50+
async def execute(self, sql: str, limit: int = 1000) -> list[dict]:
51+
return await asyncio.to_thread(self._execute_sync, sql, int(limit))
52+
53+
# --- sync workers ----------------------------------------------------
54+
55+
def _list_tables_sync(self) -> list[Table]:
56+
from sqlalchemy import inspect
57+
58+
insp = inspect(self._get_engine())
59+
schema = self._schema or insp.default_schema_name
60+
return [
61+
Table(name=t, schema=schema or "")
62+
for t in insp.get_table_names(schema=self._schema)
63+
]
64+
65+
def _describe_table_sync(self, name: str) -> Table:
66+
from sqlalchemy import inspect
67+
68+
insp = inspect(self._get_engine())
69+
cols = [
70+
Column(
71+
name=c["name"],
72+
type=str(c["type"]),
73+
nullable=bool(c.get("nullable", True)),
74+
description=c.get("comment") or "",
75+
)
76+
for c in insp.get_columns(name, schema=self._schema)
77+
]
78+
return Table(name=name, schema=self._schema or "", columns=cols)
79+
80+
def _execute_sync(self, sql: str, limit: int) -> list[dict]:
81+
from sqlalchemy import text
82+
83+
with self._get_engine().connect() as conn:
84+
result = conn.execute(text(sql))
85+
if not result.returns_rows:
86+
return []
87+
rows = result.mappings().fetchmany(limit)
88+
return [dict(r) for r in rows]

src/lang2sql/tenancy/concierge.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import os
1616

17+
from ..adapters.db.factory import explorer_from_env
1718
from ..adapters.db.postgres_explorer import PostgresExplorer
1819
from ..adapters.llm.fake import FakeLLM
1920
from ..adapters.llm.openai_ import OpenAILLM
@@ -61,7 +62,9 @@ def __init__(
6162
self._store = store if store is not None else SqliteStore(path)
6263
# Audit + session persistence both ride the one sqlite store by default.
6364
self._llm = llm if llm is not None else _default_llm()
64-
self._explorer = explorer if explorer is not None else PostgresExplorer(_DEFAULT_DSN)
65+
# Explorer precedence: explicit injection → env-configured real DB
66+
# (LANG2SQL_DB_URL / Cloudflare D1) → the canned stub for offline dev.
67+
self._explorer = explorer or explorer_from_env() or PostgresExplorer(_DEFAULT_DSN)
6568
self._safety = safety if safety is not None else SafetyPipeline()
6669
# Persistent semantic store by default so definitions survive restart.
6770
self._scope_resolver = (

0 commit comments

Comments
 (0)