94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
|
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
|