query splitter can return positions and line and column numbers

packages/query-splitter/src/options.ts

@@ -11,6 +11,8 @@ export interface SplitterOptions {
   doubleDashComments: boolean;
   multilineComments: boolean;
   javaScriptComments: boolean;
+
+  returnRichInfo: boolean;
 }
 
 export const defaultSplitterOptions: SplitterOptions = {
@@ -27,6 +29,8 @@ export const defaultSplitterOptions: SplitterOptions = {
   doubleDashComments: true,
   multilineComments: true,
   javaScriptComments: false,
+
+  returnRichInfo: false,
 };
 
 export const mysqlSplitterOptions: SplitterOptions = {

packages/query-splitter/src/splitQuery.ts

@@ -5,8 +5,16 @@ const SEMICOLON = ';';
 export interface SplitStreamContext {
   options: SplitterOptions;
   currentDelimiter: string;
-  pushOutput: (sql: string) => void;
+  pushOutput: (item: SplitResultItem) => void;
   commandPart: string;
+
+  line: number;
+  column: number;
+  streamPosition: number;
+
+  commandStartPosition: number;
+  commandStartLine: number;
+  commandStartColumn: number;
 }
 
 export interface SplitLineContext extends SplitStreamContext {
@@ -22,6 +30,41 @@ export interface SplitLineContext extends SplitStreamContext {
   //   semicolonKeyTokenRegex: RegExp;
 }
 
+export type SplitResultItem =
+  | string
+  | {
+      text: string;
+      startPosition: number;
+      endPosition: number;
+      startLine: number;
+      startColumn: number;
+      endLine: number;
+      endColumn: number;
+    };
+
+function movePosition(context: SplitLineContext, count: number) {
+  if (context.options.returnRichInfo) {
+    let { source, position, line, column, streamPosition } = context;
+    while (count > 0) {
+      if (source[position] == '\n') {
+        line += 1;
+        column = 0;
+      } else {
+        column += 1;
+      }
+      position += 1;
+      streamPosition += 1;
+      count -= 1;
+    }
+    context.position = position;
+    context.streamPosition = streamPosition;
+    context.line = line;
+    context.column = column;
+  } else {
+    context.position += count;
+  }
+}
+
 function isStringEnd(s: string, pos: number, endch: string, escapech: string) {
   if (!escapech) {
     return s[pos] == endch;
@@ -163,7 +206,31 @@ function scanToken(context: SplitLineContext): Token {
 function pushQuery(context: SplitLineContext) {
   const sql = (context.commandPart || '') + context.source.slice(context.currentCommandStart, context.position);
   const trimmed = sql.trim();
-  if (trimmed) context.pushOutput(trimmed);
+  if (trimmed) {
+    if (context.options.returnRichInfo) {
+      context.pushOutput({
+        text: trimmed,
+
+        startPosition: context.commandStartPosition,
+        startLine: context.commandStartLine,
+        startColumn: context.commandStartColumn,
+
+        endPosition: context.streamPosition,
+        endLine: context.line,
+        endColumn: context.column,
+      });
+    } else {
+      context.pushOutput(trimmed);
+    }
+  }
+}
+
+function markStartCommand(context: SplitLineContext) {
+  if (context.options.returnRichInfo) {
+    context.commandStartPosition = context.streamPosition;
+    context.commandStartLine = context.line;
+    context.commandStartColumn = context.column;
+  }
 }
 
 export function splitQueryLine(context: SplitLineContext) {
@@ -171,47 +238,50 @@ export function splitQueryLine(context: SplitLineContext) {
     const token = scanToken(context);
     if (!token) {
       // nothing special, move forward
-      context.position += 1;
+      movePosition(context, 1);
       continue;
     }
     switch (token.type) {
       case 'string':
-        context.position += token.length;
+        movePosition(context, token.length);
         context.wasDataOnLine = true;
         break;
       case 'comment':
-        context.position += token.length;
+        movePosition(context, token.length);
         context.wasDataOnLine = true;
         break;
       case 'eoln':
-        context.position += token.length;
+        movePosition(context, token.length);
         context.wasDataOnLine = false;
         break;
       case 'data':
-        context.position += token.length;
+        movePosition(context, token.length);
         context.wasDataOnLine = true;
         break;
       case 'whitespace':
-        context.position += token.length;
+        movePosition(context, token.length);
         break;
       case 'set_delimiter':
         pushQuery(context);
         context.commandPart = '';
         context.currentDelimiter = token.value;
-        context.position += token.length;
+        movePosition(context, token.length);
         context.currentCommandStart = context.position;
+        markStartCommand(context);
         break;
       case 'go_delimiter':
         pushQuery(context);
         context.commandPart = '';
-        context.position += token.length;
+        movePosition(context, token.length);
         context.currentCommandStart = context.position;
+        markStartCommand(context);
         break;
       case 'delimiter':
         pushQuery(context);
         context.commandPart = '';
-        context.position += token.length;
+        movePosition(context, token.length);
         context.currentCommandStart = context.position;
+        markStartCommand(context);
         break;
     }
   }
@@ -224,13 +294,47 @@ export function splitQueryLine(context: SplitLineContext) {
 export function getInitialDelimiter(options: SplitterOptions) {
   return options?.allowSemicolon === false ? null : SEMICOLON;
 }
-export function splitQuery(sql: string, options: SplitterOptions = null): string[] {
+
+export function finishSplitStream(context: SplitStreamContext) {
+  const trimmed = context.commandPart.trim();
+  if (trimmed) {
+    if (context.options.returnRichInfo) {
+      context.pushOutput({
+        text: trimmed,
+        startPosition: context.commandStartPosition,
+        startLine: context.commandStartLine,
+        startColumn: context.commandStartColumn,
+        endPosition: context.streamPosition,
+        endLine: context.line,
+        endColumn: context.column,
+      });
+    } else {
+      context.pushOutput(trimmed);
+    }
+  }
+}
+
+export function splitQuery(sql: string, options: SplitterOptions = null): SplitResultItem[] {
   const usedOptions = {
     ...defaultSplitterOptions,
     ...options,
   };
 
   if (usedOptions.noSplit) {
+    if (usedOptions.returnRichInfo) {
+      const lines = sql.split('\n');
+      return [
+        {
+          text: sql,
+          startLine: 0,
+          startPosition: 0,
+          startColumn: 0,
+          endLine: lines.length,
+          endColumn: lines[lines.length - 1]?.length || 0,
+          endPosition: sql.length,
+        },
+      ];
+    }
     return [sql];
   }
 
@@ -240,7 +344,13 @@ export function splitQuery(sql: string, options: SplitterOptions = null): string
     end: sql.length,
     currentDelimiter: getInitialDelimiter(options),
     position: 0,
+    column: 0,
+    line: 0,
     currentCommandStart: 0,
+    commandStartLine: 0,
+    commandStartColumn: 0,
+    commandStartPosition: 0,
+    streamPosition: 0,
     pushOutput: cmd => output.push(cmd),
     wasDataOnLine: false,
     options: usedOptions,
@@ -248,9 +358,7 @@ export function splitQuery(sql: string, options: SplitterOptions = null): string
   };
 
   splitQueryLine(context);
-
+  finishSplitStream(context);
-  const trimmed = context.commandPart.trim();
-  if (trimmed) context.pushOutput(trimmed);
 
   return output;
 }

packages/query-splitter/src/splitQueryStream.ts

@@ -1,5 +1,11 @@
 import stream from 'stream';
-import { SplitStreamContext, getInitialDelimiter, SplitLineContext, splitQueryLine } from './splitQuery';
+import {
+  SplitStreamContext,
+  getInitialDelimiter,
+  SplitLineContext,
+  splitQueryLine,
+  finishSplitStream,
+} from './splitQuery';
 import { SplitterOptions } from './options';
 
 export class SplitQueryStream extends stream.Transform {
@@ -9,6 +15,12 @@ export class SplitQueryStream extends stream.Transform {
     super({ objectMode: true });
     this.context = {
       commandPart: '',
+      commandStartLine: 0,
+      commandStartColumn: 0,
+      commandStartPosition: 0,
+      streamPosition: 0,
+      line: 0,
+      column: 0,
       options,
       currentDelimiter: getInitialDelimiter(options),
       pushOutput: cmd => this.push(cmd),
@@ -28,8 +40,7 @@ export class SplitQueryStream extends stream.Transform {
     done();
   }
   _flush(done) {
-    const trimmed = this.context.commandPart;
+    finishSplitStream(this.context);
-    if (trimmed) this.push(trimmed);
     done();
   }
 }

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

@@ -89,3 +89,69 @@ test('split mongo', () => {
   const output = splitQuery(input, mongoSplitterOptions);
   expect(output).toEqual(['db.collection.insert({x:1})', 'db.collection.insert({y:2})']);
 });
+
+test('count lines', () => {
+  const output = splitQuery('SELECT * FROM `table1`;\nSELECT * FROM `table2`;', {
+    ...mysqlSplitterOptions,
+    returnRichInfo: true,
+  });
+  expect(output).toEqual(
+    expect.arrayContaining([
+      expect.objectContaining({
+        text: 'SELECT * FROM `table1`',
+
+        startPosition: 0,
+        startLine: 0,
+        startColumn: 0,
+
+        endPosition: 22,
+        endLine: 0,
+        endColumn: 22,
+      }),
+      expect.objectContaining({
+        text: 'SELECT * FROM `table2`',
+
+        startPosition: 23,
+        startLine: 0,
+        startColumn: 23,
+
+        endPosition: 46,
+        endLine: 1,
+        endColumn: 22,
+      }),
+    ])
+  );
+});
+
+test('count lines with flush', () => {
+  const output = splitQuery('SELECT * FROM `table1`;\nSELECT * FROM `table2`', {
+    ...mysqlSplitterOptions,
+    returnRichInfo: true,
+  });
+  expect(output).toEqual(
+    expect.arrayContaining([
+      expect.objectContaining({
+        text: 'SELECT * FROM `table1`',
+
+        startPosition: 0,
+        startLine: 0,
+        startColumn: 0,
+
+        endPosition: 22,
+        endLine: 0,
+        endColumn: 22,
+      }),
+      expect.objectContaining({
+        text: 'SELECT * FROM `table2`',
+
+        startPosition: 23,
+        startLine: 0,
+        startColumn: 23,
+
+        endPosition: 46,
+        endLine: 1,
+        endColumn: 22,
+      }),
+    ])
+  );
+});