feat(streaming): lookup usage by customer id and key

Author: hekike

openmeter/streaming/clickhouse/meter_query.go

@@ -202,6 +202,26 @@ func (d *queryMeter) toSQL() (string, []interface{}, error) {
 
 		// Add the case statements for each subject to customer ID mapping
 		for _, customer := range d.FilterCustomer {
+			// Add the customer ID to the case statement
+			str := fmt.Sprintf(
+				"WHEN %s = '%s' THEN '%s' ",
+				getColumn("subject"),
+				sqlbuilder.Escape(customer.ID),
+				sqlbuilder.Escape(customer.ID),
+			)
+			caseBuilder.WriteString(str)
+
+			// Add the customer key to the case statement
+			if customer.Key != nil {
+				str := fmt.Sprintf(
+					"WHEN %s = '%s' THEN '%s' ",
+					getColumn("subject"),
+					sqlbuilder.Escape(*customer.Key),
+					sqlbuilder.Escape(customer.ID),
+				)
+				caseBuilder.WriteString(str)
+			}
+
 			for _, subjectKey := range customer.UsageAttribution.SubjectKeys {
 				str := fmt.Sprintf(
 					"WHEN %s = '%s' THEN '%s' ",
@@ -309,6 +329,14 @@ func (d *queryMeter) subjectWhere(query *sqlbuilder.SelectBuilder) *sqlbuilder.S
 
 		for _, customer := range d.FilterCustomer {
 			subjects = append(subjects, customer.UsageAttribution.SubjectKeys...)
+
+			// Add the customer ID
+			subjects = append(subjects, customer.ID)
+
+			// Add the customer key if it's set
+			if customer.Key != nil {
+				subjects = append(subjects, *customer.Key)
+			}
 		}
 
 		query = query.Where(query.Or(slicesx.Map(subjects, mapFunc)...))

openmeter/streaming/clickhouse/meter_query_test.go

@@ -290,8 +290,8 @@ func TestQueryMeter(t *testing.T) {
 				},
 				GroupBy: []string{"customer_id"},
 			},
-			wantSQL:  "SELECT tumbleStart(min(om_events.time), toIntervalMinute(1)) AS windowstart, tumbleEnd(max(om_events.time), toIntervalMinute(1)) AS windowend, sum(ifNotFinite(toFloat64OrNull(JSON_VALUE(om_events.data, '$.value')), null)) AS value, CASE WHEN om_events.subject = 'subject1' THEN 'customer1' WHEN om_events.subject = 'subject2' THEN 'customer2' ELSE '' END AS customer_id FROM openmeter.om_events WHERE om_events.namespace = ? AND om_events.type = ? AND (om_events.subject = ? OR om_events.subject = ?) GROUP BY customer_id",
-			wantArgs: []interface{}{"my_namespace", "event1", "subject1", "subject2"},
+			wantSQL:  "SELECT tumbleStart(min(om_events.time), toIntervalMinute(1)) AS windowstart, tumbleEnd(max(om_events.time), toIntervalMinute(1)) AS windowend, sum(ifNotFinite(toFloat64OrNull(JSON_VALUE(om_events.data, '$.value')), null)) AS value, CASE WHEN om_events.subject = 'customer1' THEN 'customer1' WHEN om_events.subject = 'subject1' THEN 'customer1' WHEN om_events.subject = 'customer2' THEN 'customer2' WHEN om_events.subject = 'subject2' THEN 'customer2' ELSE '' END AS customer_id FROM openmeter.om_events WHERE om_events.namespace = ? AND om_events.type = ? AND (om_events.subject = ? OR om_events.subject = ? OR om_events.subject = ? OR om_events.subject = ?) GROUP BY customer_id",
+			wantArgs: []interface{}{"my_namespace", "event1", "subject1", "customer1", "subject2", "customer2"},
 		},
 		{ // Filter by both customer and subject
 			query: queryMeter{
@@ -320,8 +320,8 @@ func TestQueryMeter(t *testing.T) {
 				FilterSubject: []string{"subject1"},
 				GroupBy:       []string{"customer_id"},
 			},
-			wantSQL:  "SELECT tumbleStart(min(om_events.time), toIntervalMinute(1)) AS windowstart, tumbleEnd(max(om_events.time), toIntervalMinute(1)) AS windowend, sum(ifNotFinite(toFloat64OrNull(JSON_VALUE(om_events.data, '$.value')), null)) AS value, CASE WHEN om_events.subject = 'subject1' THEN 'customer1' WHEN om_events.subject = 'subject2' THEN 'customer1' ELSE '' END AS customer_id FROM openmeter.om_events WHERE om_events.namespace = ? AND om_events.type = ? AND (om_events.subject = ? OR om_events.subject = ?) AND (om_events.subject = ?) GROUP BY customer_id",
-			wantArgs: []interface{}{"my_namespace", "event1", "subject1", "subject2", "subject1"},
+			wantSQL:  "SELECT tumbleStart(min(om_events.time), toIntervalMinute(1)) AS windowstart, tumbleEnd(max(om_events.time), toIntervalMinute(1)) AS windowend, sum(ifNotFinite(toFloat64OrNull(JSON_VALUE(om_events.data, '$.value')), null)) AS value, CASE WHEN om_events.subject = 'customer1' THEN 'customer1' WHEN om_events.subject = 'subject1' THEN 'customer1' WHEN om_events.subject = 'subject2' THEN 'customer1' ELSE '' END AS customer_id FROM openmeter.om_events WHERE om_events.namespace = ? AND om_events.type = ? AND (om_events.subject = ? OR om_events.subject = ? OR om_events.subject = ?) AND (om_events.subject = ?) GROUP BY customer_id",
+			wantArgs: []interface{}{"my_namespace", "event1", "subject1", "subject2", "customer1", "subject1"},
 		},
 		{ // Aggregate data with filtering for a single group and single value
 			query: queryMeter{