Add support in Gitlab::Database for checking PostgreSQL replication lagsh-add-replication-lag-query

We are starting to use this query in a bunch of places (e.g. Geo, load
balancing, etc.), so to reduce code duplication we should have it in
one central place.

2 files changed, 58 insertions, 0 deletions

diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index 8eacad078c8..fe0f8506ec2 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -57,6 +57,33 @@ module Gitlab
       !self.db_read_only?
     end
 
+    # Returns the replication lag time of this secondary in seconds as a
+    # float.
+    #
+    # This method will return nil if no lag time could be calculated.
+    def self.replication_lag_seconds
+      return unless postgresql?
+
+      # Obtain the replication lag in seconds
+      lag = ActiveRecord::Base.connection.execute(self.pg_replication_lag_query).first.fetch('lag')
+
+      lag.present? ? lag.to_f : lag
+    end
+
+    def self.pg_replication_lag_query
+      # Before we check the actual replication lag, we check if the master is actaully active.
+      # From https://www.postgresql.org/message-id/flat/CADKbJJWz9M0swPT3oqe8f9%2BtfD4-F54uE6Xtkh4nERpVsQnjnw%40mail.gmail.com.
+      <<-SQL.squish
+        SELECT CASE
+               WHEN #{Gitlab::Database.pg_last_wal_receive_lsn}() = #{Gitlab::Database.pg_last_wal_receive_lsn}()
+                 THEN 0
+               ELSE
+                 EXTRACT (EPOCH FROM now() - pg_last_xact_replay_timestamp())::float
+               END
+                 AS lag
+      SQL
+    end
+
     def self.version
       @version ||= database_version.match(/\A(?:PostgreSQL |)([^\s]+).*\z/)[1]
     end
diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index 7d76519dddd..d28303f4b1c 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -472,6 +472,37 @@ describe Gitlab::Database do
     end
   end
 
+  describe '.replication_lag_seconds' do
+    context 'when using PostgreSQL' do
+      let(:sql) { Gitlab::Database.pg_replication_lag_query }
+
+      before do
+        allow(ActiveRecord::Base.connection).to receive(:execute).and_call_original
+        allow(described_class).to receive(:postgresql?).and_return(true)
+      end
+
+      it 'reports unknown lag' do
+        allow(ActiveRecord::Base.connection).to receive(:execute).with(sql).and_return([{ "lag" => nil }])
+
+        expect(described_class.replication_lag_seconds).to be_nil
+      end
+
+      it 'reports replication lag' do
+        allow(ActiveRecord::Base.connection).to receive(:execute).with(sql).and_return([{ "lag" => "10.0" }])
+
+        expect(described_class.replication_lag_seconds).to eq(10.0)
+      end
+    end
+
+    context 'when using MySQL' do
+      before do
+        expect(described_class).to receive(:postgresql?).and_return(false)
+      end
+
+      it { expect(described_class.replication_lag_seconds).to be_falsey }
+    end
+  end
+
   describe '#sanitize_timestamp' do
     let(:max_timestamp) { Time.at((1 << 31) - 1) }