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

[hekike]

Lookup usage by customer Key and subjects in usage attribution.

• achieved by changing the default query methods both in streaming and customer/adapter

⚠️ Before merging this or otherwise deploying this change, check if conflicting subject-customer keys are present in the DB ⚠️

# e.g. scenarios like
customer1:
  key: key-1
  subject:
    key: key-2

customer2:
  key: key-2
  subject:
    key: key-1

-- aggregating query
select
  cs.namespace, count(*)
from customer_subjects cs
inner join customers c on cs.customer_id = c.id
left join customers c2 on cs.subject_key = c2.key and c2.namespace = cs.namespace and c2.deleted_at is null
where c2.id != c.id and c2.id is not null and cs.deleted_at is null
group by cs.namespace;

The effect for these conflicts after the change would be that the wrong customer / subject / usage data would be returned for a given query

Summary by CodeRabbit

• New Features

  • Customers can be looked up by either usage-attribution key or subject key for more reliable retrieval.

• Improvements

  • More accurate customer filtering in analytics queries, including cases where only a customer key is present.
  • Query generation now uses a dictionary-style mapping for subject→customer resolution, improving mapping robustness.

• Refactor

  • Internal query helper streamlined and an external dependency removed.

• Tests

  • Added descriptive subtest names and expanded tests to cover key- and subject-based lookups.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds direct customer-key matching to GetCustomerByUsageAttribution, replaces lo-based subject→customer mapping with a named subject_to_customer_id dictionary in ClickHouse query building, and updates tests to use named subtests and validate key-based customer retrieval. (49 words)

Changes

Cohort / File(s)	Summary
Customer usage attribution lookup openmeter/customer/adapter/customer.go	Extend GetCustomerByUsageAttribution to OR-match by subject-table path or direct customer key (customerdb.Or(customerdb.HasSubjectsWith(...), customerdb.Key(input.SubjectKey))); keep namespace/deleted filters and existing error handling; add explanatory comments.
ClickHouse query helper mapping logic openmeter/streaming/clickhouse/queryhelper.go	Remove lo import; add subjectToCustomerIDDictionary constant; build a WITH map(...) dictionary mapping subject keys and optional customer keys to customer_id; early-return empty customer_id when no values/customers; flatten subject keys for WHERE without lo.
ClickHouse meter query tests openmeter/streaming/clickhouse/meter_query_test.go	Add testName to table-driven tests and run named subtests; populate descriptive test names; update one test to include a customer Key and adjust expected SQL args; minor SQL formatting tweaks only.
Customer tests test/customer/customer.go	In TestGetByUsageAttribution create a customer with Key set and replace direct field assertions with retrieval-by-key checks (call GetCustomerByUsageAttribution using the customer Key twice); retain not-found case.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• refactor(streaming): resolve customer_id with map #3259 — Refactors ClickHouse subject→customer_id mapping into a WITH/map-based approach; directly overlaps queryhelper.go changes.
• fix: customer lookup by usage attribution #3292 — Also modifies GetCustomerByUsageAttribution and subject/subscription predicates; touches the same customer lookup logic.
• refactor: order of subject customer lookup #3207 — Alters customer lookup flow around subject-vs-key resolution and is related to the adapter change here.

---

📜 Recent review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 8e71591 and 82fffec.

📒 Files selected for processing (1)

• openmeter/streaming/clickhouse/queryhelper.go (3 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• openmeter/streaming/clickhouse/queryhelper.go

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (7)

• GitHub Check: Quickstart
• GitHub Check: E2E
• GitHub Check: Code Generators
• GitHub Check: Build
• GitHub Check: Lint
• GitHub Check: Test
• GitHub Check: Analyze (go)

✨ Finishing Touches

• 📝 Generate Docstrings

🧪 Generate unit tests

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch feat/streaming-customer-by-key-and-id

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 3

🧹 Nitpick comments (3)

test/customer/customer.go (1)

650-670: Remove duplicated key-based retrieval block.

The “Get the customer by key” block appears twice back-to-back with identical assertions.

Apply this diff to drop the duplicate:

@@
 	// Get the customer by key
 	cus, err = service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
 		Namespace:  s.namespace,
 		SubjectKey: TestKey,
 	})
 
 	require.NoError(t, err, "Fetching customer must not return error")
 	require.NotNil(t, cus, "Customer must not be nil")
 	require.Equal(t, s.namespace, cus.Namespace, "Customer namespace must match")
 	require.Equal(t, createdCustomer.ID, cus.ID, "Customer ID must match")
