diff --git a/src/sql/recent-query.test.ts b/src/sql/recent-query.test.ts index 96748e3..8bdd6dd 100644 --- a/src/sql/recent-query.test.ts +++ b/src/sql/recent-query.test.ts @@ -46,6 +46,30 @@ test("isSelectQuery returns false for non-SELECT statements", () => { ).toBe(false); }); +test("isSelectQuery returns false for DML containing SELECT subqueries", () => { + expect( + RecentQuery.isSelectQuery( + makeRawQuery({ + query: 'UPDATE "public"."oban_jobs" SET "state" = $1 FROM (SELECT id FROM t) s WHERE id = s.id', + }), + ), + ).toBe(false); + expect( + RecentQuery.isSelectQuery( + makeRawQuery({ + query: "INSERT INTO archive SELECT * FROM users", + }), + ), + ).toBe(false); + expect( + RecentQuery.isSelectQuery( + makeRawQuery({ + query: "DELETE FROM users WHERE EXISTS (SELECT 1 FROM banned)", + }), + ), + ).toBe(false); +}); + // --- isSystemQuery --- test("isSystemQuery returns true for pg_ tables", () => { @@ -193,3 +217,45 @@ test("analyze throws on unparseable SQL", async () => { RecentQuery.analyze(data, testHash, 3000), ).rejects.toThrow(); }); + +// --- statementType-based isSelectQuery via analyze --- + +test("analyze sets isSelectQuery=true for SELECT", async () => { + const data = makeRawQuery({ query: "SELECT * FROM users" }); + const rq = await RecentQuery.analyze(data, testHash, 1000); + expect(rq.isSelectQuery).toBe(true); +}); + +test("analyze sets isSelectQuery=true for CTE with SELECT", async () => { + const data = makeRawQuery({ + query: "WITH cte AS (SELECT id FROM users) SELECT * FROM cte", + }); + const rq = await RecentQuery.analyze(data, testHash, 1000); + expect(rq.isSelectQuery).toBe(true); +}); + +test("analyze sets isSelectQuery=false for UPDATE even with SELECT subquery", async () => { + const data = makeRawQuery({ + query: + 'UPDATE "public"."jobs" SET "state" = $1 FROM (SELECT id FROM "public"."jobs" WHERE state = $2 LIMIT 10) AS s1 WHERE "jobs".id = s1.id', + }); + const rq = await RecentQuery.analyze(data, testHash, 1000); + expect(rq.isSelectQuery).toBe(false); +}); + +test("analyze sets isSelectQuery=false for INSERT ... SELECT", async () => { + const data = makeRawQuery({ + query: "INSERT INTO archive SELECT * FROM users WHERE active = false", + }); + const rq = await RecentQuery.analyze(data, testHash, 1000); + expect(rq.isSelectQuery).toBe(false); +}); + +test("analyze sets isSelectQuery=false for DELETE with EXISTS subquery", async () => { + const data = makeRawQuery({ + query: + "DELETE FROM users WHERE EXISTS (SELECT 1 FROM banned WHERE banned.user_id = users.id)", + }); + const rq = await RecentQuery.analyze(data, testHash, 1000); + expect(rq.isSelectQuery).toBe(false); +}); diff --git a/src/sql/recent-query.ts b/src/sql/recent-query.ts index ede7d9f..877f0ca 100644 --- a/src/sql/recent-query.ts +++ b/src/sql/recent-query.ts @@ -8,6 +8,7 @@ import { PostgresQueryBuilder, PssRewriter, SQLCommenterTag, + type StatementType, type TableReference, } from "@query-doctor/core"; import { parse } from "@libpg-query/parser"; @@ -46,6 +47,7 @@ export class RecentQuery { readonly hash: QueryHash, readonly seenAt: number, analysisSkipped = false, + statementType?: StatementType, ) { this.username = data.username; this.query = data.query; @@ -57,7 +59,9 @@ export class RecentQuery { this.analysisSkipped = analysisSkipped; this.isSystemQuery = RecentQuery.isSystemQuery(tableReferences); - this.isSelectQuery = RecentQuery.isSelectQuery(data); + this.isSelectQuery = statementType !== undefined + ? statementType === "select" + : RecentQuery.isSelectQuery(data); this.isIntrospection = RecentQuery.isIntrospection(data); this.isTargetlessSelectQuery = this.isSelectQuery ? RecentQuery.isTargetlessSelectQuery(tableReferences) @@ -123,6 +127,8 @@ export class RecentQuery { analysis.nudges, hash, seenAt, + false, + analysis.statementType, ); } @@ -147,7 +153,7 @@ export class RecentQuery { } static isSelectQuery(data: RawRecentQuery): boolean { - return /select/i.test(data.query); + return /^\s*select/i.test(data.query); } static isSystemQuery(referencedTables: TableReference[]): boolean {