feat(live_agg): Enable pre-aggregation flow

Author: vincent-pochet

events-processor/processors/event_processors/enrichment_service.go

@@ -110,8 +110,7 @@ func (s *EventEnrichmentService) enrichWithSubscription(enrichedEvent *models.En
 }
 
 func (s *EventEnrichmentService) enrichWithChargeInfo(enrichedEvent *models.EnrichedEvent) utils.Result[[]*models.EnrichedEvent] {
-	// TODO(pre-aggregation): Remove the NotAPIPostProcessed condition to enable pre-aggregation
-	if !enrichedEvent.InitialEvent.NotAPIPostProcessed() || enrichedEvent.Subscription == nil {
+	if enrichedEvent.Subscription == nil {
 		return utils.SuccessResult([]*models.EnrichedEvent{enrichedEvent})
 	}
 

events-processor/processors/event_processors/enrichment_service_test.go

@@ -86,53 +86,6 @@ func TestEnrichEvent(t *testing.T) {
 		assert.Equal(t, "Error fetching billable metric", result.ErrorMessage())
 	})
 
-	t.Run("When event source is post processed on API and the result is successful", func(t *testing.T) {
-		sqlmock, delete := setupTestEnv(t)
-		defer delete()
-
-		properties := map[string]any{
-			"api_requests": "12.0",
-		}
-
-		event := models.Event{
-			OrganizationID:         "1a901a90-1a90-1a90-1a90-1a901a901a90",
-			ExternalSubscriptionID: "sub_id",
-			Code:                   "api_calls",
-			Timestamp:              1741007009,
-			Source:                 models.HTTP_RUBY,
-			Properties:             properties,
-			SourceMetadata: &models.SourceMetadata{
-				ApiPostProcess: true,
-			},
-		}
-
-		bm := models.BillableMetric{
-			ID:              "bm123",
-			OrganizationID:  event.OrganizationID,
-			Code:            event.Code,
-			AggregationType: models.AggregationTypeSum,
-			FieldName:       "api_requests",
-			Expression:      "",
-			CreatedAt:       time.Now(),
-			UpdatedAt:       time.Now(),
-		}
-		mockBmLookup(sqlmock, &bm)
-
-		sub := models.Subscription{ID: "sub123", PlanID: "plan123"}
-		mockSubscriptionLookup(sqlmock, &sub)
-
-		result := processor.EnrichEvent(&event)
-
-		assert.True(t, result.Success())
-		assert.Equal(t, 1, len(result.Value()))
-
-		eventResult := result.Value()[0]
-		assert.Equal(t, "12.0", *eventResult.Value)
-		assert.Equal(t, "sum", eventResult.AggregationType)
-		assert.Equal(t, "sub123", eventResult.SubscriptionID)
-		assert.Equal(t, "plan123", eventResult.PlanID)
-	})
-
 	t.Run("When timestamp is invalid", func(t *testing.T) {
 		sqlmock, delete := setupTestEnv(t)
 		defer delete()
@@ -195,7 +148,7 @@ func TestEnrichEvent(t *testing.T) {
 		assert.Equal(t, "Error evaluating custom expression", result.ErrorMessage())
 	})
 
-	t.Run("When event source is not post process on API", func(t *testing.T) {
+	t.Run("With a flat filter", func(t *testing.T) {
 		sqlmock, delete := setupTestEnv(t)
 		defer delete()
 
@@ -236,7 +189,7 @@ func TestEnrichEvent(t *testing.T) {
 		assert.Equal(t, "12", *eventResult.Value)
 	})
 
-	t.Run("When event source is not post process on API with multiple flat filters", func(t *testing.T) {
+	t.Run("With multiple flat filters", func(t *testing.T) {
 		sqlmock, delete := setupTestEnv(t)
 		defer delete()
 
@@ -314,7 +267,7 @@ func TestEnrichEvent(t *testing.T) {
 		assert.Equal(t, map[string]string{}, eventResult2.GroupedBy)
 	})
 
-	t.Run("When event source is not post process on API with a flat filter with pricing group keys", func(t *testing.T) {
+	t.Run("With a flat filter with pricing group keys", func(t *testing.T) {
 		sqlmock, delete := setupTestEnv(t)
 		defer delete()
 

events-processor/processors/events.go

@@ -94,10 +94,11 @@ func processEvent(event *models.Event) utils.Result[*models.EnrichedEvent] {
 
 	go processor.ProducerService.ProduceEnrichedEvent(ctx, enrichedEvent)
 
-	// TODO(pre-aggregation): Uncomment to enable the feature
-	// for _, ev := range enrichedEvents {
-	// 	go processor.ProducerService.ProduceEnrichedExpendedEvent(ctx, ev)
-	// }
+	for _, ev := range enrichedEvents {
+		if ev.ChargeID != nil {
+			go processor.ProducerService.ProduceEnrichedExpendedEvent(ctx, ev)
+		}
+	}
 
 	if enrichedEvent.Subscription != nil && event.NotAPIPostProcessed() {
 		payInAdvance := false

events-processor/processors/events_test.go

@@ -165,6 +165,17 @@ func TestProcessEvent(t *testing.T) {
 		sub := models.Subscription{ID: "sub123", PlanID: "plan123"}
 		mockSubscriptionLookup(sqlmock, &sub)
 
+		mockFlatFiltersLookup(sqlmock, []*models.FlatFilter{
+			{
+				OrganizationID:     event.OrganizationID,
+				BillableMetricCode: event.Code,
+				PlanID:             "plan_id",
+				ChargeID:           "charge_idxx",
+				ChargeUpdatedAt:    time.Now(),
+				PayInAdvance:       true,
+			},
+		})
+
 		result := processEvent(&event)
 
 		assert.True(t, result.Success())
@@ -177,7 +188,7 @@ func TestProcessEvent(t *testing.T) {
 		// TODO: Improve this by using channels in the producers methods
 		time.Sleep(50 * time.Millisecond)
 		assert.Equal(t, 1, testProducers.enrichedProducer.ExecutionCount)
-		// TODO(pre-aggregation): assert.Equal(t, 1, testProducers.enrichedExpandedProducer.ExecutionCount)
+		assert.Equal(t, 1, testProducers.enrichedExpandedProducer.ExecutionCount)
 	})
 
 	t.Run("When event source is not post process on API when timestamp is invalid", func(t *testing.T) {
@@ -367,7 +378,7 @@ func TestProcessEvent(t *testing.T) {
 		time.Sleep(50 * time.Millisecond)
 		assert.Equal(t, 1, testProducers.inAdvanceProducer.ExecutionCount)
 		assert.Equal(t, 1, testProducers.enrichedProducer.ExecutionCount)
-		// TODO(pre-aggregation): assert.Equal(t, 1, testProducers.enrichedExpandedProducer.ExecutionCount)
+		assert.Equal(t, 1, testProducers.enrichedExpandedProducer.ExecutionCount)
 
 		assert.Equal(t, 1, flagger.ExecutionCount)
 	})
@@ -441,7 +452,7 @@ func TestProcessEvent(t *testing.T) {
 		// TODO: Improve this by using channels in the producers methods
 		time.Sleep(50 * time.Millisecond)
 		assert.Equal(t, 1, testProducers.enrichedProducer.ExecutionCount)
-		// TODO(pre-aggregation): assert.Equal(t, 2, testProducers.enrichedExpandedProducer.ExecutionCount)
+		assert.Equal(t, 2, testProducers.enrichedExpandedProducer.ExecutionCount)
 	})
 
 	t.Run("When event source is not post processed on API and it matches no charges", func(t *testing.T) {