@@ -5,11 +5,12 @@ import (
55 "strings"
66
77 "github.com/huandu/go-sqlbuilder"
8- "github.com/samber/lo"
98
109 "github.com/openmeterio/openmeter/openmeter/streaming"
1110)
1211
12+ const subjectToCustomerIDDictionary = "subject_to_customer_id"
13+
1314// selectCustomerIdColumn
1415func selectCustomerIdColumn (eventsTableName string , customers []streaming.Customer , query * sqlbuilder.SelectBuilder ) * sqlbuilder.SelectBuilder {
1516 // If there are no customers, we return an empty customer id column
@@ -21,28 +22,43 @@ func selectCustomerIdColumn(eventsTableName string, customers []streaming.Custom
2122 getColumn := columnFactory (eventsTableName )
2223 subjectColumn := getColumn ("subject" )
2324
24- // Build a map of subject to customer id
25+ // Build a map of event subjects to customer ids
2526 var values []string
2627
28+ // For each customer, we map event subjects to customer ids
2729 for _ , customer := range customers {
28- // Add each subject key to the map and map it to the customer id
30+ customerIDSQL := fmt .Sprintf ("'%s'" , sqlbuilder .Escape (customer .GetUsageAttribution ().ID ))
31+
32+ // We map the customer key to the customer id if it exists
33+ if customer .GetUsageAttribution ().Key != nil {
34+ customerKeySQL := fmt .Sprintf ("'%s'" , sqlbuilder .Escape (* customer .GetUsageAttribution ().Key ))
35+ values = append (values , customerKeySQL , customerIDSQL )
36+ }
37+
38+ // We map each subject key to the customer id
2939 for _ , subjectKey := range customer .GetUsageAttribution ().SubjectKeys {
3040 subjectSQL := fmt .Sprintf ("'%s'" , sqlbuilder .Escape (subjectKey ))
31- customerIDSQL := fmt .Sprintf ("'%s'" , sqlbuilder .Escape (customer .GetUsageAttribution ().ID ))
3241
3342 values = append (values , subjectSQL , customerIDSQL )
3443 }
3544 }
3645
37- mapAs := "subject_to_customer_id"
38- mapSQL := fmt .Sprintf ("WITH map(%s) as %s" , strings .Join (values , ", " ), mapAs )
46+ // If there are no values, we return an empty customer id column
47+ // This can happen if none of the customers has key or usage attribution subjects
48+ if len (values ) == 0 {
49+ return query .SelectMore ("'' AS customer_id" )
50+ }
51+
52+ // Name of the map (dictionary)
53+
54+ mapSQL := fmt .Sprintf ("WITH map(%s) as %s" , strings .Join (values , ", " ), subjectToCustomerIDDictionary )
3955
4056 // Add the map to query via WITH clause
4157 mapQuery := sqlbuilder .ClickHouse .NewCTEBuilder ().SQL (mapSQL )
4258 query = query .With (mapQuery )
4359
4460 // Select the customer id column
45- query = query .SelectMore (fmt .Sprintf ("%s[%s] AS customer_id" , mapAs , subjectColumn ))
61+ query = query .SelectMore (fmt .Sprintf ("%s[%s] AS customer_id" , subjectToCustomerIDDictionary , subjectColumn ))
4662
4763 return query
4864}
@@ -58,12 +74,26 @@ func customersWhere(eventsTableName string, customers []streaming.Customer, quer
5874 getColumn := columnFactory (eventsTableName )
5975 subjectColumn := getColumn ("subject" )
6076
61- // If the customer filter is provided, we add all the subjects to the filter
62- subjects := lo .Map (customers , func (customer streaming.Customer , _ int ) []string {
63- return customer .GetUsageAttribution ().SubjectKeys
64- })
77+ var subjects []string
78+
79+ // Collect all the subjects from the customers
80+ for _ , customer := range customers {
81+ // Add the customer key to the filter if it exists
82+ if customer .GetUsageAttribution ().Key != nil {
83+ subjects = append (subjects , * customer .GetUsageAttribution ().Key )
84+ }
85+
86+ // Add each subject key to the filter
87+ subjects = append (subjects , customer .GetUsageAttribution ().SubjectKeys ... )
88+ }
89+
90+ // If there are no subjects, we return an empty subject filter
91+ // This can happen if none of the customers has key or usage attribution subjects
92+ if len (subjects ) == 0 {
93+ return query
94+ }
6595
66- return query .Where (query .In (subjectColumn , lo . Flatten ( subjects ) ))
96+ return query .Where (query .In (subjectColumn , subjects ))
6797}
6898
6999// subjectWhere applies the subject filter to the query.
0 commit comments