diff --git a/packages/backend/server/src/core/quota/service.ts b/packages/backend/server/src/core/quota/service.ts index e28795c00d..c8609f7c7e 100644 --- a/packages/backend/server/src/core/quota/service.ts +++ b/packages/backend/server/src/core/quota/service.ts @@ -270,6 +270,23 @@ export class QuotaService { .then(count => count > 0); } + /// check if workspaces have quota + /// return workspaces's id that have quota + async hasWorkspacesQuota( + workspaces: string[], + quota?: QuotaType + ): Promise { + const workspaceIds = await this.prisma.workspaceFeature.findMany({ + where: { + workspaceId: { in: workspaces }, + feature: { feature: quota, type: FeatureKind.Quota }, + activated: true, + }, + select: { workspaceId: true }, + }); + return Array.from(new Set(workspaceIds.map(w => w.workspaceId))); + } + async getWorkspaceConfig( workspaceId: string, type: Q diff --git a/packages/backend/server/src/core/quota/storage.ts b/packages/backend/server/src/core/quota/storage.ts index 31196b987c..310bd33d83 100644 --- a/packages/backend/server/src/core/quota/storage.ts +++ b/packages/backend/server/src/core/quota/storage.ts @@ -79,15 +79,20 @@ export class QuotaManagementService { async getUserStorageUsage(userId: string) { const workspaces = await this.permissions.getOwnedWorkspaces(userId); + const workspacesWithQuota = await this.quota.hasWorkspacesQuota(workspaces); const sizes = await Promise.allSettled( - workspaces.map(workspace => this.storage.totalSize(workspace)) + workspaces + .filter(w => !workspacesWithQuota.includes(w)) + .map(workspace => this.storage.totalSize(workspace)) ); return sizes.reduce((total, size) => { if (size.status === 'fulfilled') { - if (Number.isSafeInteger(size.value)) { - return total + size.value; + // ensure that size is within the safe range of gql + const totalSize = total + size.value; + if (Number.isSafeInteger(totalSize)) { + return totalSize; } else { this.logger.error(`Workspace size is invalid: ${size.value}`); } @@ -98,6 +103,17 @@ export class QuotaManagementService { }, 0); } + async getWorkspaceStorageUsage(workspaceId: string) { + const totalSize = this.storage.totalSize(workspaceId); + // ensure that size is within the safe range of gql + if (Number.isSafeInteger(totalSize)) { + return totalSize; + } else { + this.logger.error(`Workspace size is invalid: ${totalSize}`); + } + return 0; + } + private generateQuotaCalculator( quota: number, blobLimit: number, @@ -146,17 +162,23 @@ export class QuotaManagementService { ); } - private async getWorkspaceQuota(userId: string, workspaceId: string) { + private async getWorkspaceQuota( + userId: string, + workspaceId: string + ): Promise<{ quota: QuotaConfig; fromUser: boolean }> { const { feature: workspaceQuota } = (await this.quota.getWorkspaceQuota(workspaceId)) || {}; const { feature: userQuota } = await this.quota.getUserQuota(userId); if (workspaceQuota) { - return workspaceQuota.withOverride({ - // override user quota with workspace quota - copilotActionLimit: userQuota.copilotActionLimit, - }); + return { + quota: workspaceQuota.withOverride({ + // override user quota with workspace quota + copilotActionLimit: userQuota.copilotActionLimit, + }), + fromUser: false, + }; } - return userQuota; + return { quota: userQuota, fromUser: true }; } async checkWorkspaceSeat(workspaceId: string, excludeSelf = false) { @@ -173,17 +195,24 @@ export class QuotaManagementService { const memberCount = await this.permissions.getWorkspaceMemberCount(workspaceId); const { - name, - blobLimit, - businessBlobLimit, - historyPeriod, - memberLimit, - storageQuota, - copilotActionLimit, - humanReadable, + quota: { + name, + blobLimit, + businessBlobLimit, + historyPeriod, + memberLimit, + storageQuota, + copilotActionLimit, + humanReadable, + }, + fromUser, } = await this.getWorkspaceQuota(owner.id, workspaceId); - // get all workspaces size of owner used - const usedSize = await this.getUserStorageUsage(owner.id); + + const usedSize = fromUser + ? // get all workspaces size of owner used + await this.getUserStorageUsage(owner.id) + : // get workspace size + await this.getWorkspaceStorageUsage(workspaceId); // relax restrictions if workspace has unlimited feature // todo(@darkskygit): need a mechanism to allow feature as a middleware to edit quota const unlimited = await this.feature.hasWorkspaceFeature( diff --git a/packages/backend/server/tests/quota.spec.ts b/packages/backend/server/tests/quota.spec.ts index fd7c9a3d34..66ebcb5e9b 100644 --- a/packages/backend/server/tests/quota.spec.ts +++ b/packages/backend/server/tests/quota.spec.ts @@ -13,7 +13,7 @@ import { } from '../src/core/quota'; import { OneGB, OneMB } from '../src/core/quota/constant'; import { FreePlan, ProPlan } from '../src/core/quota/schema'; -import { StorageModule } from '../src/core/storage'; +import { StorageModule, WorkspaceBlobStorage } from '../src/core/storage'; import { WorkspaceResolver } from '../src/core/workspaces/resolvers'; import { createTestingModule } from './utils'; import { WorkspaceResolverMock } from './utils/feature'; @@ -23,6 +23,7 @@ const test = ava as TestFn<{ quota: QuotaService; quotaManager: QuotaManagementService; workspace: WorkspaceResolver; + workspaceBlob: WorkspaceBlobStorage; module: TestingModule; }>; @@ -37,16 +38,12 @@ test.beforeEach(async t => { }, }); - const quota = module.get(QuotaService); - const quotaManager = module.get(QuotaManagementService); - const workspace = module.get(WorkspaceResolver); - const auth = module.get(AuthService); - t.context.module = module; - t.context.quota = quota; - t.context.quotaManager = quotaManager; - t.context.workspace = workspace; - t.context.auth = auth; + t.context.auth = module.get(AuthService); + t.context.quota = module.get(QuotaService); + t.context.quotaManager = module.get(QuotaManagementService); + t.context.workspace = module.get(WorkspaceResolver); + t.context.workspaceBlob = module.get(WorkspaceBlobStorage); }); test.afterEach.always(async t => { @@ -165,3 +162,33 @@ test('should be able to override quota', async t => { t.is(wq3.storageQuota, 140 * OneGB, 'should be override to 120GB'); t.is(wq3.memberLimit, 2, 'should be override to 1'); }); + +test('should be able to check with workspace quota', async t => { + const { auth, quotaManager, workspace, workspaceBlob } = t.context; + + const u1 = await auth.signUp('test@affine.pro', '123456'); + const w1 = await workspace.createWorkspace(u1, null); + const w2 = await workspace.createWorkspace(u1, null); + await quotaManager.addTeamWorkspace(w2.id, 'test'); + + { + const wq = await quotaManager.getWorkspaceUsage(w1.id); + t.is(wq.usedSize, 0, 'should be 0'); + } + + { + await workspaceBlob.put(w1.id, 'test', Buffer.from([0, 0])); + const wq1 = await quotaManager.getWorkspaceUsage(w1.id); + t.is(wq1.usedSize, 2, 'should be 2'); + const wq2 = await quotaManager.getWorkspaceUsage(w2.id); + t.is(wq2.usedSize, 0, 'should be 0'); + } + + { + await workspaceBlob.put(w2.id, 'test', Buffer.from([0, 0, 0])); + const wq1 = await quotaManager.getWorkspaceUsage(w1.id); + t.is(wq1.usedSize, 2, 'should be 2'); + const wq2 = await quotaManager.getWorkspaceUsage(w2.id); + t.is(wq2.usedSize, 3, 'should be 0'); + } +});