sql-database-plugin

import { tool, Tool, ToolsProviderController } from "@lmstudio/sdk";
import { z } from "zod";
import * as db from "./db";
import { configSchematics, DatabaseConfig } from "./config";

function getDatabaseConfig(ctl: ToolsProviderController): DatabaseConfig {
  const config = ctl.getPluginConfig(configSchematics);
  return {
    dbType: config.get("dbType"),
    host: config.get("host"),
    port: config.get("port"),
    username: config.get("username"),
    password: config.get("password"),
    database: config.get("database"),
    connectionPool: config.get("connectionPool"),
  };
}

const jsonScalar = z.union([z.string(), z.number(), z.boolean(), z.null()]);
const jsonValue: z.ZodTypeAny = z.lazy(() =>
  z.union([jsonScalar, z.array(jsonValue), z.record(jsonValue)]),
);

export async function toolsProvider(ctl: ToolsProviderController) {
  const config = getDatabaseConfig(ctl);
  const tools: Tool[] = [];

  tools.push(
    tool({
      name: "query_database",
      description:
        "Execute a SELECT query against the configured database and return the resulting rows.",
      parameters: {
        sql: z.string().describe("SQL SELECT statement to execute."),
        params: z.array(jsonScalar).optional().describe("Positional query parameters."),
      },
      implementation: async ({ sql, params }) => {
        try {
          await db.initializeConnection(config);
          const results = await db.query(sql, params ?? []);
          return {
            success: true,
            rowCount: results.length,
            data: results,
          };
        } catch (error) {
          return {
            success: false,
            error: String(error),
          };
        }
      },
    }),
  );

  tools.push(
    tool({
      name: "insert_data",
      description:
        "Insert one or more rows into a table and return the affected row count.",
      parameters: {
        table: z.string().describe("Table name."),
        rows: z
          .array(z.record(jsonValue))
          .min(1)
          .describe("Array of objects representing rows to insert."),
      },
      implementation: async ({ table, rows }) => {
        try {
          await db.initializeConnection(config);

          const columns = Object.keys(rows[0]);
          const placeholders = columns.map(() => "?").join(",");
          const sql = `INSERT INTO ${table} (${columns.join(", ")}) VALUES (${placeholders})`;

          let totalAffected = 0;
          for (const row of rows) {
            const values = columns.map((col) => row[col]);
            const result = await db.execute(sql, values);
            totalAffected += result.affectedRows;
          }

          return {
            success: true,
            affectedRows: totalAffected,
          };
        } catch (error) {
          return {
            success: false,
            error: String(error),
          };
        }
      },
    }),
  );

  tools.push(
    tool({
      name: "update_data",
      description: "Update rows in a table using a WHERE clause.",
      parameters: {
        table: z.string().describe("Table name."),
        updates: z.record(jsonValue).describe("Columns and values to update."),
        whereClause: z
          .string()
          .describe("WHERE clause without the WHERE keyword, for example id = ?."),
        params: z.array(jsonScalar).optional().describe("Parameters for the WHERE clause."),
      },
      implementation: async ({ table, updates, whereClause, params }) => {
        try {
          await db.initializeConnection(config);

          const setClause = Object.keys(updates)
            .map((col) => `${col} = ?`)
            .join(", ");

          const updateValues = Object.values(updates);
          const allParams = [...updateValues, ...(params ?? [])];
          const sql = `UPDATE ${table} SET ${setClause} WHERE ${whereClause}`;
          const result = await db.execute(sql, allParams);

          return {
            success: true,
            affectedRows: result.affectedRows,
          };
        } catch (error) {
          return {
            success: false,
            error: String(error),
          };
        }
      },
    }),
  );

  tools.push(
    tool({
      name: "delete_data",
      description: "Delete rows from a table using a WHERE clause.",
      parameters: {
        table: z.string().describe("Table name."),
        whereClause: z
          .string()
          .describe("WHERE clause without the WHERE keyword, for example id = ?."),
        params: z.array(jsonScalar).optional().describe("Parameters for the WHERE clause."),
      },
      implementation: async ({ table, whereClause, params }) => {
        try {
          await db.initializeConnection(config);
          const sql = `DELETE FROM ${table} WHERE ${whereClause}`;
          const result = await db.execute(sql, params ?? []);

          return {
            success: true,
            affectedRows: result.affectedRows,
          };
        } catch (error) {
          return {
            success: false,
            error: String(error),
          };
        }
      },
    }),
  );

  tools.push(
    tool({
      name: "get_schema",
      description: "Retrieve the database schema, including tables and columns.",
      parameters: {},
      implementation: async () => {
        try {
          await db.initializeConnection(config);
          const schema = await db.getSchema();
          return {
            success: true,
            schema,
          };
        } catch (error) {
          return {
            success: false,
            error: String(error),
          };
        }
      },
    }),
  );

  tools.push(
    tool({
      name: "execute_transaction",
      description:
        "Execute a sequence of query or execute operations. Note: this currently runs sequentially and is not a real SQL transaction.",
      parameters: {
        operations: z
          .array(
            z.object({
              type: z.enum(["query", "execute"]),
              sql: z.string(),
              params: z.array(jsonScalar).optional(),
            }),
          )
          .min(1)
          .describe("Operations to execute in sequence."),
      },
      implementation: async ({ operations }) => {
        try {
          await db.initializeConnection(config);

          const results = [];
          for (const op of operations) {
            if (op.type === "query") {
              const result = await db.query(op.sql, op.params ?? []);
              results.push({ type: "query", data: result });
            } else {
              const result = await db.execute(op.sql, op.params ?? []);
              results.push({ type: "execute", affectedRows: result.affectedRows });
            }
          }

          return {
            success: true,
            results,
          };
        } catch (error) {
          return {
            success: false,
            error: String(error),
          };
        }
      },
    }),
  );

  return tools;
}