Add BatchLoader as a way to refactor the base stage code

8 files changed, 49 insertions, 45 deletions

diff --git a/app/models/cycle_analytics.rb b/app/models/cycle_analytics.rb
index c7f0e5f5cd8..52eb07ae7d6 100644
--- a/app/models/cycle_analytics.rb
+++ b/app/models/cycle_analytics.rb
@@ -6,9 +6,9 @@ class CycleAnalytics
     @options = options
   end
 
-  def self.all_medians_per_stage(projects, options)
+  def all_medians_per_stage
     STAGES.each_with_object({}) do |stage_name, hsh|
-      hsh[stage_name] = Gitlab::CycleAnalytics::Stage[stage_name].new(projects: projects, options: options).medians&.values || []
+      hsh[stage_name] = self[stage_name].median
     end
   end
 
@@ -31,7 +31,7 @@ class CycleAnalytics
   end
 
   def [](stage_name)
-    Gitlab::CycleAnalytics::Stage[stage_name].new(projects: [@project], options: @options)
+    Gitlab::CycleAnalytics::Stage[stage_name].new(project: @project, options: @options)
   end
 
   private
diff --git a/app/serializers/analytics_stage_entity.rb b/app/serializers/analytics_stage_entity.rb
index ed1103b2064..08fa8b83c24 100644
--- a/app/serializers/analytics_stage_entity.rb
+++ b/app/serializers/analytics_stage_entity.rb
@@ -6,13 +6,7 @@ class AnalyticsStageEntity < Grape::Entity
   expose :legend
   expose :description
 
-  expose :medians, as: :values do |stage|
-    medians = stage.medians
-
-    unless medians.blank?
-      medians.each do |id, median|
-        medians[id] = distance_of_time_in_words(median)
-      end
-    end
+  expose :median, as: :value do |stage|
+    stage.median && !stage.median.blank? ? distance_of_time_in_words(stage.median) : nil
   end
 end
diff --git a/lib/gitlab/cycle_analytics/base_query.rb b/lib/gitlab/cycle_analytics/base_query.rb
index 05263dd30bc..ed050287a9d 100644
--- a/lib/gitlab/cycle_analytics/base_query.rb
+++ b/lib/gitlab/cycle_analytics/base_query.rb
@@ -7,14 +7,14 @@ module Gitlab
 
       private
 
-      def base_query
-        @base_query ||= stage_query
+      def base_query(project_ids = nil)
+        stage_query(project_ids)
       end
 
-      def stage_query
+      def stage_query(project_ids = nil)
         query = mr_closing_issues_table.join(issue_table).on(issue_table[:id].eq(mr_closing_issues_table[:issue_id]))
           .join(issue_metrics_table).on(issue_table[:id].eq(issue_metrics_table[:issue_id]))
-          .where(issue_table[:project_id].in(Arel.sql(@projects.select(:id).to_sql))) # rubocop:disable Gitlab/ModuleWithInstanceVariables
+          .where(issue_table[:project_id].in(project_ids || @project.id)) # rubocop:disable Gitlab/ModuleWithInstanceVariables
           .where(issue_table[:created_at].gteq(@options[:from])) # rubocop:disable Gitlab/ModuleWithInstanceVariables
 
         # Load merge_requests
diff --git a/lib/gitlab/cycle_analytics/base_stage.rb b/lib/gitlab/cycle_analytics/base_stage.rb
index c9de27ec481..d7c98e41c63 100644
--- a/lib/gitlab/cycle_analytics/base_stage.rb
+++ b/lib/gitlab/cycle_analytics/base_stage.rb
@@ -3,8 +3,8 @@ module Gitlab
     class BaseStage
       include BaseQuery
 
-      def initialize(projects:, options:)
-        @projects = projects
+      def initialize(project:, options:)
+        @project = project
         @options = options
       end
 
@@ -20,18 +20,21 @@ module Gitlab
         raise NotImplementedError.new("Expected #{self.name} to implement title")
       end
 
-      def medians
-        cte_table = Arel::Table.new("cte_table_for_#{name}")
+      def median
+        BatchLoader.for(@project.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(base_query.dup, start_time_attrs, end_time_attrs, name.to_s))
+          # 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(base_query(project_ids), start_time_attrs, end_time_attrs, name.to_s))
 
-        median_datetimes(cte_table, interval_query, name)
+          median_datetimes(cte_table, interval_query, name, :project_id)&.each do |project_id, median|
+            loader.call(project_id, median)
+          end
+        end
       end
 
       def name
diff --git a/lib/gitlab/cycle_analytics/production_helper.rb b/lib/gitlab/cycle_analytics/production_helper.rb
index 7a889b3877f..ad1b2a007f6 100644
--- a/lib/gitlab/cycle_analytics/production_helper.rb
+++ b/lib/gitlab/cycle_analytics/production_helper.rb
@@ -1,8 +1,8 @@
 module Gitlab
   module CycleAnalytics
     module ProductionHelper
-      def stage_query
-        super
+      def stage_query(project_ids = nil)
+        super(project_ids)
           .where(mr_metrics_table[:first_deployed_to_production_at]
           .gteq(@options[:from])) # rubocop:disable Gitlab/ModuleWithInstanceVariables
       end
