diff --git a/src/mini_db_mcp/connection.py b/src/mini_db_mcp/connection.py new file mode 100644 index 0000000..eb50da1 --- /dev/null +++ b/src/mini_db_mcp/connection.py @@ -0,0 +1,110 @@ +"""Async MySQL connection manager with connection pool. + +Design mirrors DorisConnectionManager: a single shared pool, acquired via +an async context manager, with an explicit ``close()`` for clean shutdown. +""" + +from __future__ import annotations + +import logging +import re +from collections.abc import AsyncGenerator, AsyncIterator +from contextlib import asynccontextmanager +from typing import Any + +import aiomysql +from aiomysql import Pool + +from .config import Config + +logger = logging.getLogger(__name__) + +# Only these SQL keywords are permitted when write access is disabled +_READ_ONLY_KEYWORDS = frozenset( + {"SELECT", "SHOW", "DESCRIBE", "DESC", "EXPLAIN", "WITH"} +) +_FIRST_KEYWORD_RE = re.compile(r"--[^\n]*|/\*.*?\*/", re.DOTALL) + + +def _first_sql_keyword(sql: str) -> str: + stripped = _FIRST_KEYWORD_RE.sub("", sql).strip() + return stripped.split(None, 1)[0].upper() if stripped else "" + + +class ConnectionManager: + """Holds the aiomysql connection pool and enforces write policy.""" + + def __init__(self, config: Config) -> None: + self.config = config + self._pool: Pool | None = None + + async def start(self) -> None: + db = self.config.database + self._pool = await aiomysql.create_pool( + host=db.host, + port=db.port, + user=db.user, + password=db.password, + db=db.database or None, + minsize=db.pool_min, + maxsize=db.pool_max, + connect_timeout=db.connect_timeout, + autocommit=True, + charset="utf8mb4", + ) + logger.info( + "Connection pool started: %s:%d db=%s", + db.host, + db.port, + db.database or "(none)", + ) + + async def close(self) -> None: + if self._pool is not None: + self._pool.close() + await self._pool.wait_closed() + self._pool = None + logger.info("Connection pool closed") + + def _assert_write_allowed(self, sql: str) -> None: + keyword = _first_sql_keyword(sql) + if not self.config.query.allow_write and keyword not in _READ_ONLY_KEYWORDS: + raise PermissionError( + f"Write operations are disabled (first keyword: {keyword!r}). " + "Set QUERY_ALLOW_WRITE=true to enable DML/DDL." + ) + + @asynccontextmanager + async def acquire(self) -> AsyncIterator[aiomysql.Connection]: + if self._pool is None: + raise RuntimeError("Connection pool is not started") + async with self._pool.acquire() as conn: + yield conn + + async def execute( + self, + sql: str, + args: tuple[Any, ...] | None = None, + *, + database: str | None = None, + max_rows: int | None = None, + ) -> dict[str, Any]: + """Execute *sql* and return a structured result dict.""" + self._assert_write_allowed(sql) + limit = max_rows if max_rows is not None else self.config.query.max_rows + async with self.acquire() as conn: + if database: + await conn.select_db(database) + async with conn.cursor(aiomysql.DictCursor) as cur: + await cur.execute(sql, args) + rows = await cur.fetchmany(limit + 1) + truncated = len(rows) > limit + if truncated: + rows = rows[:limit] + return { + "columns": [d[0] for d in cur.description] if cur.description else [], + "rows": [dict(r) for r in rows], + "row_count": len(rows), + "truncated": truncated, + "max_rows": limit, + }