fixed postgres split query

packages/tools/src/index.ts

@@ -7,6 +7,5 @@ export * from './DatabaseAnalyser';
 export * from './driverBase';
 export * from './SqlDumper';
 export * from './testPermission';
-export * from './splitPostgresQuery';
 export * from './SqlGenerator';
 export * from './structureTools';

packages/tools/src/splitPostgresQuery.ts

@@ -1,292 +0,0 @@
-const SINGLE_QUOTE = "'";
-const DOUBLE_QUOTE = '"';
-// const BACKTICK = '`';
-const DOUBLE_DASH_COMMENT_START = '--';
-const HASH_COMMENT_START = '#';
-const C_STYLE_COMMENT_START = '/*';
-const SEMICOLON = ';';
-const LINE_FEED = '\n';
-const DELIMITER_KEYWORD = 'DELIMITER';
-
-export interface SplitOptions {
-  multipleStatements?: boolean;
-  retainComments?: boolean;
-}
-
-interface SqlStatement {
-  value: string;
-  supportMulti: boolean;
-}
-
-interface SplitExecutionContext extends Required<SplitOptions> {
-  unread: string;
-  currentDelimiter: string;
-  currentStatement: SqlStatement;
-  output: SqlStatement[];
-}
-
-interface FindExpResult {
-  expIndex: number;
-  exp: string | null;
-  nextIndex: number;
-}
-
-const regexEscapeSetRegex = /[-/\\^$*+?.()|[\]{}]/g;
-const singleQuoteStringEndRegex = /(?<!\\)'/;
-const doubleQuoteStringEndRegex = /(?<!\\)"/;
-// const backtickQuoteEndRegex = /(?<!`)`(?!`)/;
-const doubleDashCommentStartRegex = /--[ \f\n\r\t\v]/;
-const cStyleCommentStartRegex = /\/\*/;
-const cStyleCommentEndRegex = /(?<!\/)\*\//;
-const newLineRegex = /(?:[\r\n]+|$)/;
-const delimiterStartRegex = /(?:^|[\n\r]+)[ \f\t\v]*DELIMITER[ \t]+/i;
-// Best effort only, unable to find a syntax specification on delimiter
-const delimiterTokenRegex = /^(?:'(.+)'|"(.+)"|`(.+)`|([^\s]+))/;
-const semicolonKeyTokenRegex = buildKeyTokenRegex(SEMICOLON);
-const quoteEndRegexDict: Record<string, RegExp> = {
-  [SINGLE_QUOTE]: singleQuoteStringEndRegex,
-  [DOUBLE_QUOTE]: doubleQuoteStringEndRegex,
-  // [BACKTICK]: backtickQuoteEndRegex,
-};
-
-function escapeRegex(value: string): string {
-  return value.replace(regexEscapeSetRegex, '\\$&');
-}
-
-function buildKeyTokenRegex(delimiter: string): RegExp {
-  return new RegExp(
-    '(?:' +
-      [
-        escapeRegex(delimiter),
-        SINGLE_QUOTE,
-        DOUBLE_QUOTE,
-        // BACKTICK,
-        doubleDashCommentStartRegex.source,
-        HASH_COMMENT_START,
-        cStyleCommentStartRegex.source,
-        delimiterStartRegex.source,
-      ].join('|') +
-      ')',
-    'i'
-  );
-}
-
-function findExp(content: string, regex: RegExp): FindExpResult {
-  const match = content.match(regex);
-  let result: FindExpResult;
-  if (match?.index !== undefined) {
-    result = {
-      expIndex: match.index,
-      exp: match[0],
-      nextIndex: match.index + match[0].length,
-    };
-  } else {
-    result = {
-      expIndex: -1,
-      exp: null,
-      nextIndex: content.length,
-    };
-  }
-  return result;
-}
-
-function findKeyToken(content: string, currentDelimiter: string): FindExpResult {
-  let regex;
-  if (currentDelimiter === SEMICOLON) {
-    regex = semicolonKeyTokenRegex;
-  } else {
-    regex = buildKeyTokenRegex(currentDelimiter);
-  }
-  return findExp(content, regex);
-}
-
-function findEndQuote(content: string, quote: string): FindExpResult {
-  if (!(quote in quoteEndRegexDict)) {
-    throw new TypeError(`Incorrect quote ${quote} supplied`);
-  }
-  return findExp(content, quoteEndRegexDict[quote]);
-}
-
-function read(
-  context: SplitExecutionContext,
-  readToIndex: number,
-  nextUnreadIndex?: number,
-  checkSemicolon?: boolean
-): void {
-  if (checkSemicolon === undefined) {
-    checkSemicolon = true;
-  }
-  const readContent = context.unread.slice(0, readToIndex);
-  if (checkSemicolon && readContent.includes(SEMICOLON)) {
-    context.currentStatement.supportMulti = false;
-  }
-  context.currentStatement.value += readContent;
-  if (nextUnreadIndex !== undefined && nextUnreadIndex > 0) {
-    context.unread = context.unread.slice(nextUnreadIndex);
-  } else {
-    context.unread = context.unread.slice(readToIndex);
-  }
-}
-
-function readTillNewLine(context: SplitExecutionContext, checkSemicolon?: boolean): void {
-  const findResult = findExp(context.unread, newLineRegex);
-  read(context, findResult.expIndex, findResult.expIndex, checkSemicolon);
-}
-
-function discard(context: SplitExecutionContext, nextUnreadIndex: number): void {
-  if (nextUnreadIndex > 0) {
-    context.unread = context.unread.slice(nextUnreadIndex);
-  }
-}
-
-function discardTillNewLine(context: SplitExecutionContext): void {
-  const findResult = findExp(context.unread, newLineRegex);
-  discard(context, findResult.expIndex);
-}
-
-function publishStatementInMultiMode(splitOutput: SqlStatement[], currentStatement: SqlStatement): void {
-  if (splitOutput.length === 0) {
-    splitOutput.push({
-      value: '',
-      supportMulti: true,
-    });
-  }
-  const lastSplitResult = splitOutput[splitOutput.length - 1];
-  if (currentStatement.supportMulti) {
-    if (lastSplitResult.supportMulti) {
-      if (lastSplitResult.value !== '' && !lastSplitResult.value.endsWith(LINE_FEED)) {
-        lastSplitResult.value += LINE_FEED;
-      }
-      lastSplitResult.value += currentStatement.value + SEMICOLON;
-    } else {
-      splitOutput.push({
-        value: currentStatement.value + SEMICOLON,
-        supportMulti: true,
-      });
-    }
-  } else {
-    splitOutput.push({
-      value: currentStatement.value,
-      supportMulti: false,
-    });
-  }
-}
-
-function publishStatement(context: SplitExecutionContext): void {
-  const trimmed = context.currentStatement.value.trim();
-  if (trimmed !== '') {
-    if (!context.multipleStatements) {
-      context.output.push({
-        value: trimmed,
-        supportMulti: context.currentStatement.supportMulti,
-      });
-    } else {
-      context.currentStatement.value = trimmed;
-      publishStatementInMultiMode(context.output, context.currentStatement);
-    }
-  }
-  context.currentStatement.value = '';
-  context.currentStatement.supportMulti = true;
-}
-
-function handleKeyTokenFindResult(context: SplitExecutionContext, findResult: FindExpResult): void {
-  switch (findResult.exp?.trim()) {
-    case context.currentDelimiter:
-      read(context, findResult.expIndex, findResult.nextIndex);
-      publishStatement(context);
-      break;
-    // case BACKTICK:
-    case SINGLE_QUOTE:
-    case DOUBLE_QUOTE: {
-      read(context, findResult.nextIndex);
-      const findQuoteResult = findEndQuote(context.unread, findResult.exp);
-      read(context, findQuoteResult.nextIndex, undefined, false);
-      break;
-    }
-    case DOUBLE_DASH_COMMENT_START: {
-      if (context.retainComments) {
-        read(context, findResult.nextIndex);
-        readTillNewLine(context, false);
-      } else {
-        read(context, findResult.expIndex, findResult.expIndex + DOUBLE_DASH_COMMENT_START.length);
-        discardTillNewLine(context);
-      }
-      break;
-    }
-    case HASH_COMMENT_START: {
-      if (context.retainComments) {
-        read(context, findResult.nextIndex);
-        readTillNewLine(context, false);
-      } else {
-        read(context, findResult.expIndex, findResult.nextIndex);
-        discardTillNewLine(context);
-      }
-      break;
-    }
-    case C_STYLE_COMMENT_START: {
-      if (['!', '+'].includes(context.unread[findResult.nextIndex]) || context.retainComments) {
-        // Should not be skipped, see https://dev.mysql.com/doc/refman/5.7/en/comments.html
-        read(context, findResult.nextIndex);
-        const findCommentResult = findExp(context.unread, cStyleCommentEndRegex);
-        read(context, findCommentResult.nextIndex);
-      } else {
-        read(context, findResult.expIndex, findResult.nextIndex);
-        const findCommentResult = findExp(context.unread, cStyleCommentEndRegex);
-        discard(context, findCommentResult.nextIndex);
-      }
-      break;
-    }
-    case DELIMITER_KEYWORD: {
-      read(context, findResult.expIndex, findResult.nextIndex);
-      // MySQL client will return `DELIMITER cannot contain a backslash character` if backslash is used
-      // Shall we reject backslash as well?
-      const matched = context.unread.match(delimiterTokenRegex);
-      if (matched?.index !== undefined) {
-        context.currentDelimiter = matched[0].trim();
-        discard(context, matched[0].length);
-      }
-      discardTillNewLine(context);
-      break;
-    }
-    case undefined:
-    case null:
-      read(context, findResult.nextIndex);
-      publishStatement(context);
-      break;
-    default:
-      // This should never happen
-      throw new Error(`Unknown token '${findResult.exp ?? '(null)'}'`);
-  }
-}
-
-export function splitPostgresQuery(sql: string, options?: SplitOptions): string[] {
-  options = options ?? {};
-  const context: SplitExecutionContext = {
-    multipleStatements: options.multipleStatements ?? false,
-    retainComments: options.retainComments ?? false,
-    unread: sql,
-    currentDelimiter: SEMICOLON,
-    currentStatement: {
-      value: '',
-      supportMulti: true,
-    },
-    output: [],
-  };
-  let findResult: FindExpResult = {
-    expIndex: -1,
-    exp: null,
-    nextIndex: 0,
-  };
-  let lastUnreadLength;
-  do {
-    lastUnreadLength = context.unread.length;
-    findResult = findKeyToken(context.unread, context.currentDelimiter);
-    handleKeyTokenFindResult(context, findResult);
-    // Prevent infinite loop by returning incorrect result
-    if (lastUnreadLength === context.unread.length) {
-      read(context, context.unread.length);
-    }
-  } while (context.unread !== '');
-  publishStatement(context);
-  return context.output.map(v => v.value);
-}