diff --git a/lib/gitlab/cycle_analytics/test_stage.rb b/lib/gitlab/cycle_analytics/test_stage.rb
index 2b5f72bef89..f47f313d35d 100644
--- a/lib/gitlab/cycle_analytics/test_stage.rb
+++ b/lib/gitlab/cycle_analytics/test_stage.rb
@@ -25,11 +25,11 @@ module Gitlab
         _("Total test time for all commits/merges")
       end
 
-      def stage_query
+      def stage_query(project_ids = nil)
         if @options[:branch]
-          super.where(build_table[:ref].eq(@options[:branch]))
+          super(project_ids).where(build_table[:ref].eq(@options[:branch]))
         else
-          super
+          super(project_ids)
         end
       end
     end
diff --git a/lib/gitlab/cycle_analytics/usage_data.rb b/lib/gitlab/cycle_analytics/usage_data.rb
index 43ec9f9c493..0aefd434e14 100644
--- a/lib/gitlab/cycle_analytics/usage_data.rb
+++ b/lib/gitlab/cycle_analytics/usage_data.rb
@@ -15,7 +15,7 @@ module Gitlab
         values = {}
 
         medians_per_stage.each do |stage_name, medians|
-          medians = medians.compact
+          medians = medians.map(&:presence).compact
 
           stage_values = {
             average: calc_average(medians),
@@ -35,7 +35,12 @@ module Gitlab
       private
 
       def medians_per_stage
-        @medians_per_stage ||= ::CycleAnalytics.all_medians_per_stage(projects, options)
+        projects.each_with_object({}) do |project, hsh|
+          ::CycleAnalytics.new(project, options).all_medians_per_stage.each do |stage_name, median|
+            hsh[stage_name] ||= []
+            hsh[stage_name] << median
+          end
+        end
       end
 
       def calc_average(values)
@@ -61,4 +66,3 @@ module Gitlab
     end
   end
 end
-
diff --git a/lib/gitlab/database/median.rb b/lib/gitlab/database/median.rb
index 84d79cc5e19..de2cb040ad6 100644
--- a/lib/gitlab/database/median.rb
+++ b/lib/gitlab/database/median.rb
@@ -2,10 +2,10 @@
 module Gitlab
   module Database
     module Median
-      def median_datetimes(arel_table, query_so_far, column_sym)
+      def median_datetimes(arel_table, query_so_far, column_sym, partition_column)
         median_queries =
           if Gitlab::Database.postgresql?
-            pg_median_datetime_sql(arel_table, query_so_far, column_sym)
+            pg_median_datetime_sql(arel_table, query_so_far, column_sym, partition_column)
           elsif Gitlab::Database.mysql?
             mysql_median_datetime_sql(arel_table, query_so_far, column_sym)
           end
@@ -21,7 +21,7 @@ module Gitlab
 
         if Gitlab::Database.postgresql?
           result.values.map do |id, median|
-            [id, median&.to_f]
+            [id.to_i, median&.to_f]
           end.to_h
         elsif Gitlab::Database.mysql?
           result.to_a.flatten.first
@@ -53,7 +53,7 @@ module Gitlab
         ]
       end
 
-      def pg_median_datetime_sql(arel_table, query_so_far, column_sym)
+      def pg_median_datetime_sql(arel_table, query_so_far, column_sym, partition_column)
         # Create a CTE with the column we're operating on, row number (after sorting by the column
         # we're operating on), and count of the table we're operating on (duplicated across) all rows
         # of the CTE. For example, if we're looking to find the median of the `projects.star_count`
@@ -69,19 +69,22 @@ module Gitlab
           cte_table,
           arel_table
             .project(
-              arel_table[:project_id],
+              arel_table[partition_column],
               arel_table[column_sym].as(column_sym.to_s),
               Arel::Nodes::Over.new(Arel::Nodes::NamedFunction.new("rank", []),
-                                    Arel::Nodes::Window.new.partition(arel_table[:project_id])
+                                    Arel::Nodes::Window.new.partition(arel_table[partition_column])
                                       .order(arel_table[column_sym])).as('row_id'),
-              arel_table.from(arel_table.alias).project("COUNT(*)").where(arel_table[:project_id].eq(arel_table.alias[:project_id])).as('ct')).
+              arel_table.from(arel_table.alias)
+                .project("COUNT(*)")
+                .where(arel_table[partition_column].eq(arel_table.alias[partition_column])).as('ct')).
             # Disallow negative values
             where(arel_table[column_sym].gteq(zero_interval)))
 
         # From the CTE, select either the middle row or the middle two rows (this is accomplished
         # by 'where cte.row_id between cte.ct / 2.0 AND cte.ct / 2.0 + 1'). Find the average of the
         # selected rows, and this is the median value.
-        cte_table.project(cte_table[:project_id])
+        cte_table
+          .project(cte_table[partition_column])
           .project(average([extract_epoch(cte_table[column_sym])], "median"))
           .where(
             Arel::Nodes::Between.new(
@@ -93,8 +96,8 @@ module Gitlab
             )
           )
           .with(query_so_far, cte)
-          .group(cte_table[:project_id])
-          .order(cte_table[:project_id])
+          .group(cte_table[partition_column])
+          .order(cte_table[partition_column])
           .to_sql
       end