diff --git a/packages/api/src/shell/importDatabase.js b/packages/api/src/shell/importDatabase.js index 84d7b6e8c..d344527c9 100644 --- a/packages/api/src/shell/importDatabase.js +++ b/packages/api/src/shell/importDatabase.js @@ -9,14 +9,24 @@ const { getLogger } = require('dbgate-tools'); const logger = getLogger('importDb'); class ImportStream extends stream.Transform { - constructor(pool, driver) { + constructor(dbhan, driver) { super({ objectMode: true }); - this.pool = pool; + this.dbhan = dbhan; this.driver = driver; + this.writeQueryStream = null; } async _transform(chunk, encoding, cb) { try { - await this.driver.script(this.pool, chunk, { queryOptions: { importSqlDump: true } }); + if (chunk.specialMarker == 'copy_stdin_start') { + this.writeQueryStream = await this.driver.writeQueryFromStream(this.dbhan, chunk.text); + } else if (chunk.specialMarker == 'copy_stdin_data') { + this.writeQueryStream.write(chunk.text); + } else if (chunk.specialMarker == 'copy_stdin_end') { + this.writeQueryStream.end(); + this.writeQueryStream = null; + } else { + await this.driver.script(this.dbhan, chunk.text, { queryOptions: { importSqlDump: true } }); + } } catch (err) { this.emit('error', err.message); } @@ -44,7 +54,7 @@ async function importDatabase({ connection = undefined, systemConnection = undef logger.info(`Importing database`); if (!driver) driver = requireEngineDriver(connection); - const pool = systemConnection || (await connectUtility(driver, connection, 'write')); + const dbhan = systemConnection || (await connectUtility(driver, connection, 'write')); logger.info(`Connected.`); logger.info(`Input file: ${inputFile}`); @@ -52,8 +62,11 @@ async function importDatabase({ connection = undefined, systemConnection = undef logger.info(`Downloaded file: ${downloadedFile}`); const fileStream = fs.createReadStream(downloadedFile, 'utf-8'); - const splittedStream = splitQueryStream(fileStream, driver.getQuerySplitterOptions('script')); - const importStream = new ImportStream(pool, driver); + const splittedStream = splitQueryStream(fileStream, { + ...driver.getQuerySplitterOptions('import'), + returnRichInfo: true, + }); + const importStream = new ImportStream(dbhan, driver); // @ts-ignore splittedStream.pipe(importStream); await awaitStreamEnd(importStream); diff --git a/packages/types/engines.d.ts b/packages/types/engines.d.ts index 37ddc9682..e549a69ca 100644 --- a/packages/types/engines.d.ts +++ b/packages/types/engines.d.ts @@ -220,7 +220,7 @@ export interface EngineDriver extends FilterBehaviourProvider { getCollectionUpdateScript(changeSet: any, collectionInfo: CollectionInfo): string; createDatabase(dbhan: DatabaseHandle, name: string): Promise; dropDatabase(dbhan: DatabaseHandle, name: string): Promise; - getQuerySplitterOptions(usage: 'stream' | 'script' | 'editor'): any; + getQuerySplitterOptions(usage: 'stream' | 'script' | 'editor' | 'import'): any; script(dbhan: DatabaseHandle, sql: string, options?: RunScriptOptions): Promise; operation(dbhan: DatabaseHandle, operation: {}, options?: RunScriptOptions): Promise; getNewObjectTemplates(): NewObjectTemplate[]; diff --git a/plugins/dbgate-plugin-postgres/src/frontend/drivers.js b/plugins/dbgate-plugin-postgres/src/frontend/drivers.js index e9d7e5bc0..708b397be 100644 --- a/plugins/dbgate-plugin-postgres/src/frontend/drivers.js +++ b/plugins/dbgate-plugin-postgres/src/frontend/drivers.js @@ -124,6 +124,11 @@ const postgresDriverBase = { getQuerySplitterOptions: usage => usage == 'editor' ? { ...postgreSplitterOptions, ignoreComments: true, preventSingleLineSplit: true } + : usage == 'import' + ? { + ...postgreSplitterOptions, + copyFromStdin: true, + } : postgreSplitterOptions, readOnlySessions: true,