Add WHERE detection for UPDATE and DELETE

README.md

@@ -153,6 +153,23 @@ console.log(statements);
     1. `strict (bool)`: allow disable strict mode which will ignore unknown types *(default=true)*
     2. `dialect (string)`: Specify your database dialect, values: `generic`, `mysql`, `oracle`, `psql`, `sqlite`, `mssql`, `bigquery` and `dynamodb`. *(default=generic)*
 
+### Statement flags
+
+Some statement types carry flags about their own clauses. Only the statement's own clauses count, not ones inside a subquery or a CTE.
+
+- `limit (bool)`, `offset (bool)`: on `SELECT`, whether it has a `LIMIT` (or `FETCH`, or `TOP` for `mssql`) and an `OFFSET`.
+- `where (bool)`: on `UPDATE` and `DELETE`, whether the statement has its own `WHERE` clause. The condition itself isn't checked, so `WHERE 1=1` still gives `true`, and a `DELETE` that is only filtered by a `JOIN` gives `false`. A `LIMIT` or `TOP` doesn't count either.
+
+An `UPDATE` or `DELETE` inside a CTE (`WITH d AS (DELETE FROM a ...) ...`) isn't checked at all and doesn't affect the flag of the outer statement.
+
+```js
+identify('DELETE FROM users WHERE id IN (SELECT id FROM bans WHERE active)');
+// [{ type: 'DELETE', executionType: 'MODIFICATION', where: true, ... }]
+
+identify('UPDATE users SET score = (SELECT max(score) FROM games WHERE won)');
+// [{ type: 'UPDATE', executionType: 'MODIFICATION', where: false, ... }]
+```
+
 ### DynamoDB (PartiQL)
 
 DynamoDB exposes a [PartiQL](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/ql-reference.html)-compatible

src/defines.ts

@@ -153,6 +153,7 @@ export interface IdentifyResult {
   columns: ColumnReference[];
   limit?: boolean;
   offset?: boolean;
+  where?: boolean;
 }
 
 export interface Statement {
@@ -170,6 +171,7 @@ export interface Statement {
   columns: ColumnReference[];
   limit?: boolean;
   offset?: boolean;
+  where?: boolean;
   isCte?: boolean;
 }
 

src/index.ts

@@ -48,6 +48,7 @@ export function identify(query: string, options: IdentifyOptions = {}): Identify
     };
     if (statement.limit !== undefined) result.limit = statement.limit;
     if (statement.offset !== undefined) result.offset = statement.offset;
+    if (statement.where !== undefined) result.where = statement.where;
     return result;
   });
 }

src/parser.ts

