database

/**
 * Database Plugin — toolsProvider
 *
 * Tools:
 *   database(action) — query | schema | tables | sample
 */

import { text, tool, type ToolCallContext, type ToolsProvider } from "@lmstudio/sdk";
import { z } from "zod";
import { pluginConfigSchematics } from "./config";
import { sqliteQuery, sqliteWriteQuery, sqliteTables, sqliteSchema, type QueryResult } from "./sqlite";
import { postgresQuery, postgresTables, postgresSchema } from "./postgres";

function json(obj: unknown): string {
  return JSON.stringify(obj, null, 2);
}

function safe_impl<T extends Record<string, unknown>>(
  name: string,
  fn: (params: T, ctx: ToolCallContext) => Promise<string>
): (params: T, ctx: ToolCallContext) => Promise<string> {
  return async (params: T, ctx: ToolCallContext) => {
    if (ctx.signal.aborted) return JSON.stringify({ tool_error: true, tool: name, error: "cancelled" });
    try {
      return await fn(params, ctx);
    } catch (err: unknown) {
      const msg = err instanceof Error ? err.message : String(err);
      return JSON.stringify({ tool_error: true, tool: name, error: msg }, null, 2);
    }
  };
}

function toMarkdownTable(result: QueryResult): string {
  if (result.columns.length === 0) return "(no columns returned)";
  if (result.rows.length === 0) return "(0 rows)";

  const header = "| " + result.columns.join(" | ") + " |";
  const sep = "| " + result.columns.map(() => "---").join(" | ") + " |";
  const rows = result.rows.map(r => "| " + r.map(v => String(v ?? "NULL")).join(" | ") + " |");
  let table = [header, sep, ...rows].join("\n");
  if (result.truncated) {
    table += `\n\n(showing ${result.rows.length} of ${result.rowCount} rows)`;
  }
  return table;
}

const WRITE_PATTERN = /^\s*(INSERT|UPDATE|DELETE|DROP|CREATE|ALTER|TRUNCATE|REPLACE)\b/i;

export const toolsProvider: ToolsProvider = async (ctl) => {
  const cfg = ctl.getPluginConfig(pluginConfigSchematics);

  return [
    tool({
      name: "database",
      description: text`
        Query a SQLite or Postgres database. Read-only by default.

        action: "query"   — Run a SQL SELECT (or write query if allow_writes is enabled in settings)
        action: "tables"  — List all tables in the database
        action: "schema"  — Show columns and types for a specific table
        action: "sample"  — Show the first 10 rows of a table

        db_type: "sqlite" (default) or "postgres"
        db_path: Path to .db/.sqlite file (sqlite). Uses config path if blank.
        dsn: Postgres connection string (postgres). Uses config DSN if blank.
        sql: SQL statement for query action
        table: Table name for schema/sample actions
      `,
      parameters: {
        action: z.enum(["query", "tables", "schema", "sample"]),
        db_type: z.enum(["sqlite", "postgres"]).default("sqlite"),
        db_path: z.string().default(""),
        dsn: z.string().default(""),
        sql: z.string().default(""),
        table: z.string().default(""),
      },
      implementation: safe_impl("database", async ({ action, db_type, db_path, dsn, sql, table }, ctx) => {
        const isPostgres = db_type === "postgres";
        const resolvedPath = db_path || cfg.get("sqlitePath");
        const resolvedDsn = dsn || cfg.get("postgresDsn");
        const allowWrites = cfg.get("allowWrites");
        const maxRows = cfg.get("maxRows");

        if (isPostgres && !resolvedDsn) throw new Error("No Postgres DSN configured.");
        if (!isPostgres && !resolvedPath) throw new Error("No SQLite path configured.");

        if (action === "tables") {
          ctx.status("Listing tables");
          const tables = isPostgres
            ? await postgresTables(resolvedDsn)
            : sqliteTables(resolvedPath);
          return json({ db_type, tables, count: tables.length });
        }

        if (action === "schema") {
          if (!table) throw new Error("table name required for schema action");
          ctx.status(`Getting schema for ${table}`);
          const info = isPostgres
            ? await postgresSchema(resolvedDsn, table)
            : sqliteSchema(resolvedPath, table);
          return json(info);
        }

        if (action === "sample") {
          if (!table) throw new Error("table name required for sample action");
          ctx.status(`Sampling ${table}`);
          const sampleSql = `SELECT * FROM ${table} LIMIT 10`;
          const result = isPostgres
            ? await postgresQuery(resolvedDsn, sampleSql, 10)
            : sqliteQuery(resolvedPath, sampleSql, 10);
          return `## Sample: ${table}\n\n${toMarkdownTable(result)}`;
        }

        // query
        if (!sql) throw new Error("sql required for query action");
        if (WRITE_PATTERN.test(sql) && !allowWrites) {
          throw new Error("Write queries are disabled. Enable 'Allow Write Queries' in plugin settings.");
        }
        ctx.status("Running query");

        if (WRITE_PATTERN.test(sql) && !isPostgres) {
          const r = sqliteWriteQuery(resolvedPath, sql);
          return json({ changes: r.changes, lastInsertRowid: String(r.lastInsertRowid) });
        }

        const result = isPostgres
          ? await postgresQuery(resolvedDsn, sql, maxRows)
          : sqliteQuery(resolvedPath, sql, maxRows);

        return `## Query Results\n\`\`\`sql\n${sql}\n\`\`\`\n\n${toMarkdownTable(result)}`;
      }),
    }),
  ];
};