Extract smaller methods in base class of cycle analytics

1 files changed, 22 insertions, 12 deletions

diff --git a/lib/gitlab/cycle_analytics/base_stage.rb b/lib/gitlab/cycle_analytics/base_stage.rb
index 0095a7b344d..b7f52432df1 100644
--- a/lib/gitlab/cycle_analytics/base_stage.rb
+++ b/lib/gitlab/cycle_analytics/base_stage.rb
@@ -7,6 +7,7 @@ module Gitlab
 
       def initialize(projects:, options:)
         @projects = projects
+        @project = projects.first
         @options = options
       end
 
@@ -23,21 +24,12 @@ module Gitlab
       end
 
       def median
-        ids = @projects.map(&:id)
-        BatchLoader.for(@projects.first.id).batch(key: name) do |project_ids, loader|
-          cte_table = Arel::Table.new("cte_table_for_#{name}")
-          # Build a `SELECT` query. We find the first of the `end_time_attrs` that isn't `NULL` (call this end_time).
-          # Next, we find the first of the start_time_attrs that isn't `NULL` (call this start_time).
-          # We compute the (end_time - start_time) interval, and give it an alias based on the current
-          # cycle analytics stage.
-          interval_query = Arel::Nodes::As.new(cte_table,
-            subtract_datetimes(stage_query(ids), start_time_attrs, end_time_attrs, name.to_s))
-
+        BatchLoader.for(@project.id).batch(key: name) do |project_ids, loader|
           if project_ids.one?
-            loader.call(@projects.first.id, median_datetime(cte_table, interval_query, name))
+            loader.call(@project.id, median_query(project_ids))
           else
             begin
-              median_datetimes(cte_table, interval_query, name, :project_id)&.each do |project_id, median|
+              median_datetimes(cte_table, interval_query(project_ids), name, :project_id)&.each do |project_id, median|
                 loader.call(project_id, median)
               end
             rescue NotSupportedError
@@ -47,10 +39,28 @@ module Gitlab
         end
       end
 
+      def median_query(project_ids)
+        # Build a `SELECT` query. We find the first of the `end_time_attrs` that isn't `NULL` (call this end_time).
+        # Next, we find the first of the start_time_attrs that isn't `NULL` (call this start_time).
+        # We compute the (end_time - start_time) interval, and give it an alias based on the current
+        # cycle analytics stage.
+
+        median_datetime(cte_table, interval_query(project_ids), name)
+      end
+
       def name
         raise NotImplementedError.new("Expected #{self.name} to implement name")
       end
 
+      def cte_table
+        cte_table = Arel::Table.new("cte_table_for_#{name}")
+      end
+
+      def interval_query(project_ids)
+        Arel::Nodes::As.new(cte_table,
+          subtract_datetimes(stage_query(project_ids), start_time_attrs, end_time_attrs, name.to_s))
+      end
+
       private
 
       def event_fetcher