routers/_session_mgr.py

"""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()