refactor(usage): let a store answer precomputed [ING-586]

app/services/billable_metrics/aggregations/base_service.rb

@@ -129,6 +129,8 @@ def empty_results
         result
       end
 
+      delegate :precomputed?, to: :event_store
+
       protected
 
       attr_accessor :event_store,

app/services/events/stores/base_store.rb

@@ -81,6 +81,10 @@ def grouped_by_values?
         grouped_by_values.present?
       end
 
+      def precomputed?
+        false
+      end
+
       def with_grouped_by_values(grouped_by_values, &block)
         previous_grouped_by_values = @grouped_by_values
         return yield block if grouped_by_values.nil?

app/services/events/stores/usage_bucket_store.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module Events
+  module Stores
+    # Answers from the pre-aggregated usage buckets the aggregations they cover, and delegates the
+    # rest to the store it wraps.
+    class UsageBucketStore < SimpleDelegator
+      def initialize(store, usage_buckets:, charge_id:, charge_filter_id:)
+        @usage_buckets = usage_buckets
+        @charge_id = charge_id
+        @charge_filter_id = charge_filter_id
+
+        super(store)
+      end
+
+      def precomputed?
+        true
+      end
+
+      def count
+        aggregation_result
+      end
+
+      def sum(with_count: true)
+        aggregation_result
+      end
+
+      # A non-default `columns` is the presentation breakdown, which the Provider refuses, so the
+      # delegation below is a guard rather than a path.
+      def grouped_count(columns = nil)
+        return super if columns
+
+        grouped_aggregation_results
+      end
+
+      def grouped_sum(columns = nil, with_count: true)
+        return super if columns
+
+        grouped_aggregation_results
+      end
+
+      private
+
+      attr_reader :usage_buckets, :charge_id, :charge_filter_id
+
+      def aggregation_result
+        usage_buckets.aggregation_result_for(charge_id:, charge_filter_id:)
+      end
+
+      def grouped_aggregation_results
+        usage_buckets.grouped_aggregation_results_for(charge_id:, charge_filter_id:)
+      end
+    end
+  end
+end

app/services/fees/charge_service.rb

@@ -10,13 +10,15 @@ def initialize(
       billing_context:,
       cache_middleware: nil,
       filtered_aggregations: nil,
+      provider: nil,
       options: nil,
       plan: nil,
       customer: nil
     )
       @invoice = invoice
       @metered_item = metered_item
       @billing_context = billing_context
+      @provider = provider
       @options = options || Options.default
       @plan = plan
       @customer = customer
@@ -438,25 +440,47 @@ def already_billed?
       true
     end
 
+    # One instance per pricing bucket, shared by the aggregation and the zero-units hydration, so
+    # the two cannot disagree on where the units come from.
     def aggregator(selected_metered_item:)
+      @aggregators ||= {}
+      @aggregators[selected_metered_item] ||= build_aggregator(selected_metered_item)
+    end
+
+    def build_aggregator(selected_metered_item)
       aggregate = true
       aggregate = filtered_aggregations.include?(selected_metered_item.filter_id) unless filtered_aggregations.nil?
 
       BillableMetrics::AggregationFactory.new_instance(
         metered_item: selected_metered_item,
         current_usage: options.current_usage?,
         billing_context:,
-        boundaries: {
-          from_datetime: selected_metered_item.boundaries.charges_from_datetime,
-          to_datetime: selected_metered_item.boundaries.charges_to_datetime,
-          charges_duration: selected_metered_item.boundaries.charges_duration,
-          max_timestamp: selected_metered_item.boundaries.max_timestamp
-        },
+        provider:,
+        boundaries: aggregation_boundaries(selected_metered_item),
         filters: aggregation_filters(selected_metered_item:, bypass_aggregation: !aggregate),
         bypass_aggregation: !aggregate
       )
     end
 
