From 19a4d977650e48fa02a40d5c301e65e1a69818b6 Mon Sep 17 00:00:00 2001
From: "SPRINX0\\prochazka" <jan.prochazka@sprinx.com>
Date: Thu, 26 Sep 2024 15:48:29 +0200
Subject: [PATCH] postgres copystream support

---
 packages/tools/src/driverBase.ts                     |  4 ++++
 packages/types/engines.d.ts                          |  2 ++
 plugins/dbgate-plugin-postgres/package.json          |  3 +++
 .../dbgate-plugin-postgres/src/backend/drivers.js    |  6 ++++++
 yarn.lock                                            | 12 ++++++++++++
 5 files changed, 27 insertions(+)

diff --git a/packages/tools/src/driverBase.ts b/packages/tools/src/driverBase.ts
index f7c0bac18..c97f9d3ea 100644
--- a/packages/tools/src/driverBase.ts
+++ b/packages/tools/src/driverBase.ts
@@ -195,4 +195,8 @@ export const driverBase = {
   async listSchemas(pool) {
     return null;
   },
+
+  async writeQueryFromStream(dbhan, sql) {
+    return null;
+  },
 };
diff --git a/packages/types/engines.d.ts b/packages/types/engines.d.ts
index 93ed005a9..37ddc9682 100644
--- a/packages/types/engines.d.ts
+++ b/packages/types/engines.d.ts
@@ -188,6 +188,8 @@ export interface EngineDriver extends FilterBehaviourProvider {
   stream(dbhan: DatabaseHandle, sql: string, options: StreamOptions);
   readQuery(dbhan: DatabaseHandle, sql: string, structure?: TableInfo): Promise<stream.Readable>;
   readJsonQuery(dbhan: DatabaseHandle, query: any, structure?: TableInfo): Promise<stream.Readable>;
+  // eg. PostgreSQL COPY FROM stdin
+  writeQueryFromStream(dbhan: DatabaseHandle, sql: string): Promise<stream.Writable>;
   writeTable(dbhan: DatabaseHandle, name: NamedObjectInfo, options: WriteTableOptions): Promise<stream.Writable>;
   analyseSingleObject(
     dbhan: DatabaseHandle,
diff --git a/plugins/dbgate-plugin-postgres/package.json b/plugins/dbgate-plugin-postgres/package.json
index 69776b299..b8ad65873 100644
--- a/plugins/dbgate-plugin-postgres/package.json
+++ b/plugins/dbgate-plugin-postgres/package.json
@@ -37,5 +37,8 @@
     "pg": "^8.11.5",
     "webpack": "^5.91.0",
     "webpack-cli": "^5.1.4"
+  },
+  "dependencies": {
+    "pg-copy-streams": "^6.0.6"
   }
 }
diff --git a/plugins/dbgate-plugin-postgres/src/backend/drivers.js b/plugins/dbgate-plugin-postgres/src/backend/drivers.js
index eddea38a4..72d240009 100644
--- a/plugins/dbgate-plugin-postgres/src/backend/drivers.js
+++ b/plugins/dbgate-plugin-postgres/src/backend/drivers.js
@@ -4,6 +4,7 @@ const stream = require('stream');
 const driverBases = require('../frontend/drivers');
 const Analyser = require('./Analyser');
 const pg = require('pg');
+const pgCopyStreams = require('pg-copy-streams');
 const { getLogger, createBulkInsertStreamBase, makeUniqueColumnNames, extractDbNameFromComposite } =
   global.DBGATE_PACKAGES['dbgate-tools'];
 
@@ -292,6 +293,11 @@ const drivers = driverBases.map(driverBase => ({
 
     return schemas;
   },
+
+  writeQueryFromStream(dbhan, sql) {
+    const stream = dbhan.client.query(pgCopyStreams.from(sql));
+    return stream;
+  },
 }));
 
 module.exports = drivers;
diff --git a/yarn.lock b/yarn.lock
index 66c1a2057..8c7847cc9 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -7936,6 +7936,11 @@ object.pick@^1.3.0:
   dependencies:
     isobject "^3.0.1"
 
+obuf@^1.1.2:
+  version "1.1.2"
+  resolved "https://registry.yarnpkg.com/obuf/-/obuf-1.1.2.tgz#09bea3343d41859ebd446292d11c9d4db619084e"
+  integrity sha512-PX1wu0AmAdPqOL1mWhqmlOd8kOIZQwGZw6rh7uby9fTc5lhaOWFLX3I6R1hrF9k3zUY40e6igsLGkDXK92LJNg==
+
 on-exit-leak-free@^2.1.0:
   version "2.1.2"
   resolved "https://registry.yarnpkg.com/on-exit-leak-free/-/on-exit-leak-free-2.1.2.tgz#fed195c9ebddb7d9e4c3842f93f281ac8dadd3b8"
@@ -8235,6 +8240,13 @@ pg-connection-string@^2.6.4:
   resolved "https://registry.yarnpkg.com/pg-connection-string/-/pg-connection-string-2.6.4.tgz#f543862adfa49fa4e14bc8a8892d2a84d754246d"
   integrity sha512-v+Z7W/0EO707aNMaAEfiGnGL9sxxumwLl2fJvCQtMn9Fxsg+lPpPkdcyBSv/KFgpGdYkMfn+EI1Or2EHjpgLCA==
 
+pg-copy-streams@^6.0.6:
+  version "6.0.6"
+  resolved "https://registry.yarnpkg.com/pg-copy-streams/-/pg-copy-streams-6.0.6.tgz#a24dbcbd2999f4ecd4ab919f92a1b256656e4f3c"
+  integrity sha512-Z+Dd2C2NIDTsjyFKmc6a9QLlpM8tjpERx+43RSx0WmL7j3uNChERi3xSvZUL0hWJ1oRUn4S3fhyt3apdSrTyKQ==
+  dependencies:
+    obuf "^1.1.2"
+
 pg-int8@1.0.1:
   version "1.0.1"
   resolved "https://registry.yarnpkg.com/pg-int8/-/pg-int8-1.0.1.tgz#943bd463bf5b71b4170115f80f8efc9a0c0eb78c"