Add latest changes from gitlab-org/gitlab@master

22 files changed, 455 insertions, 106 deletions

diff --git a/gems/click_house-client/lib/click_house/client.rb b/gems/click_house-client/lib/click_house/client.rb
index 22c42d7be6e..579413640b1 100644
--- a/gems/click_house-client/lib/click_house/client.rb
+++ b/gems/click_house-client/lib/click_house/client.rb
@@ -3,6 +3,7 @@
 require 'addressable'
 require 'json'
 require 'active_support/time'
+require 'active_support/notifications'
 require_relative "client/database"
 require_relative "client/configuration"
 require_relative "client/formatter"
@@ -29,28 +30,41 @@ module ClickHouse
     def self.select(query, database, configuration = self.configuration)
       db = lookup_database(configuration, database)
 
-      response = configuration.http_post_proc.call(
-        db.uri.to_s,
-        db.headers,
-        "#{query} FORMAT JSON" # always return JSON
-      )
+      ActiveSupport::Notifications.instrument('sql.click_house', { query: query, database: database }) do |instrument|
+        response = configuration.http_post_proc.call(
+          db.uri.to_s,
+          db.headers,
+          "#{query} FORMAT JSON" # always return JSON
+        )
 
-      raise DatabaseError, response.body unless response.success?
+        raise DatabaseError, response.body unless response.success?
 
-      Formatter.format(configuration.json_parser.parse(response.body))
+        parsed_response = configuration.json_parser.parse(response.body)
+
+        instrument[:statistics] = parsed_response['statistics']&.symbolize_keys
+
+        Formatter.format(parsed_response)
+      end
     end
 
     # Executes any kinds of database query without returning any data (INSERT, DELETE)
     def self.execute(query, database, configuration = self.configuration)
       db = lookup_database(configuration, database)
 
-      response = configuration.http_post_proc.call(
-        db.uri.to_s,
-        db.headers,
-        query
-      )
+      ActiveSupport::Notifications.instrument('sql.click_house', { query: query, database: database }) do |instrument|
+        response = configuration.http_post_proc.call(
+          db.uri.to_s,
+          db.headers,
+          query
+        )
 
-      raise DatabaseError, response.body unless response.success?
+        raise DatabaseError, response.body unless response.success?
+
+        if response.headers['x-clickhouse-summary']
+          instrument[:statistics] =
+            Gitlab::Json.parse(response.headers['x-clickhouse-summary']).symbolize_keys
+        end
+      end
 
       true
     end
diff --git a/gems/click_house-client/lib/click_house/client/response.rb b/gems/click_house-client/lib/click_house/client/response.rb
index 898f0b0e024..8eb5eb50cc0 100644
--- a/gems/click_house-client/lib/click_house/client/response.rb
+++ b/gems/click_house-client/lib/click_house/client/response.rb
@@ -3,11 +3,12 @@
 module ClickHouse
   module Client
     class Response
-      attr_reader :body
+      attr_reader :body, :headers
 
-      def initialize(body, http_status_code)
+      def initialize(body, http_status_code, headers = {})
         @body = body
         @http_status_code = http_status_code
+        @headers = headers
       end
 
       def success?
diff --git a/gems/csv_builder/lib/csv_builder.rb b/gems/csv_builder/lib/csv_builder.rb
index 1ef38a1d6a4..86b682939dc 100644
--- a/gems/csv_builder/lib/csv_builder.rb
+++ b/gems/csv_builder/lib/csv_builder.rb
@@ -2,11 +2,13 @@
 
 require 'csv'
 require 'tempfile'
+require 'zlib'
 
 require_relative "csv_builder/version"
 require_relative "csv_builder/builder"
 require_relative "csv_builder/single_batch"
 require_relative "csv_builder/stream"
+require_relative "csv_builder/gzip"
 
 # Generates CSV when given a collection and a mapping.
 #
diff --git a/gems/csv_builder/lib/csv_builder/builder.rb b/gems/csv_builder/lib/csv_builder/builder.rb
index 3baa2155fc9..99b63153ab2 100644
--- a/gems/csv_builder/lib/csv_builder/builder.rb
+++ b/gems/csv_builder/lib/csv_builder/builder.rb
@@ -59,8 +59,10 @@ module CsvBuilder
         @collection.each_batch(order_hint: :created_at) do |relation|
           relation.preload(@associations_to_preload).order(:id).each(&block)
         end
