Fix identifying transaction statements

src/parser.ts

@@ -12,7 +12,7 @@ import type {
 } from './defines';
 
 interface StatementParser {
-  addToken: (token: Token) => void;
+  addToken: (token: Token, nextToken: Token) => void;
   getStatement: () => Statement;
 }
 
@@ -81,6 +81,15 @@ function createInitialStatement(): Statement {
   };
 }
 
+function nextNonWhitespaceToken(state: State): Token {
+  let token: Token;
+  do {
+    state = initState({ prevState: state });
+    token = scanToken(state);
+  } while (token.type === 'whitespace');
+  return token;
+}
+
 /**
  * Parser
  */
@@ -115,6 +124,7 @@ export function parse(input: string, isStrict = true, dialect: Dialect = 'generi
   while (prevState.position < topLevelState.end) {
     const tokenState = initState({ prevState });
     const token = scanToken(tokenState, dialect);
+    const nextToken = nextNonWhitespaceToken(tokenState);
 
     if (!statementParser) {
       // ignore blank tokens before the start of a CTE / not part of a statement
@@ -193,7 +203,7 @@ export function parse(input: string, isStrict = true, dialect: Dialect = 'generi
       }
     }
 
-    statementParser.addToken(token);
+    statementParser.addToken(token, nextToken);
     topLevelStatement.tokens.push(token);
     prevState = tokenState;
 
@@ -581,7 +591,6 @@ function stateMachineStatementParser(
 ): StatementParser {
   let currentStepIndex = 0;
   let prevToken: Token | undefined;
-  let prevPrevToken: Token | undefined;
   let prevNonWhitespaceToken: Token | undefined;
 
   let lastBlockOpener: Token | undefined;
@@ -606,7 +615,6 @@ function stateMachineStatementParser(
   };
 
   const setPrevToken = (token: Token) => {
-    prevPrevToken = prevToken;
     prevToken = token;
     if (token.type !== 'whitespace') {
       prevNonWhitespaceToken = token;
@@ -618,7 +626,7 @@ function stateMachineStatementParser(
       return statement;
     },
 
-    addToken(token: Token) {
+    addToken(token: Token, nextToken: Token) {
       /* eslint no-param-reassign: 0 */
       if (statement.endStatement) {
         throw new Error('This statement has already got to the end.');
@@ -650,7 +658,13 @@ function stateMachineStatementParser(
       if (
         token.type === 'keyword' &&
         blockOpeners[dialect].includes(token.value) &&
-        prevPrevToken?.value.toUpperCase() !== 'END'
+        prevNonWhitespaceToken?.value.toUpperCase() !== 'END' &&
+        (token.value.toUpperCase() !== 'BEGIN' ||
+          (token.value.toUpperCase() === 'BEGIN' &&
+            nextToken.value.toUpperCase() !== 'TRANSACTION' &&
+            (dialect !== 'sqlite' ||
+              (dialect === 'sqlite' &&
+                !['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'].includes(nextToken.value.toUpperCase())))))
       ) {
         if (
           ['oracle', 'bigquery'].includes(dialect) &&

test/identifier/multiple-statement.spec.ts

@@ -363,4 +363,75 @@ describe('identifier', () => {
       expect(sql.substring(actual[1].start, actual[1].end + 1)).to.eql(statement2);
     });
   });
+
+  describe('identifying transactions', () => {
+    it('should identify transactions', () => {
+      const statements = ['BEGIN TRANSACTION;', 'SELECT 1;', 'COMMIT;'];
+      const actual = identify(statements.join('\n'), { strict: false });
+      const expected = [
+        {
+          start: 0,
+          end: 17,
+          text: statements[0],
+          type: 'UNKNOWN',
+          executionType: 'UNKNOWN',
+          parameters: [],
+        },
+        {
+          start: 19,
+          end: 27,
+          text: statements[1],
+          type: 'SELECT',
+          executionType: 'LISTING',
+          parameters: [],
+        },
+        {
+          start: 29,
+          end: 35,
+          text: statements[2],
+          type: 'UNKNOWN',
+          executionType: 'UNKNOWN',
+          parameters: [],
+        },
+      ];
+      expect(actual).to.eql(expected);
+    });
+
+    describe('identifying keywords for sqlite transactions', () => {
+      ['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE'].forEach((type) => {
+        it(`identifies BEGIN ${type} TRANSACTION`, () => {
+          const statements = [`BEGIN ${type} TRANSACTION;`, 'SELECT 1;', 'COMMIT;'];
+          const actual = identify(statements.join('\n'), { dialect: 'sqlite', strict: false });
+          const offset = type.length + 1;
+          const expected = [
+            {
+              start: 0,
+              end: 17 + offset,
+              text: statements[0],
+              type: 'UNKNOWN',
+              executionType: 'UNKNOWN',
+              parameters: [],
+            },
+            {
+              start: 19 + offset,
+              end: 27 + offset,
+              text: statements[1],
+              type: 'SELECT',
+              executionType: 'LISTING',
+              parameters: [],
+            },
+            {
+              start: 29 + offset,
+              end: 35 + offset,
+              text: statements[2],
+              type: 'UNKNOWN',
+              executionType: 'UNKNOWN',
+              parameters: [],
+            },
+          ];
+          expect(actual).to.eql(expected);
+        });
+      });
+    });
+  });
 });