Merge pull request #6007 from Budibase/fix/postgres-connection

Postgres connection - fixing sporadic issues

packages/backend-core/package.json

@@ -41,8 +41,7 @@
   "devDependencies": {
     "ioredis-mock": "^5.5.5",
     "jest": "^26.6.3",
-    "pouchdb-adapter-memory": "^7.2.2",
-    "pouchdb-all-dbs": "^1.0.2"
+    "pouchdb-adapter-memory": "^7.2.2"
   },
   "gitHead": "d1836a898cab3f8ab80ee6d8f42be1a9eed7dcdc"
 }

packages/backend-core/src/db/index.js

@@ -3,13 +3,13 @@ const env = require("../environment")
 
 let PouchDB
 let initialised = false
+const dbList = new Set()
 
 const put =
   dbPut =>
   async (doc, options = {}) => {
-    const response = await dbPut(doc, options)
     // TODO: add created / updated
-    return response
+    return await dbPut(doc, options)
   }
 
 const checkInitialised = () => {
@@ -28,6 +28,9 @@ exports.init = opts => {
 // in situations that using the function doWithDB does not work
 exports.dangerousGetDB = (dbName, opts) => {
   checkInitialised()
+  if (env.isTest()) {
+    dbList.add(dbName)
+  }
   const db = new PouchDB(dbName, opts)
   const dbPut = db.put
   db.put = put(dbPut)
@@ -63,6 +66,9 @@ exports.doWithDB = async (dbName, cb, opts) => {
 }
 
 exports.allDbs = () => {
+  if (!env.isTest()) {
+    throw new Error("Cannot be used outside test environment.")
+  }
   checkInitialised()
-  return PouchDB.allDbs()
+  return [...dbList]
 }

packages/backend-core/src/db/pouch.js

@@ -92,11 +92,5 @@ exports.getPouch = (opts = {}) => {
     PouchDB.plugin(find)
   }
 
-  const Pouch = PouchDB.defaults(POUCH_DB_DEFAULTS)
-  if (opts.allDbs) {
-    const allDbs = require("pouchdb-all-dbs")
-    allDbs(Pouch)
-  }
-
-  return Pouch
+  return PouchDB.defaults(POUCH_DB_DEFAULTS)
 }

packages/server/__mocks__/pg.ts

@@ -14,7 +14,9 @@ module PgMock {
   function Client() {}
 
   Client.prototype.query = query
-  Client.prototype.end = jest.fn()
+  Client.prototype.end = jest.fn(cb => {
+    if (cb) cb()
+  })
   Client.prototype.connect = jest.fn()
   Client.prototype.release = jest.fn()
 

packages/server/src/integrations/postgres.ts

@@ -136,7 +136,7 @@ module PostgresModule {
           : undefined,
       }
       this.client = new Client(newConfig)
-      this.setSchema()
+      this.open = false
     }
 
     getBindingIdentifier(): string {
@@ -147,7 +147,34 @@ module PostgresModule {
       return parts.join(" || ")
     }
 
+    async openConnection() {
+      await this.client.connect()
+      if (!this.config.schema) {
+        this.config.schema = "public"
+      }
+      this.client.query(`SET search_path TO ${this.config.schema}`)
+      this.COLUMNS_SQL = `select * from information_schema.columns where table_schema = '${this.config.schema}'`
+      this.open = true
+    }
+
+    closeConnection() {
+      const pg = this
+      return new Promise<void>((resolve, reject) => {
+        this.client.end((err: any) => {
+          pg.open = false
+          if (err) {
+            reject(err)
+          } else {
+            resolve()
+          }
+        })
+      })
+    }
+
     async internalQuery(query: SqlQuery, close: boolean = true) {
+      if (!this.open) {
+        await this.openConnection()
+      }
       const client = this.client
       this.index = 1
       // need to handle a specific issue with json data types in postgres,
@@ -164,21 +191,14 @@ module PostgresModule {
       try {
         return await client.query(query.sql, query.bindings || [])
       } catch (err) {
-        await this.client.end()
+        await this.closeConnection()
         // @ts-ignore
         throw new Error(err)
       } finally {
-        if (close) await this.client.end()
+        if (close) {
+          await this.closeConnection()
         }
       }
-
-    async setSchema() {
-      await this.client.connect()
-      if (!this.config.schema) {
-        this.config.schema = "public"
-      }
-      this.client.query(`SET search_path TO ${this.config.schema}`)
-      this.COLUMNS_SQL = `select * from information_schema.columns where table_schema = '${this.config.schema}'`
     }
 
     /**
@@ -188,6 +208,7 @@ module PostgresModule {
      */
     async buildSchema(datasourceId: string, entities: Record<string, Table>) {
       let tableKeys: { [key: string]: string[] } = {}
+      await this.openConnection()
       try {
         const primaryKeysResponse = await this.client.query(
           this.PRIMARY_KEYS_SQL
@@ -251,7 +272,7 @@ module PostgresModule {
         // @ts-ignore
         throw new Error(err)
       } finally {
-        await this.client.end()
+        await this.closeConnection()
       }
     }
 
@@ -283,7 +304,7 @@ module PostgresModule {
         for (let query of input) {
           responses.push(await this.internalQuery(query, false))
         }
-        await this.client.end()
+        await this.closeConnection()
         return responses
       } else {
         const response = await this.internalQuery(input)