Move file manipulation into CustomStructure class

Move logic for manipulating the gitlab_structure.sql file into the
CustomStructure class rather than directly in the rake task, to prevent
multiple versions of the file being written inadvertently.

lib/gitlab/database/custom_structure.rb

@@ -3,10 +3,18 @@
 module Gitlab
   module Database
     class CustomStructure
-      def dump(io)
-        io << "SET search_path=public;\n\n"
+      CUSTOM_DUMP_FILE = 'db/gitlab_structure.sql'.freeze
 
-        dump_partitioned_foreign_keys(io) if partitioned_foreign_keys_exist?
+      def dump
+        File.open(self.class.custom_dump_filepath, 'wb') do |io|
+          io << "SET search_path=public;\n\n"
+
+          dump_partitioned_foreign_keys(io) if partitioned_foreign_keys_exist?
+        end
+      end
+
+      def self.custom_dump_filepath
+        Rails.root.join(CUSTOM_DUMP_FILE)
       end
 
       private

lib/tasks/gitlab/db.rake

 namespace :gitlab do
   namespace :db do
-    CUSTOM_DUMP_FILE = 'db/gitlab_structure.sql'.freeze
-
     desc 'GitLab | DB | Manually insert schema migration version'
     task :mark_migration_complete, [:version] => :environment do |_, args|
       unless args[:version]
@@ -96,15 +94,13 @@ namespace :gitlab do
 
     desc 'This dumps GitLab specific database details - it runs after db:structure:dump'
     task :dump_custom_structure do |task_name|
-      File.open(CUSTOM_DUMP_FILE, 'wb+') do |io|
-        Gitlab::Database::CustomStructure.new.dump(io)
-      end
+      Gitlab::Database::CustomStructure.new.dump
 
       # Allow this task to be called multiple times, as happens when running db:migrate:redo
       Rake::Task[task_name].reenable
     end
 
-    desc 'This loads Gitlab specific database details - runs after db:structure:dump'
+    desc 'This loads GitLab specific database details - runs after db:structure:dump'
     task :load_custom_structure do
       configuration = Rails.application.config_for(:database)
 
@@ -113,9 +109,10 @@ namespace :gitlab do
       ENV['PGPASSWORD'] = configuration['password'].to_s if configuration['password']
       ENV['PGUSER']     = configuration['username'].to_s if configuration['username']
 
-      args = ['-v', 'ON_ERROR_STOP=1', '-q', '-X', '-f', CUSTOM_DUMP_FILE]
-      args << configuration['database']
-      system('psql', *args)
+      dump_filepath = Gitlab::Database::CustomStructure.custom_dump_filepath.to_path
+      args = ['-v', 'ON_ERROR_STOP=1', '-q', '-X', '-f', dump_filepath, configuration['database']]
+
+      Kernel.system('psql', *args)
     end
 
     # Inform Rake that custom tasks should be run every time rake db:structure:dump is run

spec/lib/gitlab/database/custom_structure_spec.rb

@@ -4,11 +4,16 @@ require 'spec_helper'
 
 describe Gitlab::Database::CustomStructure do
   let_it_be(:structure) { described_class.new }
+  let_it_be(:filepath) { Rails.root.join(described_class::CUSTOM_DUMP_FILE) }
   let(:io) { StringIO.new }
 
+  before do
+    allow(File).to receive(:open).with(filepath, anything).and_yield(io)
+  end
+
   context 'when there are no partitioned_foreign_keys' do
     it 'dumps a valid structure file' do
-      structure.dump(io)
+      structure.dump
 
       expect(io.string).to eq("SET search_path=public;\n\n")
     end
@@ -24,8 +29,8 @@ describe Gitlab::Database::CustomStructure do
         cascade_delete: false, from_table: 'issues', from_column: 'moved_to_id', to_table: 'issues', to_column: 'id')
     end
 
-    it 'dumps a file with the command to restore the keys' do
-      structure.dump(io)
+    it 'dumps a file with the command to restore the current keys' do
+      structure.dump
 
       expect(io.string).to eq(<<~DATA)
         SET search_path=public;
@@ -35,6 +40,20 @@ describe Gitlab::Database::CustomStructure do
         #{second_fk.id}\tfalse\tissues\tmoved_to_id\tissues\tid
         \\.
       DATA
+
+      first_fk.destroy
+      io.truncate(0)
+      io.rewind
+
+      structure.dump
+
+      expect(io.string).to eq(<<~DATA)
+        SET search_path=public;
+
+        COPY partitioned_foreign_keys (id, cascade_delete, from_table, from_column, to_table, to_column) FROM STDIN;
+        #{second_fk.id}\tfalse\tissues\tmoved_to_id\tissues\tid
+        \\.
+      DATA
     end
   end
 end

spec/tasks/gitlab/db_rake_spec.rb

@@ -115,19 +115,43 @@ describe 'gitlab:db namespace rake task' do
     end
 
     it 'can be executed multiple times within another rake task' do
-      Rake::Task.define_task(test_task_name => :environment) do
+      expect_multiple_executions_of_task(test_task_name, clean_rake_task) do
         expect_next_instance_of(Gitlab::Database::SchemaCleaner) do |cleaner|
           expect(cleaner).to receive(:clean).with(output)
         end
-        Rake::Task[clean_rake_task].invoke
+      end
+    end
+  end
 
-        expect_next_instance_of(Gitlab::Database::SchemaCleaner) do |cleaner|
-          expect(cleaner).to receive(:clean).with(output)
+  describe 'load_custom_structure' do
+    let_it_be(:db_config) { Rails.application.config_for(:database) }
+    let_it_be(:custom_load_task) { 'gitlab:db:load_custom_structure' }
+    let_it_be(:custom_filepath) { Pathname.new('db/directory') }
+
+    it 'use the psql command to load the custom structure file' do
+      expect(Gitlab::Database::CustomStructure).to receive(:custom_dump_filepath).and_return(custom_filepath)
+
+      expect(Kernel).to receive(:system)
+        .with('psql', any_args, custom_filepath.to_path, db_config['database'])
+
+      run_rake_task(custom_load_task)
+    end
+  end
+
+  describe 'dump_custom_structure' do
+    let_it_be(:test_task_name) { 'gitlab:db:_test_multiple_task_executions' }
+    let_it_be(:custom_dump_task) { 'gitlab:db:dump_custom_structure' }
+
+    after do
+      Rake::Task[test_task_name].clear if Rake::Task.task_defined?(test_task_name)
+    end
+
+    it 'can be executed multiple times within another rake task' do
+      expect_multiple_executions_of_task(test_task_name, custom_dump_task) do
+        expect_next_instance_of(Gitlab::Database::CustomStructure) do |custom_structure|
+          expect(custom_structure).to receive(:dump)
         end
-        Rake::Task[clean_rake_task].invoke
       end
-
-      run_rake_task(test_task_name)
     end
   end
 
@@ -135,4 +159,16 @@ describe 'gitlab:db namespace rake task' do
     Rake::Task[task_name].reenable
     Rake.application.invoke_task task_name
   end
+
+  def expect_multiple_executions_of_task(test_task_name, task_to_invoke, count: 2)
+    Rake::Task.define_task(test_task_name => :environment) do
+      count.times do
+        yield
+
+        Rake::Task[task_to_invoke].invoke
+      end
+    end
+
+    run_rake_task(test_task_name)
+  end
 end