Trigger iid logic from GitHub importer for milestones.

5 files changed, 42 insertions, 6 deletions

diff --git a/app/models/internal_id.rb b/app/models/internal_id.rb
index 4eb211eff61..e7168d49db9 100644
--- a/app/models/internal_id.rb
+++ b/app/models/internal_id.rb
@@ -111,7 +111,7 @@ class InternalId < ActiveRecord::Base
 
     # Generates next internal id and returns it
     def generate
-      subject.transaction do
+      InternalId.transaction do
         # Create a record in internal_ids if one does not yet exist
         # and increment its last value
         #
@@ -125,7 +125,7 @@ class InternalId < ActiveRecord::Base
     #
     # Note this will acquire a ROW SHARE lock on the InternalId record
     def track_greatest(new_value)
-      subject.transaction do
+      InternalId.transaction do
         (lookup || create_record).track_greatest_and_save!(new_value)
       end
     end
@@ -148,7 +148,7 @@ class InternalId < ActiveRecord::Base
     # violation. We can safely roll-back the nested transaction and perform
     # a lookup instead to retrieve the record.
     def create_record
-      subject.transaction(requires_new: true) do
+      InternalId.transaction(requires_new: true) do
         InternalId.create!(
           **scope,
           usage: usage_value,
diff --git a/lib/gitlab/github_import/bulk_importing.rb b/lib/gitlab/github_import/bulk_importing.rb
index 147597289cf..da2f96b5c4b 100644
--- a/lib/gitlab/github_import/bulk_importing.rb
+++ b/lib/gitlab/github_import/bulk_importing.rb
@@ -15,10 +15,12 @@ module Gitlab
       end
 
       # Bulk inserts the given rows into the database.
-      def bulk_insert(model, rows, batch_size: 100)
+      def bulk_insert(model, rows, batch_size: 100, pre_hook: nil)
         rows.each_slice(batch_size) do |slice|
+          pre_hook.call(slice) if pre_hook
           Gitlab::Database.bulk_insert(model.table_name, slice)
         end
+        rows
       end
     end
   end
diff --git a/lib/gitlab/github_import/importer/milestones_importer.rb b/lib/gitlab/github_import/importer/milestones_importer.rb
index c53480e828a..94eb9136b9a 100644
--- a/lib/gitlab/github_import/importer/milestones_importer.rb
+++ b/lib/gitlab/github_import/importer/milestones_importer.rb
@@ -17,10 +17,20 @@ module Gitlab
         end
 
         def execute
-          bulk_insert(Milestone, build_milestones)
+          # We insert records in bulk, by-passing any standard model callbacks.
+          # The pre_hook here makes sure we track internal ids consistently.
+          # Note this has to be called before performing an insert of a batch
+          # because we're outside a transaction scope here.
+          bulk_insert(Milestone, build_milestones, pre_hook: method(:track_greatest_iid))
           build_milestones_cache
         end
 
+        def track_greatest_iid(slice)
+          greatest_iid = slice.max { |e| e[:iid] }[:iid]
+
+          InternalId.track_greatest(nil, { project: project }, :milestones, greatest_iid, ->(_) { project.milestones.maximum(:iid) })
+        end
+
         def build_milestones
           build_database_rows(each_milestone)
         end
diff --git a/spec/lib/gitlab/github_import/bulk_importing_spec.rb b/spec/lib/gitlab/github_import/bulk_importing_spec.rb
index 91229d9c7d4..861710f7e9b 100644
--- a/spec/lib/gitlab/github_import/bulk_importing_spec.rb
+++ b/spec/lib/gitlab/github_import/bulk_importing_spec.rb
@@ -58,5 +58,17 @@ describe Gitlab::GithubImport::BulkImporting do
 
       importer.bulk_insert(model, rows, batch_size: 5)
     end
+
+    it 'calls pre_hook for each slice if given' do
+      rows = [{ title: 'Foo' }] * 10
+      model = double(:model, table_name: 'kittens')
+      pre_hook = double('pre_hook', call: nil)
+      allow(Gitlab::Database).to receive(:bulk_insert)
+
+      expect(pre_hook).to receive(:call).with(rows[0..4])
+      expect(pre_hook).to receive(:call).with(rows[5..9])
+
+      importer.bulk_insert(model, rows, batch_size: 5, pre_hook: pre_hook)
+    end
   end
 end
diff --git a/spec/lib/gitlab/github_import/importer/milestones_importer_spec.rb b/spec/lib/gitlab/github_import/importer/milestones_importer_spec.rb
index b1cac3b6e46..db0be760c7b 100644
--- a/spec/lib/gitlab/github_import/importer/milestones_importer_spec.rb
+++ b/spec/lib/gitlab/github_import/importer/milestones_importer_spec.rb
@@ -29,13 +29,25 @@ describe Gitlab::GithubImport::Importer::MilestonesImporter, :clean_gitlab_redis
 
       expect(importer)
         .to receive(:bulk_insert)
-        .with(Milestone, [milestone_hash])
+        .with(Milestone, [milestone_hash], any_args)
 
       expect(importer)
         .to receive(:build_milestones_cache)
 
       importer.execute
     end
+
+    it 'tracks internal ids' do
+      milestone_hash = { iid: 1, title: '1.0', project_id: project.id }
+      allow(importer)
+        .to receive(:build_milestones)
+        .and_return([milestone_hash])
+
+      expect(InternalId).to receive(:track_greatest)
+        .with(nil, { project: project }, :milestones, 1, any_args)
+
+      importer.execute
+    end
   end
 
   describe '#build_milestones' do