Add a migration to populate fork networks

This uses the existing ForkedProjectLinks

2 files changed, 114 insertions, 0 deletions

diff --git a/lib/gitlab/background_migration/create_fork_network_memberships_range.rb b/lib/gitlab/background_migration/create_fork_network_memberships_range.rb
new file mode 100644
index 00000000000..4b468e9cd58
--- /dev/null
+++ b/lib/gitlab/background_migration/create_fork_network_memberships_range.rb
@@ -0,0 +1,60 @@
+module Gitlab
+  module BackgroundMigration
+    class CreateForkNetworkMembershipsRange
+      RESCHEDULE_DELAY = 15
+
+      class ForkedProjectLink < ActiveRecord::Base
+        self.table_name = 'forked_project_links'
+      end
+
+      def perform(start_id, end_id)
+        log("Creating memberships for forks: #{start_id} - #{end_id}")
+
+        ActiveRecord::Base.connection.execute <<~INSERT_MEMBERS
+          INSERT INTO fork_network_members (fork_network_id, project_id, forked_from_project_id)
+
+          SELECT fork_network_members.fork_network_id,
+                 forked_project_links.forked_to_project_id,
+                 forked_project_links.forked_from_project_id
+
+          FROM forked_project_links
+
+          INNER JOIN fork_network_members
+              ON forked_project_links.forked_from_project_id = fork_network_members.project_id
+
+          WHERE forked_project_links.id BETWEEN #{start_id} AND #{end_id}
+          AND NOT EXISTS (
+            SELECT true
+            FROM fork_network_members existing_members
+            WHERE existing_members.project_id = forked_project_links.forked_to_project_id
+          )
+        INSERT_MEMBERS
+
+        if missing_members?(start_id, end_id)
+          BackgroundMigrationWorker.perform_in(RESCHEDULE_DELAY, "CreateForkNetworkMembershipsRange", [start_id, end_id])
+        end
+      end
+
+      def missing_members?(start_id, end_id)
+        count_sql = <<~MISSING_MEMBERS
+          SELECT COUNT(*)
+
+          FROM forked_project_links
+
+          WHERE NOT EXISTS (
+            SELECT true
+            FROM fork_network_members
+            WHERE fork_network_members.project_id = forked_project_links.forked_to_project_id
+          )
+          AND forked_project_links.id BETWEEN #{start_id} AND #{end_id}
+        MISSING_MEMBERS
+
+        ForkNetworkMember.count_by_sql(count_sql) > 0
+      end
+
+      def log(message)
+        Rails.logger.info("#{self.class.name} - #{message}")
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/background_migration/populate_fork_networks_range.rb b/lib/gitlab/background_migration/populate_fork_networks_range.rb
new file mode 100644
index 00000000000..6c355ed1e75
--- /dev/null
+++ b/lib/gitlab/background_migration/populate_fork_networks_range.rb
@@ -0,0 +1,54 @@
+module Gitlab
+  module BackgroundMigration
+    class PopulateForkNetworksRange
+      def perform(start_id, end_id)
+        log("Creating fork networks for forked project links: #{start_id} - #{end_id}")
+
+        ActiveRecord::Base.connection.execute <<~INSERT_NETWORKS
+        INSERT INTO fork_networks (root_project_id)
+          SELECT DISTINCT forked_project_links.forked_from_project_id
+
+          FROM forked_project_links
+
+          WHERE NOT EXISTS (
+            SELECT true
+            FROM forked_project_links inner_links
+            WHERE inner_links.forked_to_project_id = forked_project_links.forked_from_project_id
+          )
+          AND NOT EXISTS (
+            SELECT true
+            FROM fork_networks
+            WHERE forked_project_links.forked_from_project_id = fork_networks.root_project_id
+          )
+          AND forked_project_links.id BETWEEN #{start_id} AND #{end_id}
+        INSERT_NETWORKS
+
+        log("Creating memberships for root projects: #{start_id} - #{end_id}")
+
+        ActiveRecord::Base.connection.execute <<~INSERT_ROOT
+          INSERT INTO fork_network_members (fork_network_id, project_id)
+          SELECT DISTINCT fork_networks.id, fork_networks.root_project_id
+
+          FROM fork_networks
+
+          INNER JOIN forked_project_links
+              ON forked_project_links.forked_from_project_id = fork_networks.root_project_id
+
+          WHERE NOT EXISTS (
+            SELECT true
+            FROM fork_network_members
+            WHERE fork_network_members.project_id = fork_networks.root_project_id
+          )
+          AND forked_project_links.id BETWEEN #{start_id} AND #{end_id}
+        INSERT_ROOT
+
+        delay = BackgroundMigration::CreateForkNetworkMembershipsRange::RESCHEDULE_DELAY
+        BackgroundMigrationWorker.perform_in(delay, "CreateForkNetworkMembershipsRange", [start_id, end_id])
+      end
+
+      def log(message)
+        Rails.logger.info("#{self.class.name} - #{message}")
+      end
+    end
+  end
+end