query splitter refactor

plugins/dbgate-plugin-mongo/package.json

@@ -32,6 +32,7 @@
   },
   "devDependencies": {
     "dbgate-plugin-tools": "^1.0.7",
+    "dbgate-query-splitter": "^4.1.1",
     "webpack": "^4.42.0",
     "webpack-cli": "^3.3.11",
     "dbgate-tools": "^4.1.1",

plugins/dbgate-plugin-mongo/src/frontend/driver.js

@@ -1,5 +1,6 @@
 const { driverBase } = global.DBGATE_TOOLS;
 const Dumper = require('./Dumper');
+const { noSplitSplitterOptions } = require('dbgate-query-splitter/lib/options');
 
 const mongoIdRegex = /^[0-9a-f]{24}$/;
 
@@ -34,6 +35,8 @@ const driver = {
   supportsDatabaseUrl: true,
   databaseUrlPlaceholder: 'e.g. mongodb://username:password@mongodb.mydomain.net/dbname',
 
+  getQuerySplitterOptions: () => noSplitSplitterOptions,
+
   showConnectionField: (field, values) => {
     if (field == 'useDatabaseUrl') return true;
     if (values.useDatabaseUrl) {

plugins/dbgate-plugin-mssql/package.json

@@ -32,6 +32,7 @@
   },
   "devDependencies": {
     "dbgate-plugin-tools": "^1.0.7",
+    "dbgate-query-splitter": "^4.1.1",
     "webpack": "^4.42.0",
     "webpack-cli": "^3.3.11",
     "dbgate-tools": "^4.1.1",

plugins/dbgate-plugin-mssql/src/frontend/driver.js

@@ -1,5 +1,6 @@
 const { driverBase } = global.DBGATE_TOOLS;
 const MsSqlDumper = require('./MsSqlDumper');
+const { mssqlSplitterOptions } = require('dbgate-query-splitter/lib/options');
 
 /** @type {import('dbgate-types').SqlDialect} */
 const dialect = {
@@ -34,6 +35,7 @@ const driver = {
   },
   showConnectionField: (field, values) =>
     ['authType', 'server', 'port', 'user', 'password', 'defaultDatabase', 'singleDatabase'].includes(field),
+  getQuerySplitterOptions: () => mssqlSplitterOptions,
 
   engine: 'mssql@dbgate-plugin-mssql',
   title: 'Microsoft SQL Server',

plugins/dbgate-plugin-mysql/package.json

@@ -32,6 +32,7 @@
   },
   "devDependencies": {
     "dbgate-plugin-tools": "^1.0.7",
+    "dbgate-query-splitter": "^4.1.1",
     "webpack": "^4.42.0",
     "webpack-cli": "^3.3.11",
     "@verycrazydog/mysql-parser": "^1.2.0",

plugins/dbgate-plugin-mysql/src/backend/drivers.js

@@ -4,7 +4,6 @@ const driverBases = require('../frontend/drivers');
 const Analyser = require('./Analyser');
 const mysql2 = require('mysql2');
 const { createBulkInsertStreamBase, makeUniqueColumnNames } = require('dbgate-tools');
-const mysqlSplitter = require('@verycrazydog/mysql-parser');
 
 function extractColumns(fields) {
   if (fields) {
@@ -24,18 +23,49 @@ function zipDataRow(rowArray, columns) {
   );
 }
 
-async function runQueryItem(connection, sql) {
-  return new Promise((resolve, reject) => {
-    connection.query(sql, function (error, results, fields) {
-      if (error) reject(error);
-      const columns = extractColumns(fields);
-      resolve({ rows: results && columns && results.map && results.map(row => zipDataRow(row, columns)), columns });
-    });
-  });
-}
+/** @type {import('dbgate-types').EngineDriver} */
+const drivers = driverBases.map(driverBase => ({
+  ...driverBase,
+  analyserClass: Analyser,
 
-async function runStreamItem(connection, sql, options) {
-  return new Promise((resolve, reject) => {
+  async connect({ server, port, user, password, database, ssl }) {
+    const connection = mysql2.createConnection({
+      host: server,
+      port,
+      user,
+      password,
+      database,
+      ssl,
+      rowsAsArray: true,
+      supportBigNumbers: true,
+      bigNumberStrings: true,
+      // TODO: test following options
+      // multipleStatements: true,
+      // dateStrings: true,
+    });
+    connection._database_name = database;
+    return connection;
+  },
+  async close(pool) {
+    return pool.close();
+  },
+  query(connection, sql) {
+    if (sql == null) {
+      return {
+        rows: [],
+        columns: [],
+      };
+    }
+
+    return new Promise((resolve, reject) => {
+      connection.query(sql, function (error, results, fields) {
+        if (error) reject(error);
+        const columns = extractColumns(fields);
+        resolve({ rows: results && columns && results.map && results.map(row => zipDataRow(row, columns)), columns });
+      });
+    });
+  },
+  async stream(connection, sql, options) {
     const query = connection.query(sql);
     let columns = [];
 
@@ -51,7 +81,7 @@ async function runStreamItem(connection, sql, options) {
     // };
 
     const handleEnd = () => {
-      resolve();
+      options.done();
     };
 
     const handleRow = row => {
@@ -86,64 +116,6 @@ async function runStreamItem(connection, sql, options) {
     };
 
     query.on('error', handleError).on('fields', handleFields).on('result', handleRow).on('end', handleEnd);
-  });
-}
-
-/** @type {import('dbgate-types').EngineDriver} */
-const drivers = driverBases.map(driverBase => ({
-  ...driverBase,
-  analyserClass: Analyser,
-
-  async connect({ server, port, user, password, database, ssl }) {
-    const connection = mysql2.createConnection({
-      host: server,
-      port,
-      user,
-      password,
-      database,
-      ssl,
-      rowsAsArray: true,
-      supportBigNumbers: true,
-      bigNumberStrings: true,
-      // TODO: test following options
-      // multipleStatements: true,
-      // dateStrings: true,
-    });
-    connection._database_name = database;
-    return connection;
-  },
-  async close(pool) {
-    return pool.close();
-  },
-  async query(connection, sql) {
-    if (sql == null) {
-      return {
-        rows: [],
-        columns: [],
-      };
-    }
-
-    const sqlSplitted = mysqlSplitter.split(sql);
-    let res = {
-      rows: [],
-      columns: [],
-    };
-    for (const sqlItem of sqlSplitted) {
-      const resultItem = await runQueryItem(connection, sqlItem);
-      if (resultItem.rows && resultItem.columns && resultItem.columns.length > 0) {
-        res = resultItem;
-      }
-    }
-    return res;
-  },
-  async stream(connection, sql, options) {
-    const sqlSplitted = mysqlSplitter.split(sql);
-
-    for (const sqlItem of sqlSplitted) {
-      await runStreamItem(connection, sqlItem, options);
-    }
-
-    options.done();
   },
   async readQuery(connection, sql, structure) {
     const query = connection.query(sql);

plugins/dbgate-plugin-mysql/src/frontend/drivers.js

@@ -1,4 +1,5 @@
 const { driverBase } = global.DBGATE_TOOLS;
+const { mysqlSplitterOptions } = require('dbgate-query-splitter/lib/options');
 const Dumper = require('./Dumper');
 
 /** @type {import('dbgate-types').SqlDialect} */
@@ -21,6 +22,7 @@ const mysqlDriverBase = {
   dumperClass: Dumper,
   dialect,
   defaultPort: 3306,
+  getQuerySplitterOptions: () => mysqlSplitterOptions,
 };
 
 /** @type {import('dbgate-types').EngineDriver} */

plugins/dbgate-plugin-postgres/package.json

@@ -31,6 +31,7 @@
   },
   "devDependencies": {
     "dbgate-plugin-tools": "^1.0.7",
+    "dbgate-query-splitter": "^4.1.1",
     "dbgate-tools": "^4.1.1",
     "lodash": "^4.17.21",
     "pg": "^7.17.0",

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

@@ -23,56 +23,6 @@ function zipDataRow(rowArray, columns) {
   );
 }
 
-async function runStreamItem(client, sql, options) {
-  return new Promise((resolve, reject) => {
-    const query = new pg.Query({
-      text: sql,
-      rowMode: 'array',
-    });
-
-    let wasHeader = false;
-
-    query.on('row', row => {
-      if (!wasHeader) {
-        columns = extractPostgresColumns(query._result);
-        if (columns && columns.length > 0) {
-          options.recordset(columns);
-        }
-        wasHeader = true;
-      }
-
-      options.row(zipDataRow(row, columns));
-    });
-
-    query.on('end', () => {
-      if (!wasHeader) {
-        columns = extractPostgresColumns(query._result);
-        if (columns && columns.length > 0) {
-          options.recordset(columns);
-        }
-        wasHeader = true;
-      }
-
-      resolve();
-    });
-
-    query.on('error', error => {
-      console.log('ERROR', error);
-      const { message, lineNumber, procName } = error;
-      options.info({
-        message,
-        line: lineNumber,
-        procedure: procName,
-        time: new Date(),
-        severity: 'error',
-      });
-      resolve();
-    });
-
-    client.query(query);
-  });
-}
-
 /** @type {import('dbgate-types').EngineDriver} */
 const drivers = driverBases.map(driverBase => ({
   ...driverBase,
@@ -127,21 +77,52 @@ const drivers = driverBases.map(driverBase => ({
     const columns = extractPostgresColumns(res);
     return { rows: res.rows.map(row => zipDataRow(row, columns)), columns };
   },
-  async stream(client, sql, options) {
-    let sqlSplitted;
-    try {
-      sqlSplitted = identify(sql, { dialect: 'psql', strict: false });
-    } catch (e) {
-      // workaround
-      sqlSplitted = [{ text: sql }];
-    }
+  stream(client, sql, options) {
+    const query = new pg.Query({
+      text: sql,
+      rowMode: 'array',
+    });
 
-    for (const sqlItem of sqlSplitted) {
-      await runStreamItem(client, sqlItem.text, options);
-    }
+    let wasHeader = false;
 
-    options.done();
-    // return stream;
+    query.on('row', row => {
+      if (!wasHeader) {
+        columns = extractPostgresColumns(query._result);
+        if (columns && columns.length > 0) {
+          options.recordset(columns);
+        }
+        wasHeader = true;
+      }
+
+      options.row(zipDataRow(row, columns));
+    });
+
+    query.on('end', () => {
+      if (!wasHeader) {
+        columns = extractPostgresColumns(query._result);
+        if (columns && columns.length > 0) {
+          options.recordset(columns);
+        }
+        wasHeader = true;
+      }
+
+      options.done();
+    });
+
+    query.on('error', error => {
+      console.log('ERROR', error);
+      const { message, lineNumber, procName } = error;
+      options.info({
+        message,
+        line: lineNumber,
+        procedure: procName,
+        time: new Date(),
+        severity: 'error',
+      });
+      options.done();
+    });
+
+    client.query(query);
   },
   async getVersion(client) {
     const { rows } = await this.query(client, 'SELECT version()');

plugins/dbgate-plugin-postgres/src/frontend/drivers.js

@@ -1,5 +1,6 @@
 const { driverBase } = global.DBGATE_TOOLS;
 const Dumper = require('./Dumper');
+const { postgreSplitterOptions } = require('dbgate-query-splitter/lib/options');
 
 /** @type {import('dbgate-types').SqlDialect} */
 const dialect = {
@@ -21,6 +22,7 @@ const postgresDriverBase = {
   dialect,
   showConnectionField: (field, values) =>
     ['server', 'port', 'user', 'password', 'defaultDatabase', 'singleDatabase'].includes(field),
+  getQuerySplitterOptions: () => postgreSplitterOptions,
 };
 
 /** @type {import('dbgate-types').EngineDriver} */

plugins/dbgate-plugin-sqlite/package.json

@@ -32,6 +32,7 @@
     "devDependencies": {
       "dbgate-tools": "^4.1.1",
       "dbgate-plugin-tools": "^1.0.4",
+      "dbgate-query-splitter": "^4.1.1",
       "byline": "^5.0.0",
       "sql-query-identifier": "^2.1.0",
       "webpack": "^4.42.0",

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

@@ -2,7 +2,7 @@ const _ = require('lodash');
 const stream = require('stream');
 const driverBase = require('../frontend/driver');
 const Analyser = require('./Analyser');
-const { identify } = require('sql-query-identifier');
+const { splitQuery, sqliteSplitterOptions } = require('dbgate-query-splitter');
 const { createBulkInsertStreamBase, makeUniqueColumnNames } = require('dbgate-tools');
 
 let Database;
@@ -82,13 +82,13 @@ const driver = {
     }
   },
   async stream(client, sql, options) {
-    const sqlSplitted = identify(sql, { dialect: 'sqlite', strict: false });
+    const sqlSplitted = splitQuery(sql, sqliteSplitterOptions);
 
     const rowCounter = { count: 0, date: null };
 
     const inTransaction = client.transaction(() => {
       for (const sqlItem of sqlSplitted) {
-        runStreamItem(client, sqlItem.text, options, rowCounter);
+        runStreamItem(client, sqlItem, options, rowCounter);
       }
 
       if (rowCounter.date) {
@@ -117,6 +117,16 @@ const driver = {
     options.done();
     // return stream;
   },
+  async script(client, sql) {
+    const inTransaction = client.transaction(() => {
+      for (const sqlItem of splitQuery(sql, this.getQuerySplitterOptions('script'))) {
+        const stmt = client.prepare(sqlItem);
+        stmt.run();
+      }
+    });
+    inTransaction();
+  },
+
   async readQueryTask(stmt, pass) {
     // let sent = 0;
     for (const row of stmt.iterate()) {

plugins/dbgate-plugin-sqlite/src/frontend/driver.js

@@ -1,5 +1,6 @@
 const { driverBase } = global.DBGATE_TOOLS;
 const Dumper = require('./Dumper');
+const { sqliteSplitterOptions, noSplitSplitterOptions } = require('dbgate-query-splitter/lib/options');
 
 function getDatabaseFileLabel(databaseFile) {
   if (!databaseFile) return databaseFile;
@@ -34,6 +35,7 @@ const driver = {
     singleDatabase: true,
     defaultDatabase: getDatabaseFileLabel(connection.databaseFile),
   }),
+  getQuerySplitterOptions: (usage) => (usage == 'stream' ? noSplitSplitterOptions : sqliteSplitterOptions),
   // isFileDatabase: true,
   isElectronOnly: true,
 };