src/mastra/tools/execute-sql.ts

import { createTool } from '@mastra/core/tools';
import { createClient } from '@libsql/client';
import { z } from 'zod';

const db = createClient({ url: 'file:./data.db' });

const MAX_QUERY_LENGTH = 2000;
const MAX_ROWS = 100;

const BLOCKED_PATTERNS = [
  /\b(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|REPLACE|UPSERT)\b/i,
  /\b(ATTACH|DETACH|PRAGMA|VACUUM|REINDEX|ANALYZE)\b/i,
  /\b(load_extension|readfile|writefile|randomblob|zeroblob)\s*\(/i,
  /--|\/\*|\*\//, // block SQL comments
  /;/, // block additional statements after optional trailing semicolon is removed
];

export const executeSql = createTool({
  id: 'execute-sql',
  description:
    'Executes one read-only SQLite SELECT query against the public demo database. Results are limited to 100 rows.',
  inputSchema: z.object({
    query: z
      .string()
      .min(1)
      .max(MAX_QUERY_LENGTH)
      .describe('One SQLite SELECT query to execute'),
  }),
  outputSchema: z.object({
    rows: z.array(z.record(z.string(), z.unknown())).describe('Query result rows'),
    rowCount: z.number().describe('Number of rows returned'),
    truncated: z.boolean().describe('Whether additional rows were excluded'),
    maxRows: z.number().describe('Maximum rows returned by the public demo'),
  }),
  execute: async ({ query }) => {
    const trimmed = query.trim().replace(/;\s*$/, '').trim();

    if (trimmed.length === 0 || trimmed.length > MAX_QUERY_LENGTH) {
      throw new Error(`Query must contain between 1 and ${MAX_QUERY_LENGTH} characters.`);
    }

    if (!/^SELECT\b/i.test(trimmed)) {
      throw new Error('Only a single SELECT query is allowed.');
    }

    for (const pattern of BLOCKED_PATTERNS) {
      if (pattern.test(trimmed)) {
        throw new Error('This query is not permitted in the public demo.');
      }
    }

    const limitedQuery = `
      SELECT *
      FROM (${trimmed}) AS public_demo_result
      LIMIT ${MAX_ROWS + 1}
    `;

    const result = await db.execute(limitedQuery);
    const allRows = result.rows as Record<string, unknown>[];
    const truncated = allRows.length > MAX_ROWS;
    const rows = allRows.slice(0, MAX_ROWS);

    return {
      rows,
      rowCount: rows.length,
      truncated,
      maxRows: MAX_ROWS,
    };
  },
});