Add upsert support to .bulk_insert

lib/gitlab/database.rb

@@ -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

spec/lib/gitlab/database_spec.rb

@@ -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