+    # Callers that run one provider for the whole computation pass theirs; the others get one
+    # scoped to this single charge.
+    def provider
+      @provider ||= Events::Stores::Provider.new(
+        organization: billing_context.organization,
+        billing_context:,
+        current_usage: options.current_usage?
+      )
+    end
+
+    def aggregation_boundaries(selected_metered_item)
+      {
+        from_datetime: selected_metered_item.boundaries.charges_from_datetime,
+        to_datetime: selected_metered_item.boundaries.charges_to_datetime,
+        charges_duration: selected_metered_item.boundaries.charges_duration,
+        max_timestamp: selected_metered_item.boundaries.max_timestamp
+      }
+    end
+
     def persist_recurring_value(aggregation_results, selected_metered_item, breakdowns_by_group)
       # TODO: Review recurring product usage persistence. CachedAggregation needs
       # product_id and product_filter_id support before segment-backed values can be persisted.

app/services/subscriptions/charge_cache_middleware.rb

@@ -13,7 +13,8 @@ def initialize(subscription:, charge:, to_datetime:, cache: true, full_usage: fa
       @last_seen_at = last_seen_at || {}
     end
 
-    def call(charge_filter:)
+    def call(charge_filter:, bypass: false)
+      return yield if bypass
       return yield unless cache
 
       # Lazily invalidate the cache when a more recent event was ingested for this charge/filter.

spec/services/events/stores/usage_bucket_store_spec.rb

@@ -0,0 +1,137 @@
+# frozen_string_literal: true
+
+require "rails_helper"
+
+RSpec.describe Events::Stores::UsageBucketStore do
+  subject(:store) do
+    described_class.new(delegated_store, usage_buckets:, charge_id: charge.id, charge_filter_id:)
+  end
+
+  let(:organization) { create(:organization) }
+  let(:subscription) { create(:subscription, organization:) }
+  let(:billing_context) { Billing::Context.from(subscription:) }
+  let(:billable_metric) { create(:sum_billable_metric, organization:) }
+  let(:charge) { create(:standard_charge, plan: subscription.plan, billable_metric:) }
+  let(:boundaries) { {from_datetime: Time.current.beginning_of_month, to_datetime: Time.current} }
+  let(:charge_filter_id) { "" }
+
+  let(:delegated_store) do
+    Events::Stores::PostgresStore.new(code: billable_metric.code, billing_context:, boundaries:)
+  end
+
+  let(:totals) { Events::Stores::UsageBucketSet::Totals.new(units: BigDecimal("42.5"), events_count: 7) }
+  let(:grouped_totals) { {} }
+  let(:usage_buckets) do
+    Events::Stores::UsageBucketSet.new(
+      totals: {[charge.id, ""] => totals},
+      grouped_totals:
+    )
+  end
+
+  describe "#precomputed?" do
+    it "is true, unlike the store it wraps" do
+      expect(store.precomputed?).to be(true)
+      expect(delegated_store.precomputed?).to be(false)
+    end
+  end
+
+  describe "#count" do
+    it "answers the units of the buckets" do
+      expect(store.count).to eq(
+        Events::Stores::BaseStore::AggregationResult.new(value: BigDecimal("42.5"), events_count: 7)
+      )
+    end
+  end
+
+  describe "#sum" do
+    it "answers the units of the buckets" do
+      expect(store.sum).to eq(
+        Events::Stores::BaseStore::AggregationResult.new(value: BigDecimal("42.5"), events_count: 7)
+      )
+    end
+
+    context "when the buckets hold no row for the charge filter" do
+      let(:charge_filter_id) { "unknown" }
+
+      it "answers zero" do
+        expect(store.sum.value).to eq(0)
+      end
+    end
+  end
+
+  describe "#grouped_count" do
+    let(:grouped_totals) do
+      {[charge.id, ""] => {{"region" => "us"} => totals}}
+    end
+
+    it "answers one result per group of the buckets" do
+      expect(store.grouped_count).to eq(
+        [
+          Events::Stores::BaseStore::GroupedAggregationResult.new(
+            groups: {"region" => "us"}, value: BigDecimal("42.5"), events_count: 7
+          )
+        ]
+      )
+    end
+
+    context "when a presentation breakdown is asked for" do
+      before { allow(delegated_store).to receive(:grouped_count).and_return([]) }
+
+      it "delegates it, as the buckets cannot answer it" do
+        store.grouped_count(["region"])
+
+        expect(delegated_store).to have_received(:grouped_count).with(["region"])
+      end
+    end
+  end
+
+  describe "#grouped_sum" do
+    let(:grouped_totals) do
+      {[charge.id, ""] => {{"region" => "us"} => totals}}
+    end
+
+    it "answers one result per group of the buckets" do
+      expect(store.grouped_sum).to eq(
+        [
+          Events::Stores::BaseStore::GroupedAggregationResult.new(
+            groups: {"region" => "us"}, value: BigDecimal("42.5"), events_count: 7
+          )
+        ]
+      )
+    end
+
+    context "when a presentation breakdown is asked for" do
+      before { allow(delegated_store).to receive(:grouped_sum).and_return([]) }
+
+      it "delegates it, as the buckets cannot answer it" do
+        store.grouped_sum(["region"], with_count: false)
+
+        expect(delegated_store).to have_received(:grouped_sum).with(["region"], with_count: false)
+      end
+    end
+  end
+
+  describe "the aggregations the buckets do not cover" do
+    context "when one of them is called" do
+      before { allow(delegated_store).to receive(:max).and_return(nil) }
+
+      it "delegates it to the store it wraps" do
+        store.max
+
+        expect(delegated_store).to have_received(:max)
+      end
+    end
+
+    it "forwards the per-charge state the aggregators write" do
+      store.aggregation_property = billable_metric.field_name
+      store.numeric_property = true
+
+      expect(delegated_store.aggregation_property).to eq(billable_metric.field_name)
+      expect(delegated_store.numeric_property).to be(true)
+    end
+
+    it "mints a plain store for a sibling window, which no bucket answers for" do
+      expect(store.for_window(**boundaries)).to be_a(Events::Stores::PostgresStore)
+    end
+  end
+end

spec/services/fees/charge_service_spec.rb

@@ -4423,7 +4423,7 @@ def create_filter(amount:, values:)
               presentation_by: ["department", "region"]
             )
           )
