feat(event): Use event stores for unique count aggregation

app/models/event.rb

@@ -8,7 +8,6 @@ class Event < EventsRecord
   include OrganizationTimezone
 
   belongs_to :organization
-  belongs_to :quantified_event, optional: true
 
   validates :transaction_id, presence: true, uniqueness: { scope: %i[organization_id external_subscription_id] }
   validates :code, presence: true

app/services/billable_metrics/aggregations/unique_count_service.rb

@@ -162,19 +162,21 @@ def base_scope
         quantified_events = QuantifiedEvent
           .where(organization_id: billable_metric.organization_id)
           .where(billable_metric_id: billable_metric.id)
+          .where(external_subscription_id: subscription.external_id)
 
-        quantified_events = if billable_metric.recurring?
-          quantified_events.where(external_subscription_id: subscription.external_id)
-        else
-          scope = Event.where(external_subscription_id: subscription.external_id)
-            .where('quantified_event_id IS NOT NULL')
-            .where(timestamp: subscription.started_at..)
-
-          scope = scope.where(timestamp: ...subscription.terminated_at) if subscription.terminated?
-
-          quantified_event_ids = scope.pluck('DISTINCT(quantified_event_id)')
+        unless billable_metric.recurring?
+          store = event_store_class.new(
+            code: billable_metric.code,
+            subscription:,
+            boundaries: {
+              from_datetime: subscription.started_at,
+              to_datetime: subscription.terminated_at,
+            },
+            group:,
+          )
+          store.aggregation_property = billable_metric.field_name
 
-          quantified_events.where(id: quantified_event_ids)
+          quantified_events = quantified_events.where(external_id: store.events_values)
         end
 
         return quantified_events unless group

app/services/events/post_process_service.rb

@@ -108,10 +108,6 @@ def should_handle_quantified_event?
     def handle_quantified_event
       service_result = quantified_event_service.call
       service_result.raise_if_error!
-
-      # TODO: Remove this relation
-      event.quantified_event_id = service_result.quantified_event&.id
-      event.save!
     end
 
     def handle_pay_in_advance

app/services/events/stores/clickhouse_store.rb

@@ -10,12 +10,12 @@ class ClickhouseStore < BaseStore
       #       and should be deduplicated depending on the aggregation logic
       def events(force_from: false)
         scope = Clickhouse::EventsRaw.where(external_subscription_id: subscription.external_id)
-          .where('events_raw.timestamp <= ?', to_datetime)
           .where(code:)
           .order(timestamp: :asc)
 
         # TODO: check how to deal with this since events are not stored forever in clickhouse
         scope = scope.where('events_raw.timestamp >= ?', from_datetime) if force_from || use_from_boundary
+        scope = scope.where('events_raw.timestamp <= ?', to_datetime) if to_datetime
         scope = scope.where(numeric_condition) if numeric_property
 
         return scope unless group

app/services/events/stores/postgres_store.rb

@@ -5,11 +5,11 @@ module Stores
     class PostgresStore < BaseStore
       def events(force_from: false)
         scope = Event.where(external_subscription_id: subscription.external_id)
-          .to_datetime(to_datetime)
           .where(code:)
           .order(timestamp: :asc)
 
         scope = scope.from_datetime(from_datetime) if force_from || use_from_boundary
+        scope = scope.to_datetime(to_datetime) if to_datetime
 
         if numeric_property
           scope = scope.where(presence_condition)

db/migrate/20231205153156_remove_events_quantified_events_relation.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+class RemoveEventsQuantifiedEventsRelation < ActiveRecord::Migration[7.0]
+  def up
+    remove_column :events, :quantified_event_id
+  end
+
+  def down
+    add_column :events, :quantified_event_id, :uuid
+    add_index :events, :quantified_event_id
+  end
+end

spec/services/billable_metrics/aggregations/unique_count_service_spec.rb

