File size: 3,121 Bytes
788e50e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """Session lifecycle — in-memory store, persistence, preloading, rate limiting."""
from __future__ import annotations
import json
import time
from pathlib import Path
from ._execute import (
execute_duckdb, execute_mssql, execute_mysql, execute_pg, execute_sqlite,
)
from ._schema import load_duckdb_schema, load_sqlite_schema
CHINOOK_PATH = Path(__file__).parent.parent / "data" / "chinook.db"
_UPLOAD_DIR = Path("/tmp/ml_sql_sessions")
_SESSION_FILE = _UPLOAD_DIR / "index.json"
_SESSION_TTL = 86_400 # 24 hours
_UPLOAD_DIR.mkdir(parents=True, exist_ok=True)
# in-memory session store: db_ref → {type, path|conn_str, schema, created_at}
_sessions: dict[str, dict] = {}
# Global rate limiter — sliding window of query timestamps
_recent_queries: list[float] = []
_RATE_LIMIT = 30 # max queries per 60 seconds globally
def get_sessions() -> dict[str, dict]:
return _sessions
def check_rate_limit() -> bool:
now = time.time()
_recent_queries[:] = [t for t in _recent_queries if t > now - 60.0]
if len(_recent_queries) >= _RATE_LIMIT:
return False
_recent_queries.append(now)
return True
def save_session_index() -> None:
index = {
ref: {"type": s["type"], "path": s.get("path", ""), "created_at": s.get("created_at", 0)}
for ref, s in _sessions.items()
if s["type"] in ("sqlite", "duckdb") and ref != "chinook"
}
try:
_SESSION_FILE.write_text(json.dumps(index))
except OSError:
pass
async def restore_sessions() -> None:
if not _SESSION_FILE.exists():
return
try:
index = json.loads(_SESSION_FILE.read_text())
except (OSError, json.JSONDecodeError):
return
cutoff = time.time() - _SESSION_TTL
for ref, meta in index.items():
if meta.get("created_at", 0) < cutoff:
continue
path = meta.get("path", "")
if path and Path(path).exists() and ref not in _sessions:
try:
stype = meta.get("type", "sqlite")
schema = await load_duckdb_schema(path) if stype == "duckdb" else await load_sqlite_schema(path)
_sessions[ref] = {"type": stype, "path": path, "schema": schema, "created_at": meta["created_at"]}
except Exception:
pass
async def exec_session(session: dict, sql: str):
stype = session["type"]
if stype == "sqlite": return await execute_sqlite(session["path"], sql)
elif stype == "duckdb": return await execute_duckdb(session["path"], sql)
elif stype == "mysql": return await execute_mysql(session["conn_str"], sql)
elif stype == "mssql": return await execute_mssql(session["conn_str"], sql)
else: return await execute_pg(session["conn_str"], sql)
async def preload_chinook() -> None:
if "chinook" not in _sessions and CHINOOK_PATH.exists():
schema = await load_sqlite_schema(str(CHINOOK_PATH))
_sessions["chinook"] = {"type": "sqlite", "path": str(CHINOOK_PATH),
"schema": schema, "created_at": time.time()}
await restore_sessions()
|