-      else
+      elsif @collection.respond_to?(:find_each)
         @collection.find_each(&block)
+      else
+        @collection.each(&block)
       end
     end
 
diff --git a/gems/csv_builder/lib/csv_builder/gzip.rb b/gems/csv_builder/lib/csv_builder/gzip.rb
new file mode 100644
index 00000000000..60875006a35
--- /dev/null
+++ b/gems/csv_builder/lib/csv_builder/gzip.rb
@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module CsvBuilder
+  class Gzip < CsvBuilder::Builder
+    # Writes the CSV file compressed and yields the written tempfile.
+    #
+    # Example:
+    # > CsvBuilder::Gzip.new(Issue, { title: -> (row) { row.title.upcase }, id: :id }).render do |tempfile|
+    # >   puts tempfile.path
+    # >   puts `zcat #{tempfile.path}`
+    # > end
+    def render
+      Tempfile.open(['csv_builder_gzip', '.csv.gz']) do |tempfile|
+        csv = CSV.new(Zlib::GzipWriter.open(tempfile.path))
+
+        write_csv csv, until_condition: -> {} # truncation must be handled outside of the CsvBuilder
+
+        csv.close
+        yield tempfile
+      end
+    end
+  end
+end
diff --git a/gems/csv_builder/spec/csv_builder/gzip_spec.rb b/gems/csv_builder/spec/csv_builder/gzip_spec.rb
new file mode 100644
index 00000000000..9d24d351247
--- /dev/null
+++ b/gems/csv_builder/spec/csv_builder/gzip_spec.rb
@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe CsvBuilder::Gzip do
+  let(:event_1) { double(title: 'Added salt', description: 'A teaspoon') }
+  let(:event_2) { double(title: 'Added sugar', description: 'Just a pinch') }
+  let(:items) { [event_1, event_2] }
+
+  subject(:builder) { described_class.new(items, 'Title' => 'title', 'Description' => 'description') }
+
+  describe '#render' do
+    it 'returns yields a tempfile' do
+      written_content = nil
+
+      builder.render do |tempfile|
+        reader = Zlib::GzipReader.new(tempfile)
+        written_content = reader.read.split("\n")
+      end
+
+      expect(written_content).to eq(
+        [
+          "Title,Description",
+          "Added salt,A teaspoon",
+          "Added sugar,Just a pinch"
+        ])
+    end
+
+    it 'requires a block' do
+      expect { builder.render }.to raise_error(LocalJumpError)
+    end
+  end
+end
diff --git a/gems/csv_builder/spec/csv_builder_spec.rb b/gems/csv_builder/spec/csv_builder_spec.rb
index 9391938f59d..9d6283b3985 100644
--- a/gems/csv_builder/spec/csv_builder_spec.rb
+++ b/gems/csv_builder/spec/csv_builder_spec.rb
@@ -2,126 +2,141 @@
 
 RSpec.describe CsvBuilder do
   let(:object) { double(question: :answer) }
-  let(:fake_relation) { described_class::FakeRelation.new([object]) }
   let(:csv_data) { subject.render }
 
   let(:subject) do
     described_class.new(
-      fake_relation, 'Q & A' => :question, 'Reversed' => ->(o) { o.question.to_s.reverse })
+      enumerable, 'Q & A' => :question, 'Reversed' => ->(o) { o.question.to_s.reverse })
   end
 
-  before do
-    stub_const("#{described_class}::FakeRelation", Array)
+  shared_examples 'csv builder examples' do
+    let(:items) { [object] }
 
-    described_class::FakeRelation.class_eval do
-      def find_each(&block)
-        each(&block)
-      end
+    it "has a version number" do
+      expect(CsvBuilder::Version::VERSION).not_to be nil
     end
-  end
 
-  it "has a version number" do
-    expect(CsvBuilder::Version::VERSION).not_to be nil
-  end
+    it 'generates a csv' do
+      expect(csv_data.scan(/(,|\n)/).join).to include ",\n,"
+    end
 
-  it 'generates a csv' do
-    expect(csv_data.scan(/(,|\n)/).join).to include ",\n,"
-  end
+    it 'uses a temporary file to reduce memory allocation' do
+      expect(CSV).to receive(:new).with(instance_of(Tempfile)).and_call_original
 
