From 2cf12abeb677bad3be81f5712cc4b93481b133c1 Mon Sep 17 00:00:00 2001
From: James <jameshy64@gmail.com>
Date: Thu, 2 Jan 2020 01:12:04 +0100
Subject: [PATCH] add support for pg_dump command arguments

---
 lib/config.js  |  1 +
 lib/handler.js |  2 +-
 lib/pgdump.js  | 29 ++++++++++++++++++-----
 test/pgdump.js | 62 ++++++++++++++++++++++++++++++++++++++++++++++----
 4 files changed, 83 insertions(+), 11 deletions(-)

diff --git a/lib/config.js b/lib/config.js
index 8dda112..42a9ed6 100644
--- a/lib/config.js
+++ b/lib/config.js
@@ -1,5 +1,6 @@
 const path = require('path')
 
+// default config that is overridden by the Lambda event
 module.exports = {
     S3_REGION: 'eu-west-1',
     PGDUMP_PATH: path.join(__dirname, '../bin/postgres-11.6'),
diff --git a/lib/handler.js b/lib/handler.js
index 06ad398..0d82ae3 100644
--- a/lib/handler.js
+++ b/lib/handler.js
@@ -39,7 +39,7 @@ async function handler(event) {
     }
     catch (error) {
         // log the error and rethrow for Lambda
-        if (process.env.NODE_ENV !== "test") {
+        if (process.env.NODE_ENV !== 'test') {
             console.error(error)
         }
         throw error
diff --git a/lib/pgdump.js b/lib/pgdump.js
index 242ee90..7f533c2 100644
--- a/lib/pgdump.js
+++ b/lib/pgdump.js
@@ -3,30 +3,47 @@ const through2 = require('through2')
 const path = require('path')
 const fs = require('fs')
 
-function spawnPgDump(config) {
+function spawnPgDump(pgdumpDir, args, env) {
     const pgDumpPath = path.join(
-        config.PGDUMP_PATH,
+        pgdumpDir,
         'pg_dump'
     )
     if (!fs.existsSync(pgDumpPath)) {
         throw new Error('pg_dump not found at ' + pgDumpPath)
     }
-    const env = { ...config, LD_LIBRARY_PATH: config.PGDUMP_PATH }
-    return spawn(pgDumpPath, ['-Fc', '-Z 1'], {
+
+    return spawn(pgDumpPath, args, {
         env
     })
 }
 
+function buildArgs(config) {
+    let args = ['-Fc', '-Z1']
+    const extraArgs = config.PGDUMP_ARGS
+
+    if (typeof extraArgs === 'string') {
+        const splitArgs = extraArgs.split(' ')
+        args = args.concat(splitArgs)
+    }
+    else if (Array.isArray(extraArgs)) {
+        args = args.concat(extraArgs)
+    }
+
+    return args
+}
+
 function pgdump(config, pgDumpSpawnFn = spawnPgDump) {
     return new Promise((resolve, reject) => {
         let headerChecked = false
         let stderr = ''
 
         // spawn pg_dump process
-        const process = pgDumpSpawnFn(config)
+        const args = buildArgs(config)
+        const env = { ...config, LD_LIBRARY_PATH: config.PGDUMP_PATH }
+        const process = pgDumpSpawnFn(config.PGDUMP_PATH, args, env)
 
         // hook into the process
-        process.stderr.on('data', (data) => {
+        process.stderr.on('data', data => {
             stderr += data.toString('utf8')
         })
 
diff --git a/test/pgdump.js b/test/pgdump.js
index b8e6bdd..d7f24b7 100644
--- a/test/pgdump.js
+++ b/test/pgdump.js
@@ -3,6 +3,7 @@ const fs = require('fs')
 const mockSpawn = require('mock-spawn')
 const chai = require('chai')
 const chaiAsPromised = require('chai-as-promised')
+const sinon = require('sinon')
 
 chai.use(chaiAsPromised)
 const { expect } = chai
@@ -27,14 +28,67 @@ describe('pgdump', () => {
         )
     })
 
+    it('should call pg_dump with some default args', async () => {
+        const pgdumpProcess = mockSpawn()()
+        const pgDumpFn = sinon.fake.returns(pgdumpProcess)
+        const config = {}
+        const p = pgdump(config, pgDumpFn)
+        pgdumpProcess.stdout.write('PGDMP - data - data')
+        pgdumpProcess.emit('close', 0)
+        await p
+
+        expect(pgDumpFn.calledOnce).to.be.true
+        const pgDumpArgs = pgDumpFn.getCall(0).args[1]
+        expect(pgDumpArgs).to.deep.equal(['-Fc', '-Z1'])
+    })
+
+    it('should call pg_dump with provided extra arguments as array', async () => {
+        const pgdumpProcess = mockSpawn()()
+        const pgDumpFn = sinon.fake.returns(pgdumpProcess)
+        const config = {
+            PGDUMP_ARGS: ['--exclude-table=ignored-table', '-N', 'public']
+        }
+        const p = pgdump(config, pgDumpFn)
+        pgdumpProcess.stdout.write('PGDMP - data - data')
+        pgdumpProcess.emit('close', 0)
+        await p
+
+        expect(pgDumpFn.calledOnce).to.be.true
+        const pgDumpArgs = pgDumpFn.getCall(0).args[1]
+
+        expect(
+            pgDumpArgs
+        ).to.deep.equal(['-Fc', '-Z1', '--exclude-table=ignored-table', '-N', 'public'])
+    })
+
+    it('should call pg_dump with provided extra arguments as string', async () => {
+        const pgdumpProcess = mockSpawn()()
+        const pgDumpFn = sinon.fake.returns(pgdumpProcess)
+        const config = {
+            PGDUMP_ARGS: '--exclude-table=ignored-table -N public'
+        }
+
+        const p = pgdump(config, pgDumpFn)
+        pgdumpProcess.stdout.write('PGDMP - data - data')
+        pgdumpProcess.emit('close', 0)
+        await p
+
+        expect(pgDumpFn.calledOnce).to.be.true
+        const pgDumpArgs = pgDumpFn.getCall(0).args[1]
+
+        expect(
+            pgDumpArgs
+        ).to.deep.equal(['-Fc', '-Z1', '--exclude-table=ignored-table', '-N', 'public'])
+    })
+
     it('should stream correctly', async () => {
-        const mySpawn = mockSpawn()()
+        const pgdumpProcess = mockSpawn()()
 
-        const pgDumpFn = () => mySpawn
+        const pgDumpFn = () => pgdumpProcess
         const config = {}
         const p = pgdump(config, pgDumpFn)
-        mySpawn.stdout.write('PGDMP - data - data')
-        mySpawn.emit('close', 0)
+        pgdumpProcess.stdout.write('PGDMP - data - data')
+        pgdumpProcess.emit('close', 0)
 
         const buffer = await p