database.py

import json
import os
import aiohttp
from config import DB_SPACE_URL, HF_TOKEN

HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}
_session: aiohttp.ClientSession | None = None


async def get_session() -> aiohttp.ClientSession:
    global _session
    if _session is None:
        _session = aiohttp.ClientSession(headers=HEADERS)
    return _session


async def close_session():
    global _session
    if _session:
        await _session.close()
        _session = None


def _escape(val):
    if val is None:
        return "NULL"
    t = type(val)
    if t is int or t is float:
        return str(val)
    if t is bool:
        return "1" if val else "0"
    s = str(val).replace("'", "''")
    return f"'{s}'"


def _build_sql(sql: str, params: list | None = None) -> str:
    # Internal use only: builds SQL from template with ? placeholders
    if not params:
        return sql
    parts = sql.split("?")
    result = parts[0]
    for i, p in enumerate(params):
        val = p
        if not isinstance(val, (int, float, bool)):
            val = str(val)
        if i + 1 < len(parts):
            result += _escape(val) + parts[i + 1]
        else:
            result += _escape(val)
    return result


async def exec(sql: str, params: list = None) -> dict:
    session = await get_session()
    full_sql = _build_sql(sql, params)
    async with session.post(f"{DB_SPACE_URL}/exec", json={"sql": full_sql}) as resp:
        return await resp.json()


async def query(sql: str, params: list = None) -> dict:
    session = await get_session()
    full_sql = _build_sql(sql, params)
    async with session.get(f"{DB_SPACE_URL}/query", params={"sql": full_sql}) as resp:
        return await resp.json()


async def fetch_all(sql: str, params: list = None) -> list[dict]:
    r = await query(sql, params)
    if not r or "rows" not in r or not r["rows"]:
        return []
    cols = r.get("cols") or r.get("columns", [])
    col_names = [c["name"] if isinstance(c, dict) else c for c in cols]
    return [dict(zip(col_names, row)) for row in r["rows"]]


async def fetch_one(sql: str, params: list = None) -> dict | None:
    rows = await fetch_all(sql, params)
    return rows[0] if rows else None


async def execute(sql: str, params: list = None) -> dict:
    return await exec(sql, params)


async def init_tables(sqls: list[str]) -> list[dict]:
    results = []
    for sql in sqls:
        r = await exec(sql)
        results.append(r)
    return results