-  it 'uses a temporary file to reduce memory allocation' do
-    expect(CSV).to receive(:new).with(instance_of(Tempfile)).and_call_original
+      subject.render
+    end
 
-    subject.render
-  end
+    it 'counts the number of rows' do
+      subject.render
 
-  it 'counts the number of rows' do
-    subject.render
+      expect(subject.rows_written).to eq 1
+    end
 
-    expect(subject.rows_written).to eq 1
-  end
+    describe 'rows_expected' do
+      it 'uses rows_written if CSV rendered successfully' do
+        subject.render
 
-  describe 'rows_expected' do
-    it 'uses rows_written if CSV rendered successfully' do
-      subject.render
+        expect(enumerable).not_to receive(:count)
+        expect(subject.rows_expected).to eq 1
+      end
 
-      expect(fake_relation).not_to receive(:count)
-      expect(subject.rows_expected).to eq 1
+      it 'falls back to calling .count before rendering begins' do
+        expect(subject.rows_expected).to eq 1
+      end
     end
 
-    it 'falls back to calling .count before rendering begins' do
-      expect(subject.rows_expected).to eq 1
-    end
-  end
+    describe 'truncation' do
+      let(:big_object) { double(question: 'Long' * 1024) }
+      let(:row_size) { big_object.question.length * 2 }
+      let(:items) { [big_object, big_object, big_object] }
+
+      it 'occurs after given number of bytes' do
+        expect(subject.render(row_size * 2).length).to be_between(row_size * 2, row_size * 3)
+        expect(subject).to be_truncated
+        expect(subject.rows_written).to eq 2
+      end
+
+      it 'is ignored by default' do
+        expect(subject.render.length).to be > row_size * 3
+        expect(subject.rows_written).to eq 3
+      end
 
-  describe 'truncation' do
-    let(:big_object) { double(question: 'Long' * 1024) }
-    let(:row_size) { big_object.question.length * 2 }
-    let(:fake_relation) { described_class::FakeRelation.new([big_object, big_object, big_object]) }
+      it 'causes rows_expected to fall back to .count' do
+        subject.render(0)
 
-    it 'occurs after given number of bytes' do
-      expect(subject.render(row_size * 2).length).to be_between(row_size * 2, row_size * 3)
-      expect(subject).to be_truncated
-      expect(subject.rows_written).to eq 2
+        expect(enumerable).to receive(:count).and_call_original
+        expect(subject.rows_expected).to eq 3
+      end
     end
 
-    it 'is ignored by default' do
-      expect(subject.render.length).to be > row_size * 3
-      expect(subject.rows_written).to eq 3
+    it 'avoids loading all data in a single query' do
+      expect(enumerable).to receive(:find_each)
+
+      subject.render
     end
 
-    it 'causes rows_expected to fall back to .count' do
-      subject.render(0)
+    it 'uses hash keys as headers' do
+      expect(csv_data).to start_with 'Q & A'
+    end
 
-      expect(fake_relation).to receive(:count).and_call_original
-      expect(subject.rows_expected).to eq 3
+    it 'gets data by calling method provided as hash value' do
+      expect(csv_data).to include 'answer'
     end
-  end
 
-  it 'avoids loading all data in a single query' do
-    expect(fake_relation).to receive(:find_each)
+    it 'allows lamdas to look up more complicated data' do
+      expect(csv_data).to include 'rewsna'
+    end
 
-    subject.render
-  end
+    describe 'excel sanitization' do
+      let(:dangerous_title) { double(title: "=cmd|' /C calc'!A0 title", description: "*safe_desc") }
+      let(:dangerous_desc) { double(title: "*safe_title", description: "=cmd|' /C calc'!A0 desc") }
+      let(:items) { [dangerous_title, dangerous_desc] }
+      let(:subject) { described_class.new(enumerable, 'Title' => 'title', 'Description' => 'description') }
+      let(:csv_data) { subject.render }
 
-  it 'uses hash keys as headers' do
-    expect(csv_data).to start_with 'Q & A'
-  end
+      it 'sanitizes dangerous characters at the beginning of a column' do
+        expect(csv_data).to include "'=cmd|' /C calc'!A0 title"
+        expect(csv_data).to include "'=cmd|' /C calc'!A0 desc"
+      end
 