@@ -60,16 +60,18 @@
   let(:unique_count_event) do
     create(
       :event,
+      organization_id: organization.id,
       code: billable_metric.code,
       external_customer_id: customer.external_id,
       external_subscription_id: subscription.external_id,
       timestamp: added_at,
-      quantified_event:,
+      properties: { unique_id: quantified_event.external_id },
     )
   end
   let(:quantified_event) do
     create(
       :quantified_event,
+      organization:,
       added_at:,
       removed_at:,
       external_subscription_id: subscription.external_id,
@@ -86,6 +88,7 @@
       let(:new_quantified_event) do
         create(
           :quantified_event,
+          organization:,
           added_at: from_datetime + 10.days,
           removed_at:,
           external_subscription_id: subscription.external_id,
@@ -95,11 +98,12 @@
       let(:new_unique_count_event) do
         create(
           :event,
+          organization_id: organization.id,
           code: billable_metric.code,
           external_customer_id: customer.external_id,
           external_subscription_id: subscription.external_id,
           timestamp: from_datetime + 10.days,
-          quantified_event: new_quantified_event,
+          properties: { unique_id: new_quantified_event.external_id },
         )
       end
 
@@ -123,16 +127,18 @@
       let(:new_unique_count_event) do
         create(
           :event,
+          organization_id: organization.id,
           code: billable_metric.code,
           external_customer_id: customer.external_id,
           external_subscription_id: subscription.external_id,
           timestamp: from_datetime + 10.days,
-          quantified_event: new_quantified_event,
+          properties: { unique_id: new_quantified_event.external_id },
         )
       end
       let(:new_quantified_event) do
         create(
           :quantified_event,
+          organization:,
           added_at: from_datetime + 10.days,
           removed_at:,
           external_subscription_id: subscription.external_id,
@@ -275,14 +281,13 @@
       let(:previous_event) do
         create(
           :event,
-          organization:,
+          organization_id: organization.id,
           code: billable_metric.code,
           external_customer_id: customer.external_id,
           external_subscription_id: subscription.external_id,
           timestamp: from_datetime + 5.days,
-          quantified_event: previous_quantified_event,
           properties: {
-            unique_id: '000',
+            unique_id: previous_quantified_event.external_id,
           },
         )
       end
@@ -335,21 +340,22 @@
     end
 
     context 'when event is given' do
-      let(:properties) { { unique_id: '111' } }
+      let(:properties) { { unique_id: new_quantified_event.external_id } }
       let(:pay_in_advance_event) do
         create(
           :event,
+          organization_id: organization.id,
           code: billable_metric.code,
           external_customer_id: customer.external_id,
           external_subscription_id: subscription.external_id,
           timestamp: from_datetime + 10.days,
           properties:,
-          quantified_event: new_quantified_event,
         )
       end
       let(:new_quantified_event) do
         create(
           :quantified_event,
+          organization:,
           added_at: from_datetime + 10.days,
           removed_at:,
           external_subscription_id: subscription.external_id,
@@ -393,20 +399,21 @@
         let(:previous_event) do
           create(
             :event,
+            organization_id: organization.id,
             code: billable_metric.code,
             external_customer_id: customer.external_id,
             external_subscription_id: subscription.external_id,
             timestamp: from_datetime + 5.days,
-            quantified_event: previous_quantified_event,
             properties: {
-              unique_id: '000',
+              unique_id: previous_quantified_event.external_id,
             },
           )
         end
 
         let(:previous_quantified_event) do
           create(
             :quantified_event,
+            organization:,
             added_at: from_datetime + 5.days,
             removed_at:,
             external_id: '000',
@@ -446,9 +453,8 @@
             external_customer_id: customer.external_id,
             external_subscription_id: subscription.external_id,
             timestamp: from_datetime + 5.days,
-            quantified_event: previous_quantified_event,
             properties: {
-              unique_id: '000',
+              unique_id: previous_quantified_event.external_id,
             },
           )
         end

spec/services/billable_metrics/prorated_aggregations/unique_count_service_spec.rb

@@ -263,13 +263,13 @@
       let(:previous_event) do
         create(
           :event,
+          organization_id: organization.id,
           code: billable_metric.code,
           external_customer_id: customer.external_id,
           external_subscription_id: subscription.external_id,
           timestamp: from_datetime + 5.days,
-          quantified_event: previous_quantified_event,
           properties: {
-            unique_id: '000',
+            unique_id: previous_quantified_event.external_id,
           },
         )
       end
@@ -318,16 +318,16 @@
     end
 
     context 'when event is given' do
-      let(:properties) { { unique_id: '111' } }
+      let(:properties) { { unique_id: new_quantified_event.external_id } }
       let(:pay_in_advance_event) do
         create(
           :event,
+          organization_id: organization.id,
           code: billable_metric.code,
           external_customer_id: customer.external_id,
           external_subscription_id: subscription.external_id,
           timestamp: from_datetime + 10.days,
           properties:,
-          quantified_event: new_quantified_event,
         )
       end
 
@@ -360,24 +360,18 @@
         let(:previous_event) do
           create(
             :event,
+            organization_id: organization.id,
             code: billable_metric.code,
             external_customer_id: customer.external_id,
             external_subscription_id: subscription.external_id,
             timestamp: from_datetime + 5.days,
-            quantified_event: previous_quantified_event,
-            properties: {
-              unique_id: '000',
-            },
-            metadata: {
-              current_aggregation: '7',
-              max_aggregation: '7',
-              max_aggregation_with_proration: '5.8',
-            },
+            properties: { unique_id: previous_quantified_event.external_id },
           )
         end
         let(:previous_quantified_event) do
           create(
             :quantified_event,
+            organization:,
             added_at: from_datetime + 5.days,
             removed_at:,
             external_id: '000',
@@ -397,20 +391,21 @@
         let(:previous_event) do
           create(
             :event,
+            organization_id: organization.id,
             code: billable_metric.code,
             external_customer_id: customer.external_id,
             external_subscription_id: subscription.external_id,
             timestamp: from_datetime + 5.days,
-            quantified_event: previous_quantified_event,
             properties: {
-              unique_id: '000',
+              unique_id: previous_quantified_event.external_id,
             },
           )
         end
 
         let(:previous_quantified_event) do
           create(
             :quantified_event,
+            organization:,
             added_at: from_datetime + 5.days,
             removed_at:,
             external_id: '000',

spec/services/events/post_process_service_spec.rb

@@ -21,7 +21,7 @@
   let(:event) do
     create(
       :event,
-      organization:,
+      organization_id: organization.id,
       external_customer_id:,
       external_subscription_id:,
       timestamp:,
@@ -108,15 +108,14 @@
         }
       end
 
-      it 'creates an association with a quantified event' do
+      it 'creates a quantified event' do
         result = nil
 
         aggregate_failures do
           expect { result = process_service.call }
             .to change(QuantifiedEvent, :count).by(1)
 
           expect(result).to be_success
-          expect(event.reload.quantified_event).to be_present
         end
       end
     end