From 88b7fc785bf83d70724433ca43140fddd258c955 Mon Sep 17 00:00:00 2001
From: Marius Bobin <mbobin@gitlab.com>
Date: Fri, 14 Jul 2023 15:01:05 +0000
Subject: [PATCH] Patch Rails 7 for CI partitioning
---
Gemfile.lock | 4 +-
gems/activerecord-gitlab/.rubocop.yml | 16 ++
gems/activerecord-gitlab/Gemfile.lock | 22 +-
.../activerecord-gitlab.gemspec | 3 +-
.../lib/active_record/gitlab_patches.rb | 1 +
.../gitlab_patches/partitioning.rb | 43 ++++
.../associations/builder/association.rb | 21 ++
.../gitlab_patches/partitioning/base.rb | 49 ++++
.../reflection/abstract_reflection.rb | 25 ++
.../reflection/association_reflection.rb | 17 ++
.../reflection/macro_reflection.rb | 19 ++
.../active_record/gitlab_patches/version.rb | 2 +-
.../associations/belongs_to_spec.rb | 52 ++++
.../associations/has_many_spec.rb | 115 +++++++++
.../partitioning/associations/has_one_spec.rb | 101 ++++++++
.../gitlab_patches/partitioning/joins_spec.rb | 41 +++
.../partitioning/preloads_spec.rb | 241 ++++++++++++++++++
.../partitioning/single_model_queries_spec.rb | 110 ++++++++
.../gitlab_patches/rescue_from_spec.rb | 6 +-
gems/activerecord-gitlab/spec/spec_helper.rb | 16 ++
.../spec/support/database.rb | 29 +++
.../spec/support/models.rb | 50 ++++
.../spec/support/query_recorder.rb | 21 ++
23 files changed, 989 insertions(+), 15 deletions(-)
create mode 100644 gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning.rb
create mode 100644 gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/associations/builder/association.rb
create mode 100644 gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/base.rb
create mode 100644 gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/abstract_reflection.rb
create mode 100644 gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/association_reflection.rb
create mode 100644 gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/macro_reflection.rb
create mode 100644 gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/belongs_to_spec.rb
create mode 100644 gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/has_many_spec.rb
create mode 100644 gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/has_one_spec.rb
create mode 100644 gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/joins_spec.rb
create mode 100644 gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/preloads_spec.rb
create mode 100644 gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/single_model_queries_spec.rb
create mode 100644 gems/activerecord-gitlab/spec/support/database.rb
create mode 100644 gems/activerecord-gitlab/spec/support/models.rb
create mode 100644 gems/activerecord-gitlab/spec/support/query_recorder.rb
diff --git a/Gemfile.lock b/Gemfile.lock
index 542591f657c84..6193a196c86a9 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -1,8 +1,8 @@
PATH
remote: gems/activerecord-gitlab
specs:
- activerecord-gitlab (0.1.0)
- activerecord (>= 6.1.7.3)
+ activerecord-gitlab (0.2.0)
+ activerecord (>= 7)
PATH
remote: gems/click_house-client
diff --git a/gems/activerecord-gitlab/.rubocop.yml b/gems/activerecord-gitlab/.rubocop.yml
index 8c670b439d3eb..5fce096769af5 100644
--- a/gems/activerecord-gitlab/.rubocop.yml
+++ b/gems/activerecord-gitlab/.rubocop.yml
@@ -1,2 +1,18 @@
inherit_from:
- ../config/rubocop.yml
+
+# FIXME
+Gitlab/RSpec/AvoidSetup:
+ Enabled: false
+
+Database/EstablishConnection:
+ Enabled: false
+
+Database/MultipleDatabases:
+ Enabled: false
+
+RSpec/EnvAssignment:
+ Enabled: false
+
+RSpec/MultipleMemoizedHelpers:
+ Enabled: false
diff --git a/gems/activerecord-gitlab/Gemfile.lock b/gems/activerecord-gitlab/Gemfile.lock
index b250ae79d4bef..ba0118eb8e397 100644
--- a/gems/activerecord-gitlab/Gemfile.lock
+++ b/gems/activerecord-gitlab/Gemfile.lock
@@ -1,23 +1,22 @@
PATH
remote: .
specs:
- activerecord-gitlab (0.1.0)
- activerecord (>= 6.1.7.3)
+ activerecord-gitlab (0.2.0)
+ activerecord (>= 7)
GEM
remote: https://rubygems.org/
specs:
- activemodel (6.1.7.3)
- activesupport (= 6.1.7.3)
- activerecord (6.1.7.3)
- activemodel (= 6.1.7.3)
- activesupport (= 6.1.7.3)
- activesupport (6.1.7.3)
+ activemodel (7.0.6)
+ activesupport (= 7.0.6)
+ activerecord (7.0.6)
+ activemodel (= 7.0.6)
+ activesupport (= 7.0.6)
+ activesupport (7.0.6)
concurrent-ruby (~> 1.0, >= 1.0.2)
i18n (>= 1.6, < 2)
minitest (>= 5.1)
tzinfo (~> 2.0)
- zeitwerk (~> 2.3)
ast (2.4.2)
concurrent-ruby (1.2.2)
diff-lcs (1.5.0)
@@ -30,6 +29,7 @@ GEM
i18n (1.13.0)
concurrent-ruby (~> 1.0)
json (2.6.3)
+ mini_portile2 (2.8.2)
minitest (5.18.0)
parallel (1.23.0)
parser (3.2.2.3)
@@ -83,10 +83,11 @@ GEM
rubocop-capybara (~> 2.17)
rubocop-factory_bot (~> 2.22)
ruby-progressbar (1.13.0)
+ sqlite3 (1.6.3)
+ mini_portile2 (~> 2.8.0)
tzinfo (2.0.6)
concurrent-ruby (~> 1.0)
unicode-display_width (2.4.2)
- zeitwerk (2.6.8)
PLATFORMS
ruby
@@ -97,6 +98,7 @@ DEPENDENCIES
rspec (~> 3.12)
rubocop (~> 1.50)
rubocop-rspec (~> 2.22)
+ sqlite3 (~> 1.6)
BUNDLED WITH
2.4.14
diff --git a/gems/activerecord-gitlab/activerecord-gitlab.gemspec b/gems/activerecord-gitlab/activerecord-gitlab.gemspec
index 36346b04602af..267938d0de332 100644
--- a/gems/activerecord-gitlab/activerecord-gitlab.gemspec
+++ b/gems/activerecord-gitlab/activerecord-gitlab.gemspec
@@ -18,10 +18,11 @@ Gem::Specification.new do |spec|
spec.files = Dir["lib/**/*.rb"]
spec.require_paths = ["lib"]
- spec.add_runtime_dependency "activerecord", ">= 6.1.7.3"
+ spec.add_runtime_dependency "activerecord", ">= 7"
spec.add_development_dependency "gitlab-styles", "~> 10.1.0"
spec.add_development_dependency "rspec", "~> 3.12"
spec.add_development_dependency "rubocop", "~> 1.50"
spec.add_development_dependency "rubocop-rspec", "~> 2.22"
+ spec.add_development_dependency "sqlite3", "~> 1.6"
end
diff --git a/gems/activerecord-gitlab/lib/active_record/gitlab_patches.rb b/gems/activerecord-gitlab/lib/active_record/gitlab_patches.rb
index 2063c272ed2b6..257602497f0b1 100644
--- a/gems/activerecord-gitlab/lib/active_record/gitlab_patches.rb
+++ b/gems/activerecord-gitlab/lib/active_record/gitlab_patches.rb
@@ -3,6 +3,7 @@
require "active_record"
require_relative "gitlab_patches/version"
require_relative "gitlab_patches/rescue_from"
+require_relative "gitlab_patches/partitioning"
module ActiveRecord
module GitlabPatches
diff --git a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning.rb b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning.rb
new file mode 100644
index 0000000000000..cf0bd4849e2d7
--- /dev/null
+++ b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning.rb
@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require_relative "partitioning/associations/builder/association"
+require_relative "partitioning/reflection/abstract_reflection"
+require_relative "partitioning/reflection/association_reflection"
+require_relative "partitioning/reflection/macro_reflection"
+require_relative "partitioning/base"
+
+module ActiveRecord
+ module GitlabPatches
+ # This allows to filter data by a dedicated column for association and joins to ActiveRecord::Base.
+ #
+ # class ApplicationRecord < ActiveRecord::Base
+ # belongs_to :pipeline,
+ # -> (build) { where(partition_id: build.partition_id) },
+ # partition_foreign_key: :partition_id
+ #
+ # To control the join filter
+ # def self.use_partition_id_filter?
+ # Feature.enabled?(...)
+ # end
+ # end
+ module Partitioning
+ ActiveSupport.on_load(:active_record) do
+ ::ActiveRecord::Associations::Builder::Association.prepend(
+ ActiveRecord::GitlabPatches::Partitioning::Associations::Builder::Association
+ )
+ ::ActiveRecord::Reflection::AbstractReflection.prepend(
+ ActiveRecord::GitlabPatches::Partitioning::Reflection::AbstractReflection
+ )
+ ::ActiveRecord::Reflection::AssociationReflection.prepend(
+ ActiveRecord::GitlabPatches::Partitioning::Reflection::AssociationReflection
+ )
+ ::ActiveRecord::Reflection::MacroReflection.prepend(
+ ActiveRecord::GitlabPatches::Partitioning::Reflection::MacroReflection
+ )
+ ::ActiveRecord::Base.prepend(
+ ActiveRecord::GitlabPatches::Partitioning::Base
+ )
+ end
+ end
+ end
+end
diff --git a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/associations/builder/association.rb b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/associations/builder/association.rb
new file mode 100644
index 0000000000000..3c92ba91c31fe
--- /dev/null
+++ b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/associations/builder/association.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+ module GitlabPatches
+ module Partitioning
+ module Associations
+ module Builder
+ module Association
+ extend ActiveSupport::Concern
+
+ class_methods do
+ def valid_options(options)
+ super + [:partition_foreign_key]
+ end
+ end
+ end
+ end
+ end
+ end
+ end
+end
diff --git a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/base.rb b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/base.rb
new file mode 100644
index 0000000000000..0c8a248b984fe
--- /dev/null
+++ b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/base.rb
@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+if ::ActiveRecord::VERSION::STRING >= "7.1"
+ raise 'New version of active-record detected, please remove or update this patch'
+end
+
+module ActiveRecord
+ module GitlabPatches
+ module Partitioning
+ module Base
+ extend ActiveSupport::Concern
+
+ def _query_constraints_hash
+ constraints_hash = super
+
+ return constraints_hash unless self.class.use_partition_id_filter?
+
+ if self.class.query_constraints_list.nil?
+ { @primary_key => id_in_database } # rubocop:disable Gitlab/ModuleWithInstanceVariables
+ else
+ self.class.query_constraints_list.index_with do |column_name|
+ attribute_in_database(column_name)
+ end
+ end
+ end
+
+ class_methods do
+ def use_partition_id_filter?
+ false
+ end
+
+ def query_constraints(*columns_list)
+ raise ArgumentError, "You must specify at least one column to be used in querying" if columns_list.empty?
+
+ @query_constraints_list = columns_list.map(&:to_s)
+ end
+
+ def query_constraints_list # :nodoc:
+ @query_constraints_list ||= if base_class? || primary_key != base_class.primary_key
+ primary_key if primary_key.is_a?(Array)
+ else
+ base_class.query_constraints_list
+ end
+ end
+ end
+ end
+ end
+ end
+end
diff --git a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/abstract_reflection.rb b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/abstract_reflection.rb
new file mode 100644
index 0000000000000..7532cd120a51e
--- /dev/null
+++ b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/abstract_reflection.rb
@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+ module GitlabPatches
+ module Partitioning
+ module Reflection
+ module AbstractReflection
+ extend ActiveSupport::Concern
+
+ def join_scope(table, foreign_table, foreign_klass)
+ klass_scope = super
+ return klass_scope unless respond_to?(:options)
+
+ partition_foreign_key = options[:partition_foreign_key]
+ if partition_foreign_key && klass.use_partition_id_filter?
+ klass_scope.where!(table[:partition_id].eq(foreign_table[partition_foreign_key]))
+ end
+
+ klass_scope
+ end
+ end
+ end
+ end
+ end
+end
diff --git a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/association_reflection.rb b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/association_reflection.rb
new file mode 100644
index 0000000000000..299ceaab973dd
--- /dev/null
+++ b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/association_reflection.rb
@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+ module GitlabPatches
+ module Partitioning
+ module Reflection
+ module AssociationReflection
+ def check_eager_loadable!
+ return if scope && scope.arity == 1 && options.key?(:partition_foreign_key)
+
+ super
+ end
+ end
+ end
+ end
+ end
+end
diff --git a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/macro_reflection.rb b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/macro_reflection.rb
new file mode 100644
index 0000000000000..7ec7da44253e8
--- /dev/null
+++ b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/partitioning/reflection/macro_reflection.rb
@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module ActiveRecord
+ module GitlabPatches
+ module Partitioning
+ module Reflection
+ module MacroReflection
+ def scope_for(relation, owner = nil)
+ if scope.arity == 1 && owner.nil? && options.key?(:partition_foreign_key)
+ relation
+ else
+ super
+ end
+ end
+ end
+ end
+ end
+ end
+end
diff --git a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/version.rb b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/version.rb
index 4045f1d5f3148..00c5f254da875 100644
--- a/gems/activerecord-gitlab/lib/active_record/gitlab_patches/version.rb
+++ b/gems/activerecord-gitlab/lib/active_record/gitlab_patches/version.rb
@@ -3,7 +3,7 @@
module ActiveRecord
module GitlabPatches
module Version
- VERSION = "0.1.0"
+ VERSION = "0.2.0"
end
end
end
diff --git a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/belongs_to_spec.rb b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/belongs_to_spec.rb
new file mode 100644
index 0000000000000..900a270c0a810
--- /dev/null
+++ b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/belongs_to_spec.rb
@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+RSpec.describe 'ActiveRecord::GitlabPatches::Partitioning::Associations::BelongsTo', :partitioning do
+ let(:pipeline) { Pipeline.create!(partition_id: 100) }
+ let(:job) { Job.create!(pipeline: pipeline, partition_id: pipeline.partition_id) }
+
+ it 'finds associated record using partition_id' do
+ find_statement = <<~SQL.squish
+ SELECT \"pipelines\".*
+ FROM \"pipelines\"
+ WHERE \"pipelines\".\"id\" = #{pipeline.id}
+ AND \"pipelines\".\"partition_id\" = #{job.partition_id}
+ LIMIT 1
+ SQL
+
+ result = QueryRecorder.log do
+ job.reset.pipeline
+ end
+
+ expect(result).to include(find_statement)
+ end
+
+ it 'builds records using partition_id' do
+ pipeline = job.build_pipeline
+
+ expect(pipeline.partition_id).to eq(job.partition_id)
+ end
+
+ it 'saves records using partition_id' do
+ create_statement = <<~SQL.squish
+ INSERT INTO \"pipelines\" (\"partition_id\") VALUES (#{job.partition_id})
+ SQL
+
+ result = QueryRecorder.log do
+ job.build_pipeline.save!
+ end
+
+ expect(result).to include(create_statement)
+ end
+
+ it 'creates records using partition_id' do
+ create_statement = <<~SQL.squish
+ INSERT INTO \"pipelines\" (\"partition_id\") VALUES (#{job.partition_id})
+ SQL
+
+ result = QueryRecorder.log do
+ job.create_pipeline!
+ end
+
+ expect(result).to include(create_statement)
+ end
+end
diff --git a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/has_many_spec.rb b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/has_many_spec.rb
new file mode 100644
index 0000000000000..3d6b24de998d9
--- /dev/null
+++ b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/has_many_spec.rb
@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+RSpec.describe 'ActiveRecord::GitlabPatches::Partitioning::Associations::HasMany', :partitioning do
+ let(:pipeline) { Pipeline.create!(partition_id: 100) }
+ let(:job) { Job.create!(pipeline: pipeline, partition_id: pipeline.partition_id) }
+
+ it 'finds individual records using partition_id' do
+ find_statement = <<~SQL.squish
+ SELECT \"jobs\".*
+ FROM \"jobs\"
+ WHERE \"jobs\".\"pipeline_id\" = #{pipeline.id}
+ AND \"jobs\".\"partition_id\" = #{pipeline.partition_id}
+ AND \"jobs\".\"id\" = #{job.id}
+ LIMIT 1
+ SQL
+
+ result = QueryRecorder.log do
+ pipeline.jobs.find(job.id)
+ end
+
+ expect(result).to include(find_statement)
+ end
+
+ it 'finds all records using partition_id' do
+ find_statement = <<~SQL.squish
+ SELECT \"jobs\".*
+ FROM \"jobs\"
+ WHERE \"jobs\".\"pipeline_id\" = #{pipeline.id}
+ AND \"jobs\".\"partition_id\" = #{pipeline.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ pipeline.jobs.all.to_a
+ end
+
+ expect(result).to include(find_statement)
+ end
+
+ it 'jobs records using partition_id' do
+ build = pipeline.jobs.new(name: 'test job')
+
+ expect(build.pipeline_id).to eq(pipeline.id)
+ expect(build.partition_id).to eq(pipeline.partition_id)
+ end
+
+ it 'saves records using partition_id' do
+ create_statement = <<~SQL.squish
+ INSERT INTO \"jobs\" (\"pipeline_id\", \"partition_id\", \"name\")
+ VALUES (#{pipeline.id}, #{pipeline.partition_id}, 'test job')
+ SQL
+
+ result = QueryRecorder.log do
+ build = pipeline.jobs.new(name: 'test job')
+ build.save!
+ end
+
+ expect(result).to include(create_statement)
+ end
+
+ it 'creates records using partition_id' do
+ create_statement = <<~SQL.squish
+ INSERT INTO \"jobs\" (\"pipeline_id\", \"partition_id\", \"name\")
+ VALUES (#{pipeline.id}, #{pipeline.partition_id}, 'test job')
+ SQL
+
+ result = QueryRecorder.log do
+ pipeline.jobs.create!(name: 'test job')
+ end
+
+ expect(result).to include(create_statement)
+ end
+
+ it 'deletes_all records using partition_id' do
+ delete_statement = <<~SQL.squish
+ DELETE FROM \"jobs\"
+ WHERE \"jobs\".\"pipeline_id\" = #{pipeline.id}
+ AND \"jobs\".\"partition_id\" = #{pipeline.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ pipeline.jobs.delete_all
+ end
+
+ expect(result).to include(delete_statement)
+ end
+
+ it 'destroy_all records using partition_id' do
+ destroy_statement = <<~SQL.squish
+ DELETE FROM \"jobs\"
+ WHERE \"jobs\".\"id\" = #{job.id}
+ AND \"jobs\".\"partition_id\" = #{pipeline.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ pipeline.jobs.destroy_all # rubocop: disable Cop/DestroyAll
+ end
+
+ expect(result).to include(destroy_statement)
+ end
+
+ it 'counts records using partition_id' do
+ destroy_statement = <<~SQL.squish
+ SELECT COUNT(*)
+ FROM \"jobs\"
+ WHERE \"jobs\".\"pipeline_id\" = #{pipeline.id}
+ AND \"jobs\".\"partition_id\" = #{pipeline.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ pipeline.jobs.count
+ end
+
+ expect(result).to include(destroy_statement)
+ end
+end
diff --git a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/has_one_spec.rb b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/has_one_spec.rb
new file mode 100644
index 0000000000000..aeb565c6dad22
--- /dev/null
+++ b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/associations/has_one_spec.rb
@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+RSpec.describe 'ActiveRecord::GitlabPatches::Partitioning::Associations::HasOne', :partitioning do
+ let(:pipeline) { Pipeline.create!(partition_id: 100) }
+ let(:job) { Job.create!(pipeline: pipeline, partition_id: pipeline.partition_id) }
+
+ it 'finds associated record using partition_id' do
+ find_statement = <<~SQL.squish
+ SELECT \"metadata\".*
+ FROM \"metadata\"
+ WHERE \"metadata\".\"job_id\" = #{job.id}
+ AND \"metadata\".\"partition_id\" = #{job.partition_id}
+ LIMIT 1
+ SQL
+
+ result = QueryRecorder.log do
+ job.reset.metadata
+ end
+
+ expect(result).to include(find_statement)
+ end
+
+ it 'builds records using partition_id' do
+ metadata = job.build_metadata
+
+ expect(metadata.job_id).to eq(job.id)
+ expect(metadata.partition_id).to eq(job.partition_id)
+ end
+
+ it 'saves records using partition_id' do
+ create_statement = <<~SQL.squish
+ INSERT INTO \"metadata\" (\"job_id\", \"partition_id\") VALUES (#{job.id}, #{job.partition_id})
+ SQL
+
+ result = QueryRecorder.log do
+ job.build_metadata.save!
+ end
+
+ expect(result).to include(create_statement)
+ end
+
+ it 'creates records using partition_id' do
+ create_statement = <<~SQL.squish
+ INSERT INTO \"metadata\" (\"job_id\", \"partition_id\") VALUES (#{job.id}, #{job.partition_id})
+ SQL
+
+ result = QueryRecorder.log do
+ job.create_metadata
+ end
+
+ expect(result).to include(create_statement)
+ end
+
+ it 'uses nested attributes on create' do
+ skip '`partitionable` will assign the `partition_id` value in this case.'
+
+ statement1 = <<~SQL.squish
+ INSERT INTO \"jobs\" (\"pipeline_id\", \"partition_id\", \"name\")
+ VALUES (#{pipeline.id}, #{pipeline.partition_id}, 'test')
+ SQL
+
+ statement2 = <<~SQL.squish
+ INSERT INTO \"metadata\" (\"job_id\", \"partition_id\", \"test_flag\")
+ VALUES (#{job.id}, #{job.partition_id}, 1)
+ SQL
+
+ insert_statements = [statement1, statement2]
+
+ result = QueryRecorder.log do
+ pipeline.jobs.create!(name: 'test', metadata_attributes: { test_flag: true })
+ end
+
+ insert_statements.each do |statement|
+ expect(result).to include(statement)
+ end
+ end
+
+ it 'uses nested attributes on update' do
+ statement1 = <<~SQL.squish
+ UPDATE \"jobs\" SET \"name\" = 'other test'
+ WHERE \"jobs\".\"id\" = #{job.id} AND \"jobs\".\"partition_id\" = #{job.partition_id}
+ SQL
+
+ statement2 = <<~SQL.squish
+ INSERT INTO \"metadata\" (\"job_id\", \"partition_id\", \"test_flag\") VALUES (#{job.id}, #{job.partition_id}, 1)
+ SQL
+
+ update_statements = [statement1, statement2]
+
+ job.name = 'other test'
+ job.metadata_attributes = { test_flag: true }
+
+ result = QueryRecorder.log do
+ job.save!
+ end
+
+ update_statements.each do |statement|
+ expect(result).to include(statement)
+ end
+ end
+end
diff --git a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/joins_spec.rb b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/joins_spec.rb
new file mode 100644
index 0000000000000..038fae4364482
--- /dev/null
+++ b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/joins_spec.rb
@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+RSpec.describe 'ActiveRecord::GitlabPatches::Partitioning::Associations::Joins', :partitioning do
+ let!(:pipeline) { Pipeline.create!(partition_id: 100) }
+ let!(:job) { Job.create!(pipeline: pipeline, partition_id: pipeline.partition_id) }
+ let!(:metadata) { Metadata.create!(job: job, partition_id: job.partition_id) }
+
+ it 'joins using partition_id' do
+ join_statement = <<~SQL.squish
+ SELECT \"pipelines\".*
+ FROM \"pipelines\"
+ INNER JOIN \"jobs\" ON \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
+ AND \"jobs\".\"partition_id\" = \"pipelines\".\"partition_id\"
+ WHERE \"pipelines\".\"partition_id\" = #{pipeline.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ Pipeline.where(partition_id: pipeline.partition_id).joins(:jobs).to_a
+ end
+
+ expect(result).to include(join_statement)
+ end
+
+ it 'joins other models using partition_id' do
+ join_statement = <<~SQL.squish
+ SELECT \"pipelines\".*
+ FROM \"pipelines\"
+ INNER JOIN \"jobs\" ON \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
+ AND \"jobs\".\"partition_id\" = \"pipelines\".\"partition_id\"
+ INNER JOIN \"metadata\" ON \"metadata\".\"job_id\" = \"jobs\".\"id\"
+ AND \"metadata\".\"partition_id\" = \"jobs\".\"partition_id\"
+ WHERE \"pipelines\".\"partition_id\" = #{pipeline.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ Pipeline.where(partition_id: pipeline.partition_id).joins(jobs: :metadata).to_a
+ end
+
+ expect(result).to include(join_statement)
+ end
+end
diff --git a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/preloads_spec.rb b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/preloads_spec.rb
new file mode 100644
index 0000000000000..f37a563fe9e84
--- /dev/null
+++ b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/preloads_spec.rb
@@ -0,0 +1,241 @@
+# frozen_string_literal: true
+
+RSpec.describe 'ActiveRecord::GitlabPatches::Partitioning::Associations::Preloads', :partitioning do
+ let(:project) { Project.create! }
+
+ let!(:pipeline) { Pipeline.create!(project: project, partition_id: 100) }
+ let!(:other_pipeline) { Pipeline.create!(project: project, partition_id: 100) }
+
+ let!(:job) { Job.create!(pipeline: pipeline, partition_id: pipeline.partition_id) }
+ let!(:other_job) { Job.create!(pipeline: pipeline, partition_id: pipeline.partition_id) }
+
+ describe 'preload queries with single partition' do
+ it 'preloads metadata for jobs' do
+ statement1 = <<~SQL.squish
+ SELECT \"jobs\".* FROM \"jobs\" WHERE \"jobs\".\"partition_id\" = 100
+ SQL
+
+ statement2 = <<~SQL.squish
+ SELECT \"metadata\".* FROM \"metadata\"
+ WHERE \"metadata\".\"partition_id\" = 100 AND \"metadata\".\"job_id\" IN (#{job.id}, #{other_job.id})
+ SQL
+
+ preload_statements = [statement1, statement2]
+
+ result = QueryRecorder.log do
+ Job.where(partition_id: 100).preload(:metadata).to_a
+ end
+
+ preload_statements.each do |statement|
+ expect(result).to include(statement)
+ end
+ end
+
+ it 'preloads jobs for pipelines' do
+ statement1 = <<~SQL.squish
+ SELECT \"pipelines\".* FROM \"pipelines\" WHERE \"pipelines\".\"partition_id\" = 100
+ SQL
+
+ statement2 = <<~SQL.squish
+ SELECT \"jobs\".* FROM \"jobs\"
+ WHERE \"jobs\".\"partition_id\" = 100 AND \"jobs\".\"pipeline_id\" IN (#{pipeline.id}, #{other_pipeline.id})
+ SQL
+
+ preload_statements = [statement1, statement2]
+
+ result = QueryRecorder.log do
+ Pipeline.where(partition_id: 100).preload(:jobs).to_a
+ end
+
+ preload_statements.each do |statement|
+ expect(result).to include(statement)
+ end
+ end
+
+ it 'preloads jobs and metadata for pipelines' do
+ statement1 = <<~SQL.squish
+ SELECT \"pipelines\".* FROM \"pipelines\" WHERE \"pipelines\".\"partition_id\" = 100
+ SQL
+
+ statement2 = <<~SQL.squish
+ SELECT \"jobs\".* FROM \"jobs\"
+ WHERE \"jobs\".\"partition_id\" = 100 AND \"jobs\".\"pipeline_id\" IN (#{pipeline.id}, #{other_pipeline.id})
+ SQL
+
+ statement3 = <<~SQL.squish
+ SELECT \"metadata\".* FROM \"metadata\"
+ WHERE \"metadata\".\"partition_id\" = 100 AND \"metadata\".\"job_id\" IN (#{job.id}, #{other_job.id})
+ SQL
+
+ preload_statements = [statement1, statement2, statement3]
+
+ result = QueryRecorder.log do
+ Pipeline.where(partition_id: 100).preload(jobs: :metadata).to_a
+ end
+
+ preload_statements.each do |statement|
+ expect(result).to include(statement)
+ end
+ end
+ end
+
+ describe 'preload queries with multiple partitions' do
+ let!(:recent_pipeline) { Pipeline.create!(project: project, partition_id: 200) }
+ let!(:test_job) { Job.create!(pipeline: recent_pipeline, partition_id: recent_pipeline.partition_id) }
+ let!(:deploy_job) { Job.create!(pipeline: recent_pipeline, partition_id: recent_pipeline.partition_id) }
+
+ it 'preloads metadata for jobs' do
+ statement1 = <<~SQL.squish
+ SELECT \"jobs\".* FROM \"jobs\" WHERE \"jobs\".\"partition_id\" IN (100, 200)
+ SQL
+
+ statement2 = <<~SQL.squish
+ SELECT \"metadata\".* FROM \"metadata\"
+ WHERE \"metadata\".\"partition_id\" = 100 AND \"metadata\".\"job_id\" IN (#{job.id}, #{other_job.id})
+ SQL
+
+ statement3 = <<~SQL.squish
+ SELECT \"jobs\".* FROM \"jobs\" WHERE \"jobs\".\"partition_id\" IN (100, 200)
+ SQL
+
+ preload_statements = [statement1, statement2, statement3]
+
+ result = QueryRecorder.log do
+ Job.where(partition_id: [100, 200]).preload(:metadata).to_a
+ end
+
+ preload_statements.each do |statement|
+ expect(result).to include(statement)
+ end
+ end
+
+ it 'preloads jobs for pipelines' do
+ statement1 = <<~SQL.squish
+ SELECT \"pipelines\".* FROM \"pipelines\" WHERE \"pipelines\".\"partition_id\" IN (100, 200)
+ SQL
+
+ statement2 = <<~SQL.squish
+ SELECT \"jobs\".* FROM \"jobs\"
+ WHERE \"jobs\".\"partition_id\" = 100 AND \"jobs\".\"pipeline_id\" IN (#{pipeline.id}, #{other_pipeline.id})
+ SQL
+
+ statement3 = <<~SQL.squish
+ SELECT \"jobs\".* FROM \"jobs\"
+ WHERE \"jobs\".\"partition_id\" = 200 AND \"jobs\".\"pipeline_id\" = #{recent_pipeline.id}
+ SQL
+
+ preload_statements = [statement1, statement2, statement3]
+
+ result = QueryRecorder.log do
+ Pipeline.where(partition_id: [100, 200]).preload(:jobs).to_a
+ end
+
+ preload_statements.each do |statement|
+ expect(result).to include(statement)
+ end
+ end
+
+ it 'preloads jobs and metadata for pipelines' do
+ statement1 = <<~SQL.squish
+ SELECT \"pipelines\".* FROM \"pipelines\" WHERE \"pipelines\".\"partition_id\" IN (100, 200)
+ SQL
+
+ statement2 = <<~SQL.squish
+ SELECT \"jobs\".* FROM \"jobs\"
+ WHERE \"jobs\".\"partition_id\" = 100 AND \"jobs\".\"pipeline_id\" IN (#{pipeline.id}, #{other_pipeline.id})
+ SQL
+
+ statement3 = <<~SQL.squish
+ SELECT \"jobs\".* FROM \"jobs\"
+ WHERE \"jobs\".\"partition_id\" = 200 AND \"jobs\".\"pipeline_id\" = #{recent_pipeline.id}
+ SQL
+
+ statement4 = <<~SQL.squish
+ SELECT \"metadata\".* FROM \"metadata\"
+ WHERE \"metadata\".\"partition_id\" = 100 AND \"metadata\".\"job_id\" IN (#{job.id}, #{other_job.id})
+ SQL
+
+ statement5 = <<~SQL.squish
+ SELECT \"metadata\".* FROM \"metadata\"
+ WHERE \"metadata\".\"partition_id\" = 200 AND \"metadata\".\"job_id\" IN (#{test_job.id}, #{deploy_job.id})
+ SQL
+
+ preload_statements = [statement1, statement2, statement3, statement4, statement5]
+
+ result = QueryRecorder.log do
+ Pipeline.where(partition_id: [100, 200]).preload(jobs: :metadata).to_a
+ end
+
+ preload_statements.each do |statement|
+ expect(result).to include(statement)
+ end
+ end
+ end
+
+ describe 'includes queries' do
+ it 'preloads data for pipeline with multiple queries' do
+ statement1 = <<~SQL.squish
+ SELECT \"pipelines\".* FROM \"pipelines\"
+ WHERE \"pipelines\".\"project_id\" = 1 AND \"pipelines\".\"id\"
+ IN (#{pipeline.id}, #{other_pipeline.id}) AND \"pipelines\".\"partition_id\" = 100
+ SQL
+
+ statement2 = <<~SQL.squish
+ SELECT \"jobs\".* FROM \"jobs\"
+ WHERE \"jobs\".\"partition_id\" = 100 AND \"jobs\".\"pipeline_id\" IN (#{pipeline.id}, #{other_pipeline.id})
+ SQL
+
+ statement3 = <<~SQL.squish
+ SELECT \"metadata\".* FROM \"metadata\"
+ WHERE \"metadata\".\"partition_id\" = 100 AND \"metadata\".\"job_id\" IN (#{job.id}, #{other_job.id})
+ SQL
+
+ preload_statements = [statement1, statement2, statement3]
+
+ result = QueryRecorder.log do
+ project.pipelines.includes(jobs: :metadata).where(id: [pipeline.id, other_pipeline.id], partition_id: 100).to_a
+ end
+
+ preload_statements.each do |statement|
+ expect(result).to include(statement)
+ end
+ end
+
+ it 'preloads data for pipeline with join query' do
+ preload_statement = <<~SQL.squish
+ SELECT \"pipelines\".\"id\"
+ AS t0_r0, \"pipelines\".\"project_id\"
+ AS t0_r1, \"pipelines\".\"partition_id\"
+ AS t0_r2, \"jobs\".\"id\"
+ AS t1_r0, \"jobs\".\"pipeline_id\"
+ AS t1_r1, \"jobs\".\"partition_id\"
+ AS t1_r2, \"jobs\".\"name\"
+ AS t1_r3, \"metadata\".\"id\"
+ AS t2_r0, \"metadata\".\"job_id\"
+ AS t2_r1, \"metadata\".\"partition_id\"
+ AS t2_r2, \"metadata\".\"test_flag\"
+ AS t2_r3
+ FROM \"pipelines\"
+ LEFT OUTER JOIN \"jobs\" ON \"jobs\".\"pipeline_id\" = \"pipelines\".\"id\"
+ AND \"jobs\".\"partition_id\" = \"pipelines\".\"partition_id\"
+ LEFT OUTER JOIN \"metadata\" ON \"metadata\".\"job_id\" = \"jobs\".\"id\"
+ AND \"metadata\".\"partition_id\" = \"jobs\".\"partition_id\"
+ WHERE \"pipelines\".\"project_id\" = 1
+ AND \"pipelines\".\"id\"
+ IN (#{pipeline.id}, #{other_pipeline.id})
+ AND \"pipelines\".\"partition_id\" = 100
+ SQL
+
+ result = QueryRecorder.log do
+ project
+ .pipelines
+ .includes(jobs: :metadata)
+ .references(:jobs, :metadata)
+ .where(id: [1, 2], partition_id: 100)
+ .to_a
+ end
+
+ expect(result).to include(preload_statement)
+ end
+ end
+end
diff --git a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/single_model_queries_spec.rb b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/single_model_queries_spec.rb
new file mode 100644
index 0000000000000..b035d7a627767
--- /dev/null
+++ b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/partitioning/single_model_queries_spec.rb
@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+RSpec.describe 'ActiveRecord::GitlabPatches::Partitioning::Associations::SingleModelQueries', :partitioning do
+ let(:project) { Project.create! }
+ let(:pipeline) { Pipeline.create!(project: project, partition_id: 100) }
+ let(:job) { Job.create!(pipeline: pipeline, partition_id: pipeline.partition_id) }
+
+ it 'creates using id and partition_id' do
+ create_statement = <<~SQL.squish
+ INSERT INTO \"jobs\" (\"pipeline_id\", \"partition_id\")
+ VALUES (#{pipeline.id}, #{pipeline.partition_id})
+ SQL
+
+ result = QueryRecorder.log do
+ Job.create!(pipeline_id: pipeline.id, partition_id: pipeline.partition_id)
+ end
+
+ expect(result).to include(create_statement)
+ end
+
+ it 'finds with id and partition_id' do
+ find_statement = <<~SQL.squish
+ SELECT \"jobs\".*
+ FROM \"jobs\"
+ WHERE \"jobs\".\"id\" = #{job.id}
+ AND \"jobs\".\"partition_id\" = #{job.partition_id}
+ LIMIT 1
+ SQL
+
+ result = QueryRecorder.log do
+ Job.find_by!(id: job.id, partition_id: job.partition_id)
+ end
+
+ expect(result).to include(find_statement)
+ end
+
+ it 'saves using id and partition_id' do
+ update_statement = <<~SQL.squish
+ UPDATE \"jobs\"
+ SET \"name\" = 'test'
+ WHERE \"jobs\".\"id\" = #{job.id}
+ AND \"jobs\".\"partition_id\" = #{job.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ job.name = 'test'
+
+ job.save!
+ end
+
+ expect(result).to include(update_statement)
+ end
+
+ it 'updates using id and partition_id' do
+ update_statement = <<~SQL.squish
+ UPDATE \"jobs\"
+ SET \"name\" = 'test2'
+ WHERE \"jobs\".\"id\" = #{job.id}
+ AND \"jobs\".\"partition_id\" = #{job.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ job.update!(name: 'test2')
+ end
+
+ expect(result).to include(update_statement)
+ end
+
+ it 'deletes using id and partition_id' do
+ delete_statement = <<~SQL.squish
+ DELETE FROM \"jobs\"
+ WHERE \"jobs\".\"id\" = #{job.id}
+ AND \"jobs\".\"partition_id\" = #{job.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ job.delete
+ end
+
+ expect(result).to include(delete_statement)
+ end
+
+ it 'destroys using id and partition_id' do
+ destroy_statement = <<~SQL.squish
+ DELETE FROM \"jobs\"
+ WHERE \"jobs\".\"id\" = #{job.id}
+ AND \"jobs\".\"partition_id\" = #{job.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ job.destroy
+ end
+
+ expect(result).to include(destroy_statement)
+ end
+
+ it 'destroy_all using partition_id' do
+ destroy_statement = <<~SQL.squish
+ DELETE FROM \"jobs\"
+ WHERE \"jobs\".\"id\" = #{job.id}
+ AND \"jobs\".\"partition_id\" = #{job.partition_id}
+ SQL
+
+ result = QueryRecorder.log do
+ Job.where(id: job.id).destroy_all # rubocop: disable Cop/DestroyAll
+ end
+
+ expect(result).to include(destroy_statement)
+ end
+end
diff --git a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/rescue_from_spec.rb b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/rescue_from_spec.rb
index 22c8db7b17418..c1537c3bd909c 100644
--- a/gems/activerecord-gitlab/spec/active_record/gitlab_patches/rescue_from_spec.rb
+++ b/gems/activerecord-gitlab/spec/active_record/gitlab_patches/rescue_from_spec.rb
@@ -1,6 +1,6 @@
# frozen_string_literal: true
-RSpec.describe ActiveRecord::GitlabPatches::RescueFrom do
+RSpec.describe ActiveRecord::GitlabPatches::RescueFrom, :without_sqlite3 do
let(:model_with_rescue_from) do
Class.new(ActiveRecord::Base) do
rescue_from ActiveRecord::ConnectionNotEstablished, with: :handle_exception
@@ -16,12 +16,16 @@ def handle_exception(exception); end
end
it 'triggers rescue_from' do
+ stub_const('ModelWithRescueFrom', model_with_rescue_from)
+
expect(model_with_rescue_from).to receive(:handle_exception)
expect { model_with_rescue_from.all.load }.not_to raise_error
end
it 'does not trigger rescue_from' do
+ stub_const('ModelWithoutRescueFrom', model_without_rescue_from)
+
expect { model_without_rescue_from.all.load }.to raise_error(ActiveRecord::ConnectionNotEstablished)
end
end
diff --git a/gems/activerecord-gitlab/spec/spec_helper.rb b/gems/activerecord-gitlab/spec/spec_helper.rb
index 4250023b898ca..548295b3f2cf7 100644
--- a/gems/activerecord-gitlab/spec/spec_helper.rb
+++ b/gems/activerecord-gitlab/spec/spec_helper.rb
@@ -1,5 +1,7 @@
# frozen_string_literal: true
+ENV["RAILS_ENV"] = "test"
+
require "active_record/gitlab_patches"
RSpec.configure do |config|
@@ -12,4 +14,18 @@
config.expect_with :rspec do |c|
c.syntax = :expect
end
+
+ Dir[File.expand_path("spec/support/**/*.rb")].each { |f| require f }
+
+ config.around(:all, :partitioning) do |example|
+ ActiveRecord::Base.transaction do
+ example.run
+
+ raise ActiveRecord::Rollback
+ end
+ end
+
+ config.before(:all, :without_sqlite3) do
+ ActiveRecord::Base.remove_connection
+ end
end
diff --git a/gems/activerecord-gitlab/spec/support/database.rb b/gems/activerecord-gitlab/spec/support/database.rb
new file mode 100644
index 0000000000000..998d945c31150
--- /dev/null
+++ b/gems/activerecord-gitlab/spec/support/database.rb
@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+RSpec.configure do |config|
+ config.before(:all) do
+ ActiveRecord::Base.establish_connection(adapter: "sqlite3", database: ":memory:")
+ ActiveRecord::Base.logger = Logger.new('/dev/null')
+
+ ActiveRecord::Schema.define do
+ create_table :projects, force: true
+
+ create_table :pipelines, force: true do |t|
+ t.integer :project_id
+ t.integer :partition_id
+ end
+
+ create_table :jobs, force: true do |t|
+ t.integer :pipeline_id
+ t.integer :partition_id
+ t.string :name
+ end
+
+ create_table :metadata, force: true do |t|
+ t.integer :job_id
+ t.integer :partition_id
+ t.boolean :test_flag, default: false
+ end
+ end
+ end
+end
diff --git a/gems/activerecord-gitlab/spec/support/models.rb b/gems/activerecord-gitlab/spec/support/models.rb
new file mode 100644
index 0000000000000..c0017656ea814
--- /dev/null
+++ b/gems/activerecord-gitlab/spec/support/models.rb
@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+class PartitionedRecord < ActiveRecord::Base
+ self.abstract_class = true
+
+ def self.use_partition_id_filter?
+ true
+ end
+
+ alias_method :reset, :reload
+end
+
+class Project < ActiveRecord::Base
+ has_many :pipelines
+end
+
+class Pipeline < PartitionedRecord
+ belongs_to :project
+ query_constraints :id, :partition_id
+
+ has_many :jobs,
+ ->(pipeline) { where(partition_id: pipeline.partition_id) },
+ partition_foreign_key: :partition_id,
+ dependent: :destroy
+end
+
+class Job < PartitionedRecord
+ query_constraints :id, :partition_id
+
+ belongs_to :pipeline,
+ ->(build) { where(partition_id: build.partition_id) },
+ partition_foreign_key: :partition_id
+
+ has_one :metadata,
+ ->(build) { where(partition_id: build.partition_id) },
+ foreign_key: :job_id,
+ partition_foreign_key: :partition_id,
+ inverse_of: :job,
+ autosave: true
+
+ accepts_nested_attributes_for :metadata
+end
+
+class Metadata < PartitionedRecord
+ self.table_name = :metadata
+ query_constraints :id, :partition_id
+
+ belongs_to :job,
+ ->(metadata) { where(partition_id: metadata.partition_id) }
+end
diff --git a/gems/activerecord-gitlab/spec/support/query_recorder.rb b/gems/activerecord-gitlab/spec/support/query_recorder.rb
new file mode 100644
index 0000000000000..5129bae92407a
--- /dev/null
+++ b/gems/activerecord-gitlab/spec/support/query_recorder.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+class QueryRecorder
+ attr_reader :log
+
+ def initialize(&block)
+ @log = []
+
+ ActiveRecord::Base.connection.unprepared_statement do
+ ActiveSupport::Notifications.subscribed(method(:callback), 'sql.active_record', &block)
+ end
+ end
+
+ def callback(_name, _start, _finish, _message_id, values)
+ @log << values[:sql]
+ end
+
+ def self.log(&block)
+ new(&block).log
+ end
+end
--
GitLab