-  it 'gets data by calling method provided as hash value' do
-    expect(csv_data).to include 'answer'
-  end
+      it 'does not sanitize safe symbols at the beginning of a column' do
+        expect(csv_data).not_to include "'*safe_desc"
+        expect(csv_data).not_to include "'*safe_title"
+      end
 
-  it 'allows lamdas to look up more complicated data' do
-    expect(csv_data).to include 'rewsna'
-  end
+      context 'when dangerous characters are after a line break' do
+        let(:items) { [double(title: "Safe title", description: "With task list\n-[x] todo 1")] }
 
-  describe 'excel sanitization' do
-    let(:dangerous_title) { double(title: "=cmd|' /C calc'!A0 title", description: "*safe_desc") }
-    let(:dangerous_desc) { double(title: "*safe_title", description: "=cmd|' /C calc'!A0 desc") }
-    let(:fake_relation) { described_class::FakeRelation.new([dangerous_title, dangerous_desc]) }
-    let(:subject) { described_class.new(fake_relation, 'Title' => 'title', 'Description' => 'description') }
-    let(:csv_data) { subject.render }
+        it 'does not append single quote to description' do
+          builder = described_class.new(enumerable, 'Title' => 'title', 'Description' => 'description')
 
-    it 'sanitizes dangerous characters at the beginning of a column' do
-      expect(csv_data).to include "'=cmd|' /C calc'!A0 title"
-      expect(csv_data).to include "'=cmd|' /C calc'!A0 desc"
-    end
+          csv_data = builder.render
 
-    it 'does not sanitize safe symbols at the beginning of a column' do
-      expect(csv_data).not_to include "'*safe_desc"
-      expect(csv_data).not_to include "'*safe_title"
+          expect(csv_data).to eq("Title,Description\nSafe title,\"With task list\n-[x] todo 1\"\n")
+        end
+      end
     end
+  end
 
-    context 'when dangerous characters are after a line break' do
-      it 'does not append single quote to description' do
-        fake_object = double(title: "Safe title", description: "With task list\n-[x] todo 1")
-        fake_relation = described_class::FakeRelation.new([fake_object])
-        builder = described_class.new(fake_relation, 'Title' => 'title', 'Description' => 'description')
+  context 'when ActiveRecord::Relation like object is given' do
+    let(:enumerable) { described_class::FakeRelation.new(items) }
 
-        csv_data = builder.render
+    before do
+      stub_const("#{described_class}::FakeRelation", Array)
 
-        expect(csv_data).to eq("Title,Description\nSafe title,\"With task list\n-[x] todo 1\"\n")
+      described_class::FakeRelation.class_eval do
+        def find_each(&block)
+          each(&block)
+        end
       end
     end
+
+    it_behaves_like 'csv builder examples'
+  end
+
+  context 'when Enumerable like object is given' do
+    let(:enumerable) { items }
+
+    it_behaves_like 'csv builder examples'
   end
 end
diff --git a/gems/gitlab-schema-validation/Gemfile.lock b/gems/gitlab-schema-validation/Gemfile.lock
index 5ad804d3660..124e03422e9 100644
--- a/gems/gitlab-schema-validation/Gemfile.lock
+++ b/gems/gitlab-schema-validation/Gemfile.lock
@@ -40,6 +40,7 @@ GEM
     parser (3.2.2.3)
       ast (~> 2.4.1)
       racc
+    pg (1.5.3)
     pg_query (4.2.1)
       google-protobuf (>= 3.22.3)
     proc_to_ast (0.1.0)
@@ -126,6 +127,7 @@ PLATFORMS
 DEPENDENCIES
   gitlab-schema-validation!
   gitlab-styles (~> 10.1.0)
+  pg (~> 1.5.3)
   pry
   rspec (~> 3.0)
   rspec-benchmark (~> 0.6.0)
diff --git a/gems/gitlab-schema-validation/gitlab-schema-validation.gemspec b/gems/gitlab-schema-validation/gitlab-schema-validation.gemspec
index 47ca8b65b5d..513035dac4b 100644
--- a/gems/gitlab-schema-validation/gitlab-schema-validation.gemspec
+++ b/gems/gitlab-schema-validation/gitlab-schema-validation.gemspec
@@ -23,6 +23,7 @@ Gem::Specification.new do |spec|
   spec.add_runtime_dependency "pg_query"
 
   spec.add_development_dependency "gitlab-styles", "~> 10.1.0"
+  spec.add_development_dependency "pg", "~> 1.5.3"
   spec.add_development_dependency "pry"
   spec.add_development_dependency "rspec", "~> 3.0"
   spec.add_development_dependency "rspec-benchmark", "~> 0.6.0"
diff --git a/gems/gitlab-schema-validation/lib/gitlab/schema/validation.rb b/gems/gitlab-schema-validation/lib/gitlab/schema/validation.rb
index 5211358a197..98da37cf16b 100644
--- a/gems/gitlab-schema-validation/lib/gitlab/schema/validation.rb
+++ b/gems/gitlab-schema-validation/lib/gitlab/schema/validation.rb
@@ -28,8 +28,12 @@ require_relative 'validation/validators/different_definition_triggers'
 require_relative 'validation/validators/extra_triggers'
 require_relative 'validation/validators/missing_triggers'
 
+require_relative 'validation/sources/connection_adapters/base'
+require_relative 'validation/sources/connection_adapters/active_record_adapter'
+require_relative 'validation/sources/connection_adapters/pg_adapter'
 require_relative 'validation/sources/structure_sql'
 require_relative 'validation/sources/database'
+require_relative 'validation/sources/connection'
 
 require_relative 'validation/schema_objects/base'
 require_relative 'validation/schema_objects/column'
diff --git a/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection.rb b/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection.rb
new file mode 100644
index 00000000000..83f282ef56d
--- /dev/null
+++ b/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection.rb
@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Schema
+    module Validation
+      module Sources
+        class AdapterNotSupportedError < StandardError
+          def initialize(adapter)
+            @adapter = adapter
+          end
+
+          def message
+            "#{adapter} is not supported"
+          end
+
+          private
+
+          attr_reader :adapter
+        end
+
+        class Connection
+          CONNECTION_ADAPTERS = {
+            'Gitlab::Database::LoadBalancing::ConnectionProxy' => ConnectionAdapters::ActiveRecordAdapter,
+            'ActiveRecord::ConnectionAdapters::PostgreSQLAdapter' => ConnectionAdapters::ActiveRecordAdapter,
+            'PG::Connection' => ConnectionAdapters::PgAdapter
+          }.freeze
+
+          def initialize(connection)
+            @connection_adapter = fetch_adapter(connection)
+          end
+
+          def current_schema
+            connection_adapter.current_schema
+          end
+
+          def select_rows(sql, schemas = [])
+            connection_adapter.select_rows(sql, schemas)
+          end
+
+          def exec_query(sql, schemas = [])
+            connection_adapter.exec_query(sql, schemas)
+          end
+
+          private
+
+          attr_reader :connection_adapter
+
+          def fetch_adapter(connection)
+            CONNECTION_ADAPTERS.fetch(connection.class.name).new(connection)
+          rescue KeyError => e
+            raise AdapterNotSupportedError, e.key
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection_adapters/active_record_adapter.rb b/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection_adapters/active_record_adapter.rb
new file mode 100644
index 00000000000..10af0fc3647
--- /dev/null
+++ b/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection_adapters/active_record_adapter.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Schema
+    module Validation
+      module Sources
+        module ConnectionAdapters
+          class ActiveRecordAdapter < Base
+            extend Forwardable
+
+            def_delegators :@connection, :current_schema
+
+            def exec_query(sql, schemas)
+              connection.exec_query(sql, nil, schemas)
+            end
+
+            def select_rows(sql, schemas)
+              connection.select_rows(sql, nil, schemas)
+            end
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection_adapters/base.rb b/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection_adapters/base.rb
new file mode 100644
index 00000000000..713cbbe0a6d
--- /dev/null
+++ b/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection_adapters/base.rb
@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Schema
+    module Validation
+      module Sources
+        module ConnectionAdapters
+          class Base
+            def initialize(connection)
+              @connection = connection
+            end
+
+            def current_schema
+              raise NotImplementedError, "#{self.class} does not implement #{__method__}"
+            end
+
+            def select_rows(sql, schemas = [])
+              raise NotImplementedError, "#{self.class} does not implement #{__method__}"
+            end
+
+            def exec_query(sql, schemas = [])
+              raise NotImplementedError, "#{self.class} does not implement #{__method__}"
+            end
+
+            private
+
+            attr_reader :connection
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection_adapters/pg_adapter.rb b/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection_adapters/pg_adapter.rb
new file mode 100644
index 00000000000..4dcaf15be71
--- /dev/null
+++ b/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/connection_adapters/pg_adapter.rb
@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module Schema
+    module Validation
+      module Sources
+        module ConnectionAdapters
+          class PgAdapter < Base
+            def initialize(connection)
+              @connection = connection
+              @connection.type_map_for_results = PG::BasicTypeMapForResults.new(connection)
+            end
+
+            def current_schema
+              connection.exec('SELECT current_schema').first['current_schema']
+            end
+
+            def exec_query(sql, schemas)
+              connection.exec(sql, schemas)
+            end
+
+            def select_rows(sql, schemas)
+              exec_query(sql, schemas).values
+            end
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/database.rb b/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/database.rb
index 8505d1f149a..45ce4d8ebfe 100644
--- a/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/database.rb
+++ b/gems/gitlab-schema-validation/lib/gitlab/schema/validation/sources/database.rb
@@ -8,7 +8,7 @@ module Gitlab
           STATIC_PARTITIONS_SCHEMA = 'gitlab_partitions_static'
 
           def initialize(connection)
