From f36ab98d52603fe266ad3ae1a2575d7071c89ff5 Mon Sep 17 00:00:00 2001
From: Hugo Sales <hugo@visrez.com>
Date: Wed, 20 Dec 2023 19:11:28 +0000
Subject: [PATCH] Add support for specifying extra command line options during
 database dump and load

---
 CHANGELOG.md                        |  1 +
 lib/capistrano-db-tasks/database.rb | 20 ++++++++++++++++----
 lib/capistrano-db-tasks/dbtasks.rb  |  2 ++
 3 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0e92f70..3ee7657 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ Reverse Chronological Order:
 https://github.com/sgruhier/capistrano-db-tasks/compare/v0.6...HEAD
 
 * Your contribution here!
+* Add support for passing extra options to the database dump and load commands (@someonewithpc)
 * Added support for [zstd compressor](http://facebook.github.io/zstd/) (@ocha)
 
 # 0.6 (Dec 14 2016)
diff --git a/lib/capistrano-db-tasks/database.rb b/lib/capistrano-db-tasks/database.rb
index 08d044c..1825492 100644
--- a/lib/capistrano-db-tasks/database.rb
+++ b/lib/capistrano-db-tasks/database.rb
@@ -72,18 +72,18 @@ def dump_cmd
 
     def import_cmd(file)
       if mysql?
-        "mysql #{credentials} -D #{database} < #{file}"
+        "mysql #{credentials} -D #{database} #{import_cmd_opts} < #{file}"
       elsif postgresql?
         terminate_connection_sql = "SELECT pg_terminate_backend(pg_stat_activity.pid) FROM pg_stat_activity WHERE pg_stat_activity.datname = '#{database}' AND pid <> pg_backend_pid();"
-        "#{pgpass} psql -c \"#{terminate_connection_sql};\" #{credentials} #{database}; #{pgpass} dropdb #{credentials} #{database}; #{pgpass} createdb #{credentials} #{database}; #{pgpass} psql #{credentials} -d #{database} < #{file}"
+        "#{pgpass} psql -c \"#{terminate_connection_sql};\" #{credentials} #{database}; #{pgpass} dropdb #{credentials} #{database}; #{pgpass} createdb #{credentials} #{database}; #{pgpass} psql #{credentials} -d #{database} #{import_cmd_opts} < #{file}"
       end
     end
 
     def dump_cmd_opts
       if mysql?
-        "--lock-tables=false #{dump_cmd_ignore_tables_opts} #{dump_cmd_ignore_data_tables_opts}"
+        "--lock-tables=false #{dump_cmd_ignore_tables_opts} #{dump_cmd_ignore_data_tables_opts} #{dump_cmd_extra_opts}"
       elsif postgresql?
-        "--no-acl --no-owner #{dump_cmd_ignore_tables_opts} #{dump_cmd_ignore_data_tables_opts}"
+        "--no-acl --no-owner #{dump_cmd_ignore_tables_opts} #{dump_cmd_ignore_data_tables_opts} #{dump_cmd_extra_opts}"
       end
     end
 
@@ -100,6 +100,18 @@ def dump_cmd_ignore_data_tables_opts
       ignore_tables = @cap.fetch(:db_ignore_data_tables, [])
       ignore_tables.map { |t| "--exclude-table-data=#{t}" }.join(" ") if postgresql?
     end
+
+    def dump_cmd_extra_opts
+      @cap.fetch(:db_dump_extra_opts, "")
+    end
+
+    def import_cmd_opts
+      import_cmd_extra_opts
+    end
+
+    def import_cmd_extra_opts
+      @cap.fetch(:db_import_extra_opts, "")
+    end
   end
 
   class Remote < Base
diff --git a/lib/capistrano-db-tasks/dbtasks.rb b/lib/capistrano-db-tasks/dbtasks.rb
index 15891be..8e0b143 100644
--- a/lib/capistrano-db-tasks/dbtasks.rb
+++ b/lib/capistrano-db-tasks/dbtasks.rb
@@ -14,6 +14,8 @@
 set :skip_data_sync_confirm, ENV['SKIP_DATA_SYNC_CONFIRM'].to_s.casecmp('true').zero?
 set :disallow_pushing, false unless fetch(:disallow_pushing)
 set :compressor, :gzip unless fetch(:compressor)
+set :db_dump_extra_opts, '' unless fetch(:db_dump_extra_opts)
+set :db_import_extra_opts, '' unless fetch(:db_import_extra_opts)
 
 namespace :capistrano_db_tasks do
   task :check_can_push do