database

"use strict";
/**
 * Database Plugin — toolsProvider
 *
 * Tools:
 *   database(action) — query | schema | tables | sample
 */
Object.defineProperty(exports, "__esModule", { value: true });
exports.toolsProvider = void 0;
const sdk_1 = require("@lmstudio/sdk");
const zod_1 = require("zod");
const config_1 = require("./config");
const sqlite_1 = require("./sqlite");
const postgres_1 = require("./postgres");
function json(obj) {
    return JSON.stringify(obj, null, 2);
}
function safe_impl(name, fn) {
    return async (params, ctx) => {
        if (ctx.signal.aborted)
            return JSON.stringify({ tool_error: true, tool: name, error: "cancelled" });
        try {
            return await fn(params, ctx);
        }
        catch (err) {
            const msg = err instanceof Error ? err.message : String(err);
            return JSON.stringify({ tool_error: true, tool: name, error: msg }, null, 2);
        }
    };
}
function toMarkdownTable(result) {
    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;
const toolsProvider = async (ctl) => {
    const cfg = ctl.getPluginConfig(config_1.pluginConfigSchematics);
    return [
        (0, sdk_1.tool)({
            name: "database",
            description: (0, sdk_1.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: zod_1.z.enum(["query", "tables", "schema", "sample"]),
                db_type: zod_1.z.enum(["sqlite", "postgres"]).default("sqlite"),
                db_path: zod_1.z.string().default(""),
                dsn: zod_1.z.string().default(""),
                sql: zod_1.z.string().default(""),
                table: zod_1.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 (0, postgres_1.postgresTables)(resolvedDsn)
                        : (0, sqlite_1.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 (0, postgres_1.postgresSchema)(resolvedDsn, table)
                        : (0, sqlite_1.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 (0, postgres_1.postgresQuery)(resolvedDsn, sampleSql, 10)
                        : (0, sqlite_1.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 = (0, sqlite_1.sqliteWriteQuery)(resolvedPath, sql);
                    return json({ changes: r.changes, lastInsertRowid: String(r.lastInsertRowid) });
                }
                const result = isPostgres
                    ? await (0, postgres_1.postgresQuery)(resolvedDsn, sql, maxRows)
                    : (0, sqlite_1.sqliteQuery)(resolvedPath, sql, maxRows);
                return `## Query Results\n\`\`\`sql\n${sql}\n\`\`\`\n\n${toMarkdownTable(result)}`;
            }),
        }),
    ];
};
exports.toolsProvider = toolsProvider;