From ccd0d95d574c6761e99640048df83e03079366fa Mon Sep 17 00:00:00 2001
From: Thomas Heartman <thomas@getunleash.io>
Date: Wed, 6 Nov 2024 13:51:27 +0100
Subject: [PATCH] refactor: add `countProjectTokens` method to api token store

---
 src/lib/db/api-token-store.ts                 |  8 ++++
 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 ++
 4 files changed, 57 insertions(+)

diff --git a/src/lib/db/api-token-store.ts b/src/lib/db/api-token-store.ts
index 338e293c2b56..c430fbdf0c80 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/types/stores/api-token-store.ts b/src/lib/types/stores/api-token-store.ts
index 08fd584e2d0f..9da92337b6d4 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 c7edd1ef7524..dbda4bb080d8 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 6b90064fceb6..d24f920e83e8 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;
+    }
 }