@@ -14,6 +14,22 @@ import type {
 import { ColumnParser } from './column-parser';
 import { TableParser } from './table-parser';
 
+// statement types whose LIMIT/OFFSET/WHERE flags only count outside parentheses
+const PARENS_TRACKED_TYPES = new Set(['SELECT', 'UPDATE', 'DELETE']);
+
+// keywords that can't appear at the top level of an UPDATE or DELETE, only start a new statement
+const NEXT_STATEMENT_KEYWORDS = new Set([
+  'SELECT',
+  'INSERT',
+  'UPDATE',
+  'DELETE',
+  'MERGE',
+  'TRUNCATE',
+  'CREATE',
+  'ALTER',
+  'DROP',
+]);
+
 interface StatementParser {
   addToken: (token: Token, nextToken: Token) => void;
   getStatement: () => Statement;
@@ -572,6 +588,7 @@ function createUpdateStatementParser(options: ParseOptions) {
       },
       add: (token) => {
         statement.type = 'UPDATE';
+        statement.where = false;
         if (statement.start < 0) {
           statement.start = token.start;
         }
@@ -595,6 +612,7 @@ function createDeleteStatementParser(options: ParseOptions) {
       },
       add: (token) => {
         statement.type = 'DELETE';
+        statement.where = false;
         if (statement.start < 0) {
           statement.start = token.start;
         }
@@ -1024,6 +1042,7 @@ function stateMachineStatementParser(
 
   let openBlocks = 0;
   let parensDepth = 0;
+  let sawNextStatement = false;
 
   const columnParser = new ColumnParser(dialect);
   const tableParser = new TableParser(dialect);
@@ -1173,12 +1192,32 @@ function stateMachineStatementParser(
         statement.parameters.push(token.value);
       }
 
-      if (statement.type === 'SELECT' && token.value === '(') {
-        parensDepth++;
+      if (statement.type && PARENS_TRACKED_TYPES.has(statement.type)) {
+        if (token.value === '(') {
+          parensDepth++;
+        } else if (token.value === ')') {
+          parensDepth--;
+        }
       }
 
-      if (statement.type === 'SELECT' && token.value === ')') {
-        parensDepth--;
+      if (
+        (statement.type === 'UPDATE' || statement.type === 'DELETE') &&
+        token.type === 'keyword' &&
+        parensDepth === 0
+      ) {
+        const keyword = token.value.toUpperCase();
+        // Without a semicolon between them (fine in mssql), the next statement ends up in
+        // this one, and its WHERE would be counted as ours. Don't trust the flag then.
+        // INSERT( and TRUNCATE( are mysql functions, not statements.
+        if (NEXT_STATEMENT_KEYWORDS.has(keyword) && nextToken.value !== '(') {
+          sawNextStatement = true;
+          statement.where = false;
+        }
+        // Only a WHERE of the statement itself limits the rows it changes, not one in a
+        // subquery. A CTE's never reaches this parser.
+        if (keyword === 'WHERE' && !sawNextStatement) {
+          statement.where = true;
+        }
       }
 
       if (

src/tokenizer.ts

@@ -93,6 +93,7 @@ const KEYWORDS = [
   'FETCH',
   'TOP',
   'ALL',
+  'WHERE',
 ];
 
 const SNOWFLAKE_KEYWORDS = [
@@ -115,6 +116,13 @@ const INDIVIDUALS: Record<string, Token['type']> = {
   ';': 'semicolon',
 };
 
+const ORACLE_QUOTE_CLOSINGS: Record<string, string> = {
+  '[': ']',
+  '{': '}',
+  '<': '>',
+  '(': ')',
+};
+
 const ENDTOKENS: Record<string, Char> = {
   '"': '"',
   "'": "'",
@@ -138,11 +146,21 @@ export function scanToken(
   }
 
   if (isCommentBlock(ch, state)) {
-    return scanCommentBlock(state);
+    return scanCommentBlock(state, dialect);
   }
 
   if (isString(ch, dialect) && ch !== null) {
-    return scanString(state, ENDTOKENS[ch]);
+    return scanString(state, ENDTOKENS[ch], hasBackslashEscapes(dialect));
+  }
+
+  // Postgres escape string: E'It\'s'
+  if (dialect === 'psql' && (ch === 'E' || ch === 'e') && peek(state) === "'") {
+    read(state);
+    return scanString(state, "'", true);
+  }
+
+  if (dialect === 'oracle' && isOracleQuotedString(ch, state)) {
+    return scanOracleQuotedString(state);
   }
 
   if (isParameter(ch, state, paramTypes)) {
@@ -296,10 +314,15 @@ function scanDollarQuotedString(state: State): Token {
   };
 }
 
-function scanString(state: State, endToken: Char): Token {
+function scanString(state: State, endToken: Char, backslashEscapes = false): Token {
   let nextChar: Char;
   do {
     nextChar = read(state);
+    // supporting backslash escaping: 'str\'ing'
+    if (backslashEscapes && nextChar === '\\') {
+      read(state);
+      continue;
+    }
     // supporting double quote escaping: 'str''ing'
     if (nextChar === endToken) {
       if (peek(state) === endToken) {
@@ -332,6 +355,23 @@ function getCustomParam(state: State, paramTypes: ParamTypes): string | null | u
   return matches ? matches[0] : null;
 }
 
+// Oracle alternative quoting: q'[It's]', nq'{It's}', q'!It's!'
+function scanOracleQuotedString(state: State): Token {
+  const delimiterPos = state.input.indexOf("'", state.start) + 1;
+  const delimiter = state.input[delimiterPos];
+  const closing = delimiter ? (ORACLE_QUOTE_CLOSINGS[delimiter] ?? delimiter) + "'" : null;
+  const end = closing ? state.input.indexOf(closing, delimiterPos + 1) : -1;
+  state.position = end === -1 ? state.input.length - 1 : end + 1;
+
+  const value = state.input.slice(state.start, state.position + 1);
+  return {
+    type: 'string',
+    value,
+    start: state.start,
+    end: state.start + value.length - 1,
+  };
+}
+
 function scanParameter(state: State, dialect: Dialect, paramTypes: ParamTypes): Token {
   const curCh = state.input[state.start];
   const nextChar = peek(state);
@@ -410,14 +450,25 @@ function scanParameter(state: State, dialect: Dialect, paramTypes: ParamTypes):
   };
 }
 
-function scanCommentBlock(state: State): Token {
+function scanCommentBlock(state: State, dialect: Dialect): Token {
   let nextChar: Char = '';
   let prevChar: Char;
+  // Postgres and SQL Server allow nested block comments: /* a /* b */ c */
+  const nested = dialect === 'psql' || dialect === 'mssql';
+  let depth = 1;
 
   do {
     prevChar = nextChar;
     nextChar = read(state);
-  } while ((prevChar || '') + (nextChar || '') !== '*/' && nextChar !== null);
+    const pair = (prevChar || '') + (nextChar || '');
+    if (pair === '*/') {
+      depth--;
+      nextChar = depth > 0 ? '' : nextChar;
+    } else if (nested && pair === '/*') {
+      depth++;
+      nextChar = '';
+    }
+  } while (depth > 0 && nextChar !== null);
 
   if (nextChar !== null && nextChar !== '/') {
     unread(state);
@@ -476,7 +527,7 @@ function scanWord(state: State, dialect: Dialect): Token {
   }
 
   const value = state.input.slice(state.start, state.position + 1);
-  if (!isKeyword(value, dialect)) {
+  if (!isKeyword(value, dialect) || isQualifiedName(state)) {
     return skipWord(state, value);
   }
 
@@ -525,10 +576,39 @@ function isWhitespace(ch: Char): boolean {
   return ch === ' ' || ch === '\t' || ch === '\n' || ch === '\r';
 }
 
+function isInlineSpace(ch: Char): boolean {
+  return ch === ' ' || ch === '\t';
+}
+
 function isAlphaNumeric(ch: Char): boolean {
   return ch !== null && /\w/.test(ch);
 }
 
+// a column or variable named like a keyword: t.where, t . where, @where
+function isQualifiedName(state: State): boolean {
+  const { input } = state;
+  let i = state.start - 1;
+  if (input[i] === '@') {
+    return true;
+  }
+  // stay on the same line, so a line comment ending in a dot (-- done.\nDELETE ...) doesn't count
+  while (i >= 0 && isInlineSpace(input[i])) i--;
+  if (input[i] !== '.') {
+    return false;
+  }
+
+  // a number with a trailing dot (1. FROM t) is not a qualifier
+  let end = i - 1;
+  while (end >= 0 && isInlineSpace(input[end])) end--;
+  let begin = end;
+  while (begin >= 0 && isAlphaNumeric(input[begin])) begin--;
+  return !/^\d+$/.test(input.slice(begin + 1, end + 1));
+}
+
+function hasBackslashEscapes(dialect: Dialect): boolean {
+  return ['mysql', 'bigquery', 'snowflake'].includes(dialect);
+}
+
 function isString(ch: Char, dialect: Dialect): boolean {
   const stringStart: Char[] = dialect === 'mysql' ? ["'", '"'] : ["'"];
   return stringStart.includes(ch);
@@ -579,6 +659,13 @@ function isDollarQuotedString(state: State): boolean {
   return /^\$[\w]*\$/.exec(state.input.slice(state.start)) !== null;
 }
 
+function isOracleQuotedString(ch: Char, state: State): boolean {
+  if (ch !== 'q' && ch !== 'Q' && ch !== 'n' && ch !== 'N') {
+    return false;
+  }
+  return /^n?q'/i.test(state.input.slice(state.start, state.start + 3));
+}
+
 function isQuotedIdentifier(ch: Char, dialect: Dialect): boolean {
   const startQuoteChars: Char[] = getStartQuotes(dialect);
   return startQuoteChars.includes(ch);