Wait until DB is ready to list data on legacy storage

Required to prevent Docker upgrade to 14.0 if there
data on legacy storage.

See: https://gitlab.com/gitlab-org/omnibus-gitlab/-/merge_requests/5311#note_590454698

Changelog: changed

lib/tasks/gitlab/storage.rake

@@ -96,8 +96,12 @@ namespace :gitlab do
 
     desc 'Gitlab | Storage | Summary of existing projects using Legacy Storage'
     task legacy_projects: :environment do
-      helper = Gitlab::HashedStorage::RakeHelper
-      helper.relation_summary('projects using Legacy Storage', Project.without_storage_feature(:repository))
+      # Required to prevent Docker upgrade to 14.0 if there data on legacy storage
+      # See: https://gitlab.com/gitlab-org/omnibus-gitlab/-/merge_requests/5311#note_590454698
+      wait_until_database_is_ready do
+        helper = Gitlab::HashedStorage::RakeHelper
+        helper.relation_summary('projects using Legacy Storage', Project.without_storage_feature(:repository))
+      end
     end
 
     desc 'Gitlab | Storage | List existing projects using Legacy Storage'
@@ -135,8 +139,12 @@ namespace :gitlab do
 
     desc 'Gitlab | Storage | Summary of project attachments using Legacy Storage'
     task legacy_attachments: :environment do
-      helper = Gitlab::HashedStorage::RakeHelper
-      helper.relation_summary('attachments using Legacy Storage', helper.legacy_attachments_relation)
+      # Required to prevent Docker upgrade to 14.0 if there data on legacy storage
+      # See: https://gitlab.com/gitlab-org/omnibus-gitlab/-/merge_requests/5311#note_590454698
+      wait_until_database_is_ready do
+        helper = Gitlab::HashedStorage::RakeHelper
+        helper.relation_summary('attachments using Legacy Storage', helper.legacy_attachments_relation)
+      end
     end
 
     desc 'Gitlab | Storage | List existing project attachments using Legacy Storage'
@@ -156,5 +164,23 @@ namespace :gitlab do
       helper = Gitlab::HashedStorage::RakeHelper
       helper.attachments_list('attachments using Hashed Storage', helper.hashed_attachments_relation)
     end
+
+    def wait_until_database_is_ready
+      attempts = (ENV['MAX_DATABASE_CONNECTION_CHECKS'] || 1).to_i
+      inverval = (ENV['MAX_DATABASE_CONNECTION_CHECK_INTERVAL'] || 10).to_f
+
+      attempts.to_i.times do
+        unless Gitlab::Database.exists?
+          puts "Waiting until database is ready before continuing...".color(:yellow)
+          sleep inverval
+        end
+      end
+
+      yield
+    rescue ActiveRecord::ConnectionNotEstablished => ex
+      puts "Failed to connect to the database...".color(:red)
+      puts "Error: #{ex}"
+      exit 1
+    end
   end
 end

spec/tasks/gitlab/storage_rake_spec.rb

@@ -88,6 +88,27 @@ RSpec.describe 'rake gitlab:storage:*', :silence_stdout do
     end
   end
 
+  shared_examples 'wait until database is ready' do
+    it 'checks if the database is ready once' do
+      expect(Gitlab::Database).to receive(:exists?).once
+
+      run_rake_task(task)
+    end
+
+    context 'handles custom env vars' do
+      before do
+        stub_env('MAX_DATABASE_CONNECTION_CHECKS' => 3)
+        stub_env('MAX_DATABASE_CONNECTION_INTERVAL' => 0.1)
+      end
+
+      it 'tries for 3 times, polling every 0.1 seconds' do
+        expect(Gitlab::Database).to receive(:exists?).exactly(3).times.and_return(false)
+
+        run_rake_task(task)
+      end
+    end
+  end
+
   describe 'gitlab:storage:migrate_to_hashed' do
     let(:task) { 'gitlab:storage:migrate_to_hashed' }
 
@@ -198,6 +219,10 @@ RSpec.describe 'rake gitlab:storage:*', :silence_stdout do
       let(:task) { 'gitlab:storage:legacy_projects' }
       let(:create_collection) { create_list(:project, 3, :legacy_storage) }
     end
+
+    it_behaves_like 'wait until database is ready' do
+      let(:task) { 'gitlab:storage:legacy_projects' }
+    end
   end
 
   describe 'gitlab:storage:list_legacy_projects' do
@@ -227,6 +252,10 @@ RSpec.describe 'rake gitlab:storage:*', :silence_stdout do
       let(:project) { create(:project, storage_version: 1) }
       let(:create_collection) { create_list(:upload, 3, model: project) }
     end
+
+    it_behaves_like 'wait until database is ready' do
+      let(:task) { 'gitlab:storage:legacy_attachments' }
+    end
   end
 
   describe 'gitlab:storage:list_legacy_attachments' do