diff --git a/packages/tools/src/index.ts b/packages/tools/src/index.ts index 956184ec4..7ab818b54 100644 --- a/packages/tools/src/index.ts +++ b/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'; diff --git a/packages/tools/src/splitPostgresQuery.ts b/packages/tools/src/splitPostgresQuery.ts deleted file mode 100644 index 2986def1c..000000000 --- a/packages/tools/src/splitPostgresQuery.ts +++ /dev/null @@ -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 { - unread: string; - currentDelimiter: string; - currentStatement: SqlStatement; - output: SqlStatement[]; -} - -interface FindExpResult { - expIndex: number; - exp: string | null; - nextIndex: number; -} - -const regexEscapeSetRegex = /[-/\\^$*+?.()|[\]{}]/g; -const singleQuoteStringEndRegex = /(? = { - [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); -} diff --git a/plugins/dbgate-plugin-postgres/package.json b/plugins/dbgate-plugin-postgres/package.json index a3c14e8a8..6fc72f3c8 100644 --- a/plugins/dbgate-plugin-postgres/package.json +++ b/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" } -} \ No newline at end of file +} diff --git a/plugins/dbgate-plugin-postgres/src/backend/driver.js b/plugins/dbgate-plugin-postgres/src/backend/driver.js index 0373c8dda..5caec35ac 100644 --- a/plugins/dbgate-plugin-postgres/src/backend/driver.js +++ b/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(); diff --git a/yarn.lock b/yarn.lock index 41aeb66b1..d4391016a 100644 --- a/yarn.lock +++ b/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"