apply_migration.ts•3.61 kB
import { z } from "zod";
import * as pg from "pg";
import type { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
export interface ToolProps {
pool: pg.Pool;
}
export function ApplyMigration(server: McpServer, props: ToolProps) {
server.tool(
"apply_migration",
"Apply a SQL migration (limited to DDL/schema changes that are tracked)",
{
name: z.string().describe("The name of the migration in snake_case"),
query: z.string().describe("The SQL migration query to apply (DDL operations only)"),
},
async ({ name, query }) => {
if (!props.pool) {
throw new Error("Database pool not initialized");
}
// Basic validation to ensure this is likely a DDL operation
const ddlKeywords = [
"CREATE",
"ALTER",
"DROP",
"RENAME",
"TRUNCATE",
"COMMENT",
"GRANT",
"REVOKE",
"SET",
"RESET",
];
const upperQuery = query.trim().toUpperCase();
const isDDL = ddlKeywords.some((keyword) => upperQuery.startsWith(keyword));
if (!isDDL) {
throw new Error(
"Migration queries should be DDL operations (CREATE, ALTER, DROP, etc.). Use execute_sql for data operations.",
);
}
const client = await props.pool.connect();
try {
await client.query("BEGIN");
// Execute the migration
const result = await client.query(query);
// Try to record the migration in a migrations table if one exists
const migrationTables = [
"supabase_migrations.schema_migrations",
"public.schema_migrations",
"public.migrations",
];
let migrationRecorded = false;
for (const table of migrationTables) {
try {
const [schema, tableName] = table.includes(".")
? table.split(".")
: ["public", table];
// Check if table exists
const tableExistsQuery = `
SELECT EXISTS (
SELECT FROM information_schema.tables
WHERE table_schema = $1 AND table_name = $2
);
`;
const tableExists = await client.query(tableExistsQuery, [
schema,
tableName,
]);
if (tableExists.rows[0].exists) {
// Try to insert migration record
if (table.includes("supabase_migrations")) {
await client.query(
`INSERT INTO ${table} (version, name, executed_at) VALUES ($1, $2, NOW())`,
[Date.now().toString(), name],
);
} else {
// Try common migration table structure
await client.query(
`INSERT INTO ${table} (version, name, executed_at) VALUES ($1, $2, NOW())`,
[Date.now().toString(), name],
);
}
migrationRecorded = true;
break;
}
} catch (error) {
// Continue to next table if this one fails
console.warn(`Failed to record migration in ${table}:`, error);
}
}
await client.query("COMMIT");
return {
content: [
{
type: "text",
text: JSON.stringify(
{
success: true,
migration_name: name,
rows_affected: result.rowCount || 0,
migration_recorded: migrationRecorded,
message: migrationRecorded
? "Migration applied and recorded successfully"
: "Migration applied successfully (no migration table found for recording)",
},
null,
2,
),
},
],
};
} catch (error) {
await client.query("ROLLBACK").catch(() => {});
throw new Error(
`Migration failed: ${error instanceof Error ? error.message : "Unknown error"}`,
);
} finally {
client.release();
}
},
);
}