connection

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,
+                }