Tools reference

src/mini_db_mcp/db.py

@@ -0,0 +1,93 @@
+from typing import Any
+
+from sqlalchemy import URL, inspect
+from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
+
+from mini_db_mcp.config import Config
+
+
+class Database:
+    def __init__(self, settings: Config):
+        self.settings = settings
+        self.database = settings.database
+        url = URL.create(
+            "mysql+aiomysql",
+            username=self.database.user,
+            password=self.database.password,
+            host=self.database.host,
+            port=self.database.port,
+            database=self.database.schema or None,
+        )
+        self.engine: AsyncEngine = create_async_engine(
+            url,
+            pool_size=self.database.pool_size,
+            max_overflow=self.database.max_overflow,
+            pool_timeout=self.database.pool_timeout,
+            pool_recycle=self.database.pool_recycle,
+            pool_pre_ping=self.database.pool_pre_ping,
+            connect_args={"connect_timeout": self.database.connect_timeout},
+        )
+
+    async def close(self) -> None:
+        await self.engine.dispose()
+
+
+    async def list_tables(
+        self,
+        *,
+        keyword: str | None,
+        limit: int,
+        cursor: str | None,
+    ) -> tuple[list[dict[str, Any]], str | None]:
+        """Return one name-sorted page of tables and views.
+
+        ``keyword`` filters object names case-insensitively. ``cursor`` must be
+        the last object name returned by the previous page. The returned cursor
+        is ``None`` when no additional objects remain.
+
+        Raises:
+            ValueError: If the cursor is not present in the filtered result.
+        """
+        schema = self.settings.database.schema or None
+
+        def inspect_tables(sync_connection):
+            inspector = inspect(sync_connection)
+            table_names = set(inspector.get_table_names(schema=schema))
+            try:
+                view_names = set(inspector.get_view_names(schema=schema))
+            except NotImplementedError:
+                view_names = set()
+            return table_names, view_names
+
+        async with self.engine.connect() as connection:
+            table_names, view_names = await connection.run_sync(inspect_tables)
+
+        keyword_cf = keyword.casefold() if keyword else None
+        names = sorted(
+            name
+            for name in table_names | view_names
+            if keyword_cf is None or keyword_cf in name.casefold()
+        )
+
+        start = 0
+        if cursor is not None:
+            try:
+                start = names.index(cursor) + 1
+            except ValueError as exc:
+                raise ValueError("invalid pagination cursor") from exc
+
+        page_names = names[start : start + limit]
+        next_cursor = page_names[-1] if start + len(page_names) < len(names) else None
+
+        objects: list[dict[str, Any]] = []
+        for name in page_names:
+            object_type = "view" if name in view_names else "table"
+            objects.append(
+                {
+                    "schema_name": schema,
+                    "object_name": name,
+                    "object_type": object_type,
+                    "description": None,
+                }
+            )
+        return objects, next_cursor