feat: skipe separator in begin-end closure option

Author: Nybkox

src/options.ts

@@ -15,6 +15,7 @@ export interface SplitterOptions {
   doubleDashComments: boolean;
   multilineComments: boolean;
   javaScriptComments: boolean;
+  skipSeparatorBeginEnd: boolean;
   // comments willl be not part of output
   ignoreComments: boolean;
   // if more commands are on single line, they are not splitted
@@ -42,6 +43,7 @@ export const defaultSplitterOptions: SplitterOptions = {
   allowSlashDelimiter: false,
   allowDollarDollarString: false,
   noSplit: false,
+  skipSeparatorBeginEnd: false,
 
   doubleDashComments: true,
   multilineComments: true,
@@ -89,6 +91,7 @@ export const postgreSplitterOptions: SplitterOptions = {
 export const sqliteSplitterOptions: SplitterOptions = {
   ...defaultSplitterOptions,
 
+  skipSeparatorBeginEnd: true,
   stringsBegins: ["'", '"'],
   stringsEnds: { "'": "'", '"': '"' },
   stringEscapes: { "'": "'", '"': '"' },

src/queryParamHandler.ts

@@ -11,6 +11,7 @@ function createParameterizerContext(sql: string, options: SplitterOptions) {
     wasDataOnLine: false,
     isCopyFromStdin: false,
     isCopyFromStdinCandidate: false,
+    beginEndIdentLevel: 0,
   };
 }
 

src/splitQuery.ts

@@ -1,6 +1,10 @@
 import { SplitterOptions, defaultSplitterOptions } from './options';
 
 const SEMICOLON = ';';
+const BEGIN_EXTRA_KEYWORDS = ['DEFERRED', 'IMMEDIATE', 'EXCLUSIVE', 'TRANSACTION'];
+const BEGIN_EXTRA_KEYWORDS_REGEX = new RegExp(`^(?:${BEGIN_EXTRA_KEYWORDS.join('|')})`, 'i');
+const END_EXTRA_KEYWORDS = ['TRANSACTION', 'IF'];
+const END_EXTRA_KEYWORDS_REGEX = new RegExp(`^(?:${END_EXTRA_KEYWORDS.join('|')})`, 'i');
 
 type SplitterSpecialMarkerType = 'copy_stdin_start' | 'copy_stdin_end' | 'copy_stdin_line';
 export interface SplitStreamContext {
@@ -40,6 +44,7 @@ export interface ScannerContext {
   readonly wasDataOnLine: boolean;
   readonly isCopyFromStdin: boolean;
   readonly isCopyFromStdinCandidate: boolean;
+  readonly beginEndIdentLevel: number;
 }
 
 export interface SplitLineContext extends SplitStreamContext {
@@ -49,6 +54,7 @@ export interface SplitLineContext extends SplitStreamContext {
   end: number;
   wasDataOnLine: boolean;
   currentCommandStart: number;
+  beginEndIdentLevel: number;
 
   //   unread: string;
   //   currentStatement: string;
@@ -111,6 +117,8 @@ interface Token {
   type:
     | 'string'
     | 'delimiter'
+    | 'end'
+    | 'begin'
     | 'whitespace'
     | 'eoln'
     | 'data'
@@ -228,7 +236,8 @@ export function scanToken(context: ScannerContext): Token {
     };
   }
 
-  if (context.currentDelimiter && s.slice(pos).startsWith(context.currentDelimiter)) {
+  const isInBeginEnd = context.options.skipSeparatorBeginEnd && context.beginEndIdentLevel > 0;
+  if (context.currentDelimiter && s.slice(pos).startsWith(context.currentDelimiter) && !isInBeginEnd) {
     return {
       type: 'delimiter',
       length: context.currentDelimiter.length,
@@ -350,6 +359,35 @@ export function scanToken(context: ScannerContext): Token {
     };
   }
 
+  if (context.options.skipSeparatorBeginEnd && s.slice(pos).match(/^begin/i)) {
+    let pos2 = pos + 'BEGIN'.length;
+    let pos0 = pos2;
+
+    while (pos0 < context.end && /[^a-zA-Z0-9]/.test(s[pos0])) pos0++;
+
+    if (!BEGIN_EXTRA_KEYWORDS_REGEX.test(s.slice(pos0))) {
+      return {
+        type: 'begin',
+        length: pos2 - pos,
+        lengthWithoutWhitespace: pos0 - pos,
+      };
+    }
+  }
+
+  if (context.options.skipSeparatorBeginEnd && s.slice(pos).match(/^end/i)) {
+    let pos2 = pos + 'END'.length;
+    let pos0 = pos2;
+
+    while (pos0 < context.end && /[^a-zA-Z0-9]/.test(s[pos0])) pos0++;
+
+    if (!END_EXTRA_KEYWORDS_REGEX.test(s.slice(pos0))) {
+      return {
+        type: 'end',
+        length: pos2 - pos,
+      };
+    }
+  }
+
   const dollarString = scanDollarQuotedString(context);
   if (dollarString) return dollarString;
 
@@ -366,6 +404,7 @@ function containsDataAfterDelimiterOnLine(context: ScannerContext, delimiter: To
     wasDataOnLine: context.wasDataOnLine,
     isCopyFromStdinCandidate: context.isCopyFromStdinCandidate,
     isCopyFromStdin: context.isCopyFromStdin,
+    beginEndIdentLevel: context.beginEndIdentLevel,
   };
 
   cloned.position += delimiter.length;
@@ -594,6 +633,18 @@ export function splitQueryLine(context: SplitLineContext) {
         markStartCommand(context);
         context.isCopyFromStdinCandidate = false;
         break;
+      case 'begin':
+        if (context.options.skipSeparatorBeginEnd) {
+          context.beginEndIdentLevel++;
+        }
+        movePosition(context, token.length, false);
+        break;
+      case 'end':
+        if (context.options.skipSeparatorBeginEnd && context.beginEndIdentLevel > 0) {
+          context.beginEndIdentLevel--;
+        }
+        movePosition(context, token.length, false);
+        break;
     }
   }
 
@@ -660,6 +711,7 @@ export function splitQuery(sql: string, options: SplitterOptions = null): SplitR
     trimCommandStartPosition: 0,
     trimCommandStartLine: 0,
     trimCommandStartColumn: 0,
+    beginEndIdentLevel: 0,
 
     wasDataInCommand: false,
     isCopyFromStdin: false,

src/splitQueryStream.ts

@@ -45,6 +45,7 @@ export class SplitQueryStream extends stream.Transform {
   flushBuffer() {
     const lineContext: SplitLineContext = {
       ...this.context,
+      beginEndIdentLevel: 0,
       position: 0,
       currentCommandStart: 0,
       wasDataOnLine: false,

src/splitter.test.ts

@@ -6,6 +6,7 @@ import {
   noSplitSplitterOptions,
   redisSplitterOptions,
   oracleSplitterOptions,
+  sqliteSplitterOptions,
 } from './options';
 import { splitQuery } from './splitQuery';
 
@@ -33,14 +34,10 @@ test('correct split 2 queries - with escaped string', () => {
 });
 
 test('incorrect used escape', () => {
-  const output = splitQuery(
-    "query1\\",
-    mysqlSplitterOptions
-  );
-  expect(output).toEqual(["query1\\"]);
+  const output = splitQuery('query1\\', mysqlSplitterOptions);
+  expect(output).toEqual(['query1\\']);
 });
 
-
 test('delete empty query', () => {
   const output = splitQuery(';;;\n;;SELECT * FROM `table1`;;;;;SELECT * FROM `table2`;;; ;;;', mysqlSplitterOptions);
   expect(output).toEqual(['SELECT * FROM `table1`', 'SELECT * FROM `table2`']);
@@ -302,3 +299,75 @@ test('postgres copy from stdin', () => {
     ])
   );
 });
+
+test('sqlite skipSeparatorBeginEnd', () => {
+  const input = 'CREATE TRIGGER obj1 AFTER INSERT ON t1 FOR EACH ROW BEGIN SELECT * FROM t1; END';
+  const output = splitQuery(input, {
+    ...sqliteSplitterOptions,
+  });
+
+  expect(output.length).toBe(1);
+  expect(output[0]).toEqual(input);
+});
+
+test('sqlite skipSeparatorBeginEnd in lowercase', () => {
+  const input = 'create trigger obj1 after insert on t1 for each row begin select * from t1; end';
+  const output = splitQuery(input, {
+    ...sqliteSplitterOptions,
+  });
+
+  expect(output.length).toBe(1);
+  expect(output[0]).toEqual(input);
+});
+
+test('sqlite skipSeparatorBeginEnd in TRANSACTION', () => {
+  const input = `
+START TRANSACTION;
+
+UPDATE accounts
+SET balance = balance - 100
+WHERE account_id = 1;
+
+IF ROW_COUNT() > 0 THEN
+    COMMIT;
+ELSE
+    ROLLBACK;
+END IF;
+`;
+  const output = splitQuery(input, {
+    ...sqliteSplitterOptions,
+  });
+
+  expect(output.length).toBe(5);
+  expect(output[0]).toEqual('START TRANSACTION');
+  expect(output[1]).toEqual('UPDATE accounts\nSET balance = balance - 100\nWHERE account_id = 1');
+  expect(output[2]).toEqual('IF ROW_COUNT() > 0 THEN\n    COMMIT');
+  expect(output[3]).toEqual('ELSE\n    ROLLBACK');
+  expect(output[4]).toEqual('END IF');
+});
+
+test('sqlite skipSeparatorBeginEnd in TRANSACTION in lowercase', () => {
+  const input = `
+start transaction;
+
+update accounts
+set balance = balance - 100
+where account_id = 1;
+
+if row_count() > 0 then
+    commit;
+else
+    rollback;
+end if;
+`;
+  const output = splitQuery(input, {
+    ...sqliteSplitterOptions,
+  });
+
+  expect(output.length).toBe(5);
+  expect(output[0]).toEqual('start transaction');
+  expect(output[1]).toEqual('update accounts\nset balance = balance - 100\nwhere account_id = 1');
+  expect(output[2]).toEqual('if row_count() > 0 then\n    commit');
+  expect(output[3]).toEqual('else\n    rollback');
+  expect(output[4]).toEqual('end if');
+});