From c70c023143406a31bee02cf76c66b0e2b85303af Mon Sep 17 00:00:00 2001
From: Thomas Heartman <thomas@getunleash.io>
Date: Wed, 6 Nov 2024 14:13:44 +0100
Subject: [PATCH] refactor: introduce `countProjectTokens` method on
 ApiTokenStore (#8674)

This change introduces a new method `countProjectTokens` on the
`IApiTokenStore` interface. It also swaps out the manual filtering for
api tokens belonging to a project in the project status service.
---
 src/lib/db/api-token-store.ts                 |  8 ++++
 .../project-status/project-status-service.ts  | 12 +----
 src/lib/types/stores/api-token-store.ts       |  1 +
 .../e2e/stores/api-token-store.e2e.test.ts    | 44 +++++++++++++++++++
 src/test/fixtures/fake-api-token-store.ts     |  4 ++
 5 files changed, 58 insertions(+), 11 deletions(-)

diff --git a/src/lib/db/api-token-store.ts b/src/lib/db/api-token-store.ts
index 338e293c2b..75b5b53542 100644
--- a/src/lib/db/api-token-store.ts
+++ b/src/lib/db/api-token-store.ts
@@ -326,4 +326,12 @@ export class ApiTokenStore implements IApiTokenStore {
             activeLegacyTokens,
         };
     }
+
+    async countProjectTokens(projectId: string): Promise<number> {
+        const count = await this.db(API_LINK_TABLE)
+            .where({ project: projectId })
+            .count()
+            .first();
+        return Number(count?.count ?? 0);
+    }
 }
diff --git a/src/lib/features/project-status/project-status-service.ts b/src/lib/features/project-status/project-status-service.ts
index ae45732625..763c71a558 100644
--- a/src/lib/features/project-status/project-status-service.ts
+++ b/src/lib/features/project-status/project-status-service.ts
@@ -38,17 +38,7 @@ export class ProjectStatusService {
         ] = await Promise.all([
             this.projectStore.getConnectedEnvironmentCountForProject(projectId),
             this.projectStore.getMembersCountByProject(projectId),
-            this.apiTokenStore
-                .getAll()
-                .then(
-                    (tokens) =>
-                        tokens.filter(
-                            (token) =>
-                                token.project === projectId ||
-                                token.projects.includes(projectId),
-                        ).length,
-                ),
-
+            this.apiTokenStore.countProjectTokens(projectId),
             this.segmentStore.getProjectSegmentCount(projectId),
             this.eventStore.getProjectRecentEventActivity(projectId),
         ]);
diff --git a/src/lib/types/stores/api-token-store.ts b/src/lib/types/stores/api-token-store.ts
index 08fd584e2d..9da92337b6 100644
--- a/src/lib/types/stores/api-token-store.ts
+++ b/src/lib/types/stores/api-token-store.ts
@@ -14,4 +14,5 @@ export interface IApiTokenStore extends Store<IApiToken, string> {
         legacyTokens: number;
         activeLegacyTokens: number;
     }>;
+    countProjectTokens(projectId: string): Promise<number>;
 }
diff --git a/src/test/e2e/stores/api-token-store.e2e.test.ts b/src/test/e2e/stores/api-token-store.e2e.test.ts
index c7edd1ef75..dbda4bb080 100644
--- a/src/test/e2e/stores/api-token-store.e2e.test.ts
+++ b/src/test/e2e/stores/api-token-store.e2e.test.ts
@@ -2,6 +2,7 @@ import dbInit, { type ITestDb } from '../helpers/database-init';
 import getLogger from '../../fixtures/no-logger';
 import type { IUnleashStores } from '../../../lib/types';
 import { ApiTokenType } from '../../../lib/types/models/api-token';
+import { randomId } from '../../../lib/util';
 
 let stores: IUnleashStores;
 let db: ITestDb;
@@ -182,3 +183,46 @@ describe('count deprecated tokens', () => {
         });
     });
 });
+
+describe('count project tokens', () => {
+    test('counts only tokens belonging to the specified project', async () => {
+        const project = await stores.projectStore.create({
+            id: randomId(),
+            name: 'project A',
+        });
+
+        const store = stores.apiTokenStore;
+        await store.insert({
+            secret: `default:default.${randomId()}`,
+            environment: 'default',
+            type: ApiTokenType.CLIENT,
+            projects: ['default'],
+            tokenName: 'token1',
+        });
+        await store.insert({
+            secret: `*:*.${randomId()}`,
+            environment: 'default',
+            type: ApiTokenType.CLIENT,
+            projects: ['*'],
+            tokenName: 'token2',
+        });
+
+        await store.insert({
+            secret: `${project.id}:default.${randomId()}`,
+            environment: 'default',
+            type: ApiTokenType.CLIENT,
+            projects: [project.id],
+            tokenName: 'token3',
+        });
+
+        await store.insert({
+            secret: `[]:default.${randomId()}`,
+            environment: 'default',
+            type: ApiTokenType.CLIENT,
+            projects: [project.id, 'default'],
+            tokenName: 'token4',
+        });
+
+        expect(await store.countProjectTokens(project.id)).toBe(2);
+    });
+});
diff --git a/src/test/fixtures/fake-api-token-store.ts b/src/test/fixtures/fake-api-token-store.ts
index 6b90064fce..d24f920e83 100644
--- a/src/test/fixtures/fake-api-token-store.ts
+++ b/src/test/fixtures/fake-api-token-store.ts
@@ -92,4 +92,8 @@ export default class FakeApiTokenStore
             activeLegacyTokens: 0,
         };
     }
+
+    async countProjectTokens(): Promise<number> {
+        return 0;
+    }
 }