-            @connection = connection
+            @connection = Connection.new(connection)
           end
 
           def fetch_index_by_name(index_name)
@@ -102,7 +102,7 @@ module Gitlab
             SQL
             # rubocop:enable Rails/SquishedSQLHeredocs
 
-            connection.select_rows(sql, nil, schemas).to_h
+            connection.select_rows(sql, schemas).to_h
           end
 
           def table_map
@@ -136,7 +136,7 @@ module Gitlab
             SQL
             # rubocop:enable Rails/SquishedSQLHeredocs
 
-            connection.exec_query(sql, nil, schemas).group_by { |row| row['table_name'] }
+            connection.exec_query(sql, schemas).group_by { |row| row['table_name'] }
           end
 
           def fetch_indexes
@@ -148,7 +148,7 @@ module Gitlab
             SQL
             # rubocop:enable Rails/SquishedSQLHeredocs
 
-            connection.select_rows(sql, nil, schemas).to_h
+            connection.select_rows(sql, schemas).to_h
           end
 
           def index_map
@@ -183,7 +183,7 @@ module Gitlab
             SQL
             # rubocop:enable Rails/SquishedSQLHeredocs
 
-            connection.exec_query(sql, nil, [connection.current_schema])
+            connection.exec_query(sql, [connection.current_schema])
           end
         end
       end
