From 0505809cde9e879bd67764cb221a08e93f279ada Mon Sep 17 00:00:00 2001
From: Nick LaMuro <nicklamuro@gmail.com>
Date: Wed, 31 Oct 2018 18:38:45 -0500
Subject: [PATCH] [PostgresAdmin] Add backup_type option to .restore

Allows pre-fetching a backup_type when calling .restore, so no extra
file type checking is done on the `opt[:local_file]` (which could be a
streamed IO where that would break things down the line).

This behaves as it had if this option isn't passed.
---
 lib/gems/pending/util/postgres_admin.rb |  7 ++++---
 spec/util/postgres_admin_spec.rb        | 22 ++++++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/lib/gems/pending/util/postgres_admin.rb b/lib/gems/pending/util/postgres_admin.rb
index 55abbafbf..75ad5e4fc 100644
--- a/lib/gems/pending/util/postgres_admin.rb
+++ b/lib/gems/pending/util/postgres_admin.rb
@@ -97,10 +97,11 @@ def self.backup(opts)
   end
 
   def self.restore(opts)
-    file = opts[:local_file]
-    if pg_dump_file?(file)
+    file        = opts[:local_file]
+    backup_type = opts.delete(:backup_type)
+    if backup_type == :pgdump || pg_dump_file?(file)
       restore_pg_dump(opts)
-    elsif base_backup_file?(file)
+    elsif backup_type == :basebackup || base_backup_file?(file)
       restore_pg_basebackup(file)
     else
       raise "#{file} is not a database backup"
diff --git a/spec/util/postgres_admin_spec.rb b/spec/util/postgres_admin_spec.rb
index 60a1962d8..0032cde49 100644
--- a/spec/util/postgres_admin_spec.rb
+++ b/spec/util/postgres_admin_spec.rb
@@ -59,6 +59,28 @@
         expect(book_count).to   eq(3)
       end
     end
+
+    # Note, we aren't actually prefetching the magic here, but this is mean to
+    # simulate that an override works as expected.  We are stubbing the the
+    # restore calls here, so just making sure the logic works.
+    context "'pre-fetching' magic number" do
+      let(:dummy_base_opts) { { :local_file => "foo" } }
+
+      before do
+        allow(PostgresAdmin).to receive(:pg_dump_file?).and_return(false)
+        allow(PostgresAdmin).to receive(:base_backup_file?).and_return(false)
+      end
+
+      it "calls `.restore_pg_dump` with :backup_type => :pgdump" do
+        expect(PostgresAdmin).to receive(:restore_pg_dump).with(dummy_base_opts)
+        PostgresAdmin.restore(dummy_base_opts.merge(:backup_type => :pgdump))
+      end
+
+      it "calls `.restore_pg_basebackup` with :backup_type => :basebackup" do
+        expect(PostgresAdmin).to receive(:restore_pg_basebackup).with("foo")
+        PostgresAdmin.restore(dummy_base_opts.merge(:backup_type => :basebackup))
+      end
+    end
   end
 
   describe ".base_backup_file?" do