Merge branch 'mwaw/refactor_aggregates_class' into 'master'

Clean up Gitlab::Usage::Metrics::Aggregates::Aggregate class

See merge request https://gitlab.com/gitlab-org/gitlab/-/merge_requests/98202



Merged-by: charlie ablett <cablett@gitlab.com>
Approved-by: Tarun Vellishetty <tvellishetty@gitlab.com>
Approved-by: charlie ablett <cablett@gitlab.com>
Co-authored-by: Mikolaj Wawrzyniak <mwawrzyniak@gitlab.com>

lib/gitlab/usage/metrics/aggregates/aggregate.rb

@@ -13,62 +13,72 @@ def initialize(recorded_at)
           end
 
           def all_time_data
-            aggregated_metrics_data(start_date: nil, end_date: nil, time_frame: Gitlab::Usage::TimeFrame::ALL_TIME_TIME_FRAME_NAME)
+            aggregated_metrics_data(Gitlab::Usage::TimeFrame::ALL_TIME_TIME_FRAME_NAME)
           end
 
           def monthly_data
-            aggregated_metrics_data(**monthly_time_range.merge(time_frame: Gitlab::Usage::TimeFrame::TWENTY_EIGHT_DAYS_TIME_FRAME_NAME))
+            aggregated_metrics_data(Gitlab::Usage::TimeFrame::TWENTY_EIGHT_DAYS_TIME_FRAME_NAME)
           end
 
           def weekly_data
-            aggregated_metrics_data(**weekly_time_range.merge(time_frame: Gitlab::Usage::TimeFrame::SEVEN_DAYS_TIME_FRAME_NAME))
+            aggregated_metrics_data(Gitlab::Usage::TimeFrame::SEVEN_DAYS_TIME_FRAME_NAME)
           end
 
           private
 
           attr_accessor :aggregated_metrics, :recorded_at
 
-          def aggregated_metrics_data(start_date:, end_date:, time_frame:)
+          def aggregated_metrics_data(time_frame)
             aggregated_metrics.each_with_object({}) do |aggregation, data|
               next if aggregation[:feature_flag] && Feature.disabled?(aggregation[:feature_flag], type: :development)
               next unless aggregation[:time_frame].include?(time_frame)
 
-              case aggregation[:source]
-              when REDIS_SOURCE
-                if time_frame == Gitlab::Usage::TimeFrame::ALL_TIME_TIME_FRAME_NAME
-                  data[aggregation[:name]] = Gitlab::Utils::UsageData::FALLBACK
-                  Gitlab::ErrorTracking
-                    .track_and_raise_for_dev_exception(
-                      DisallowedAggregationTimeFrame.new("Aggregation time frame: 'all' is not allowed for aggregation with source: '#{REDIS_SOURCE}'")
-                    )
-                else
-                  data[aggregation[:name]] = calculate_count_for_aggregation(aggregation: aggregation, start_date: start_date, end_date: end_date)
-                end
-              when DATABASE_SOURCE
-                data[aggregation[:name]] = calculate_count_for_aggregation(aggregation: aggregation, start_date: start_date, end_date: end_date)
-              else
-                Gitlab::ErrorTracking
-                  .track_and_raise_for_dev_exception(UnknownAggregationSource.new("Aggregation source: '#{aggregation[:source]}' must be included in #{SOURCES.keys}"))
+              data[aggregation[:name]] = calculate_count_for_aggregation(aggregation: aggregation, time_frame: time_frame)
+            end
+          end
 
-                data[aggregation[:name]] = Gitlab::Utils::UsageData::FALLBACK
+          def calculate_count_for_aggregation(aggregation:, time_frame:)
+            with_validate_configuration(aggregation, time_frame) do
+              source = SOURCES[aggregation[:source]]
+
+              if aggregation[:operator] == UNION_OF_AGGREGATED_METRICS
+                source.calculate_metrics_union(**time_constraints(time_frame).merge(metric_names: aggregation[:events], recorded_at: recorded_at))
+              else
+                source.calculate_metrics_intersections(**time_constraints(time_frame).merge(metric_names: aggregation[:events], recorded_at: recorded_at))
               end
             end
+          rescue Gitlab::UsageDataCounters::HLLRedisCounter::EventError, AggregatedMetricError => error
+            failure(error)
           end
 
-          def calculate_count_for_aggregation(aggregation:, start_date:, end_date:)
-            source = SOURCES[aggregation[:source]]
-
-            case aggregation[:operator]
-            when UNION_OF_AGGREGATED_METRICS
-              source.calculate_metrics_union(metric_names: aggregation[:events], start_date: start_date, end_date: end_date, recorded_at: recorded_at)
-            when INTERSECTION_OF_AGGREGATED_METRICS
-              source.calculate_metrics_intersections(metric_names: aggregation[:events], start_date: start_date, end_date: end_date, recorded_at: recorded_at)
-            else
-              Gitlab::ErrorTracking
-                .track_and_raise_for_dev_exception(UnknownAggregationOperator.new("Events should be aggregated with one of operators #{ALLOWED_METRICS_AGGREGATIONS}"))
-              Gitlab::Utils::UsageData::FALLBACK
+          def with_validate_configuration(aggregation, time_frame)
+            source = aggregation[:source]
+
+            unless ALLOWED_METRICS_AGGREGATIONS.include?(aggregation[:operator])
+              return failure(
+                UnknownAggregationOperator
+                  .new("Events should be aggregated with one of operators #{ALLOWED_METRICS_AGGREGATIONS}")
+              )
             end
-          rescue Gitlab::UsageDataCounters::HLLRedisCounter::EventError, AggregatedMetricError => error
+
+            unless SOURCES[source]
+              return failure(
+                UnknownAggregationSource
+                  .new("Aggregation source: '#{source}' must be included in #{SOURCES.keys}")
+              )
+            end
+
+            if time_frame == Gitlab::Usage::TimeFrame::ALL_TIME_TIME_FRAME_NAME && source == REDIS_SOURCE
+              return failure(
+                DisallowedAggregationTimeFrame
+                  .new("Aggregation time frame: 'all' is not allowed for aggregation with source: '#{REDIS_SOURCE}'")
+              )
+            end
+
+            yield
+          end
+
+          def failure(error)
             Gitlab::ErrorTracking.track_and_raise_for_dev_exception(error)
             Gitlab::Utils::UsageData::FALLBACK
           end
@@ -82,6 +92,17 @@ def load_metrics(wildcard)
           def load_yaml_from_path(path)
             YAML.safe_load(File.read(path), aliases: true)&.map(&:with_indifferent_access)
           end
+
+          def time_constraints(time_frame)
+            case time_frame
+            when Gitlab::Usage::TimeFrame::TWENTY_EIGHT_DAYS_TIME_FRAME_NAME
+              monthly_time_range
+            when Gitlab::Usage::TimeFrame::SEVEN_DAYS_TIME_FRAME_NAME
+              weekly_time_range
+            when Gitlab::Usage::TimeFrame::ALL_TIME_TIME_FRAME_NAME
+              { start_date: nil, end_date: nil }
+            end
+          end
         end
       end
     end