diff --git a/packages/query-splitter/src/splitQuery.ts b/packages/query-splitter/src/splitQuery.ts index 219dfdad2..8e6e04974 100644 --- a/packages/query-splitter/src/splitQuery.ts +++ b/packages/query-splitter/src/splitQuery.ts @@ -47,6 +47,30 @@ const DATA_TOKEN: Token = { length: 1, }; +function scanDollarQuotedString(context: SplitExecutionContext): Token { + if (!context.options.allowDollarDollarString) return null; + + let pos = context.position; + const s = context.source; + + const match = /^(\$[a-zA-Z0-9_]*\$)/.exec(s.slice(pos)); + if (!match) return null; + const label = match[1]; + pos += label.length; + + while (pos < context.end) { + if (s.slice(pos).startsWith(label)) { + return { + type: 'string', + length: pos + label.length - context.position, + }; + } + pos++; + } + + return null; +} + function scanToken(context: SplitExecutionContext): Token { let pos = context.position; const s = context.source; @@ -115,6 +139,9 @@ function scanToken(context: SplitExecutionContext): Token { } } + const dollarString = scanDollarQuotedString(context); + if (dollarString) return dollarString; + return DATA_TOKEN; } diff --git a/packages/query-splitter/src/splitter.test.ts b/packages/query-splitter/src/splitter.test.ts index 9198a9a41..fcba2acc9 100644 --- a/packages/query-splitter/src/splitter.test.ts +++ b/packages/query-splitter/src/splitter.test.ts @@ -1,4 +1,4 @@ -import { mysqlSplitterOptions, mssqlSplitterOptions } from './options'; +import { mysqlSplitterOptions, mssqlSplitterOptions, postgreSplitterOptions } from './options'; import { splitQuery } from './splitQuery'; test('simple query', () => { @@ -59,3 +59,9 @@ test('multi line comment test', () => { const output = splitQuery(input, mysqlSplitterOptions); expect(output).toEqual(['SELECT 1 /* comment1;comment2\ncomment3*/', 'SELECT 2']); }); + +test('dollar string', () => { + const input = 'CREATE PROC $$ SELECT 1; SELECT 2; $$ ; SELECT 3'; + const output = splitQuery(input, postgreSplitterOptions); + expect(output).toEqual(['CREATE PROC $$ SELECT 1; SELECT 2; $$', 'SELECT 3']); +});