feat(pre-aggregation): Extract enrich in a dedicated processor (#561)

Author: vincent-pochet

events-processor/models/event.go

@@ -27,7 +27,9 @@ type SourceMetadata struct {
 }
 
 type EnrichedEvent struct {
-	IntialEvent *Event `json:"-"`
+	IntialEvent    *Event          `json:"-"`
+	BillableMetric *BillableMetric `json:"-"`
+	Subscription   *Subscription   `json:"-"`
 
 	OrganizationID          string         `json:"organization_id"`
 	ExternalSubscriptionID  string         `json:"external_subscription_id"`

events-processor/processors/event_processors/base_service.go

@@ -0,0 +1,13 @@
+package event_processors
+
+import (
+	"github.com/getlago/lago/events-processor/models"
+	"github.com/getlago/lago/events-processor/utils"
+)
+
+func failedResult(r utils.AnyResult, code string, message string) utils.Result[*models.EnrichedEvent] {
+	result := utils.FailedResult[*models.EnrichedEvent](r.Error()).AddErrorDetails(code, message)
+	result.Retryable = r.IsRetryable()
+	result.Capture = r.IsCapturable()
+	return result
+}

events-processor/processors/event_processors/enrichment_service.go

@@ -0,0 +1,105 @@
+package event_processors
+
+import (
+	"encoding/json"
+	"fmt"
+
+	"github.com/getlago/lago-expression/expression-go"
+	"github.com/getlago/lago/events-processor/models"
+	"github.com/getlago/lago/events-processor/utils"
+)
+
+type EventEnrichmentService struct {
+	apiStore *models.ApiStore
+}
+
+func NewEventEnrichmentService(apiStore *models.ApiStore) *EventEnrichmentService {
+	return &EventEnrichmentService{
+		apiStore: apiStore,
+	}
+}
+
+func (s *EventEnrichmentService) EnrichEvent(event *models.Event) utils.Result[*models.EnrichedEvent] {
+	enrichedEventResult := event.ToEnrichedEvent()
+	if enrichedEventResult.Failure() {
+		return failedResult(enrichedEventResult, "build_enriched_event", "Error while converting event to enriched event")
+	}
+	enrichedEvent := enrichedEventResult.Value()
+
+	bmResult := s.apiStore.FetchBillableMetric(event.OrganizationID, event.Code)
+	if bmResult.Failure() {
+		return failedResult(bmResult, "fetch_billable_metric", "Error fetching billable metric")
+	}
+
+	bm := bmResult.Value()
+	if bm != nil {
+		enrichBmResult := s.enrichWithBillableMetric(enrichedEvent, bm)
+		if enrichBmResult.Failure() {
+			return enrichBmResult
+		}
+	}
+
+	subResult := s.apiStore.FetchSubscription(event.OrganizationID, event.ExternalSubscriptionID, enrichedEvent.Time)
+	if subResult.Failure() && subResult.IsCapturable() {
+		// We want to keep processing the event even if the subscription is not found
+		return failedResult(subResult, "fetch_subscription", "Error fetching subscription")
+	}
+
+	sub := subResult.Value()
+	if sub != nil {
+		enrichSubResult := s.enrichWithSubscription(enrichedEvent, sub)
+		if enrichSubResult.Failure() {
+			return enrichSubResult
+		}
+	}
+
+	return utils.SuccessResult(enrichedEvent)
+}
+
+func (s *EventEnrichmentService) enrichWithBillableMetric(enrichedEvent *models.EnrichedEvent, bm *models.BillableMetric) utils.Result[*models.EnrichedEvent] {
+	enrichedEvent.BillableMetric = bm
+	enrichedEvent.AggregationType = bm.AggregationType.String()
+
+	if enrichedEvent.Source != models.HTTP_RUBY {
+		expressionResult := s.evaluateExpression(enrichedEvent, bm)
+		if expressionResult.Failure() {
+			return failedResult(expressionResult, "evaluate_expression", "Error evaluating custom expression")
+		}
+	}
+
+	var value = fmt.Sprintf("%v", enrichedEvent.Properties[bm.FieldName])
+	enrichedEvent.Value = &value
+
+	return utils.SuccessResult(enrichedEvent)
+}
+
+func (s *EventEnrichmentService) evaluateExpression(ev *models.EnrichedEvent, bm *models.BillableMetric) utils.Result[bool] {
+	if bm.Expression == "" {
+		return utils.SuccessResult(false)
+	}
+
+	eventJson, err := json.Marshal(ev)
+	if err != nil {
+		return utils.FailedBoolResult(err).NonRetryable()
+	}
+	eventJsonString := string(eventJson[:])
+
+	result := expression.Evaluate(bm.Expression, eventJsonString)
+	if result != nil {
+		ev.Properties[bm.FieldName] = *result
+	} else {
+		return utils.
+			FailedBoolResult(fmt.Errorf("Failed to evaluate expr: %s with json: %s", bm.Expression, eventJsonString)).
+			NonRetryable()
+	}
+
+	return utils.SuccessResult(true)
+}
+
+func (s *EventEnrichmentService) enrichWithSubscription(enrichedEvent *models.EnrichedEvent, sub *models.Subscription) utils.Result[*models.EnrichedEvent] {
+	enrichedEvent.Subscription = sub
+	enrichedEvent.SubscriptionID = sub.ID
+	enrichedEvent.PlanID = sub.PlanID
+
+	return utils.SuccessResult(enrichedEvent)
+}

events-processor/processors/event_processors/enrichment_service_test.go

@@ -0,0 +1,218 @@
+package event_processors
+
+import (
+	"testing"
+	"time"
+
+	"github.com/DATA-DOG/go-sqlmock"
+	"github.com/getlago/lago/events-processor/models"
+	"github.com/getlago/lago/events-processor/tests"
+	"github.com/getlago/lago/events-processor/utils"
+	"github.com/stretchr/testify/assert"
+	"gorm.io/gorm"
+)
+
+var processor *EventEnrichmentService
+
+func setupTestEnv(t *testing.T) (sqlmock.Sqlmock, func()) {
+	db, mock, delete := tests.SetupMockStore(t)
+	apiStore := models.NewApiStore(db)
+
+	processor = &EventEnrichmentService{
+		apiStore: apiStore,
+	}
+
+	return mock, delete
+}
+
+func mockBmLookup(sqlmock sqlmock.Sqlmock, bm *models.BillableMetric) {
+	columns := []string{"id", "organization_id", "code", "aggregation_type", "field_name", "expression", "created_at", "updated_at", "deleted_at"}
+
+	rows := sqlmock.NewRows(columns).
+		AddRow(bm.ID, bm.OrganizationID, bm.Code, bm.AggregationType, bm.FieldName, bm.Expression, bm.CreatedAt, bm.UpdatedAt, bm.DeletedAt)
+
+	sqlmock.ExpectQuery("SELECT \\* FROM \"billable_metrics\".*").WillReturnRows(rows)
+}
+
+func mockSubscriptionLookup(sqlmock sqlmock.Sqlmock, sub *models.Subscription) {
+	columns := []string{"id", "external_id", "plan_id", "created_at", "updated_at", "terminated_at"}
+
+	rows := sqlmock.NewRows(columns).
+		AddRow(sub.ID, sub.ExternalID, sub.PlanID, sub.CreatedAt, sub.UpdatedAt, sub.TerminatedAt)
+
+	sqlmock.ExpectQuery(".* FROM \"subscriptions\".*").WillReturnRows(rows)
+}
+
+func TestEnrichEvent(t *testing.T) {
+	t.Run("Without Billable Metric", func(t *testing.T) {
+		sqlmock, delete := setupTestEnv(t)
+		defer delete()
+
+		event := models.Event{
+			OrganizationID:         "1a901a90-1a90-1a90-1a90-1a901a901a90",
+			ExternalSubscriptionID: "sub_id",
+			Code:                   "api_calls",
+			Timestamp:              1741007009,
+		}
+
+		sqlmock.ExpectQuery(".*").WillReturnError(gorm.ErrRecordNotFound)
+
+		result := processor.EnrichEvent(&event)
+		assert.False(t, result.Success())
+		assert.Equal(t, "record not found", result.ErrorMsg())
+		assert.Equal(t, "fetch_billable_metric", result.ErrorCode())
+		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, "12.0", *result.Value().Value)
+		assert.Equal(t, "sum", result.Value().AggregationType)
+		assert.Equal(t, "sub123", result.Value().SubscriptionID)
+		assert.Equal(t, "plan123", result.Value().PlanID)
+	})
+
+	t.Run("When timestamp is invalid", func(t *testing.T) {
+		sqlmock, delete := setupTestEnv(t)
+		defer delete()
+
+		event := models.Event{
+			OrganizationID:         "1a901a90-1a90-1a90-1a90-1a901a901a90",
+			ExternalSubscriptionID: "sub_id",
+			Code:                   "api_calls",
+			Timestamp:              "2025-03-06T12:00:00Z",
+			Source:                 "SQS",
+		}
+
+		bm := models.BillableMetric{
+			ID:              "bm123",
+			OrganizationID:  event.OrganizationID,
+			Code:            event.Code,
+			AggregationType: models.AggregationTypeWeightedSum,
+			FieldName:       "api_requests",
+			Expression:      "",
+			CreatedAt:       time.Now(),
+			UpdatedAt:       time.Now(),
+		}
+		mockBmLookup(sqlmock, &bm)
+
+		result := processor.EnrichEvent(&event)
+		assert.False(t, result.Success())
+		assert.Equal(t, "strconv.ParseFloat: parsing \"2025-03-06T12:00:00Z\": invalid syntax", result.ErrorMsg())
+		assert.Equal(t, "build_enriched_event", result.ErrorCode())
+		assert.Equal(t, "Error while converting event to enriched event", result.ErrorMessage())
+	})
+
+	t.Run("When expression failed to evaluate", func(t *testing.T) {
+		sqlmock, delete := setupTestEnv(t)
+		defer delete()
+
+		event := models.Event{
+			OrganizationID:         "1a901a90-1a90-1a90-1a90-1a901a901a90",
+			ExternalSubscriptionID: "sub_id",
+			Code:                   "api_calls",
+			Timestamp:              "1741007009.123",
+			Source:                 "SQS",
+		}
+
+		bm := models.BillableMetric{
+			ID:              "bm123",
+			OrganizationID:  event.OrganizationID,
+			Code:            event.Code,
+			AggregationType: models.AggregationTypeWeightedSum,
+			FieldName:       "api_requests",
+			Expression:      "round(event.properties.value)",
+			CreatedAt:       time.Now(),
+			UpdatedAt:       time.Now(),
+		}
+		mockBmLookup(sqlmock, &bm)
+
+		result := processor.EnrichEvent(&event)
+		assert.False(t, result.Success())
+		assert.Contains(t, result.ErrorMsg(), "Failed to evaluate expr: round(event.properties.value)")
+		assert.Equal(t, "evaluate_expression", result.ErrorCode())
+		assert.Equal(t, "Error evaluating custom expression", result.ErrorMessage())
+	})
+}
+
+func TestEvaluateExpression(t *testing.T) {
+	_, delete := setupTestEnv(t)
+	defer delete()
+
+	bm := models.BillableMetric{}
+	event := models.EnrichedEvent{Timestamp: 1741007009.0, Code: "foo"}
+	var result utils.Result[bool]
+
+	t.Run("Without expression", func(t *testing.T) {
+		result = processor.evaluateExpression(&event, &bm)
+		assert.True(t, result.Success(), "It should succeed when Billable metric does not have a custom expression")
+	})
+
+	t.Run("With an expression but witout required fields", func(t *testing.T) {
+		bm.Expression = "round(event.properties.value * event.properties.units)"
+		bm.FieldName = "total_value"
+		result = processor.evaluateExpression(&event, &bm)
+		assert.False(t, result.Success())
+		assert.Contains(
+			t,
+			result.ErrorMsg(),
+			"Failed to evaluate expr:",
+			"It should fail when the event does not hold the required fields",
+		)
+	})
+
+	t.Run("With an expression and with required fields", func(t *testing.T) {
+		properties := map[string]any{
+			"value": "12.0",
+			"units": 3,
+		}
+		event.Properties = properties
+		result = processor.evaluateExpression(&event, &bm)
+		assert.True(t, result.Success())
+		assert.Equal(t, "36", event.Properties["total_value"])
+	})
+
+	t.Run("With a float timestamp", func(t *testing.T) {
+		event.Timestamp = 1741007009.123
+
+		result = processor.evaluateExpression(&event, &bm)
+		assert.True(t, result.Success())
+		assert.Equal(t, "36", event.Properties["total_value"])
+	})
+}