diff --git a/app/lib/__tests__/db.server.test.ts b/app/lib/__tests__/db.server.test.ts index 1ce1fd8..4f2193f 100644 --- a/app/lib/__tests__/db.server.test.ts +++ b/app/lib/__tests__/db.server.test.ts @@ -1,5 +1,324 @@ import { describe, test, expect } from "vitest"; -import { needsTransaction } from "../db.server"; +import { needsTransaction, analyzeSQL } from "../db.server"; + +describe("analyzeSQL - RETURNING clause detection", () => { + describe("SELECT queries", () => { + test("detects basic SELECT query", () => { + const sql = "SELECT * FROM users;"; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(true); + expect(result.hasReturning).toBe(false); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects SELECT with WHERE clause", () => { + const sql = "SELECT id, name FROM users WHERE active = true;"; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(true); + expect(result.hasReturning).toBe(false); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects multiline SELECT query", () => { + const sql = ` + SELECT + id, + name, + email + FROM users + WHERE active = true; + `; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + }); + + describe("UPDATE with RETURNING", () => { + test("detects UPDATE with RETURNING *", () => { + const sql = "UPDATE charges SET amount = 500 WHERE id = 123 RETURNING *;"; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(false); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects UPDATE with RETURNING specific columns", () => { + const sql = + "UPDATE charges SET amount = 500 WHERE id = 123 RETURNING id, amount, updated_at;"; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(false); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects multiline UPDATE with RETURNING", () => { + const sql = ` + UPDATE charges + SET + amount = 500, + campus_id = 'abc-123' + WHERE stripe_charge_id = 'ch_xyz' + RETURNING id, amount, campus_id, updated_at; + `; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(false); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects Cameron's backfill query with RETURNING", () => { + const sql = ` + UPDATE charges + SET + amount = (se.stripe_event_data -> 'data' -> 'object' -> 'amount')::INTEGER, + campus_id = CASE + WHEN se.stripe_event_data -> 'data' -> 'object' -> 'metadata' -> 'campusId' IS NOT NULL + THEN SPLIT_PART(se.stripe_event_data -> 'data' -> 'object' -> 'metadata' ->> 'campusId', ':', 2)::UUID + ELSE NULL + END + FROM stripe_events se + WHERE se.stripe_event_data -> 'data' -> 'object' ->> 'id' = charges.stripe_charge_id + AND charges.amount IS NULL + RETURNING charges.id, charges.stripe_charge_id, charges.amount; + `; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(false); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + }); + + describe("INSERT with RETURNING", () => { + test("detects INSERT with RETURNING", () => { + const sql = + "INSERT INTO users (name, email) VALUES ('John', 'john@example.com') RETURNING id, created_at;"; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(false); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects multiline INSERT with RETURNING *", () => { + const sql = ` + INSERT INTO users (name, email, active) + VALUES ('Jane Doe', 'jane@example.com', true) + RETURNING *; + `; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(false); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects bulk INSERT with RETURNING", () => { + const sql = ` + INSERT INTO test_table (name, value) + VALUES + ('test1', 100), + ('test2', 200), + ('test3', 300) + RETURNING id, name, value; + `; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + }); + + describe("DELETE with RETURNING", () => { + test("detects DELETE with RETURNING", () => { + const sql = + "DELETE FROM old_records WHERE created_at < '2020-01-01' RETURNING id, created_at;"; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(false); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects DELETE with RETURNING *", () => { + const sql = "DELETE FROM temp_data WHERE processed = true RETURNING *;"; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + }); + + describe("DML without RETURNING", () => { + test("UPDATE without RETURNING should not capture rows", () => { + const sql = "UPDATE charges SET amount = 500 WHERE id = 123;"; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(false); + expect(result.hasReturning).toBe(false); + expect(result.shouldCaptureRows).toBe(false); + }); + + test("INSERT without RETURNING should not capture rows", () => { + const sql = "INSERT INTO users (name) VALUES ('test');"; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(false); + expect(result.hasReturning).toBe(false); + expect(result.shouldCaptureRows).toBe(false); + }); + + test("DELETE without RETURNING should not capture rows", () => { + const sql = "DELETE FROM users WHERE id = 123;"; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(false); + expect(result.hasReturning).toBe(false); + expect(result.shouldCaptureRows).toBe(false); + }); + }); + + describe("Edge cases with comments", () => { + test("ignores RETURNING in single-line comment", () => { + const sql = ` + -- This query is not returning anything + UPDATE charges SET amount = 500 WHERE id = 123; + `; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(false); + expect(result.shouldCaptureRows).toBe(false); + }); + + test("ignores RETURNING in multi-line comment", () => { + const sql = ` + /* + * Note: We could use RETURNING here but chose not to + * for performance reasons + */ + UPDATE charges SET amount = 500 WHERE id = 123; + `; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(false); + expect(result.shouldCaptureRows).toBe(false); + }); + + test("detects RETURNING when also in comments", () => { + const sql = ` + -- Update with RETURNING clause + UPDATE charges + SET amount = 500 + WHERE id = 123 + RETURNING id, amount; -- Return the updated row + `; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("handles mixed comments and RETURNING", () => { + const sql = ` + /* Block comment */ + UPDATE charges SET amount = 500 -- inline comment + WHERE id = 123 + RETURNING *; -- return everything + `; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + }); + + describe("Case insensitivity", () => { + test("detects lowercase returning", () => { + const sql = "update charges set amount = 500 returning id;"; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects mixed case ReTuRnInG", () => { + const sql = "UPDATE charges SET amount = 500 ReTuRnInG id;"; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects lowercase select", () => { + const sql = "select * from users;"; + const result = analyzeSQL(sql); + expect(result.isSelectQuery).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + }); + + describe("Multi-statement queries", () => { + test("detects RETURNING in first statement", () => { + const sql = ` + UPDATE table1 SET x = 1 RETURNING *; + UPDATE table2 SET y = 2; + `; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects RETURNING in second statement", () => { + const sql = ` + UPDATE table1 SET x = 1; + UPDATE table2 SET y = 2 RETURNING *; + `; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("detects RETURNING in both statements", () => { + const sql = ` + UPDATE table1 SET x = 1 RETURNING id; + UPDATE table2 SET y = 2 RETURNING id; + `; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + }); + + describe("Complex real-world queries", () => { + test("handles subqueries with RETURNING", () => { + const sql = ` + UPDATE charges c + SET campus_id = (SELECT id FROM campuses WHERE is_default = true) + WHERE c.campus_id IS NULL + RETURNING c.id, c.campus_id; + `; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("handles CTE (WITH clause) with RETURNING", () => { + const sql = ` + WITH updated_charges AS ( + UPDATE charges SET processed = true WHERE amount > 100 + RETURNING id, amount + ) + SELECT * FROM updated_charges; + `; + const result = analyzeSQL(sql); + // This starts with WITH but contains RETURNING and SELECT + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + + test("handles JOIN in UPDATE with RETURNING", () => { + const sql = ` + UPDATE charges c + SET + amount = se.amount, + campus_id = se.campus_id + FROM stripe_events se + WHERE se.charge_id = c.stripe_charge_id + AND c.amount IS NULL + RETURNING c.id, c.amount, c.campus_id; + `; + const result = analyzeSQL(sql); + expect(result.hasReturning).toBe(true); + expect(result.shouldCaptureRows).toBe(true); + }); + }); +}); describe("needsTransaction", () => { test("single statement with semicolon returns false", () => { diff --git a/app/lib/db.server.ts b/app/lib/db.server.ts index 5f37f6c..5690c9b 100644 --- a/app/lib/db.server.ts +++ b/app/lib/db.server.ts @@ -54,6 +54,30 @@ export function needsTransaction(sql: string): boolean { return semicolons >= 2; } +// Helper function to clean SQL and detect query characteristics +export function analyzeSQL(sql: string): { + isSelectQuery: boolean; + hasReturning: boolean; + shouldCaptureRows: boolean; +} { + // Remove comments and whitespace to detect query type properly + const sqlClean = sql + .replace(/--.*$/gm, "") // Remove single-line comments + .replace(/\/\*[\s\S]*?\*\//g, "") // Remove multi-line comments + .trim() + .toUpperCase(); + + const isSelectQuery = sqlClean.startsWith("SELECT"); + const hasReturning = sqlClean.includes("RETURNING"); + const shouldCaptureRows = isSelectQuery || hasReturning; + + return { + isSelectQuery, + hasReturning, + shouldCaptureRows, + }; +} + // Execute SQL against target database export async function executeSQL( target: "staging" | "production", @@ -87,18 +111,11 @@ export async function executeSQL( const executionTime = Date.now() - start; - // Check if this is a SELECT query (has result rows) - // Remove comments and whitespace to detect SELECT queries properly - const sqlClean = sql - .replace(/--.*$/gm, "") // Remove single-line comments - .replace(/\/\*[\s\S]*?\*\//g, "") // Remove multi-line comments - .trim() - .toUpperCase(); - const isSelectQuery = sqlClean.startsWith("SELECT"); - const resultRows = - isSelectQuery && result.rows && result.rows.length > 0 - ? result.rows - : undefined; + // Analyze SQL to determine if we should capture result rows + const { shouldCaptureRows: shouldCapture } = analyzeSQL(sql); + const shouldCaptureRows = + shouldCapture && result.rows && result.rows.length > 0; + const resultRows = shouldCaptureRows ? result.rows : undefined; // Limit result rows to prevent storing huge result sets (max 100 rows) const limitedRows = @@ -106,11 +123,9 @@ export async function executeSQL( ? resultRows.slice(0, 100) : resultRows; - // For SELECT queries, use rows.length; for DML/DDL, use rowCount - const rowsAffected = isSelectQuery - ? result.rows - ? result.rows.length - : 0 + // For SELECT/RETURNING queries, use rows.length; for DML/DDL, use rowCount + const rowsAffected = shouldCaptureRows + ? result.rows.length : result.rowCount || 0; return {