Commit e7d94c80 authored by Kamil Trzciński's avatar Kamil Trzciński Committed by Grzegorz Bizon

Support Rails 6 `insert_all!`

This fixes Load Balancing implementation
to support and test `insert_all!`
parent 3294c889
---
title: Support Rails 6 `insert_all!`
merge_request: 26595
author:
type: fixed
......@@ -23,15 +23,16 @@ module Gitlab
update_all
).freeze
NON_STICKY_READS = %i(
select
quote_column_name
).freeze
# hosts - The hosts to use for load balancing.
def initialize(hosts = [])
@load_balancer = LoadBalancer.new(hosts)
end
def select(*args)
read_using_load_balancer(:select, args)
end
def select_all(arel, name = nil, binds = [], preparable: nil)
if arel.respond_to?(:locked) && arel.locked
# SELECT ... FOR UPDATE queries should be sent to the primary.
......@@ -42,6 +43,12 @@ module Gitlab
end
end
NON_STICKY_READS.each do |name|
define_method(name) do |*args, &block|
read_using_load_balancer(name, args, &block)
end
end
STICKY_WRITES.each do |name|
define_method(name) do |*args, &block|
write_using_load_balancer(name, args, sticky: true, &block)
......
......@@ -44,6 +44,17 @@ describe Gitlab::Database::LoadBalancing::ConnectionProxy do
end
end
Gitlab::Database::LoadBalancing::ConnectionProxy::NON_STICKY_READS.each do |name|
describe "#{name}" do
it 'runs the query on the replica' do
expect(proxy).to receive(:read_using_load_balancer)
.with(name, ['foo'])
proxy.send(name, 'foo')
end
end
end
Gitlab::Database::LoadBalancing::ConnectionProxy::STICKY_WRITES.each do |name|
describe "#{name}" do
it 'runs the query on the primary and sticks to it' do
......@@ -55,6 +66,45 @@ describe Gitlab::Database::LoadBalancing::ConnectionProxy do
end
end
describe '.insert_all!' do
before do
ActiveRecord::Schema.define do
create_table :connection_proxy_bulk_insert, force: true do |t|
t.string :name, null: true
end
end
end
after do
ActiveRecord::Schema.define do
drop_table :connection_proxy_bulk_insert, force: true
end
end
let(:model_class) do
Class.new(ApplicationRecord) do
self.table_name = "connection_proxy_bulk_insert"
end
end
it 'inserts data in bulk' do
expect(model_class).to receive(:connection)
.at_least(:once)
.and_return(proxy)
expect(proxy).to receive(:write_using_load_balancer)
.at_least(:once)
.and_call_original
expect do
model_class.insert_all! [
{ name: "item1" },
{ name: "item2" }
]
end.to change { model_class.count }.by(2)
end
end
# We have an extra test for #transaction here to make sure that nested queries
# are also sent to a primary.
describe '#transaction' do
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment