From 2b201d5993c9a46c955a9b2b7624430bc96cebbb Mon Sep 17 00:00:00 2001
From: Pagan Gazzard <page@balena.io>
Date: Fri, 13 Dec 2024 16:13:48 +0000
Subject: [PATCH] Loki: split tenants by organization rather than fleet

Change-type: major
---
 src/features/device-logs/lib/backends/loki.ts | 35 +++++++++++--------
 src/features/device-logs/lib/struct.ts        |  1 +
 test/13_loki-backend.ts                       |  1 +
 3 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/src/features/device-logs/lib/backends/loki.ts b/src/features/device-logs/lib/backends/loki.ts
index e1100596b..85545d794 100644
--- a/src/features/device-logs/lib/backends/loki.ts
+++ b/src/features/device-logs/lib/backends/loki.ts
@@ -100,8 +100,8 @@ function backoff<T extends (...args: any[]) => any>(
 async function assertLokiLogContext(
 	ctx: LogContext & Partial<LokiLogContext>,
 ): Promise<LokiLogContext> {
-	if ('appId' in ctx) {
-		return ctx as types.RequiredField<typeof ctx, 'appId'>;
+	if ('appId' in ctx && 'orgId' in ctx) {
+		return ctx as types.RequiredField<typeof ctx, 'appId' | 'orgId'>;
 	}
 
 	const device = await sbvrUtils.api.resin.get({
@@ -110,6 +110,11 @@ async function assertLokiLogContext(
 		passthrough: { req: permissions.root },
 		options: {
 			$select: ['belongs_to__application'],
+			$expand: {
+				belongs_to__application: {
+					$select: ['id', 'organization'],
+				},
+			},
 		},
 	});
 
@@ -117,15 +122,17 @@ async function assertLokiLogContext(
 		throw new Error(`Device '${ctx.id}' not found`);
 	}
 
-	// Mutate so that we don't have to repeatedly amend the same context and instead cache it
-	(ctx as Writable<typeof ctx>).appId =
-		`${device.belongs_to__application?.__id}`;
-
-	if (ctx.appId == null) {
+	if (device.belongs_to__application[0] == null) {
 		throw new Error(`Device '${ctx.id}' app not found`);
 	}
 
-	return ctx as types.RequiredField<typeof ctx, 'appId'>;
+	// Mutate so that we don't have to repeatedly amend the same context and instead cache it
+	(ctx as Writable<typeof ctx>).appId =
+		`${device.belongs_to__application[0].id}`;
+	(ctx as Writable<typeof ctx>).orgId =
+		`${device.belongs_to__application[0].organization.__id}`;
+
+	return ctx as types.RequiredField<typeof ctx, 'appId' | 'orgId'>;
 }
 
 export class LokiBackend implements DeviceLogsBackend {
@@ -189,7 +196,7 @@ export class LokiBackend implements DeviceLogsBackend {
 		const [, body] = await requestAsync({
 			url: `http://${lokiQueryAddress}/loki/api/v1/query_range`,
 			headers: {
-				'X-Scope-OrgID': ctx.appId,
+				'X-Scope-OrgID': ctx.orgId,
 			},
 			qs: {
 				query: this.getDeviceQuery(ctx),
@@ -262,7 +269,7 @@ export class LokiBackend implements DeviceLogsBackend {
 			await new Promise<loki.PushResponse>((resolve, reject) => {
 				this.pusher.push(
 					pushRequest,
-					loki.createOrgIdMetadata(ctx.appId),
+					loki.createOrgIdMetadata(ctx.orgId),
 					{
 						deadline: startAt + PUSH_TIMEOUT,
 					},
@@ -290,7 +297,7 @@ export class LokiBackend implements DeviceLogsBackend {
 
 			const call = this.querier.tail(
 				request,
-				loki.createOrgIdMetadata(ctx.appId),
+				loki.createOrgIdMetadata(ctx.orgId),
 			);
 			call.on('data', (response: loki.TailResponse) => {
 				const stream = response.getStream();
@@ -329,11 +336,11 @@ export class LokiBackend implements DeviceLogsBackend {
 	}
 
 	private getDeviceQuery(ctx: LokiLogContext) {
-		return `{application_id="${ctx.appId}"} | device_id="${ctx.id}"`;
+		return `{fleet_id="${ctx.appId}"} | device_id="${ctx.id}"`;
 	}
 
 	private getKey(ctx: LokiLogContext) {
-		return `a${ctx.appId}:d${ctx.id}`;
+		return `o${ctx.orgId}:a${ctx.appId}:d${ctx.id}`;
 	}
 
 	private getStructuredMetadata(ctx: LogContext): loki.LabelPairAdapter[] {
@@ -343,7 +350,7 @@ export class LokiBackend implements DeviceLogsBackend {
 	}
 
 	private getLabels(ctx: LokiLogContext): string {
-		return `{application_id="${ctx.appId}"}`;
+		return `{fleet_id="${ctx.appId}"}`;
 	}
 
 	private validateLog(log: DeviceLog): asserts log is DeviceLog {
diff --git a/src/features/device-logs/lib/struct.ts b/src/features/device-logs/lib/struct.ts
index bdbc300cb..4a3938feb 100644
--- a/src/features/device-logs/lib/struct.ts
+++ b/src/features/device-logs/lib/struct.ts
@@ -4,6 +4,7 @@ export interface LogContext {
 	readonly retention_limit: number;
 }
 export interface LokiLogContext extends LogContext {
+	readonly orgId: string;
 	readonly appId: string;
 }
 
diff --git a/test/13_loki-backend.ts b/test/13_loki-backend.ts
index 72a0abb87..3d792ddfc 100644
--- a/test/13_loki-backend.ts
+++ b/test/13_loki-backend.ts
@@ -23,6 +23,7 @@ const createContext = (extra = {}): LokiLogContext => {
 		id: 1,
 		uuid: '1',
 		appId: '1',
+		orgId: '1',
 		retention_limit: 100,
 		...extra,
 	};