Handle failure of rake call to psql command

In the rake task to load a custom database structure.sql, check for the
result of the call to the psql command which loads the file, and display
a message to the user on error.

db/gitlab_structure.sql

+-- this file tracks custom GitLab data, such as foreign keys referencing partitioned tables
+-- more details can be found in the issue: https://gitlab.com/gitlab-org/gitlab/-/issues/201872
 SET search_path=public;
 

lib/gitlab/database/custom_structure.rb

@@ -3,10 +3,12 @@
 module Gitlab
   module Database
     class CustomStructure
-      CUSTOM_DUMP_FILE = 'db/gitlab_structure.sql'.freeze
+      CUSTOM_DUMP_FILE = 'db/gitlab_structure.sql'
 
       def dump
         File.open(self.class.custom_dump_filepath, 'wb') do |io|
+          io << "-- this file tracks custom GitLab data, such as foreign keys referencing partitioned tables\n"
+          io << "-- more details can be found in the issue: https://gitlab.com/gitlab-org/gitlab/-/issues/201872\n"
           io << "SET search_path=public;\n\n"
 
           dump_partitioned_foreign_keys(io) if partitioned_foreign_keys_exist?

lib/tasks/gitlab/db.rake

@@ -109,10 +109,14 @@ namespace :gitlab do
       ENV['PGPASSWORD'] = configuration['password'].to_s if configuration['password']
       ENV['PGUSER']     = configuration['username'].to_s if configuration['username']
 
+      command = 'psql'
       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)
+      unless Kernel.system(command, *args)
+        raise "failed to execute:\n#{command} #{args.join(' ')}\n\n" \
+          "Please ensure `#{command}` is installed in your PATH and has proper permissions.\n\n"
+      end
     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

@@ -5,6 +5,14 @@ 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_it_be(:file_header) do
+    <<~DATA
+      -- this file tracks custom GitLab data, such as foreign keys referencing partitioned tables
+      -- more details can be found in the issue: https://gitlab.com/gitlab-org/gitlab/-/issues/201872
+      SET search_path=public;
+    DATA
+  end
+
   let(:io) { StringIO.new }
 
   before do
@@ -15,7 +23,7 @@ describe Gitlab::Database::CustomStructure do
     it 'dumps a valid structure file' do
       structure.dump
 
-      expect(io.string).to eq("SET search_path=public;\n\n")
+      expect(io.string).to eq("#{file_header}\n")
     end
   end
 
@@ -33,8 +41,7 @@ describe Gitlab::Database::CustomStructure do
       structure.dump
 
       expect(io.string).to eq(<<~DATA)
-        SET search_path=public;
-
+        #{file_header}
         COPY partitioned_foreign_keys (id, cascade_delete, from_table, from_column, to_table, to_column) FROM STDIN;
         #{first_fk.id}\ttrue\tissues\tproject_id\tprojects\tid
         #{second_fk.id}\tfalse\tissues\tmoved_to_id\tissues\tid
@@ -48,8 +55,7 @@ describe Gitlab::Database::CustomStructure do
       structure.dump
 
       expect(io.string).to eq(<<~DATA)
-        SET search_path=public;
-
+        #{file_header}
         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
         \\.

spec/tasks/gitlab/db_rake_spec.rb

@@ -128,14 +128,23 @@ describe 'gitlab:db namespace rake task' do
     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
+    it 'uses 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'])
+        .with('psql', any_args, custom_filepath.to_path, db_config['database']).and_return(true)
 
       run_rake_task(custom_load_task)
     end
+
+    it 'raises an error when the call to the psql command fails' 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']).and_return(nil)
+
+      expect { run_rake_task(custom_load_task) }.to raise_error(/failed to execute:\s*psql/)
+    end
   end
 
   describe 'dump_custom_structure' do