diff --git a/gems/gitlab-schema-validation/spec/lib/gitlab/schema/validation/sources/connection_spec.rb b/gems/gitlab-schema-validation/spec/lib/gitlab/schema/validation/sources/connection_spec.rb
new file mode 100644
index 00000000000..f3469665d57
--- /dev/null
+++ b/gems/gitlab-schema-validation/spec/lib/gitlab/schema/validation/sources/connection_spec.rb
@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::Schema::Validation::Sources::Connection, feature_category: :database do
+  let(:sql) { 'SELECT column_one, column_two FROM my_table WHERE schema_name IN ($1);' }
+  let(:schemas) { ['public'] }
+  let(:query_result) do
+    [
+      { 'name' => 'Person one', 'email' => 'person.one@gitlab.com' },
+      { 'name' => 'Person two', 'email' => 'person.two@gitlab.com' }
+    ]
+  end
+
+  let(:rows) { query_result.map(&:values) }
+
+  context 'when using active record for postgres adapter' do
+    let(:schema) { 'public' }
+    let(:connection_class_name) { 'ActiveRecord::ConnectionAdapters::PostgreSQLAdapter' }
+    let(:adapter_class) { Gitlab::Schema::Validation::Sources::ConnectionAdapters::ActiveRecordAdapter }
+
+    it_behaves_like 'connection adapter'
+  end
+
+  context 'when using gitlab active record adapter' do
+    let(:schema) { 'gitlab_main' }
+    let(:connection_class_name) { 'Gitlab::Database::LoadBalancing::ConnectionProxy' }
+    let(:adapter_class) { Gitlab::Schema::Validation::Sources::ConnectionAdapters::ActiveRecordAdapter }
+
+    it_behaves_like 'connection adapter'
+  end
+
+  context 'when using postgres adapter' do
+    let(:schema) { 'public' }
+    let(:connection_class_name) { 'PG::Connection' }
+    let(:adapter_class) { Gitlab::Schema::Validation::Sources::ConnectionAdapters::PgAdapter }
+
+    before do
+      allow(connection_object).to receive(:exec)
+      allow(connection_object).to receive(:type_map_for_results=)
+    end
+
+    it_behaves_like 'connection adapter'
+  end
+
+  context 'when using an unsupported connection adapter' do
+    subject(:connection) { described_class.new(connection_object) }
+
+    let(:connection_class_name) { 'ActiveRecord::ConnectionAdapters::InvalidAdapter' }
+    let(:connection_class) { class_double(Class, name: connection_class_name) }
+    let(:connection_object) { instance_double(connection_class_name, class: connection_class) }
+    let(:error_class) { Gitlab::Schema::Validation::Sources::AdapterNotSupportedError }
+    let(:error_message) { 'ActiveRecord::ConnectionAdapters::InvalidAdapter is not supported' }
+
+    it { expect { connection }.to raise_error(error_class, error_message) }
+  end
+end
diff --git a/gems/gitlab-schema-validation/spec/spec_helper.rb b/gems/gitlab-schema-validation/spec/spec_helper.rb
index c11c5021e3b..4eee709dc67 100644
--- a/gems/gitlab-schema-validation/spec/spec_helper.rb
+++ b/gems/gitlab-schema-validation/spec/spec_helper.rb
@@ -2,6 +2,7 @@
 
 require "gitlab/schema/validation"
 require 'rspec-parameterized'
+require 'pg'
 
 RSpec.configure do |config|
   # Enable flags like --only-failures and --next-failure
diff --git a/gems/gitlab-schema-validation/spec/support/shared_examples/connection_adapter_shared_examples.rb b/gems/gitlab-schema-validation/spec/support/shared_examples/connection_adapter_shared_examples.rb
new file mode 100644
index 00000000000..f2b1ee7a934
--- /dev/null
+++ b/gems/gitlab-schema-validation/spec/support/shared_examples/connection_adapter_shared_examples.rb
@@ -0,0 +1,36 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.shared_examples 'connection adapter' do
+  subject(:connection) { described_class.new(connection_object) }
+
+  let(:connection_class) { class_double(Class, name: connection_class_name) }
+  let(:connection_object) { instance_double(connection_class_name, class: connection_class) }
+  let(:adapter) do
+    instance_double(
+      described_class::CONNECTION_ADAPTERS[connection_class_name],
+      current_schema: schema,
+      exec_query: query_result,
+      select_rows: rows
+    )
+  end
+
+  before do
+    allow(connection).to receive(:connection_adapter).and_return(adapter)
+  end
+
+  context 'when using a valid connection adapter' do
+    describe '#current_schema' do
+      it { expect(connection.current_schema).to eq(schema) }
+    end
+
+    describe '#select_rows' do
+      it { expect(connection.select_rows(sql, schemas)).to eq(rows) }
+    end
+
+    describe '#exec_query' do
+      it { expect(connection.exec_query(sql, schemas)).to eq(query_result) }
+    end
+  end
+end
diff --git a/gems/gitlab-schema-validation/spec/support/shared_examples/foreign_key_validators_shared_examples.rb b/gems/gitlab-schema-validation/spec/support/shared_examples/foreign_key_validators_shared_examples.rb
index 1f33c8bd760..8de9c165f8a 100644
--- a/gems/gitlab-schema-validation/spec/support/shared_examples/foreign_key_validators_shared_examples.rb
+++ b/gems/gitlab-schema-validation/spec/support/shared_examples/foreign_key_validators_shared_examples.rb
@@ -10,8 +10,11 @@ RSpec.shared_examples 'foreign key validators' do |validator, expected_result|
   let(:inconsistency_type) { validator.to_s }
   let(:database_name) { 'main' }
   let(:schema) { 'public' }
+  let(:connection_class) { class_double(Class, name: 'ActiveRecord::ConnectionAdapters::PostgreSQLAdapter') }
   # rubocop:disable RSpec/VerifiedDoubleReference
