lan_sqlite_api.py

from __future__ import annotations

import json
import re
import sqlite3
import sys
from contextlib import closing
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
from urllib.parse import parse_qs, unquote, urlparse

from app.store import DEFAULT_DB_PATH


DEFAULT_HOST = "0.0.0.0"
DEFAULT_PORT = 8770
MAX_LIMIT = 500
TABLE_NAME_PATTERN = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")


class ReadOnlySqliteApi:
    def __init__(self, db_path: str | Path = DEFAULT_DB_PATH) -> None:
        self.db_path = Path(db_path)

    def list_tables(self) -> list[str]:
        with closing(self._connect()) as conn:
            rows = conn.execute(
                """
                select name
                from sqlite_master
                where type = 'table'
                  and name not like 'sqlite_%'
                order by name asc
                """
            ).fetchall()
        return [str(row["name"]) for row in rows]

    def list_columns(self, table: str) -> list[dict[str, str]]:
        table = self._validated_table(table)
        with closing(self._connect()) as conn:
            rows = conn.execute(f'pragma table_info("{table}")').fetchall()
        return [{"name": str(row["name"]), "type": str(row["type"] or "")} for row in rows]

    def list_rows(self, table: str, limit: int = 100, offset: int = 0) -> dict:
        table = self._validated_table(table)
        limit = _bounded_int(limit, default=100, minimum=1, maximum=MAX_LIMIT)
        offset = _bounded_int(offset, default=0, minimum=0, maximum=10_000_000)
        with closing(self._connect()) as conn:
            rows = conn.execute(f'select * from "{table}" limit ? offset ?', (limit, offset)).fetchall()
        return {"table": table, "limit": limit, "offset": offset, "rows": [_row_to_dict(row) for row in rows]}

    def query(self, sql: str, params: list | None = None, limit: int = 100) -> dict:
        sql = str(sql or "").strip()
        if not _is_read_only_sql(sql):
            raise ValueError("仅允许只读查询")
        limit = _bounded_int(limit, default=100, minimum=1, maximum=MAX_LIMIT)
        params = params if isinstance(params, list) else []
        with closing(self._connect()) as conn:
            rows = conn.execute(sql, params).fetchmany(limit)
        return {"limit": limit, "rows": [_row_to_dict(row) for row in rows]}

    def _connect(self) -> sqlite3.Connection:
        uri = f"file:{self.db_path.as_posix()}?mode=ro"
        conn = sqlite3.connect(uri, uri=True)
        conn.row_factory = sqlite3.Row
        return conn

    def _validated_table(self, table: str) -> str:
        table = str(table or "").strip()
        if not TABLE_NAME_PATTERN.match(table):
            raise ValueError("表名不合法")
        if table not in self.list_tables():
            raise ValueError("表不存在")
        return table


class LanSqliteApiHandler(BaseHTTPRequestHandler):
    api = ReadOnlySqliteApi()

    def do_OPTIONS(self) -> None:
        self.send_response(204)
        self._send_cors_headers()
        self.end_headers()

    def do_GET(self) -> None:
        parsed = urlparse(self.path)
        try:
            if parsed.path == "/health":
                self._send_health()
                return
            if parsed.path == "/api/tables":
                self._send_tables()
                return
            if parsed.path.startswith("/api/tables/") and parsed.path.endswith("/columns"):
                table = unquote(parsed.path.removeprefix("/api/tables/").removesuffix("/columns").strip("/"))
                self._send_columns(table)
                return
            if parsed.path.startswith("/api/tables/") and parsed.path.endswith("/rows"):
                table = unquote(parsed.path.removeprefix("/api/tables/").removesuffix("/rows").strip("/"))
                self._send_rows(table, parsed.query)
                return
        except ValueError as exc:
            self._send_json({"error": str(exc)}, status=400)
            return
        except Exception as exc:
            self._send_json({"error": str(exc)}, status=500)
            return
        self.send_error(404)

    def do_POST(self) -> None:
        if urlparse(self.path).path != "/api/query":
            self.send_error(404)
            return
        try:
            self._handle_query()
        except ValueError as exc:
            self._send_json({"error": str(exc)}, status=400)
        except Exception as exc:
            self._send_json({"error": str(exc)}, status=500)

    def _send_health(self) -> None:
        self._send_json({"status": "ok", "dbPath": str(self.api.db_path)})

    def _send_tables(self) -> None:
        self._send_json({"tables": self.api.list_tables()})

    def _send_columns(self, table: str) -> None:
        self._send_json({"table": table, "columns": self.api.list_columns(table)})

    def _send_rows(self, table: str, query: str) -> None:
        params = parse_qs(query)
        limit = int((params.get("limit") or ["100"])[0])
        offset = int((params.get("offset") or ["0"])[0])
        self._send_json(self.api.list_rows(table, limit=limit, offset=offset))

    def _handle_query(self) -> None:
        length = int(self.headers.get("Content-Length") or "0")
        body = self.rfile.read(length).decode("utf-8") if length else "{}"
        payload = json.loads(body, strict=False)
        if not isinstance(payload, dict):
            raise ValueError("请求体必须是 JSON 对象")
        self._send_json(
            self.api.query(
                str(payload.get("sql") or ""),
                payload.get("params") if isinstance(payload.get("params"), list) else [],
                limit=int(payload.get("limit") or 100),
            )
        )

    def _send_json(self, payload: dict, status: int = 200) -> None:
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Cache-Control", "no-store")
        self._send_cors_headers()
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _send_cors_headers(self) -> None:
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type")

    def log_message(self, fmt: str, *args: object) -> None:
        sys.stderr.write("%s - %s\n" % (self.log_date_time_string(), fmt % args))


def parse_server_args(args: list[str]) -> tuple[str, int, Path]:
    port = int(args[0]) if args else DEFAULT_PORT
    host = args[1] if len(args) > 1 else DEFAULT_HOST
    db_path = Path(args[2]) if len(args) > 2 else DEFAULT_DB_PATH
    return host, port, db_path


def main() -> None:
    host, port, db_path = parse_server_args(sys.argv[1:])
    handler = type("ConfiguredLanSqliteApiHandler", (LanSqliteApiHandler,), {"api": ReadOnlySqliteApi(db_path)})
    server = ThreadingHTTPServer((host, port), handler)
    print(f"SQLite 只读局域网 API 已启动: http://{host}:{port}")
    print(f"数据库: {db_path}")
    server.serve_forever()


def _is_read_only_sql(sql: str) -> bool:
    normalized = sql.lstrip().lower()
    if not normalized.startswith(("select", "with")):
        return False
    return not any(token in normalized for token in (";", " insert ", " update ", " delete ", " drop ", " alter ", " create ", " attach ", " detach ", " replace ", " vacuum ", " pragma "))


def _bounded_int(value: int, default: int, minimum: int, maximum: int) -> int:
    try:
        parsed = int(value)
    except (TypeError, ValueError):
        parsed = default
    return min(max(parsed, minimum), maximum)


def _row_to_dict(row: sqlite3.Row) -> dict:
    return {key: row[key] for key in row.keys()}


if __name__ == "__main__":
    main()