Commit 19f5184e authored by David Fernandez's avatar David Fernandez
Browse files

Add upsert support to .bulk_insert

parent 899e36f7
......@@ -95,6 +95,10 @@ module Gitlab
version.to_f >= 9.6
end
 
def self.upsert_supported?
version.to_f >= 9.5
end
# map some of the function names that changed between PostgreSQL 9 and 10
# https://wiki.postgresql.org/wiki/New_in_postgres_10
def self.pg_wal_lsn_diff
......@@ -158,7 +162,9 @@ module Gitlab
# disable_quote - A key or an Array of keys to exclude from quoting (You
# become responsible for protection from SQL injection for
# these keys!)
def self.bulk_insert(table, rows, return_ids: false, disable_quote: [])
# on_conflict - Defines an upsert. Values can be: :disabled (default) or
# :do_nothing
def self.bulk_insert(table, rows, return_ids: false, disable_quote: [], on_conflict: nil)
return if rows.empty?
 
keys = rows.first.keys
......@@ -176,10 +182,12 @@ module Gitlab
VALUES #{tuples.map { |tuple| "(#{tuple.join(', ')})" }.join(', ')}
EOF
 
if return_ids
sql = "#{sql}RETURNING id"
if upsert_supported? && on_conflict == :do_nothing
sql = "#{sql} ON CONFLICT DO NOTHING"
end
 
sql = "#{sql} RETURNING id" if return_ids
result = connection.execute(sql)
 
if return_ids
......
......@@ -228,6 +228,7 @@ describe Gitlab::Database do
describe '.bulk_insert' do
before do
allow(described_class).to receive(:connection).and_return(connection)
allow(described_class).to receive(:version).and_return(version)
allow(connection).to receive(:quote_column_name, &:itself)
allow(connection).to receive(:quote, &:itself)
allow(connection).to receive(:execute)
......@@ -242,6 +243,8 @@ describe Gitlab::Database do
]
end
 
let_it_be(:version) { 9.6 }
it 'does nothing with empty rows' do
expect(connection).not_to receive(:execute)
 
......@@ -307,6 +310,29 @@ describe Gitlab::Database do
 
expect(ids).to eq([10])
end
context 'with version >= 9.5' do
it 'allows setting the upsert to do nothing' do
expect(connection)
.to receive(:execute)
.with(/ON CONFLICT DO NOTHING/)
described_class
.bulk_insert('test', [{ number: 10 }], on_conflict: :do_nothing)
end
end
context 'with version < 9.5' do
let(:version) { 9.4 }
it 'refuses setting the upsert' do
expect(connection)
.not_to receive(:execute)
.with(/ON CONFLICT/)
described_class
.bulk_insert('test', [{ number: 10 }], on_conflict: :do_nothing)
end
end
end
end
 
......
Supports Markdown
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