diff --git a/packages/query-splitter/src/splitQuery.ts b/packages/query-splitter/src/splitQuery.ts index 8e6e04974..cf593e007 100644 --- a/packages/query-splitter/src/splitQuery.ts +++ b/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; diff --git a/packages/query-splitter/src/splitter.test.ts b/packages/query-splitter/src/splitter.test.ts index fcba2acc9..422d8de50 100644 --- a/packages/query-splitter/src/splitter.test.ts +++ b/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']); +});