query splitter - postgre dollar string

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;
 }
 

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']);
+});