-  let(:connection) { instance_double('connection', exec_query: database_query, current_schema: 'public') }
+  let(:connection) do
+    instance_double('connection', class: connection_class, exec_query: database_query, current_schema: 'public')
+  end
   # rubocop:enable RSpec/VerifiedDoubleReference
 
   let(:database) { Gitlab::Schema::Validation::Sources::Database.new(connection) }
diff --git a/gems/gitlab-schema-validation/spec/support/shared_examples/index_validators_shared_examples.rb b/gems/gitlab-schema-validation/spec/support/shared_examples/index_validators_shared_examples.rb
index cc20c0dc765..36a621ffb35 100644
--- a/gems/gitlab-schema-validation/spec/support/shared_examples/index_validators_shared_examples.rb
+++ b/gems/gitlab-schema-validation/spec/support/shared_examples/index_validators_shared_examples.rb
@@ -13,9 +13,12 @@ RSpec.shared_examples 'index validators' do |validator, expected_result|
   end
 
   let(:inconsistency_type) { validator.name }
+  let(:connection_class) { class_double(Class, name: 'ActiveRecord::ConnectionAdapters::PostgreSQLAdapter') }
 
   # rubocop:disable RSpec/VerifiedDoubleReference
-  let(:connection) { instance_double('connection', select_rows: database_indexes, current_schema: 'public') }
+  let(:connection) do
+    instance_double('connection', class: connection_class, select_rows: database_indexes, current_schema: 'public')
+  end
   # rubocop:enable RSpec/VerifiedDoubleReference
 
   let(:schema) { 'public' }
diff --git a/gems/gitlab-schema-validation/spec/support/shared_examples/table_validators_shared_examples.rb b/gems/gitlab-schema-validation/spec/support/shared_examples/table_validators_shared_examples.rb
index d2a51a9b202..5c5a67c4628 100644
--- a/gems/gitlab-schema-validation/spec/support/shared_examples/table_validators_shared_examples.rb
+++ b/gems/gitlab-schema-validation/spec/support/shared_examples/table_validators_shared_examples.rb
@@ -7,9 +7,13 @@ RSpec.shared_examples "table validators" do |validator, expected_result|
 
   let(:structure_file_path) { 'spec/fixtures/structure.sql' }
   let(:inconsistency_type) { validator.to_s }
+  let(:connection_class) { class_double(Class, name: 'ActiveRecord::ConnectionAdapters::PostgreSQLAdapter') }
   # rubocop:disable RSpec/VerifiedDoubleReference
-  let(:connection) { instance_double('connection', exec_query: database_tables, current_schema: 'public') }
+  let(:connection) do
+    instance_double('connection', class: connection_class, exec_query: database_tables, current_schema: 'public')
+  end
   # rubocop:enable RSpec/VerifiedDoubleReference
+
   let(:schema) { 'public' }
   let(:database) { Gitlab::Schema::Validation::Sources::Database.new(connection) }
   let(:structure_file) { Gitlab::Schema::Validation::Sources::StructureSql.new(structure_file_path, schema) }
diff --git a/gems/gitlab-schema-validation/spec/support/shared_examples/trigger_validators_shared_examples.rb b/gems/gitlab-schema-validation/spec/support/shared_examples/trigger_validators_shared_examples.rb
index 45ed87082bb..36eedcde280 100644
--- a/gems/gitlab-schema-validation/spec/support/shared_examples/trigger_validators_shared_examples.rb
+++ b/gems/gitlab-schema-validation/spec/support/shared_examples/trigger_validators_shared_examples.rb
@@ -11,9 +11,12 @@ RSpec.shared_examples 'trigger validators' do |validator, expected_result|
   let(:database_name) { 'main' }
   let(:schema) { 'public' }
   let(:database) { Gitlab::Schema::Validation::Sources::Database.new(connection) }
+  let(:connection_class) { class_double(Class, name: 'ActiveRecord::ConnectionAdapters::PostgreSQLAdapter') }
 
   # rubocop:disable RSpec/VerifiedDoubleReference
-  let(:connection) { instance_double('connection', select_rows: database_triggers, current_schema: 'public') }
+  let(:connection) do
+    instance_double('connection', class: connection_class, select_rows: database_triggers, current_schema: 'public')
+  end
   # rubocop:enable RSpec/VerifiedDoubleReference
 
   let(:database_triggers) do