Migration observer to report size change

Relates to
https://gitlab.com/gitlab-org/database-team/team-tasks/-/issues/145

lib/gitlab/database/migrations/instrumentation.rb

@@ -3,16 +3,11 @@
 module Gitlab
   module Database
     module Migrations
-      Observation = Struct.new(
-        :migration,
-        :walltime,
-        :success
-      )
-
       class Instrumentation
         attr_reader :observations
 
-        def initialize
+        def initialize(observers = ::Gitlab::Database::Migrations::Observers.all_observers)
+          @observers = observers
           @observations = []
         end
 
@@ -22,6 +17,8 @@ module Gitlab
 
           exception = nil
 
+          on_each_observer { |observer| observer.before }
+
           observation.walltime = Benchmark.realtime do
             yield
           rescue => e
@@ -29,6 +26,9 @@ module Gitlab
             observation.success = false
           end
 
+          on_each_observer { |observer| observer.after }
+          on_each_observer { |observer| observer.record(observation) }
+
           record_observation(observation)
 
           raise exception if exception
@@ -38,9 +38,19 @@ module Gitlab
 
         private
 
+        attr_reader :observers
+
         def record_observation(observation)
           @observations << observation
         end
+
+        def on_each_observer(&block)
+          observers.each do |observer|
+            yield observer
+          rescue => e
+            Gitlab::AppLogger.error("Migration observer #{observer.class} failed with: #{e}")
+          end
+        end
       end
     end
   end

lib/gitlab/database/migrations/observation.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      Observation = Struct.new(
+        :migration,
+        :walltime,
+        :success,
+        :total_database_size_change
+      )
+    end
+  end
+end

lib/gitlab/database/migrations/observers.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      module Observers
+        def self.all_observers
+          [
+            TotalDatabaseSizeChange.new
+          ]
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/migrations/observers/migration_observer.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      module Observers
+        class MigrationObserver
+          attr_reader :connection
+
+          def initialize
+            @connection = ActiveRecord::Base.connection
+          end
+
+          def before
+            # implement in subclass
+          end
+
+          def after
+            # implement in subclass
+          end
+
+          def record(observation)
+            raise NotImplementedError, 'implement in subclass'
+          end
+        end
+      end
+    end
+  end
+end

lib/gitlab/database/migrations/observers/total_database_size_change.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module Database
+    module Migrations
+      module Observers
+        class TotalDatabaseSizeChange < MigrationObserver
+          def before
+            @size_before = get_total_database_size
+          end
+
+          def after
+            @size_after = get_total_database_size
+          end
+
+          def record(observation)
+            return unless @size_after && @size_before
+
+            observation.total_database_size_change = @size_after - @size_before
+          end
+
+          private
+
+          def get_total_database_size
+            connection.execute("select pg_database_size(current_database())").first['pg_database_size']
+          end
+        end
+      end
+    end
+  end
+end

spec/lib/gitlab/database/migrations/instrumentation_spec.rb

@@ -11,6 +11,41 @@ RSpec.describe Gitlab::Database::Migrations::Instrumentation do
       expect { |b| subject.observe(migration, &b) }.to yield_control
     end
 
+    context 'behavior with observers' do
+      subject { described_class.new(observers).observe(migration) {} }
+
+      let(:observers) { [observer] }
+      let(:observer) { instance_double('Gitlab::Database::Migrations::Observers::MigrationObserver', before: nil, after: nil, record: nil) }
+
+      it 'calls #before, #after, #record on given observers' do
+        expect(observer).to receive(:before).ordered
+        expect(observer).to receive(:after).ordered
+        expect(observer).to receive(:record).ordered do |observation|
+          expect(observation.migration).to eq(migration)
+        end
+
+        subject
+      end
+
+      it 'ignores errors coming from observers #before' do
+        expect(observer).to receive(:before).and_raise('some error')
+
+        subject
+      end
+
+      it 'ignores errors coming from observers #after' do
+        expect(observer).to receive(:after).and_raise('some error')
+
+        subject
+      end
+
+      it 'ignores errors coming from observers #record' do
+        expect(observer).to receive(:record).and_raise('some error')
+
+        subject
+      end
+    end
+
     context 'on successful execution' do
       subject { described_class.new.observe(migration) {} }
 

spec/lib/gitlab/database/migrations/observers/total_database_size_change_spec.rb

+# frozen_string_literal: true
+require 'spec_helper'
+
+RSpec.describe Gitlab::Database::Migrations::Observers::TotalDatabaseSizeChange do
+  subject { described_class.new }
+
+  let(:observation) { Gitlab::Database::Migrations::Observation.new }
+  let(:connection) { ActiveRecord::Base.connection }
+  let(:query) { 'select pg_database_size(current_database())' }
+
+  it 'records the size change' do
+    expect(connection).to receive(:execute).with(query).once.and_return([{ 'pg_database_size' => 1024 }])
+    expect(connection).to receive(:execute).with(query).once.and_return([{ 'pg_database_size' => 256 }])
+
+    subject.before
+    subject.after
+    subject.record(observation)
+
+    expect(observation.total_database_size_change).to eq(256 - 1024)
+  end
+
+  context 'out of order calls' do
+    before do
+      allow(connection).to receive(:execute).with(query).and_return([{ 'pg_database_size' => 1024 }])
+    end
+
+    it 'does not record anything if before size is unknown' do
+      subject.after
+
+      expect { subject.record(observation) }.not_to change { observation.total_database_size_change }
+    end
+
+    it 'does not record anything if after size is unknown' do
+      subject.before
+
+      expect { subject.record(observation) }.not_to change { observation.total_database_size_change }
+    end
+  end
+end