dsh-db-mcp 0.6.4

package/server/main.py

@@ -0,0 +1,260 @@
+"""db-mcp: the MCP server that lets the agent read the databases registered in
+the dsh-db-mcp plugin panel.
+
+Run over stdio (the plugin Host half spawns it):
+
+    python main.py
+
+Every tool re-reads `datasources.json` and `secrets.json`, so sources added or
+edited in the panel are visible without a restart. Tools never return passwords.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import time
+from typing import Any, Optional
+
+from mcp.server.fastmcp import FastMCP
+
+import engines
+import store
+
+DEFAULT_MAX_ROWS = int(os.environ.get("DB_MCP_MAX_ROWS") or 200)
+HARD_MAX_ROWS = int(os.environ.get("DB_MCP_MAX_ROWS_HARD") or 2000)
+
+mcp = FastMCP(
+    "db-mcp",
+    instructions=(
+        "Query the databases the user registered in the DSH Database panel. "
+        "Start with list_datasources, then list_tables, describe_table and finally "
+        "run_query. Prefer these tools over guessing table names."
+    ),
+)
+
+
+def _clamp_rows(limit: Optional[int]) -> int:
+    try:
+        value = int(limit) if limit is not None else DEFAULT_MAX_ROWS
+    except (TypeError, ValueError):
+        value = DEFAULT_MAX_ROWS
+    if value <= 0:
+        value = DEFAULT_MAX_ROWS
+    return min(value, HARD_MAX_ROWS)
+
+
+def _error(error: Exception) -> dict:
+    if isinstance(error, (engines.DbError, store.StoreError)):
+        payload = {"ok": False, "error": {"code": error.code, "message": error.message}}
+        if error.hint:
+            payload["error"]["hint"] = error.hint
+        return payload
+    return {
+        "ok": False,
+        "error": {"code": "infra_error", "message": f"{type(error).__name__}: {error}"},
+    }
+
+
+def _open(datasource: Optional[str]):
+    profile, secrets = store.resolve(datasource)
+    password = store.password_of(profile, secrets)
+    return profile, engines.connect(profile, password)
+
+
+def _source(profile: dict) -> dict:
+    config = profile.get("config") or {}
+    kind = profile.get("kind")
+    if kind == "sqlite":
+        target = config.get("file") or ""
+    elif kind == "oracle":
+        target = f"{config.get('host')}:{config.get('port') or 1521}/{config.get('serviceName') or config.get('sid') or ''}"
+    else:
+        target = f"{config.get('host')}:{config.get('port') or ''}/{config.get('database') or ''}"
+    return {
+        "id": profile.get("id"),
+        "name": profile.get("name"),
+        "kind": kind,
+        "readOnly": bool(profile.get("readOnly")),
+        "isDefault": bool(profile.get("isDefault")),
+        "target": target,
+    }
+
+
+def _inline_profile(kind: str, config: Any, read_only: Optional[bool]) -> dict:
+    if not isinstance(config, dict):
+        raise store.StoreError("user_error", "inline test needs a config object", "pass kind plus config")
+    if kind not in ("sqlite", "postgres", "mysql", "oracle"):
+        raise store.StoreError("user_error", f"unsupported datasource kind: {kind}")
+    return {
+        "id": None,
+        "name": "inline test",
+        "kind": kind,
+        "config": config,
+        "readOnly": True if read_only is None else bool(read_only),
+        "isDefault": False,
+        "tags": [],
+    }
+
+
+@mcp.tool()
+def test_connection(
+    datasource: Optional[str] = None,
+    kind: Optional[str] = None,
+    config: Optional[dict] = None,
+    password: Optional[str] = None,
+    read_only: Optional[bool] = None,
+) -> dict:
+    """Log in to a database and run one trivial query, without returning data.
+
+    Use a stored `datasource` id, or pass `kind` plus `config` (and optionally
+    `password`, `read_only`) to test a connection that has not been saved yet.
+    Returns the server banner, the connected user, the database/schema and how
+    long the login took.
+    """
+    try:
+        if kind and config is not None:
+            profile = _inline_profile(kind, config, read_only)
+            secret = password
+        else:
+            profile, secrets = store.resolve(datasource)
+            secret = password if password is not None else store.password_of(profile, secrets)
+        result = engines.probe(profile, secret)
+        return {**result, "datasource": _source(profile)}
+    except Exception as error:  # noqa: BLE001 - reported as structured data
+        payload = _error(error)
+        payload["latencyMs"] = 0
+        return payload
+
+
+@mcp.tool()
+def list_datasources() -> dict:
+    """List the database connections the user registered in DSH.
+
+    Call this first. Use the returned `id` as the `datasource` argument of the
+    other tools. The default datasource may be omitted everywhere.
+    """
+    try:
+        sources = store.all_sources()
+        return {
+            "ok": True,
+            "datasources": [_source(source) for source in sources],
+            "count": len(sources),
+        }
+    except Exception as error:  # noqa: BLE001 - reported as structured data
+        return _error(error)
+
+
+@mcp.tool()
+def list_tables(datasource: Optional[str] = None, schema: Optional[str] = None) -> dict:
+    """List the tables and views of one datasource.
+
+    `datasource` is an id from list_datasources; omit it to use the default
+    source. `schema` narrows the listing when the engine has several schemas.
+    """
+    try:
+        profile, (connection, close) = _open(datasource)
+        try:
+            result = engines.list_tables(profile, connection, schema)
+        finally:
+            close()
+        return {"ok": True, "datasource": _source(profile), **result}
+    except Exception as error:  # noqa: BLE001
+        return _error(error)
+
+
+@mcp.tool()
+def describe_table(table: str, datasource: Optional[str] = None, schema: Optional[str] = None) -> dict:
+    """Describe one table: column names, types, nullability, defaults and primary key.
+
+    Call this before writing SQL so the column names are exact.
+    """
+    try:
+        profile, (connection, close) = _open(datasource)
+        try:
+            result = engines.describe_table(profile, connection, table, schema)
+        finally:
+            close()
+        return {"ok": True, "datasource": _source(profile), **result}
+    except Exception as error:  # noqa: BLE001
+        return _error(error)
+
+
+@mcp.tool()
+def sample_rows(
+    table: str,
+    datasource: Optional[str] = None,
+    limit: int = 20,
+    schema: Optional[str] = None,
+) -> dict:
+    """Return a few rows from one table, to see what the data looks like."""
+    try:
+        rows = _clamp_rows(limit)
+        profile, (connection, close) = _open(datasource)
+        try:
+            result = engines.sample_rows(profile, connection, table, rows, schema)
+        finally:
+            close()
+        return {"ok": True, "datasource": _source(profile), **result}
+    except Exception as error:  # noqa: BLE001
+        return _error(error)
+
+
+@mcp.tool()
+def run_query(sql: str, datasource: Optional[str] = None, max_rows: Optional[int] = None) -> dict:
+    """Run one SQL statement and return columns plus rows.
+
+    Read-only datasources accept SELECT/WITH/SHOW/EXPLAIN statements only.
+    Results are capped by `max_rows` (default 200); `truncated` says whether more
+    rows existed.
+    """
+    try:
+        limit = _clamp_rows(max_rows)
+        profile, (connection, close) = _open(datasource)
+        try:
+            result = engines.run_query(profile, connection, sql, limit)
+        finally:
+            close()
+        return {"ok": True, "datasource": _source(profile), **result}
+    except Exception as error:  # noqa: BLE001
+        return _error(error)
+
+
+@mcp.tool()
+def explain_query(sql: str, datasource: Optional[str] = None) -> dict:
+    """Show the query plan of one statement, without running it."""
+    try:
+        profile, (connection, close) = _open(datasource)
+        try:
+            result = engines.explain_query(profile, connection, sql)
+        finally:
+            close()
+        return {"ok": True, "datasource": _source(profile), **result}
+    except Exception as error:  # noqa: BLE001
+        return _error(error)
+
+
+@mcp.tool()
+def get_ddl(table: str, datasource: Optional[str] = None, schema: Optional[str] = None) -> dict:
+    """Show the CREATE TABLE statement of one table.
+
+    SQLite returns the statement recorded at creation time; PostgreSQL, MySQL
+    and Oracle have theirs rebuilt from the column metadata.
+    """
+    try:
+        profile, (connection, close) = _open(datasource)
+        try:
+            result = engines.table_ddl(profile, connection, table, schema)
+        finally:
+            close()
+        return {"ok": True, "datasource": _source(profile), **result}
+    except Exception as error:  # noqa: BLE001
+        return _error(error)
+
+
+def main() -> None:
+    mcp.run(transport="stdio")
+
+
+if __name__ == "__main__":
+    main()

package/server/requirements.txt

@@ -0,0 +1,10 @@
+# Installed into <DSH_HOME>/db-mcp/.venv by the Database panel (P5 will do this
+# automatically). Until then:
+#
+#   <DSH_HOME>\db-mcp\.venv\Scripts\python.exe -m pip install -r server\requirements.txt
+#
+# mcp 1.x is required: 2.x renamed FastMCP to MCPServer and moved the module.
+mcp>=1.9,<2
+psycopg[binary]>=3.2
+PyMySQL>=1.1
+oracledb>=2.0

package/server/store.py

@@ -0,0 +1,98 @@
+"""Registry access for the db-mcp MCP server.
+
+The plugin Host half owns the files; this module only reads them. Every tool call
+re-reads the registry, so a source added or edited in the panel is visible to the
+agent without restarting the server process.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+SCHEMA_VERSION = 1
+
+
+class StoreError(Exception):
+    def __init__(self, code: str, message: str, hint: str | None = None):
+        super().__init__(message)
+        self.code = code
+        self.message = message
+        self.hint = hint
+
+
+def runtime_dir() -> Path:
+    override = os.environ.get("DB_MCP_RUNTIME")
+    if override:
+        return Path(override)
+    home = os.environ.get("DSH_HOME")
+    if not home:
+        home = str(Path.home() / ".dsh")
+    return Path(home) / "db-mcp"
+
+
+def _read_json(path: Path, default):
+    try:
+        with path.open("r", encoding="utf-8") as handle:
+            return json.load(handle)
+    except FileNotFoundError:
+        return default
+    except json.JSONDecodeError as error:
+        raise StoreError("infra_error", f"{path.name} is not valid JSON: {error}")
+
+
+def load_registry() -> dict:
+    data = _read_json(runtime_dir() / "datasources.json", {"schemaVersion": SCHEMA_VERSION, "datasources": []})
+    sources = data.get("datasources")
+    if not isinstance(sources, list):
+        raise StoreError("infra_error", "datasources.json has no datasources array")
+    return data
+
+
+def load_secrets() -> dict:
+    data = _read_json(runtime_dir() / "secrets.json", {"secrets": {}})
+    secrets = data.get("secrets")
+    return secrets if isinstance(secrets, dict) else {}
+
+
+def password_of(profile: dict, secrets: dict | None = None) -> str | None:
+    secrets = secrets if secrets is not None else load_secrets()
+    entry = secrets.get(profile.get("id"))
+    if not isinstance(entry, dict):
+        return None
+    for key in ("password", "#password"):
+        value = entry.get(key)
+        if isinstance(value, str) and value != "":
+            return value
+    return None
+
+
+def all_sources() -> list[dict]:
+    return list(load_registry().get("datasources") or [])
+
+
+def resolve(name: str | None) -> tuple[dict, dict]:
+    """Resolve a datasource by id or name; fall back to the default source."""
+    sources = all_sources()
+    if not sources:
+        raise StoreError(
+            "user_error",
+            "no datasource is configured yet",
+            "open the Database panel in DSH and add one, or import from DBeaver",
+        )
+    if name:
+        wanted = name.strip()
+        for source in sources:
+            if source.get("id") == wanted:
+                return source, load_secrets()
+        folded = wanted.casefold()
+        for source in sources:
+            if str(source.get("name") or "").casefold() == folded:
+                return source, load_secrets()
+        known = ", ".join(str(source.get("id")) for source in sources)
+        raise StoreError("user_error", f"unknown datasource: {name}", f"known datasources: {known}")
+    for source in sources:
+        if source.get("isDefault"):
+            return source, load_secrets()
+    return sources[0], load_secrets()