plugins/dbgate-plugin-postgres/package.json

@@ -31,11 +31,12 @@
   },
   "devDependencies": {
     "dbgate-plugin-tools": "^1.0.7",
-    "webpack": "^4.42.0",
-    "webpack-cli": "^3.3.11",
     "dbgate-tools": "^4.1.1",
     "lodash": "^4.17.15",
     "pg": "^7.17.0",
-    "pg-query-stream": "^3.1.1"
+    "pg-query-stream": "^3.1.1",
+    "webpack": "^4.42.0",
+    "webpack-cli": "^3.3.11",
+    "sql-query-identifier": "^2.1.0"
   }
-}
+}

plugins/dbgate-plugin-postgres/src/backend/driver.js

@@ -1,10 +1,12 @@
 const _ = require('lodash');
 const stream = require('stream');
+const { identify } = require('sql-query-identifier');
+
 const driverBase = require('../frontend/driver');
 const Analyser = require('./Analyser');
 const pg = require('pg');
 const pgQueryStream = require('pg-query-stream');
-const { createBulkInsertStreamBase, splitPostgresQuery, makeUniqueColumnNames } = require('dbgate-tools');
+const { createBulkInsertStreamBase, makeUniqueColumnNames } = require('dbgate-tools');
 
 function extractPostgresColumns(result) {
   if (!result || !result.fields) return [];
@@ -119,10 +121,10 @@ const driver = {
     return { rows: res.rows.map(row => zipDataRow(row, columns)), columns };
   },
   async stream(client, sql, options) {
-    const sqlSplitted = splitPostgresQuery(sql);
+    const sqlSplitted = identify(sql, { dialect: 'psql' });
 
     for (const sqlItem of sqlSplitted) {
-      await runStreamItem(client, sqlItem, options);
+      await runStreamItem(client, sqlItem.text, options);
     }
 
     options.done();

yarn.lock

@@ -8215,6 +8215,11 @@ sql-formatter@^2.3.3:
   dependencies:
     lodash "^4.16.0"
 
+sql-query-identifier@^2.1.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/sql-query-identifier/-/sql-query-identifier-2.1.0.tgz#dbf0f34b11bc14c8ade44de13350271047eb566e"
+  integrity sha512-DcC+inWZvN6fiTyxv+9uhFoTRC9C8LTeApVl1N7JJTTCzto6yhuaI423DzPPqDk10z4naL2mF9g/eNhUfxuMpA==
+
 sqlstring@^2.3.2:
   version "2.3.2"
   resolved "https://registry.yarnpkg.com/sqlstring/-/sqlstring-2.3.2.tgz#cdae7169389a1375b18e885f2e60b3e460809514"