Database 1.0.0
db_tables.py
1.8 KB · raw
"""List the tables in a saved database."""
from __future__ import annotations
import asyncio
from aworg.tools.base import ToolContext, ToolError, ToolResult
from ._db import kind_of, open_connection, resolve
NAME = "db_tables"
DESCRIPTION = "List the tables and views in a saved database connection."
INPUT_SCHEMA = {
"type": "object",
"properties": {
"connection": {"type": "string", "description": "The saved connection's name."},
},
"required": ["connection"],
}
SQL = {
"sqlite": "SELECT name, type FROM sqlite_master WHERE type IN ('table','view') "
"AND name NOT LIKE 'sqlite_%' ORDER BY name",
"postgresql": "SELECT table_schema || '.' || table_name, table_type FROM information_schema.tables "
"WHERE table_schema NOT IN ('pg_catalog','information_schema') ORDER BY 1",
"mysql": "SELECT table_name, table_type FROM information_schema.tables "
"WHERE table_schema = DATABASE() ORDER BY 1",
}
async def run(context: ToolContext, connection: str = "") -> ToolResult:
url = resolve(str(connection).strip())
kind = kind_of(url)
workspace = getattr(getattr(context, "paths", None), "workspace", None)
def go():
conn = open_connection(url, workspace, read_only=True)
try:
cur = conn.cursor()
cur.execute(SQL[kind])
return cur.fetchall()
finally:
conn.close()
try:
rows = await asyncio.to_thread(go)
except ToolError:
raise
except Exception as exc: # noqa: BLE001
raise ToolError(f"{exc}") from None
if not rows:
return ToolResult(text="No tables.", summary="0 tables")
return ToolResult(
text="\n".join(f"{name} ({str(kind_).lower()})" for name, kind_ in rows),
summary=f"{len(rows)} tables",
)