query splitter - mssql go delimiter

packages/query-splitter/src/splitQuery.ts

@@ -29,7 +29,7 @@ function isStringEnd(s: string, pos: number, endch: string, escapech: string) {
 }
 
 interface Token {
-  type: 'string' | 'delimiter' | 'whitespace' | 'eoln' | 'data' | 'set_delimiter' | 'comment';
+  type: 'string' | 'delimiter' | 'whitespace' | 'eoln' | 'data' | 'set_delimiter' | 'comment' | 'go_delimiter';
   length: number;
   value?: string;
 }
@@ -139,6 +139,16 @@ function scanToken(context: SplitExecutionContext): Token {
     }
   }
 
+  if (context.options.allowGoDelimiter && !context.wasDataOnLine) {
+    const m = s.slice(pos).match(/^GO[\t\r ]*(\n|$)/i);
+    if (m) {
+      return {
+        type: 'go_delimiter',
+        length: m[0].length - 1,
+      };
+    }
+  }
+
   const dollarString = scanDollarQuotedString(context);
   if (dollarString) return dollarString;
 
@@ -199,6 +209,11 @@ export function splitQuery(sql: string, options: SplitterOptions = null): string
         context.position += token.length;
         context.currentCommandStart = context.position;
         break;
+      case 'go_delimiter':
+        pushQuery(context);
+        context.position += token.length;
+        context.currentCommandStart = context.position;
+        break;
       case 'delimiter':
         pushQuery(context);
         context.position += token.length;

packages/query-splitter/src/splitter.test.ts

@@ -65,3 +65,9 @@ test('dollar string', () => {
   const output = splitQuery(input, postgreSplitterOptions);
   expect(output).toEqual(['CREATE PROC $$ SELECT 1; SELECT 2; $$', 'SELECT 3']);
 });
+
+test('go delimiter', () => {
+  const input = 'SELECT 1\ngo\nSELECT 2';
+  const output = splitQuery(input, mssqlSplitterOptions);
+  expect(output).toEqual(['SELECT 1', 'SELECT 2']);
+});