diff --git a/packages/backend/server/schema.prisma b/packages/backend/server/schema.prisma index 6d5a16bb57c06..2ddb93c15a288 100644 --- a/packages/backend/server/schema.prisma +++ b/packages/backend/server/schema.prisma @@ -214,7 +214,7 @@ model WorkspaceFeature { workspaceId String @map("workspace_id") @db.VarChar featureId Int @map("feature_id") @db.Integer - // override feature's configs + // override quota's configs configs Json @default("{}") @db.Json // we will record the reason why the feature is enabled/disabled // for example: diff --git a/packages/backend/server/src/core/quota/index.ts b/packages/backend/server/src/core/quota/index.ts index 9fea2c4da6ae0..65645661216fa 100644 --- a/packages/backend/server/src/core/quota/index.ts +++ b/packages/backend/server/src/core/quota/index.ts @@ -3,7 +3,6 @@ import { Module } from '@nestjs/common'; import { FeatureModule } from '../features'; import { PermissionModule } from '../permission'; import { StorageModule } from '../storage'; -import { QuotaOverrideService } from './override'; import { QuotaManagementResolver } from './resolver'; import { QuotaService } from './service'; import { QuotaManagementService } from './storage'; @@ -16,19 +15,12 @@ import { QuotaManagementService } from './storage'; */ @Module({ imports: [FeatureModule, StorageModule, PermissionModule], - providers: [ - QuotaService, - QuotaOverrideService, - QuotaManagementResolver, - QuotaManagementService, - ], - exports: [QuotaService, QuotaOverrideService, QuotaManagementService], + providers: [QuotaService, QuotaManagementResolver, QuotaManagementService], + exports: [QuotaService, QuotaManagementService], }) export class QuotaModule {} export { QuotaManagementService, QuotaService }; -export { OneDay, OneGB, OneMB } from './constant'; -export { QuotaOverride, QuotaOverrideService } from './override'; export { Quota_FreePlanV1_1, Quota_ProPlanV1 } from './schema'; export { formatDate, diff --git a/packages/backend/server/src/core/quota/override.ts b/packages/backend/server/src/core/quota/override.ts deleted file mode 100644 index 4d0a3b111edf8..0000000000000 --- a/packages/backend/server/src/core/quota/override.ts +++ /dev/null @@ -1,53 +0,0 @@ -import { Injectable, Logger } from '@nestjs/common'; - -import type { QuotaBusinessType } from './types'; - -export abstract class QuotaOverride { - abstract readonly name: string; - abstract overrideQuota( - ownerId: string, - workspaceId: string, - quota: QuotaBusinessType - ): Promise; -} - -@Injectable() -export class QuotaOverrideService { - private readonly logger = new Logger(QuotaOverrideService.name); - private readonly overrides: QuotaOverride[] = []; - - registerOverride(override: QuotaOverride) { - if ( - !this.overrides.includes(override) && - typeof override.overrideQuota === 'function' - ) { - this.overrides.push(override); - } - } - - async overrideQuota( - ownerId: string, - workspaceId: string, - quota: QuotaBusinessType - ): Promise { - let lastQuota = quota; - for (const override of this.overrides) { - try { - const quota = await override.overrideQuota( - ownerId, - workspaceId, - lastQuota - ); - if (quota) { - lastQuota = quota; - } - } catch (e) { - this.logger.error( - `Failed to override quota ${override.name} for workspace ${workspaceId}`, - e - ); - } - } - return lastQuota; - } -} diff --git a/packages/backend/server/src/core/quota/quota.ts b/packages/backend/server/src/core/quota/quota.ts index b7ebddf1fdd31..7d6a73167916f 100644 --- a/packages/backend/server/src/core/quota/quota.ts +++ b/packages/backend/server/src/core/quota/quota.ts @@ -61,11 +61,17 @@ export class QuotaConfig { } withOverride(override: any) { - return new QuotaConfig(this.config, override); + if (override) { + return new QuotaConfig(this.config, override); + } + return this; } checkOverride(override: any) { - return QuotaSchema.safeParse(Object.assign({}, this.config, override)); + return QuotaSchema.safeParse({ + ...this.config, + configs: Object.assign({}, this.config.configs, override), + }); } get version() { @@ -84,8 +90,8 @@ export class QuotaConfig { get businessBlobLimit() { return ( this.override?.businessBlobLimit || - this.override?.blobLimit || this.config.configs.businessBlobLimit || + this.override?.blobLimit || this.config.configs.blobLimit ); } diff --git a/packages/backend/server/src/core/quota/schema.ts b/packages/backend/server/src/core/quota/schema.ts index 24657157afb0f..c46fcd0e105fd 100644 --- a/packages/backend/server/src/core/quota/schema.ts +++ b/packages/backend/server/src/core/quota/schema.ts @@ -174,7 +174,6 @@ export const Quotas: Quota[] = [ copilotActionLimit: 10, }, }, - { feature: QuotaType.TeamPlanV1, type: FeatureKind.Quota, @@ -183,7 +182,7 @@ export const Quotas: Quota[] = [ // quota name name: 'Team Workspace', // single blob limit 100MB - blobLimit: 100 * OneMB, + blobLimit: 500 * OneMB, // total blob limit 100GB storageQuota: 100 * OneGB, // seat quota 20GB per seat diff --git a/packages/backend/server/src/core/quota/service.ts b/packages/backend/server/src/core/quota/service.ts index 29163931e715c..0cbc0a0128873 100644 --- a/packages/backend/server/src/core/quota/service.ts +++ b/packages/backend/server/src/core/quota/service.ts @@ -1,6 +1,5 @@ import { Injectable } from '@nestjs/common'; import { PrismaClient } from '@prisma/client'; -import { JsonObject } from '@prisma/client/runtime/library'; import type { EventPayload } from '../../fundamentals'; import { OnEvent, PrismaTransaction } from '../../fundamentals'; @@ -17,9 +16,12 @@ export class QuotaService { ) {} async getQuota( - quota: Q + quota: Q, + tx?: PrismaTransaction ): Promise { - const data = await this.prisma.feature.findFirst({ + const executor = tx ?? this.prisma; + + const data = await executor.feature.findFirst({ where: { feature: quota, type: FeatureKind.Quota }, select: { id: true }, orderBy: { version: 'desc' }, @@ -38,9 +40,7 @@ export class QuotaService { const quota = await this.prisma.userFeature.findFirst({ where: { userId, - feature: { - type: FeatureKind.Quota, - }, + feature: { type: FeatureKind.Quota }, activated: true, }, select: { @@ -65,9 +65,7 @@ export class QuotaService { const quotas = await this.prisma.userFeature.findMany({ where: { userId, - feature: { - type: FeatureKind.Quota, - }, + feature: { type: FeatureKind.Quota }, }, select: { activated: true, @@ -76,9 +74,7 @@ export class QuotaService { expiredAt: true, featureId: true, }, - orderBy: { - id: 'asc', - }, + orderBy: { id: 'asc' }, }); const configs = await Promise.all( quotas.map(async quota => { @@ -107,11 +103,7 @@ export class QuotaService { ) { await this.prisma.$transaction(async tx => { const hasSameActivatedQuota = await this.hasUserQuota(userId, quota, tx); - - if (hasSameActivatedQuota) { - // don't need to switch - return; - } + if (hasSameActivatedQuota) return; // don't need to switch const featureId = await tx.feature .findFirst({ @@ -175,12 +167,11 @@ export class QuotaService { const quota = await this.prisma.workspaceFeature.findFirst({ where: { workspaceId, - feature: { - type: FeatureKind.Quota, - }, + feature: { type: FeatureKind.Quota }, activated: true, }, select: { + configs: true, reason: true, createdAt: true, expiredAt: true, @@ -190,7 +181,7 @@ export class QuotaService { if (quota) { const feature = await QuotaConfig.get(this.prisma, quota.featureId); - return { ...quota, feature }; + return { ...quota, feature: feature.withOverride(quota.configs) }; } return null; } @@ -209,24 +200,13 @@ export class QuotaService { quota, tx ); - - if (hasSameActivatedQuota) { - // don't need to switch - return; - } + if (hasSameActivatedQuota) return; // don't need to switch const featureId = await tx.feature .findFirst({ - where: { - feature: quota, - type: FeatureKind.Quota, - }, - select: { - id: true, - }, - orderBy: { - version: 'desc', - }, + where: { feature: quota, type: FeatureKind.Quota }, + select: { id: true }, + orderBy: { version: 'desc' }, }) .then(f => f?.id); @@ -294,18 +274,17 @@ export class QuotaService { type: Q ): Promise { const quota = await this.getQuota(type); - const configs = await this.prisma.workspaceFeature - .findFirst({ - where: { - workspaceId, - feature: { feature: type, type: FeatureKind.Feature }, - activated: true, - }, - select: { configs: true }, - }) - .then(q => q?.configs); - - if (quota && configs) { + if (quota) { + const configs = await this.prisma.workspaceFeature + .findFirst({ + where: { + workspaceId, + feature: { feature: type, type: FeatureKind.Feature }, + activated: true, + }, + select: { configs: true }, + }) + .then(q => q?.configs); return quota.withOverride(configs); } return undefined; @@ -314,7 +293,7 @@ export class QuotaService { async updateWorkspaceConfig( workspaceId: string, quota: QuotaType, - configs: JsonObject + configs: any ) { const current = await this.getWorkspaceConfig(workspaceId, quota); @@ -327,7 +306,7 @@ export class QuotaService { const r = await this.prisma.workspaceFeature.updateMany({ where: { workspaceId, - feature: { feature: quota, type: FeatureKind.Feature }, + feature: { feature: quota, type: FeatureKind.Quota }, activated: true, }, data: { configs }, diff --git a/packages/backend/server/src/core/quota/storage.ts b/packages/backend/server/src/core/quota/storage.ts index 2e5d0276ad675..78999de6541ae 100644 --- a/packages/backend/server/src/core/quota/storage.ts +++ b/packages/backend/server/src/core/quota/storage.ts @@ -4,7 +4,6 @@ import { FeatureService, FeatureType } from '../features'; import { PermissionService } from '../permission'; import { WorkspaceBlobStorage } from '../storage'; import { OneGB } from './constant'; -import { QuotaOverrideService } from './override'; import { QuotaConfig } from './quota'; import { QuotaService } from './service'; import { formatSize, Quota, type QuotaBusinessType, QuotaType } from './types'; @@ -17,8 +16,7 @@ export class QuotaManagementService { private readonly feature: FeatureService, private readonly quota: QuotaService, private readonly permissions: PermissionService, - private readonly storage: WorkspaceBlobStorage, - private readonly override: QuotaOverrideService + private readonly storage: WorkspaceBlobStorage ) {} async getUserQuota(userId: string) { @@ -147,6 +145,12 @@ export class QuotaManagementService { ); } + private async getWorkspaceQuota(userId: string, workspaceId: string) { + const workspaceQuota = await this.quota.getWorkspaceQuota(workspaceId); + if (workspaceQuota) return workspaceQuota; + return await this.quota.getUserQuota(userId); + } + // get workspace's owner quota and total size of used // quota was apply to owner's account async getWorkspaceUsage(workspaceId: string): Promise { @@ -164,7 +168,7 @@ export class QuotaManagementService { copilotActionLimit, humanReadable, }, - } = await this.quota.getUserQuota(owner.id); + } = await this.getWorkspaceQuota(owner.id, workspaceId); // get all workspaces size of owner used const usedSize = await this.getUserUsage(owner.id); // relax restrictions if workspace has unlimited feature @@ -192,7 +196,7 @@ export class QuotaManagementService { return this.mergeUnlimitedQuota(quota); } - return await this.override.overrideQuota(owner.id, workspaceId, quota); + return quota; } private mergeUnlimitedQuota(orig: QuotaBusinessType): QuotaBusinessType { diff --git a/packages/backend/server/src/data/migrations/1732786991577-team-workspace-feature.ts b/packages/backend/server/src/data/migrations/1732786991577-team-workspace-feature.ts deleted file mode 100644 index 4e6ee32565215..0000000000000 --- a/packages/backend/server/src/data/migrations/1732786991577-team-workspace-feature.ts +++ /dev/null @@ -1,14 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { FeatureType } from '../../core/features'; -import { upsertLatestFeatureVersion } from './utils/user-features'; - -export class TeamWorkspaceFeature1732786991577 { - // do the migration - static async up(db: PrismaClient) { - await upsertLatestFeatureVersion(db, FeatureType.TeamWorkspace); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index d26dafe401784..1afa8eab943d9 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -298,7 +298,6 @@ enum FeatureType { Admin Copilot EarlyAccess - TeamWorkspace UnlimitedCopilot UnlimitedWorkspace } @@ -547,7 +546,6 @@ type Mutation { """Update workspace""" updateWorkspace(input: UpdateWorkspaceInput!): WorkspaceType! - updateWorkspaceTeamConfig(input: UpdateTeamWorkspaceConfigInput!): Boolean! """Upload user avatar""" uploadAvatar(avatar: Upload!): UserType! @@ -835,11 +833,6 @@ enum SubscriptionVariant { Onetime } -type TeamWorkspaceConfigType { - enableAi: Boolean! - enableShare: Boolean! -} - type UnknownOauthProviderDataType { name: String! } @@ -848,12 +841,6 @@ type UnsupportedSubscriptionPlanDataType { plan: String! } -input UpdateTeamWorkspaceConfigInput { - enableAi: Boolean - enableShare: Boolean - id: ID! -} - input UpdateUserInput { """User name""" name: String @@ -1013,9 +1000,6 @@ type WorkspaceType { """The team subscription of the workspace, if exists.""" subscription: SubscriptionType - - """Team workspace config""" - teamConfig: TeamWorkspaceConfigType } type tokenType { diff --git a/packages/backend/server/tests/quota.spec.ts b/packages/backend/server/tests/quota.spec.ts index 18289f6ac4266..fd7c9a3d34733 100644 --- a/packages/backend/server/tests/quota.spec.ts +++ b/packages/backend/server/tests/quota.spec.ts @@ -6,19 +6,12 @@ import ava from 'ava'; import { AuthService } from '../src/core/auth'; import { - FeatureModule, - FeatureService, - FeatureType, -} from '../src/core/features'; -import { - QuotaBusinessType, QuotaManagementService, QuotaModule, - QuotaOverride, - QuotaOverrideService, QuotaService, QuotaType, } 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 { WorkspaceResolver } from '../src/core/workspaces/resolvers'; @@ -27,9 +20,7 @@ import { WorkspaceResolverMock } from './utils/feature'; const test = ava as TestFn<{ auth: AuthService; - feature: FeatureService; quota: QuotaService; - quotaOverride: QuotaOverrideService; quotaManager: QuotaManagementService; workspace: WorkspaceResolver; module: TestingModule; @@ -37,7 +28,7 @@ const test = ava as TestFn<{ test.beforeEach(async t => { const module = await createTestingModule({ - imports: [StorageModule, FeatureModule, QuotaModule], + imports: [StorageModule, QuotaModule], providers: [WorkspaceResolver], tapModule: module => { module @@ -46,17 +37,13 @@ test.beforeEach(async t => { }, }); - const feature = module.get(FeatureService); const quota = module.get(QuotaService); - const quotaOverride = module.get(QuotaOverrideService); const quotaManager = module.get(QuotaManagementService); const workspace = module.get(WorkspaceResolver); const auth = module.get(AuthService); t.context.module = module; - t.context.feature = feature; t.context.quota = quota; - t.context.quotaOverride = quotaOverride; t.context.quotaManager = quotaManager; t.context.workspace = workspace; t.context.auth = auth; @@ -155,41 +142,26 @@ test('should be able to check quota', async t => { }); test('should be able to override quota', async t => { - class TestQuotaOverride implements QuotaOverride { - get name(): string { - return TestQuotaOverride.name; - } - async overrideQuota( - _: string, - __: string, - features: FeatureType[], - quota: QuotaBusinessType - ): Promise { - if (features.includes(FeatureType.TeamWorkspace)) { - return { - ...quota, - blobLimit: 1024, - businessBlobLimit: 1024, - memberLimit: 1, - }; - } - return quota; - } - } - const { auth, feature, quotaOverride, quotaManager, workspace } = t.context; - quotaOverride.registerOverride(new TestQuotaOverride()); + const { auth, quotaManager, workspace } = t.context; const u1 = await auth.signUp('test@affine.pro', '123456'); const w1 = await workspace.createWorkspace(u1, null); const wq1 = await quotaManager.getWorkspaceUsage(w1.id); - t.is(wq1.blobLimit, 1024 * 1024 * 10, 'should be 10MB'); - t.is(wq1.businessBlobLimit, 1024 * 1024 * 100, 'should be 100MB'); + t.is(wq1.blobLimit, 10 * OneMB, 'should be 10MB'); + t.is(wq1.businessBlobLimit, 100 * OneMB, 'should be 100MB'); t.is(wq1.memberLimit, 3, 'should be 3'); - await feature.addWorkspaceFeature(w1.id, FeatureType.TeamWorkspace, 'test'); + await quotaManager.addTeamWorkspace(w1.id, 'test'); const wq2 = await quotaManager.getWorkspaceUsage(w1.id); - t.is(wq2.blobLimit, 1024, 'should be override to 1KB'); - t.is(wq2.businessBlobLimit, 1024, 'should be override to 1KB'); + t.is(wq2.storageQuota, 120 * OneGB, 'should be override to 100GB'); + t.is(wq2.businessBlobLimit, 500 * OneMB, 'should be override to 500MB'); t.is(wq2.memberLimit, 1, 'should be override to 1'); + + await quotaManager.updateWorkspaceConfig(w1.id, QuotaType.TeamPlanV1, { + memberLimit: 2, + }); + const wq3 = await quotaManager.getWorkspaceUsage(w1.id); + t.is(wq3.storageQuota, 140 * OneGB, 'should be override to 120GB'); + t.is(wq3.memberLimit, 2, 'should be override to 1'); });