Add returning IDs to Gitlab::Database.bulk_insert

This adds the keyword argument "return_ids" to
Gitlab::Database.bulk_insert. When set to `true` (and PostgreSQL is
used) this method will return an Array of the IDs of the inserted rows,
otherwise it will return an empty Array.

1 files changed, 20 insertions, 0 deletions

diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index 7aeb85b8f5a..fcddfad3f9f 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -202,6 +202,26 @@ describe Gitlab::Database do
     it 'handles non-UTF-8 data' do
       expect { described_class.bulk_insert('test', [{ a: "\255" }]) }.not_to raise_error
     end
+
+    context 'when using PostgreSQL' do
+      before do
+        allow(described_class).to receive(:mysql?).and_return(false)
+      end
+
+      it 'allows the returning of the IDs of the inserted rows' do
+        result = double(:result, values: [['10']])
+
+        expect(connection)
+          .to receive(:execute)
+          .with(/RETURNING id/)
+          .and_return(result)
+
+        ids = described_class
+          .bulk_insert('test', [{ number: 10 }], return_ids: true)
+
+        expect(ids).to eq([10])
+      end
+    end
   end
 
   describe '.create_connection_pool' do