Bring in line with EE needs

7 files changed, 84 insertions, 51 deletions

diff --git a/lib/gitlab/metrics_dashboard/common_metrics_inserter.rb b/lib/gitlab/metrics_dashboard/common_metrics_inserter.rb
deleted file mode 100644
index 5bc83e6266e..00000000000
--- a/lib/gitlab/metrics_dashboard/common_metrics_inserter.rb
+++ /dev/null
@@ -1,32 +0,0 @@
-# frozen_string_literal: true
-
-module Gitlab
-  module MetricsDashboard
-    class CommonMetricsInserter
-      class << self
-        # For each metric in the dashboard config, attempts to find a corresponding
-        # database record. If found, includes the record's id in the dashboard config.
-        def transform!(dashboard, _project)
-          common_metrics = ::PrometheusMetric.common
-
-          for_metrics(dashboard) do |metric|
-            metric_record = common_metrics.find { |m| m.identifier == metric[:id] }
-            metric[:metric_id] = metric_record.id if metric_record
-          end
-        end
-
-        private
-
-        def for_metrics(dashboard)
-          dashboard[:panel_groups].each do |panel_group|
-            panel_group[:panels].each do |panel|
-              panel[:metrics].each do |metric|
-                yield metric
-              end
-            end
-          end
-        end
-      end
-    end
-  end
-end
diff --git a/lib/gitlab/metrics_dashboard/processor.rb b/lib/gitlab/metrics_dashboard/processor.rb
index 518e0123220..14ee12049d0 100644
--- a/lib/gitlab/metrics_dashboard/processor.rb
+++ b/lib/gitlab/metrics_dashboard/processor.rb
@@ -3,17 +3,24 @@
 module Gitlab
   module MetricsDashboard
     class Processor
-      STAGES = [CommonMetricsInserter, ProjectMetricsInserter, Sorter].freeze
-
-      def initialize(dashboard, project)
+      def initialize(dashboard, project, environment)
         @dashboard = dashboard.deep_transform_keys(&:to_sym)
         @project = project
+        @environment = environment
+      end
+
+      def stages
+        @stages ||= [
+          Stages::CommonMetricsInserter,
+          Stages::ProjectMetricsInserter,
+          Stages::Sorter
+        ].freeze
       end
 
       def process
-        STAGES.each { |stage| stage.transform!(@dashboard, @project) }
+        stages.each { |stage| stage.new(@dashboard, @project, @environment).transform! }
 
-        @dashboard.to_json
+        @dashboard
       end
     end
   end
diff --git a/lib/gitlab/metrics_dashboard/service.rb b/lib/gitlab/metrics_dashboard/service.rb
index 521a3914e9e..01e61b257e2 100644
--- a/lib/gitlab/metrics_dashboard/service.rb
+++ b/lib/gitlab/metrics_dashboard/service.rb
@@ -11,7 +11,7 @@ module Gitlab
       def get_dashboard
         dashboard_string = Rails.cache.fetch(cache_key) { system_dashboard }
 
-        dashboard = JSON.parse(process_dashboard(dashboard_string))
+        dashboard = process_dashboard(dashboard_string)
 
         success(dashboard: dashboard)
       end
@@ -28,7 +28,7 @@ module Gitlab
       end
 
       def process_dashboard(dashboard)
-        Processor.new(dashboard, @project).process
+        Processor.new(dashboard, project, params[:environment]).process
       end
     end
   end
diff --git a/lib/gitlab/metrics_dashboard/stages/base_stage.rb b/lib/gitlab/metrics_dashboard/stages/base_stage.rb
new file mode 100644
index 00000000000..72085e0c09e
--- /dev/null
+++ b/lib/gitlab/metrics_dashboard/stages/base_stage.rb
@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module MetricsDashboard
+    module Stages
+      class BaseStage
+        DEFAULT_PANEL_TYPE = 'area-chart'
+
+        attr_reader :dashboard, :project, :environment
+
+        def initialize(dashboard, project, environment)
+          @dashboard = dashboard
+          @project = project
+          @environment = environment
+        end
+
+        # Entry-point to the stage
+        # @param dashboard [Hash]
+        # @param project [Project]
+        # @param environment [Environment]
+        def transform!
+          raise NotImplementedError
+        end
+
+        protected
+
+        def for_metrics
+          dashboard[:panel_groups].each do |panel_group|
+            panel_group[:panels].each do |panel|
+              panel[:metrics].each do |metric|
+                yield metric
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/metrics_dashboard/stages/common_metrics_inserter.rb b/lib/gitlab/metrics_dashboard/stages/common_metrics_inserter.rb
new file mode 100644
index 00000000000..e85bdb2700b
--- /dev/null
+++ b/lib/gitlab/metrics_dashboard/stages/common_metrics_inserter.rb
@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module MetricsDashboard
+    module Stages
+      class CommonMetricsInserter < BaseStage
+        # For each metric in the dashboard config, attempts to
+        # find a corresponding database record. If found,
+        # includes the record's id in the dashboard config.
+        def transform!
+          common_metrics = ::PrometheusMetric.common
+
+          for_metrics do |metric|
+            metric_record = common_metrics.find { |m| m.identifier == metric[:id] }
+            metric[:metric_id] = metric_record.id if metric_record
+          end
+        end
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/metrics_dashboard/project_metrics_inserter.rb b/lib/gitlab/metrics_dashboard/stages/project_metrics_inserter.rb
index f4463645d2c..7b694c6e5fa 100644
--- a/lib/gitlab/metrics_dashboard/project_metrics_inserter.rb
+++ b/lib/gitlab/metrics_dashboard/stages/project_metrics_inserter.rb
@@ -2,13 +2,11 @@
 
 module Gitlab
   module MetricsDashboard
-    class ProjectMetricsInserter
-      DEFAULT_PANEL_TYPE = 'area-chart'
-
-      class << self
+    module Stages
+      class ProjectMetricsInserter < BaseStage
         # Inserts project-specific metrics into the dashboard config.
         # If there are no project-specific metrics, this will have no effect.
-        def transform!(dashboard, project)
+        def transform!
           project.prometheus_metrics.each do |project_metric|
             group = find_or_create_panel_group(dashboard[:panel_groups], project_metric)
             panel = find_or_create_panel(group[:panels], project_metric)
diff --git a/lib/gitlab/metrics_dashboard/sorter.rb b/lib/gitlab/metrics_dashboard/stages/sorter.rb
index 59b21d207e1..74b596038fe 100644
--- a/lib/gitlab/metrics_dashboard/sorter.rb
+++ b/lib/gitlab/metrics_dashboard/stages/sorter.rb
@@ -2,22 +2,22 @@
 
 module Gitlab
   module MetricsDashboard
-    class Sorter
-      class << self
-        def transform!(dashboard, _project)
-          sort_groups!(dashboard)
-          sort_panels!(dashboard)
+    module Stages
+      class Sorter < BaseStage
+        def transform!
+          sort_groups!
+          sort_panels!
         end
 
         private
 
         # Sorts the groups in the dashboard by the :priority key
-        def sort_groups!(dashboard)
+        def sort_groups!
           dashboard[:panel_groups] = dashboard[:panel_groups].sort_by { |group| -group[:priority].to_i }
         end
 
         # Sorts the panels in the dashboard by the :weight key
-        def sort_panels!(dashboard)
+        def sort_panels!
           dashboard[:panel_groups].each do |group|
             group[:panels] = group[:panels].sort_by { |panel| -panel[:weight].to_i }
           end