-
-	// Get the customer by key
-	cus, err = service.GetCustomerByUsageAttribution(ctx, customer.GetCustomerByUsageAttributionInput{
-		Namespace:  s.namespace,
-		SubjectKey: TestKey,
-	})
-
-	require.NoError(t, err, "Fetching customer must not return error")
-	require.NotNil(t, cus, "Customer must not be nil")
-	require.Equal(t, s.namespace, cus.Namespace, "Customer namespace must match")
-	require.Equal(t, createdCustomer.ID, cus.ID, "Customer ID must match")

openmeter/customer/adapter/customer.go (1)

385-399: OR-lookup by subject or key is correct; clarify not-found message.

The OR between HasSubjectsWith(...) and Key(...) matches the intended behavior and still respects DeletedAt filters. Consider adjusting the not-found error to reflect both lookup modes.

-					fmt.Errorf("customer with subject key %s not found in %s namespace", input.SubjectKey, input.Namespace),
+					fmt.Errorf("customer with subject key or key %q not found in %s namespace", input.SubjectKey, input.Namespace),

Also applies to: 406-410

openmeter/streaming/clickhouse/meter_query_test.go (1)

29-52: Name the first table case for consistency.

tt.testName is empty for the first case, resulting in an unnamed subtest.

@@
-		{
+		{
+			testName: "aggregate by window with subject and extra group-by",
@@
-	for _, tt := range tests {
-		t.Run(tt.testName, func(t *testing.T) {
+	for _, tt := range tests {
+		t.Run(tt.testName, func(t *testing.T) {

Also applies to: 452-454

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 8eb2506 and 8dfc561.

📒 Files selected for processing (4)

• openmeter/customer/adapter/customer.go (1 hunks)
• openmeter/streaming/clickhouse/meter_query_test.go (18 hunks)
• openmeter/streaming/clickhouse/queryhelper.go (3 hunks)
• test/customer/customer.go (2 hunks)

🧰 Additional context used

🧠 Learnings (1)

📚 Learning: 2025-03-07T12:17:43.129Z

Learnt from: GAlexIHU
PR: openmeterio/openmeter#2383
File: openmeter/entitlement/metered/lateevents_test.go:37-45
Timestamp: 2025-03-07T12:17:43.129Z
Learning: In the OpenMeter codebase, test files like `openmeter/entitlement/metered/lateevents_test.go` may use variables like `meterSlug` and `namespace` without explicit declarations visible in the same file. This appears to be an accepted pattern in their test structure.

Applied to files:

• openmeter/streaming/clickhouse/meter_query_test.go

🧬 Code graph analysis (2)

test/customer/customer.go (2)

openmeter/ent/db/customer/where.go (2)

• Key (137-139)
• Namespace (71-73)

openmeter/customer/customer.go (1)

• GetCustomerByUsageAttributionInput (162-165)

openmeter/customer/adapter/customer.go (2)

openmeter/ent/db/customer/where.go (5)

• Or (1439-1441)
• HasSubjectsWith (1330-1339)
• DeletedAtIsNil (348-350)
• DeletedAtGT (328-330)
• Key (137-139)

openmeter/ent/db/customersubjects/where.go (4)

• Or (397-399)
• SubjectKey (69-71)
• DeletedAtIsNil (359-361)
• DeletedAtGT (339-341)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (7)

• GitHub Check: Quickstart
• GitHub Check: E2E
• GitHub Check: Lint
• GitHub Check: Code Generators
• GitHub Check: Test
• GitHub Check: Build
• GitHub Check: Analyze (go)

🔇 Additional comments (4)

test/customer/customer.go (2)

78-89: Setting key on creation aligns tests with new lookup semantics — good.
Covers the key-based retrieval path added in the adapter.

---

629-633: Also setting a key here is correct for exercising key-based attribution.
No issues.

openmeter/streaming/clickhouse/meter_query_test.go (2)

24-24: Named subtests improve diagnostics — good change.

---

309-356: Customer-key-inclusive filter case is well covered.
Captures inclusion order (key + subjects) in args.

[coderabbitai]

Actionable comments posted: 0

♻️ Duplicate comments (3)

openmeter/streaming/clickhouse/queryhelper.go (3)

46-50: Good guard for empty dictionary.

Early return prevents an untyped empty map and downstream indexing issues.

---

29-36: Wrong source for customer ID and missing escaping.

Use the actual customer ID and escape it before inlining; the current code pulls the usage-attribution ID and leaves values unescaped.

Apply:

-        customerIDSQL := fmt.Sprintf("'%s'", customer.GetUsageAttribution().ID)
+        // Map to the customer's ID; escape because values are inlined into the CTE.
+        customerIDSQL := fmt.Sprintf("'%s'", sqlbuilder.Escape(customer.GetID()))

If the accessor is named differently (e.g., GetId or ID()), adjust accordingly.

---

77-90: Do not pre-escape values passed to IN(); let the builder bind raw params.

Pre-escaping corrupts values (e.g., quotes doubled) and is inconsistent with how Key is handled.

Apply:

-        for _, subjectKey := range customer.GetUsageAttribution().SubjectKeys {
-            subjects = append(subjects, sqlbuilder.Escape(subjectKey))
-        }
+        for _, subjectKey := range customer.GetUsageAttribution().SubjectKeys {
+            subjects = append(subjects, subjectKey)
+        }

🧹 Nitpick comments (4)

openmeter/streaming/clickhouse/queryhelper.go (4)

98-99: Ensure IN() expands slice arguments correctly.

Pass a list sentinel so go-sqlbuilder expands the slice; otherwise it may bind the entire slice as a single arg.

Apply:

-return query.Where(query.In(subjectColumn, subjects))
+return query.Where(query.In(subjectColumn, sqlbuilder.List(subjects)))

Optionally mirror the same in subjectWhere (Line 115) for consistency.

---

32-44: Minor: centralize literal-quoting to avoid drift.

You repeatedly wrap escaped strings with fmt.Sprintf("'%s'", ...). Consider a tiny helper to keep quoting/escaping consistent.

Example:

sqlLit := func(s string) string { return fmt.Sprintf("'%s'", sqlbuilder.Escape(s)) }

Then use sqlLit for customerKeySQL/subjectSQL/customerIDSQL.

---

29-44: Optional: deduplicate keys before building the map.

If two customers share a subject/key, duplicate keys in map(...) can yield surprising precedence. Dedup to the intended winner or fail fast.

---

114-116: Consistency: use List() here too.

Align subjectWhere with customersWhere to avoid differing IN() binding behavior.

Apply:

-        query = query.Where(query.In(subjectColumn, subjects))
+        query = query.Where(query.In(subjectColumn, sqlbuilder.List(subjects)))

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 8dfc561 and 6eb8b95.

📒 Files selected for processing (1)

• openmeter/streaming/clickhouse/queryhelper.go (3 hunks)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Analyze (go)

🔇 Additional comments (2)

openmeter/streaming/clickhouse/queryhelper.go (2)

12-13: Nice: stable alias for the map.

Defining a constant alias improves readability and avoids magic strings.

---

29-36: No changes needed for GetUsageAttribution accessor or IN() usage
Verified that streaming.Customer’s GetUsageAttribution().ID is the intended field and that query.In(subjectColumn, subjects) matches existing patterns in this file.