-        ).twice
+        ).once
       end
 
       context "when presentation_group_keys is empty" do
@@ -4438,7 +4438,7 @@ def create_filter(amount:, values:)
                 charge_id: charge.id
               )
             )
-          ).twice
+          ).once
         end
       end
 
@@ -4454,7 +4454,7 @@ def create_filter(amount:, values:)
                 presentation_by: []
               )
             )
-          ).twice
+          ).once
         end
 
         context "when presentation_group_keys is empty" do
@@ -4469,7 +4469,7 @@ def create_filter(amount:, values:)
                   charge_id: charge.id
                 )
               )
-            ).twice
+            ).once
           end
         end
       end
@@ -4486,7 +4486,7 @@ def create_filter(amount:, values:)
                 presentation_by: ["region"]
               )
             )
-          ).twice
+          ).once
         end
 
         context "when presentation_group_keys is empty" do
@@ -4501,7 +4501,7 @@ def create_filter(amount:, values:)
                   charge_id: charge.id
                 )
               )
-            ).twice
+            ).once
           end
         end
       end
@@ -4518,7 +4518,7 @@ def create_filter(amount:, values:)
                 presentation_by: []
               )
             )
-          ).twice
+          ).once
         end
       end
     end

spec/services/subscriptions/charge_cache_middleware_spec.rb

@@ -37,6 +37,17 @@
       end
     end
 
+    context "when the filter is bypassed" do
+      it "yields and returns the block result without touching the cache" do
+        fees = [build(:charge_fee, subscription:, charge:)]
+
+        result = middleware.call(charge_filter:, bypass: true) { fees }
+
+        expect(result).to eq(fees)
+        expect(Rails.cache.exist?(cache_key)).to be(false)
+      end
+    end
+
     context "when the cache is empty" do
       let(:fee) { build(:charge_fee, subscription:, charge:, amount_cents: 999) }