streaming/clickhouse/queryhelper.go

package clickhouse

import (
	"fmt"
	"strings"

	"github.com/huandu/go-sqlbuilder"

	"github.com/openmeterio/openmeter/openmeter/streaming"
)

const subjectToCustomerIDDictionary = "subject_to_customer_id"

// selectCustomerIdColumn
func selectCustomerIdColumn(eventsTableName string, customers []streaming.Customer, query *sqlbuilder.SelectBuilder) *sqlbuilder.SelectBuilder {
	// If there are no customers, we return an empty customer id column
	if len(customers) == 0 {
		return query.SelectMore("'' AS customer_id")
	}

	// Helper function to get the subject column
	getColumn := columnFactory(eventsTableName)
	subjectColumn := getColumn("subject")

	// Build a map of event subjects to customer ids
	var values []string

	// For each customer, we map event subjects to customer ids
	for _, customer := range customers {
		customerIDSQL := fmt.Sprintf("'%s'", sqlbuilder.Escape(customer.GetUsageAttribution().ID))

		// We map the customer key to the customer id if it exists
		if customer.GetUsageAttribution().Key != nil {
			customerKeySQL := fmt.Sprintf("'%s'", sqlbuilder.Escape(*customer.GetUsageAttribution().Key))
			values = append(values, customerKeySQL, customerIDSQL)
		}

		// We map each subject key to the customer id
		for _, subjectKey := range customer.GetUsageAttribution().SubjectKeys {
			subjectSQL := fmt.Sprintf("'%s'", sqlbuilder.Escape(subjectKey))

			values = append(values, subjectSQL, customerIDSQL)
		}
	}

	// If there are no values, we return an empty customer id column
	// This can happen if none of the customers has key or usage attribution subjects
	if len(values) == 0 {
		return query.SelectMore("'' AS customer_id")
	}

	// Name of the map (dictionary)

	mapSQL := fmt.Sprintf("WITH map(%s) as %s", strings.Join(values, ", "), subjectToCustomerIDDictionary)

	// Add the map to query via WITH clause
	mapQuery := sqlbuilder.ClickHouse.NewCTEBuilder().SQL(mapSQL)
	query = query.With(mapQuery)

	// Select the customer id column
	query = query.SelectMore(fmt.Sprintf("%s[%s] AS customer_id", subjectToCustomerIDDictionary, subjectColumn))

	return query
}

// customersWhere applies the customer filter to the query.
func customersWhere(eventsTableName string, customers []streaming.Customer, query *sqlbuilder.SelectBuilder) *sqlbuilder.SelectBuilder {
	// If there are no customers, we return an empty subject filter
	if len(customers) == 0 {
		return query
	}

	// Helper function to filter by subject
	getColumn := columnFactory(eventsTableName)
	subjectColumn := getColumn("subject")

	var subjects []string

	// Collect all the subjects from the customers
	for _, customer := range customers {
		subjects = append(subjects, customer.GetUsageAttribution().GetValues()...)
	}

	// If there are no subjects, we return an empty subject filter
	// This can happen if none of the customers has key or usage attribution subjects
	if len(subjects) == 0 {
		return query
	}

	return query.Where(query.In(subjectColumn, subjects))
}

// subjectWhere applies the subject filter to the query.
func subjectWhere(
	eventsTableName string,
	subjects []string,
	query *sqlbuilder.SelectBuilder,
) *sqlbuilder.SelectBuilder {
	// Helper function to filter by subject
	getColumn := columnFactory(eventsTableName)
	subjectColumn := getColumn("subject")

	// If we have a subject filter, we add it to the query
	// If we have both a customer filter and a subject filter,
	// this is an AND between the two filters
	if len(subjects) > 0 {
		query = query.Where(query.In(subjectColumn, subjects))
	}

	return query
}

func columnFactory(alias string) func(string) string {
	return func(column string) string {
		return fmt.Sprintf("%s.%s", alias, column)
	}
}