Use pg-pool, add timeout to fetching total number of rows in content …

…tab, update to electron 15
Paxa · Nov 13, 2021 · 12b725c · 12b725c
1 parent 6cd7767
commit 12b725c
Show file tree

Hide file tree

Showing 26 changed files with 322 additions and 220 deletions.
diff --git a/app.js b/app.js
@@ -1,4 +1,4 @@
-var remote = require('electron').remote;
+var remote = require('@electron/remote');
 var events = require('events');
 var path = require('path');
 

diff --git a/app/connection.js b/app/connection.js
@@ -9,7 +9,7 @@ colors.enabled = true;
 var vsprintf = require("sprintf-js").vsprintf;
 var ConnectionString = require('connection-string').ConnectionString;
 var SshConnect = require('../lib/ssh_connect');
-
+var electronRemote = require('@electron/remote');
 
 // Change postgres' type parser to use moment instance instead of js Date instance
 // because momentjs object has timezone information as in original string (js Date object always use local timezone)
@@ -198,7 +198,9 @@ class Connection {
     logger.info('Connecting to', this.connectString);
 
     if (this.connection) {
-      this.connection.end();
+      if (!this.connection.ended) {
+        this.connection.end();
+      }
       delete this.connection; // = null;
     }
 
@@ -232,55 +234,73 @@ class Connection {
   }
 
   _startConnection(options, callback) {
-    return this.connection.connect().then(() => {
-      this.connection.on('notification', (msg) => {
-        this.notificationCallbacks.forEach((fn) => {
-          fn(msg);
+    this.connection.on('error', (e) => { this.onConnectionLost(e) });
+
+    this.connection.on('connect', client => {
+      if (!client.postbirdEventsInit) {
+        client.on('notification', (msg) => {
+          this.notificationCallbacks.forEach((fn) => {
+            fn(msg);
+          });
+          App.log("notification.recieved", msg);
         });
-        App.log("notification.recieved", msg);
-      });
 
-      this.connection.on('error', (e) => { this.onConnectionLost(e) });
+        client.postbirdEventsInit = true;
+      }
+    })
 
-      this.serverVersion().then(version => {
-        if (this.logging) {
-          console.log("Server version is", version);
-        }
-        if (this.startQuery) {
-          this.query(this.startQuery, (res, error) => {
-            if (error) {
-              var formattedError = new Error(`Start query: ${this.startQuery}\n\nError: ${error.message}${error.hint ? "\n----\n" + error.hint : ''}`);
-              callback && callback(false, formattedError);
-              Promise.reject(error);
-            } else {
-              callback && callback(true);
-              Promise.resolve(true);
-            }
-          });
-        } else {
-          callback && callback(true);
-          Promise.resolve(true);
-        }
-      });
+    return this.serverVersion().then(version => {
       App.log("connect.success", JSON.parse(JSON.stringify(options)));
+
+      if (this.logging) {
+        console.log("Server version is", version);
+      }
+
+      if (this.startQuery) {
+        this.query(this.startQuery, (res, error) => {
+          if (error) {
+            var formattedError = new Error(`Start query: ${this.startQuery}\n\nError: ${error.message}${error.hint ? "\n----\n" + error.hint : ''}`);
+            callback && callback(false, formattedError);
+            Promise.reject(error);
+          } else {
+            callback && callback(true);
+            Promise.resolve(true);
+          }
+        });
+      } else {
+        callback && callback(true);
+        Promise.resolve(true);
+      }
     }).catch(error => {
       callback && callback(false, error);
       App.log("connect.error", this, JSON.parse(JSON.stringify(options)), error);
     });
   }
 
   _initConnection(connectString) /*: pg.ClientExt */ {
+    delete this._serverVersion;
+    delete this._serverVersionFull;
+
     // @ts-ignore
     var clientConfig = Connection.parseConnectionString(connectString) /*:: as pg.ClientConfig */;
-    // clientConfig.connectionTimeoutMillis = 10000; // 10 sec
-    return new pg.Client(clientConfig) /*:: as pg.ClientExt */;
+
+    clientConfig.idleTimeoutMillis = 600000;
+    clientConfig.max = clientConfig.max || 5;
+    clientConfig.log = (a, b, c) => { console.log('Pool:', a, b, c); }
+    var pool = new pg.Pool(clientConfig);
+
+    // pool.on('remove', client => {
+    //   console.log('pool removed');
+    // })
+
+    return pool;
   }
 
   onConnectionLost(error) {
     if (this.onDisconnect) {
       this.onDisconnect(error);
     } else {
-      var dialog = electron.remote.dialog;
+      var dialog = electronRemote.dialog;
       var message = error.message.replace(/\n\s+/g, "\n") + "\nTo re-open connection, use File -> Reconnect";
       dialog.showErrorBox("Server Connection Error", message);
     }
@@ -306,7 +326,7 @@ class Connection {
     var startStack = new Error().stack;
 
     return new Promise((resolve, reject) => {
-      this.connection.query(options, (error, result) => {
+      var queryCallback = (error, result) => {
         historyRecord.time = Date.now() - time;
         if (this.logging) logger.print("SQL:" + colors.green(" Done ") + historyRecord.time + "\n");
 
@@ -316,19 +336,22 @@ class Connection {
         }
 
         if (error) {
+          var customTimeoutError = options.query_timeout && error.message == 'Query read timeout';
           error.stack = error.stack + "\n" + startStack.substring(startStack.indexOf("\n") + 1);
           historyRecord.error = error;
           historyRecord.state = 'failed';
           App.log("sql.failed", historyRecord);
           // @ts-ignore
           error.query = sql;
-          if (this.logging) {
+          if (this.logging && !customTimeoutError) {
             console.error("SQL failed", sql);
             console.error(error);
           }
           //if (query) query.state = 'error';
           if (callback) callback(result, error);
-          this.onConnectionError(error);
+          if (!customTimeoutError) {
+            this.onConnectionError(error);
+          }
           reject(error);
         } else {
           historyRecord.state = 'success';
@@ -339,7 +362,9 @@ class Connection {
           if (callback) callback(result);
           resolve(result);
         }
-      });
+      };
+
+      this.connection.query(options, queryCallback);
     });
   }
 
@@ -381,7 +406,7 @@ class Connection {
       return Promise.resolve(this._serverVersion);
     }
 
-    console.log("Client don't support serverVersion, getting it with sql");
+    // console.log("Client don't support serverVersion, getting it with sql");
 
     return this.server.fetchServerVersion().then(version => {
       if (version.match(/CockroachDB/i)) {
@@ -617,10 +642,11 @@ class Connection {
     });
   }
 
-  close(callback /*:: ?: Function */) {
+  async close(callback /*:: ?: Function */) {
     if (this.connection) {
-      this.connection.end();
+      await this.connection.end()
     }
+
     if (this.sshConnection) {
       this.sshConnection.disconnect();
     }
@@ -629,12 +655,13 @@ class Connection {
     if (index != -1) {
       global.Connection.instances.splice(index, 1);
     }
+
     callback && callback();
   }
 
   reconnect(callback /*: (success: boolean, error?: Error) => void */) {
-    this.close(() => {
-      this.connectToServer(this.options, callback);
+    return this.close().then(() => {
+      return this.connectToServer(this.options, callback);
     });
   }
 
@@ -691,27 +718,16 @@ class Connection {
         console.log('no running query');
       }
     } else {
-      query = this.connection.activeQuery;
-      if (query) {
-        var otherConn = this._initConnection(this.connectString);
-        otherConn.connect((error) => {
-          if (error) {
-            console.log(error);
-            return;
-          }
-
-          console.log("Stopping query via sql. PID:", this.connection.processID);
-          var sql = `select pg_cancel_backend(${this.connection.processID})`;
-          otherConn.query(sql).then(() => {
-            otherConn.end();
+      this.connection._clients.forEach(poolClient => {
+        if (poolClient.activeQuery) {
+          var sql = `select pg_cancel_backend(${poolClient.processID})`;
+          this.connection.query(sql).then(res => {
+            console.log('query canceled');
           }).catch(err => {
             console.error(err);
-            otherConn.end();
           });
-        });
-      } else {
-        console.log('no running query');
-      }
+        }
+      })
     }
   }
 }

diff --git a/app/controllers/updates_controller.js b/app/controllers/updates_controller.js
@@ -1,4 +1,4 @@
-var remote = require('electron').remote;
+var remote = require('@electron/remote');
 var semver = require('semver');
 var needle = require('needle');
 var strftime = require('strftime');

diff --git a/app/db_screen.js b/app/db_screen.js
@@ -1,3 +1,5 @@
+var electronRemote = require('@electron/remote');
+
 class DbScreen {
   /*::
   type: string
@@ -39,13 +41,18 @@ class DbScreen {
 
     this.currentTab = null;
 
-    if (this.options.fetchDbList) this.fetchDbList();
-
     this.database = this.connection.options.database;
     App.emit('database.changed', this.database);
-    this.fetchTablesAndSchemas(() => {
-      this.view.showDatabaseContent();
-    });
+
+    (async () => {
+      if (this.options.fetchDbList) {
+        await this.fetchDbList();
+      }
+
+      await this.fetchTablesAndSchemas(() => {
+        this.view.showDatabaseContent();
+      });
+    })();
 
     this.connection.onNotification((message) => {
       window.alertify.alert("Recieve Message:<br>" + JSON.stringify(message));
@@ -722,7 +729,7 @@ class DbScreen {
   }
 
   async showConnectionLostDialog () {
-    var dialog = electron.remote.dialog;
+    var dialog = electronRemote.dialog;
     var message = this.connectionLostError.message.replace(/\n\s+/g, "\n");
     var res = await dialog.showMessageBox({
       type: "error",

diff --git a/app/help_screen.js b/app/help_screen.js
@@ -1,3 +1,5 @@
+var electronRemote = require('@electron/remote');
+
 class HelpScreen {
   /*::
   type: string
@@ -18,7 +20,7 @@ class HelpScreen {
     this.content.find('.page a.external').bind('click', (e) => {
       $u.stopEvent(e);
       var url = e.target.href;
-      electron.remote.shell.openExternal(url);
+      electronRemote.shell.openExternal(url);
     });
 
     new PgDumpRunner().version().then(version => {

diff --git a/app/logger.js b/app/logger.js
@@ -1,7 +1,5 @@
 var sprintf = require("sprintf-js").sprintf;
 var util = require('util');
-var remote = require('electron').remote;
-var topProcess = remote ? remote.process : process;
 
 var LOG_LEVELS = ['emergency', 'alert', 'critical', 'error', 'warning', 'notice', 'info', 'debug'];
 
@@ -11,6 +9,9 @@ class Logger {
   */
 
   constructor (level) {
+    var isRenderer = require('is-electron-renderer');
+    var remote = isRenderer ? require('@electron/remote') : require("@electron/remote/main");
+    this.topProcess = remote ? remote.process : process;
     this.logLevel = level;
   }
 
@@ -19,7 +20,7 @@ class Logger {
   }
 
   print (string) {
-    topProcess.stdout.write(string);
+    this.topProcess.stdout.write(string);
   }
 
   emergency(...messageArgs) { this.write('emergency', messageArgs); }

diff --git a/app/login_components/heroku_client.js b/app/login_components/heroku_client.js
@@ -5,6 +5,7 @@ var url = require('url');
 var querystring = require('querystring');
 var needle = require('needle');
 var http = require('http');
+var electronRemote = require('@electron/remote');
 
 /*::
 interface HerokuClientOptions {
@@ -122,12 +123,12 @@ class HerokuClient {
       if (this.catcher) this.catcher.stop();
       this.catcher = new HerokuCatcher(requestToken => {
         this.setRequestToken(requestToken);
-        electron.remote.app.mainWindow.focus();
+        electronRemote.app.mainWindow.focus();
         callback();
       });
       this.catcher.start();
       console.log("Opening url " + url);
-      electron.remote.shell.openExternal(url);
+      electronRemote.shell.openExternal(url);
       //child_process.spawn('open', [url]);
     }
   }

diff --git a/app/models/table.js b/app/models/table.js
@@ -1,5 +1,6 @@
 var SqlExporter = require('../../lib/sql_exporter');
 var pgEscape = require('pg-escape');
+var pg = require('pg');
 
 var TABLE_TYPES = {
   VIEW: 'VIEW',
@@ -505,9 +506,19 @@ class Table extends ModelBase {
   }
 
   async getTotalRows () /*: Promise<number> */ {
-    var sql = `SELECT count(*) AS rows_count FROM ${this.sqlTable()}`;
-    var data = await this.q(sql);
-    return parseInt(data.rows[0].rows_count, 10);
+    return new Promise ((resolve, reject) => {
+      var sql = `SELECT count(*) AS rows_count FROM ${this.sqlTable()}`;
+      var queryOptions = {
+        query_timeout: 5000,
+      };
+      this.connection().queryWithOptions(sql, queryOptions, (data, error) => {
+        if (error) {
+          reject(error)
+        } else {
+          resolve(parseInt(data.rows[0].rows_count, 10));
+        }
+      }).catch(err => {});
+    });
   }
 
   insertRow (values) {
@@ -560,7 +571,7 @@ class Table extends ModelBase {
     var tableType = await this.getTableType();
 
     if (tableType == TABLE_TYPES.VIEW) {
-      var sql = `SELECT pg_get_viewdef('${this.table}', true) AS view_def`;
+      var sql = `SELECT pg_get_viewdef('${this.sqlTable()}', true) AS view_def`;
       return new Promise((resolve, reject) => {
         this.q(sql).then(res => {
           var viewSource = res.rows[0] && res.rows[0].view_def;