From e3db57efd084164c8c73e3d257e442977dfcc9a6 Mon Sep 17 00:00:00 2001
From: Varjitt Jeeva <varjitt.jeeva@plangrid.com>
Date: Mon, 9 Oct 2017 09:47:38 -0700
Subject: [PATCH] Fixing #11 (#12)

* Fixing #11 added extra param for which restore command to use, can't generalize based on file type.

* Need to catch empty logs oops

* psql command reverted to forked original and fixed grabbing of parameter for restore type

* Missed an equals sign... oops

* @ sign added by autocomplete, oops...
---
 README.mdown              | 3 ++-
 lib/engines/postgresql.rb | 7 ++++---
 ui.rb                     | 8 +++++---
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/README.mdown b/README.mdown
index 3840bd1..486aea9 100644
--- a/README.mdown
+++ b/README.mdown
@@ -311,7 +311,7 @@ Unless you need to change any of the defaults, a standard configuration for Mong
 
 ###### Full `postgresql` Engine Example
 
-You can specify the username for accessing the postgres server (other than the default `postgres`) for when you run the tests. Furthermore, you can specify the file type to be `sql` or `dump`. If `sql`, then the restore will be done with the `psql` command, which takes a `.sql` file. If `dump`, then the restore will be done with the `pg_restore` command, which takes a `.pg_dump` file (usually, or whatever file is obtained using the pg_dump command).
+You can specify the username for accessing the postgres server (other than the default `postgres`) for when you run the tests. Furthermore, you can specify the file type. Since the file type can be anything, the `restore_type` must be set to know which command to restore with. If `psql`, then the restore will be done with the `psql` command, which takes a `.sql` file. If `pg_restore`, then the restore will be done with the `pg_restore` command, which takes a dump file.
 
     engine:
       type: postgresql
@@ -320,6 +320,7 @@ You can specify the username for accessing the postgres server (other than the d
         docker_image: postgres
         docker_tag: 9.5.0
         extension: .sql
+        restore_type: psql
         port: 5432
         sql_file: PostgreSQL.sql
         username: postgres
diff --git a/lib/engines/postgresql.rb b/lib/engines/postgresql.rb
index 50fb735..6ffc4d8 100644
--- a/lib/engines/postgresql.rb
+++ b/lib/engines/postgresql.rb
@@ -7,6 +7,7 @@ def initialize(params)
         @docker_image ||= 'postgres'
         @docker_tag ||= '9.5'
         @extension ||= '.sql'
+        @restore_type = params[:restore_type] || 'psql'
         @port ||= 5432
         @username = params[:username] || 'postgres'
         @docker_env ||= ["POSTGRES_USER=#{@username}", "POSTGRES_DB=#{@database}"]
@@ -24,9 +25,9 @@ def initialize(params)
 
       def load_backup(path)
         Dir.chdir(path) do
-          if @extension.include?(".sql")
-            run_command("psql --no-password --exit-on-error --username=#{@username} --host=#{container_ip_address} --port=#{@port} --dbname=#{@database} < #{@sql_file}")
-          elsif @extension.include?(".pg_dump")
+          if @restore_type.include?("psql")
+            run_command("psql --no-password --set ON_ERROR_STOP=on --username=#{@username} --host=#{container_ip_address} --port=#{@port} --dbname=#{@database} < #{@sql_file}")
+          elsif @restore_type.include?("pg_restore")
              run_command("pg_restore --no-password --exit-on-error --username=#{@username} --host=#{container_ip_address} --port=#{@port} --dbname=#{@database} < #{@sql_file}")
           else
             raise 'File Type parameter in Postgresql Driver is invalid.'
diff --git a/ui.rb b/ui.rb
index 473bd94..fc30f9e 100644
--- a/ui.rb
+++ b/ui.rb
@@ -153,9 +153,11 @@ def backups_with_logs
           backup = {}
           backup[:date] = Files.extract_datetime_from_backup_path(config, path)
           backup[:backup] = path
-          logss.each do |log|
-            if log[:epoch].to_s == backup[:date].strftime('%s').to_s
-              backup[:log] = log
+          if logss != nil
+            logss.each do |log|
+              if log[:epoch].to_s == backup[:date].strftime('%s').to_s
+                backup[:log] = log
+              end
             end
           end
           backups << backup