diff --git a/packages/backend/server/migrations/20250203112209_data_migration_table_unique_name/migration.sql b/packages/backend/server/migrations/20250203112209_data_migration_table_unique_name/migration.sql new file mode 100644 index 0000000000..5f93e192f9 --- /dev/null +++ b/packages/backend/server/migrations/20250203112209_data_migration_table_unique_name/migration.sql @@ -0,0 +1,2 @@ +-- CreateIndex +CREATE UNIQUE INDEX "_data_migrations_name_key" ON "_data_migrations"("name"); diff --git a/packages/backend/server/migrations/20250203142831_standardize_features/migration.sql b/packages/backend/server/migrations/20250203142831_standardize_features/migration.sql new file mode 100644 index 0000000000..9a89b687dc --- /dev/null +++ b/packages/backend/server/migrations/20250203142831_standardize_features/migration.sql @@ -0,0 +1,24 @@ +-- AlterTable +ALTER TABLE "features" ADD COLUMN "updated_at" TIMESTAMPTZ(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, +ALTER COLUMN "type" SET DEFAULT 0, +ALTER COLUMN "configs" SET DEFAULT '{}'; + +-- AlterTable +ALTER TABLE "user_features" ADD COLUMN "name" VARCHAR NOT NULL DEFAULT '', +ADD COLUMN "type" INTEGER NOT NULL DEFAULT 0; + +-- AlterTable +ALTER TABLE "workspace_features" ADD COLUMN "name" VARCHAR NOT NULL DEFAULT '', +ADD COLUMN "type" INTEGER NOT NULL DEFAULT 0; + +-- CreateIndex +CREATE INDEX "user_features_name_idx" ON "user_features"("name"); + +-- CreateIndex +CREATE INDEX "user_features_feature_id_idx" ON "user_features"("feature_id"); + +-- CreateIndex +CREATE INDEX "workspace_features_name_idx" ON "workspace_features"("name"); + +-- CreateIndex +CREATE INDEX "workspace_features_feature_id_idx" ON "workspace_features"("feature_id"); diff --git a/packages/backend/server/schema.prisma b/packages/backend/server/schema.prisma index 7dcffd8057..0453735a04 100644 --- a/packages/backend/server/schema.prisma +++ b/packages/backend/server/schema.prisma @@ -1,7 +1,7 @@ generator client { provider = "prisma-client-js" binaryTargets = ["native", "debian-openssl-3.0.x", "linux-arm64-openssl-3.0.x"] - previewFeatures = ["metrics", "tracing", "relationJoins", "nativeDistinct"] + previewFeatures = ["metrics", "relationJoins", "nativeDistinct"] } datasource db { @@ -102,10 +102,10 @@ model Workspace { enableAi Boolean @default(true) @map("enable_ai") enableUrlPreview Boolean @default(false) @map("enable_url_preview") + features WorkspaceFeature[] pages WorkspacePage[] permissions WorkspaceUserPermission[] pagePermissions WorkspacePageUserPermission[] - features WorkspaceFeature[] blobs Blob[] @@map("workspaces") @@ -178,82 +178,76 @@ model WorkspacePageUserPermission { @@map("workspace_page_user_permissions") } -// feature gates is a way to enable/disable features for a user -// for example: -// - early access is a feature that allow some users to access the insider version -// - pro plan is a quota that allow some users access to more resources after they pay -model UserFeature { - id Int @id @default(autoincrement()) - userId String @map("user_id") @db.VarChar - featureId Int @map("feature_id") @db.Integer +model Feature { + id Int @id @default(autoincrement()) + name String @map("feature") @db.VarChar + configs Json @default("{}") @db.Json + createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3) + updatedAt DateTime @default(now()) @updatedAt @map("updated_at") @db.Timestamptz(3) + /// TODO(@forehalo): remove in the coming version + /// @deprecated + /// we don't need to record all the historical version of features + deprecatedVersion Int @default(0) @map("version") @db.Integer + /// @deprecated + /// we don't need to record type of features any more, there are always static, + /// but set it in `WorkspaceFeature` and `UserFeature` for fast query with just a little redundant. + deprecatedType Int @default(0) @map("type") @db.Integer - // we will record the reason why the feature is enabled/disabled - // for example: - // - pro_plan_v1: "user buy the pro plan" + userFeatures UserFeature[] + workspaceFeatures WorkspaceFeature[] + + @@unique([name, deprecatedVersion]) + @@map("features") +} + +model UserFeature { + id Int @id @default(autoincrement()) + userId String @map("user_id") @db.VarChar + featureId Int @map("feature_id") @db.Integer + // it should be typed as `optional` in the codebase, but we would keep all values exists during data migration. + // so it's safe to assert it a non-null value. + name String @default("") @map("name") @db.VarChar + // a little redundant, but fast the queries + type Int @default(0) @map("type") @db.Integer reason String @db.VarChar - // record the quota enabled time createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3) - // record the quota expired time, pay plan is a subscription, so it will expired expiredAt DateTime? @map("expired_at") @db.Timestamptz(3) - // whether the feature is activated - // for example: - // - if we switch the user to another plan, we will set the old plan to deactivated, but dont delete it activated Boolean @default(false) feature Feature @relation(fields: [featureId], references: [id], onDelete: Cascade) user User @relation(fields: [userId], references: [id], onDelete: Cascade) @@index([userId]) + @@index([name]) + @@index([featureId]) @@map("user_features") } -// feature gates is a way to enable/disable features for a workspace -// for example: -// - copilet is a feature that allow some users in a workspace to access the copilet feature model WorkspaceFeature { - id Int @id @default(autoincrement()) - workspaceId String @map("workspace_id") @db.VarChar - featureId Int @map("feature_id") @db.Integer - - // override quota's configs - configs Json @default("{}") @db.Json - // we will record the reason why the feature is enabled/disabled - // for example: - // - copilet_v1: "owner buy the copilet feature package" - reason String @db.VarChar - // record the feature enabled time - createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3) - // record the quota expired time, pay plan is a subscription, so it will expired - expiredAt DateTime? @map("expired_at") @db.Timestamptz(3) - // whether the feature is activated - // for example: - // - if owner unsubscribe a feature package, we will set the feature to deactivated, but dont delete it - activated Boolean @default(false) + id Int @id @default(autoincrement()) + workspaceId String @map("workspace_id") @db.VarChar + featureId Int @map("feature_id") @db.Integer + // it should be typed as `optional` in the codebase, but we would keep all values exists during data migration. + // so it's safe to assert it a non-null value. + name String @default("") @map("name") @db.VarChar + // a little redundant, but fast the queries + type Int @default(0) @map("type") @db.Integer + /// overrides for the default feature configs + configs Json @default("{}") @db.Json + reason String @db.VarChar + createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3) + activated Boolean @default(false) + expiredAt DateTime? @map("expired_at") @db.Timestamptz(3) feature Feature @relation(fields: [featureId], references: [id], onDelete: Cascade) workspace Workspace @relation(fields: [workspaceId], references: [id], onDelete: Cascade) @@index([workspaceId]) + @@index([name]) + @@index([featureId]) @@map("workspace_features") } -model Feature { - id Int @id @default(autoincrement()) - feature String @db.VarChar - version Int @default(0) @db.Integer - // 0: feature, 1: quota - type Int @db.Integer - // configs, define by feature controller - configs Json @db.Json - createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3) - - UserFeatureGates UserFeature[] - WorkspaceFeatures WorkspaceFeature[] - - @@unique([feature, version]) - @@map("features") -} - // the latest snapshot of each doc that we've seen // Snapshot + Updates are the latest state of the doc model Snapshot { @@ -417,7 +411,7 @@ model AiSession { model DataMigration { id String @id @default(uuid()) @db.VarChar - name String @db.VarChar + name String @unique @db.VarChar startedAt DateTime @default(now()) @map("started_at") @db.Timestamptz(3) finishedAt DateTime? @map("finished_at") @db.Timestamptz(3) @@ -552,21 +546,21 @@ model Subscription { } model Invoice { - stripeInvoiceId String @id @map("stripe_invoice_id") - targetId String @map("target_id") @db.VarChar - currency String @db.VarChar(3) + stripeInvoiceId String @id @map("stripe_invoice_id") + targetId String @map("target_id") @db.VarChar + currency String @db.VarChar(3) // CNY 12.50 stored as 1250 - amount Int @db.Integer - status String @db.VarChar(20) - createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3) - updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3) + amount Int @db.Integer + status String @db.VarChar(20) + createdAt DateTime @default(now()) @map("created_at") @db.Timestamptz(3) + updatedAt DateTime @updatedAt @map("updated_at") @db.Timestamptz(3) // billing reason - reason String? @db.VarChar - lastPaymentError String? @map("last_payment_error") @db.Text + reason String? @db.VarChar + lastPaymentError String? @map("last_payment_error") @db.Text // stripe hosted invoice link - link String? @db.Text + link String? @db.Text // whether the onetime subscription has been redeemed - onetimeSubscriptionRedeemed Boolean @map("onetime_subscription_redeemed") @default(false) + onetimeSubscriptionRedeemed Boolean @default(false) @map("onetime_subscription_redeemed") @@index([targetId]) @@map("invoices") diff --git a/packages/backend/server/src/__tests__/app/selfhost.e2e.ts b/packages/backend/server/src/__tests__/app/selfhost.e2e.ts index 4133fa7704..7533beeb4d 100644 --- a/packages/backend/server/src/__tests__/app/selfhost.e2e.ts +++ b/packages/backend/server/src/__tests__/app/selfhost.e2e.ts @@ -1,7 +1,6 @@ import { mkdirSync, writeFileSync } from 'node:fs'; import path from 'node:path'; -import type { INestApplication } from '@nestjs/common'; import { PrismaClient } from '@prisma/client'; import type { TestFn } from 'ava'; import ava from 'ava'; @@ -10,10 +9,10 @@ import request from 'supertest'; import { buildAppModule } from '../../app.module'; import { Config } from '../../base'; import { ServerService } from '../../core/config'; -import { createTestingApp, initTestingDB } from '../utils'; +import { createTestingApp, type TestingApp } from '../utils'; const test = ava as TestFn<{ - app: INestApplication; + app: TestingApp; db: PrismaClient; }>; @@ -54,7 +53,7 @@ test.before('init selfhost server', async t => { }); test.beforeEach(async t => { - await initTestingDB(t.context.db); + await t.context.app.initTestingDB(); const server = t.context.app.get(ServerService); // @ts-expect-error disable cache server._initialized = false; @@ -188,7 +187,8 @@ test('should redirect to admin if initialized', async t => { t.is(res.header.location, '/admin'); }); -test('should return mobile assets if visited by mobile', async t => { +// TODO(@forehalo): return mobile when it's ready +test.skip('should return web assets if visited by mobile', async t => { await t.context.db.user.create({ data: { name: 'test', diff --git a/packages/backend/server/src/__tests__/auth/service.spec.ts b/packages/backend/server/src/__tests__/auth/service.spec.ts index 2f1a7c21c4..6f5cd11902 100644 --- a/packages/backend/server/src/__tests__/auth/service.spec.ts +++ b/packages/backend/server/src/__tests__/auth/service.spec.ts @@ -1,4 +1,3 @@ -import { TestingModule } from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import ava, { TestFn } from 'ava'; @@ -8,7 +7,7 @@ import { FeatureModule } from '../../core/features'; import { QuotaModule } from '../../core/quota'; import { UserModule } from '../../core/user'; import { Models } from '../../models'; -import { createTestingModule, initTestingDB } from '../utils'; +import { createTestingModule, type TestingModule } from '../utils'; const test = ava as TestFn<{ auth: AuthService; @@ -31,7 +30,7 @@ test.before(async t => { }); test.beforeEach(async t => { - await initTestingDB(t.context.db); + await t.context.m.initTestingDB(); t.context.u1 = await t.context.auth.signUp('u1@affine.pro', '1'); }); diff --git a/packages/backend/server/src/__tests__/copilot-provider.spec.ts b/packages/backend/server/src/__tests__/copilot-provider.spec.ts index 70344ac26a..72bdee3e1b 100644 --- a/packages/backend/server/src/__tests__/copilot-provider.spec.ts +++ b/packages/backend/server/src/__tests__/copilot-provider.spec.ts @@ -1,6 +1,5 @@ /// -import { TestingModule } from '@nestjs/testing'; import type { ExecutionContext, TestFn } from 'ava'; import ava from 'ava'; @@ -27,7 +26,7 @@ import { CopilotCheckHtmlExecutor, CopilotCheckJsonExecutor, } from '../plugins/copilot/workflow/executor'; -import { createTestingModule } from './utils'; +import { createTestingModule, TestingModule } from './utils'; import { TestAssets } from './utils/copilot'; type Tester = { diff --git a/packages/backend/server/src/__tests__/copilot.e2e.ts b/packages/backend/server/src/__tests__/copilot.e2e.ts index f2a698221a..1f6e11ce17 100644 --- a/packages/backend/server/src/__tests__/copilot.e2e.ts +++ b/packages/backend/server/src/__tests__/copilot.e2e.ts @@ -2,7 +2,6 @@ import { randomUUID } from 'node:crypto'; -import { INestApplication } from '@nestjs/common'; import type { TestFn } from 'ava'; import ava from 'ava'; import Sinon from 'sinon'; @@ -27,6 +26,7 @@ import { createWorkspace, inviteUser, signUp, + TestingApp, } from './utils'; import { array2sse, @@ -47,7 +47,7 @@ import { const test = ava as TestFn<{ auth: AuthService; - app: INestApplication; + app: TestingApp; prompt: PromptService; provider: CopilotProviderService; storage: CopilotStorage; diff --git a/packages/backend/server/src/__tests__/copilot.spec.ts b/packages/backend/server/src/__tests__/copilot.spec.ts index a0920a1b9b..7b9c5d2507 100644 --- a/packages/backend/server/src/__tests__/copilot.spec.ts +++ b/packages/backend/server/src/__tests__/copilot.spec.ts @@ -1,6 +1,3 @@ -/// - -import { TestingModule } from '@nestjs/testing'; import type { TestFn } from 'ava'; import ava from 'ava'; import Sinon from 'sinon'; @@ -41,7 +38,7 @@ import { } from '../plugins/copilot/workflow/executor'; import { AutoRegisteredWorkflowExecutor } from '../plugins/copilot/workflow/executor/utils'; import { WorkflowGraphList } from '../plugins/copilot/workflow/graph'; -import { createTestingModule } from './utils'; +import { createTestingModule, TestingModule } from './utils'; import { MockCopilotTestProvider, WorkflowTestCases } from './utils/copilot'; const test = ava as TestFn<{ diff --git a/packages/backend/server/src/__tests__/doc/history.spec.ts b/packages/backend/server/src/__tests__/doc/history.spec.ts index e53f09fceb..d9534b2051 100644 --- a/packages/backend/server/src/__tests__/doc/history.spec.ts +++ b/packages/backend/server/src/__tests__/doc/history.spec.ts @@ -1,4 +1,3 @@ -import { TestingModule } from '@nestjs/testing'; import type { Snapshot } from '@prisma/client'; import { PrismaClient } from '@prisma/client'; import test from 'ava'; @@ -7,7 +6,7 @@ import * as Sinon from 'sinon'; import { DocStorageModule, PgWorkspaceDocStorageAdapter } from '../../core/doc'; import { DocStorageOptions } from '../../core/doc/options'; import { DocRecord } from '../../core/doc/storage'; -import { createTestingModule, initTestingDB } from '../utils'; +import { createTestingModule, type TestingModule } from '../utils'; let m: TestingModule; let adapter: PgWorkspaceDocStorageAdapter; @@ -24,7 +23,7 @@ test.before(async () => { }); test.beforeEach(async () => { - await initTestingDB(db); + await m.initTestingDB(); const options = m.get(DocStorageOptions); Sinon.stub(options, 'historyMaxAge').resolves(1000); }); diff --git a/packages/backend/server/src/__tests__/doc/renderer.spec.ts b/packages/backend/server/src/__tests__/doc/renderer.spec.ts index d8840f9134..6638f7b4b1 100644 --- a/packages/backend/server/src/__tests__/doc/renderer.spec.ts +++ b/packages/backend/server/src/__tests__/doc/renderer.spec.ts @@ -75,7 +75,8 @@ test('should render correct html', async t => { ); }); -test('should render correct mobile html', async t => { +// TODO(@forehalo): enable it when mobile version is ready +test.skip('should render correct mobile html', async t => { const res = await request(t.context.app.getHttpServer()) .get('/workspace/xxxx/xxx') .set('user-agent', mobileUAString) diff --git a/packages/backend/server/src/__tests__/doc/workspace.spec.ts b/packages/backend/server/src/__tests__/doc/workspace.spec.ts index 47e6ec7855..4be1ffb5b3 100644 --- a/packages/backend/server/src/__tests__/doc/workspace.spec.ts +++ b/packages/backend/server/src/__tests__/doc/workspace.spec.ts @@ -1,4 +1,3 @@ -import { TestingModule } from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import test from 'ava'; import * as Sinon from 'sinon'; @@ -9,7 +8,7 @@ import { DocStorageModule, PgWorkspaceDocStorageAdapter as Adapter, } from '../../core/doc'; -import { createTestingModule, initTestingDB } from '../utils'; +import { createTestingModule, type TestingModule } from '../utils'; let m: TestingModule; let db: PrismaClient; @@ -35,7 +34,7 @@ test.before('init testing module', async () => { }); test.beforeEach(async () => { - await initTestingDB(db); + await m.initTestingDB(); }); test.after.always(async () => { diff --git a/packages/backend/server/src/__tests__/feature.spec.ts b/packages/backend/server/src/__tests__/feature.spec.ts deleted file mode 100644 index bb3bbac61f..0000000000 --- a/packages/backend/server/src/__tests__/feature.spec.ts +++ /dev/null @@ -1,176 +0,0 @@ -/// - -import { INestApplication } from '@nestjs/common'; -import type { TestFn } from 'ava'; -import ava from 'ava'; - -import { Runtime } from '../base'; -import { AuthService } from '../core/auth/service'; -import { - FeatureManagementService, - FeatureModule, - FeatureService, - FeatureType, -} from '../core/features'; -import { WorkspaceResolver } from '../core/workspaces/resolvers'; -import { createTestingApp } from './utils'; -import { WorkspaceResolverMock } from './utils/feature'; - -const test = ava as TestFn<{ - auth: AuthService; - feature: FeatureService; - workspace: WorkspaceResolver; - management: FeatureManagementService; - app: INestApplication; -}>; - -test.beforeEach(async t => { - const { app } = await createTestingApp({ - imports: [FeatureModule], - providers: [WorkspaceResolver], - tapModule: module => { - module - .overrideProvider(WorkspaceResolver) - .useClass(WorkspaceResolverMock); - }, - }); - - const runtime = app.get(Runtime); - await runtime.set('flags/earlyAccessControl', true); - t.context.app = app; - t.context.auth = app.get(AuthService); - t.context.feature = app.get(FeatureService); - t.context.workspace = app.get(WorkspaceResolver); - t.context.management = app.get(FeatureManagementService); -}); - -test.afterEach.always(async t => { - await t.context.app.close(); -}); - -test('should be able to set user feature', async t => { - const { auth, feature } = t.context; - - const u1 = await auth.signUp('test@test.com', '123456'); - - const f1 = await feature.getUserFeatures(u1.id); - t.is(f1.length, 0, 'should be empty'); - - await feature.addUserFeature(u1.id, FeatureType.EarlyAccess, 'test'); - - const f2 = await feature.getUserFeatures(u1.id); - t.is(f2.length, 1, 'should have 1 feature'); - t.is(f2[0].feature.name, FeatureType.EarlyAccess, 'should be early access'); -}); - -test('should be able to check early access', async t => { - const { auth, feature, management } = t.context; - const u1 = await auth.signUp('test@test.com', '123456'); - - const f1 = await management.canEarlyAccess(u1.email); - t.false(f1, 'should not have early access'); - - await management.addEarlyAccess(u1.id); - const f2 = await management.canEarlyAccess(u1.email); - t.true(f2, 'should have early access'); - - const f3 = await feature.listUsersByFeature(FeatureType.EarlyAccess); - t.is(f3.length, 1, 'should have 1 user'); - t.is(f3[0].id, u1.id, 'should be the same user'); -}); - -test('should be able revert user feature', async t => { - const { auth, feature, management } = t.context; - const u1 = await auth.signUp('test@test.com', '123456'); - - const f1 = await management.canEarlyAccess(u1.email); - t.false(f1, 'should not have early access'); - - await management.addEarlyAccess(u1.id); - const f2 = await management.canEarlyAccess(u1.email); - t.true(f2, 'should have early access'); - const q1 = await management.listEarlyAccess(); - t.is(q1.length, 1, 'should have 1 user'); - t.is(q1[0].id, u1.id, 'should be the same user'); - - await management.removeEarlyAccess(u1.id); - const f3 = await management.canEarlyAccess(u1.email); - t.false(f3, 'should not have early access'); - const q2 = await management.listEarlyAccess(); - t.is(q2.length, 0, 'should have no user'); - - const q3 = await feature.getUserFeatures(u1.id); - t.is(q3.length, 1, 'should have 1 feature'); - t.is(q3[0].feature.name, FeatureType.EarlyAccess, 'should be early access'); - t.is(q3[0].activated, false, 'should be deactivated'); -}); - -test('should be same instance after reset the user feature', async t => { - const { auth, feature, management } = t.context; - const u1 = await auth.signUp('test@test.com', '123456'); - - await management.addEarlyAccess(u1.id); - const f1 = (await feature.getUserFeatures(u1.id))[0]; - - await management.removeEarlyAccess(u1.id); - - await management.addEarlyAccess(u1.id); - const f2 = (await feature.getUserFeatures(u1.id))[1]; - - t.is(f1.feature, f2.feature, 'should be same instance'); -}); - -test('should be able to set workspace feature', async t => { - const { auth, feature, workspace } = t.context; - - const u1 = await auth.signUp('test@test.com', '123456'); - const w1 = await workspace.createWorkspace(u1, null); - - const f1 = await feature.getWorkspaceFeatures(w1.id); - t.is(f1.length, 0, 'should be empty'); - - await feature.addWorkspaceFeature(w1.id, FeatureType.Copilot, 'test'); - - const f2 = await feature.getWorkspaceFeatures(w1.id); - t.is(f2.length, 1, 'should have 1 feature'); - t.is(f2[0].feature.name, FeatureType.Copilot, 'should be copilot'); -}); - -test('should be able to check workspace feature', async t => { - const { auth, feature, workspace, management } = t.context; - const u1 = await auth.signUp('test@test.com', '123456'); - const w1 = await workspace.createWorkspace(u1, null); - - const f1 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot); - t.false(f1, 'should not have copilot'); - - await management.addWorkspaceFeatures(w1.id, FeatureType.Copilot, 'test'); - const f2 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot); - t.true(f2, 'should have copilot'); - - const f3 = await feature.listWorkspacesByFeature(FeatureType.Copilot); - t.is(f3.length, 1, 'should have 1 workspace'); - t.is(f3[0].id, w1.id, 'should be the same workspace'); -}); - -test('should be able revert workspace feature', async t => { - const { auth, feature, workspace, management } = t.context; - const u1 = await auth.signUp('test@test.com', '123456'); - const w1 = await workspace.createWorkspace(u1, null); - - const f1 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot); - t.false(f1, 'should not have feature'); - - await management.addWorkspaceFeatures(w1.id, FeatureType.Copilot, 'test'); - const f2 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot); - t.true(f2, 'should have feature'); - - await management.removeWorkspaceFeature(w1.id, FeatureType.Copilot); - const f3 = await management.hasWorkspaceFeature(w1.id, FeatureType.Copilot); - t.false(f3, 'should not have feature'); - - const q3 = await feature.getWorkspaceFeatures(w1.id); - t.is(q3.length, 1, 'should have 1 feature'); - t.is(q3[0].feature.name, FeatureType.Copilot, 'should be copilot'); - t.is(q3[0].activated, false, 'should be deactivated'); -}); diff --git a/packages/backend/server/src/__tests__/mailer.spec.ts b/packages/backend/server/src/__tests__/mailer.spec.ts index cde65207f0..5d24495eab 100644 --- a/packages/backend/server/src/__tests__/mailer.spec.ts +++ b/packages/backend/server/src/__tests__/mailer.spec.ts @@ -5,7 +5,6 @@ import Sinon from 'sinon'; import { AppModule } from '../app.module'; import { MailService } from '../base/mailer'; -import { FeatureManagementService } from '../core/features'; import { createTestingApp, createWorkspace, inviteUser, signUp } from './utils'; const test = ava as TestFn<{ app: INestApplication; @@ -16,13 +15,6 @@ import * as renderers from '../mails'; test.beforeEach(async t => { const { module, app } = await createTestingApp({ imports: [AppModule], - tapModule: module => { - module.overrideProvider(FeatureManagementService).useValue({ - hasWorkspaceFeature() { - return false; - }, - }); - }, }); const mail = module.get(MailService); diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/feature-user.spec.ts.md b/packages/backend/server/src/__tests__/models/__snapshots__/feature-user.spec.ts.md new file mode 100644 index 0000000000..9221205f5d --- /dev/null +++ b/packages/backend/server/src/__tests__/models/__snapshots__/feature-user.spec.ts.md @@ -0,0 +1,44 @@ +# Snapshot report for `src/__tests__/models/feature-user.spec.ts` + +The actual snapshot is saved in `feature-user.spec.ts.snap`. + +Generated by [AVA](https://avajs.dev). + +## should get user quota + +> free plan + + { + blobLimit: 10485760, + businessBlobLimit: 104857600, + copilotActionLimit: 10, + historyPeriod: 604800000, + memberLimit: 3, + name: 'Free', + storageQuota: 10737418240, + } + +## should switch user quota + +> switch to pro plan + + { + blobLimit: 104857600, + copilotActionLimit: 10, + historyPeriod: 2592000000, + memberLimit: 10, + name: 'Pro', + storageQuota: 107374182400, + } + +> switch to free plan + + { + blobLimit: 10485760, + businessBlobLimit: 104857600, + copilotActionLimit: 10, + historyPeriod: 604800000, + memberLimit: 3, + name: 'Free', + storageQuota: 10737418240, + } diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/feature-user.spec.ts.snap b/packages/backend/server/src/__tests__/models/__snapshots__/feature-user.spec.ts.snap new file mode 100644 index 0000000000..faecfcd623 Binary files /dev/null and b/packages/backend/server/src/__tests__/models/__snapshots__/feature-user.spec.ts.snap differ diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/feature-workspace.spec.ts.md b/packages/backend/server/src/__tests__/models/__snapshots__/feature-workspace.spec.ts.md new file mode 100644 index 0000000000..f3c67c3588 --- /dev/null +++ b/packages/backend/server/src/__tests__/models/__snapshots__/feature-workspace.spec.ts.md @@ -0,0 +1,18 @@ +# Snapshot report for `src/__tests__/models/feature-workspace.spec.ts` + +The actual snapshot is saved in `feature-workspace.spec.ts.snap`. + +Generated by [AVA](https://avajs.dev). + +## should get workspace quota + +> team plan + + { + blobLimit: 524288000, + historyPeriod: 2592000000, + memberLimit: 100, + name: 'Team Workspace', + seatQuota: 21474836480, + storageQuota: 2254857830400, + } diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/feature-workspace.spec.ts.snap b/packages/backend/server/src/__tests__/models/__snapshots__/feature-workspace.spec.ts.snap new file mode 100644 index 0000000000..df1c07a0a0 Binary files /dev/null and b/packages/backend/server/src/__tests__/models/__snapshots__/feature-workspace.spec.ts.snap differ diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.md b/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.md index aa90298228..a8c4f83458 100644 --- a/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.md +++ b/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.md @@ -10,6 +10,7 @@ Generated by [AVA](https://avajs.dev). { blobLimit: 10485760, + businessBlobLimit: 104857600, copilotActionLimit: 10, historyPeriod: 604800000, memberLimit: 3, @@ -23,6 +24,7 @@ Generated by [AVA](https://avajs.dev). { blobLimit: 10485760, + businessBlobLimit: 104857600, copilotActionLimit: 10, historyPeriod: 604800000, memberLimit: 3, diff --git a/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.snap b/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.snap index 296805ddc2..fca8a58a00 100644 Binary files a/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.snap and b/packages/backend/server/src/__tests__/models/__snapshots__/feature.spec.ts.snap differ diff --git a/packages/backend/server/src/__tests__/models/feature-user.spec.ts b/packages/backend/server/src/__tests__/models/feature-user.spec.ts index f6ebe26ee1..4b54096a0b 100644 --- a/packages/backend/server/src/__tests__/models/feature-user.spec.ts +++ b/packages/backend/server/src/__tests__/models/feature-user.spec.ts @@ -1,9 +1,8 @@ -import { TestingModule } from '@nestjs/testing'; -import { PrismaClient, User } from '@prisma/client'; +import { User } from '@prisma/client'; import ava, { TestFn } from 'ava'; -import { UserFeatureModel, UserModel } from '../../models'; -import { createTestingModule, initTestingDB } from '../utils'; +import { FeatureType, UserFeatureModel, UserModel } from '../../models'; +import { createTestingModule, TestingModule } from '../utils'; interface Context { module: TestingModule; @@ -21,7 +20,7 @@ test.before(async t => { }); test.beforeEach(async t => { - await initTestingDB(t.context.module.get(PrismaClient)); + await t.context.module.initTestingDB(); t.context.u1 = await t.context.module.get(UserModel).create({ email: 'u1@affine.pro', registered: true, @@ -41,7 +40,13 @@ test('should get null if user feature not found', async t => { test('should get user feature', async t => { const { model, u1 } = t.context; const userFeature = await model.get(u1.id, 'free_plan_v1'); - t.is(userFeature?.feature, 'free_plan_v1'); + t.is(userFeature?.name, 'free_plan_v1'); +}); + +test('should get user quota', async t => { + const { model, u1 } = t.context; + const userQuota = await model.getQuota(u1.id); + t.snapshot(userQuota?.configs, 'free plan'); }); test('should list user features', async t => { @@ -50,6 +55,16 @@ test('should list user features', async t => { t.like(await model.list(u1.id), ['free_plan_v1']); }); +test('should list user features by type', async t => { + const { model, u1 } = t.context; + + await model.add(u1.id, 'free_plan_v1', 'test'); + await model.add(u1.id, 'unlimited_copilot', 'test'); + + t.like(await model.list(u1.id, FeatureType.Quota), ['free_plan_v1']); + t.like(await model.list(u1.id, FeatureType.Feature), ['unlimited_copilot']); +}); + test('should directly test user feature existence', async t => { const { model, u1 } = t.context; @@ -82,14 +97,29 @@ test('should remove user feature', async t => { t.false((await model.list(u1.id)).includes('free_plan_v1')); }); -test('should switch user feature', async t => { +test('should switch user quota', async t => { const { model, u1 } = t.context; - await model.switch(u1.id, 'free_plan_v1', 'pro_plan_v1', 'test'); + await model.switchQuota(u1.id, 'pro_plan_v1', 'test'); + const quota = await model.getQuota(u1.id); + t.snapshot(quota?.configs, 'switch to pro plan'); - t.false(await model.has(u1.id, 'free_plan_v1')); - t.true(await model.has(u1.id, 'pro_plan_v1')); - - t.false((await model.list(u1.id)).includes('free_plan_v1')); - t.true((await model.list(u1.id)).includes('pro_plan_v1')); + await model.switchQuota(u1.id, 'free_plan_v1', 'test'); + const quota2 = await model.getQuota(u1.id); + t.snapshot(quota2?.configs, 'switch to free plan'); +}); + +test('should not switch user quota if the new quota is the same as the current one', async t => { + const { model, u1 } = t.context; + + await model.switchQuota(u1.id, 'free_plan_v1', 'test not switch'); + + // @ts-expect-error private + const quota = await model.db.userFeature.findFirst({ + where: { + userId: u1.id, + }, + }); + + t.not(quota?.reason, 'test not switch'); }); diff --git a/packages/backend/server/src/__tests__/models/feature-workspace.spec.ts b/packages/backend/server/src/__tests__/models/feature-workspace.spec.ts index a3726ee141..4e03b369c0 100644 --- a/packages/backend/server/src/__tests__/models/feature-workspace.spec.ts +++ b/packages/backend/server/src/__tests__/models/feature-workspace.spec.ts @@ -1,9 +1,13 @@ -import { TestingModule } from '@nestjs/testing'; -import { PrismaClient, Workspace } from '@prisma/client'; +import { Workspace } from '@prisma/client'; import ava, { TestFn } from 'ava'; -import { UserModel, WorkspaceFeatureModel, WorkspaceModel } from '../../models'; -import { createTestingModule, initTestingDB } from '../utils'; +import { + FeatureType, + UserModel, + WorkspaceFeatureModel, + WorkspaceModel, +} from '../../models'; +import { createTestingModule, type TestingModule } from '../utils'; interface Context { module: TestingModule; @@ -21,7 +25,7 @@ test.before(async t => { }); test.beforeEach(async t => { - await initTestingDB(t.context.module.get(PrismaClient)); + await t.context.module.initTestingDB(); const u1 = await t.context.module.get(UserModel).create({ email: 'u1@affine.pro', registered: true, @@ -46,18 +50,52 @@ test('should directly test workspace feature existence', async t => { t.false(await model.has(ws.id, 'unlimited_workspace')); }); +test('should get workspace quota', async t => { + const { model, ws } = t.context; + + await model.add(ws.id, 'team_plan_v1', 'test', { + memberLimit: 100, + }); + + const quota = await model.getQuota(ws.id); + t.snapshot(quota?.configs, 'team plan'); +}); + +test('should return null if quota removed', async t => { + const { model, ws } = t.context; + + await model.add(ws.id, 'team_plan_v1', 'test', { + memberLimit: 100, + }); + + await model.remove(ws.id, 'team_plan_v1'); + + const quota = await model.getQuota(ws.id); + t.is(quota, null); +}); + test('should list empty workspace features', async t => { const { model, ws } = t.context; t.deepEqual(await model.list(ws.id), []); }); +test('should list workspace features by type', async t => { + const { model, ws } = t.context; + + await model.add(ws.id, 'unlimited_workspace', 'test'); + await model.add(ws.id, 'team_plan_v1', 'test'); + + t.like(await model.list(ws.id, FeatureType.Quota), ['team_plan_v1']); + t.like(await model.list(ws.id, FeatureType.Feature), ['unlimited_workspace']); +}); + test('should add workspace feature', async t => { const { model, ws } = t.context; await model.add(ws.id, 'unlimited_workspace', 'test'); t.is( - (await model.get(ws.id, 'unlimited_workspace'))?.feature, + (await model.get(ws.id, 'unlimited_workspace'))?.name, 'unlimited_workspace' ); t.true(await model.has(ws.id, 'unlimited_workspace')); @@ -103,25 +141,3 @@ test('should remove workspace feature', async t => { t.false(await model.has(ws.id, 'team_plan_v1')); t.false((await model.list(ws.id)).includes('team_plan_v1')); }); - -test('should switch workspace feature', async t => { - const { model, ws } = t.context; - - await model.switch(ws.id, 'team_plan_v1', 'unlimited_workspace', 'test'); - - t.false(await model.has(ws.id, 'team_plan_v1')); - t.true(await model.has(ws.id, 'unlimited_workspace')); - - t.false((await model.list(ws.id)).includes('team_plan_v1')); - t.true((await model.list(ws.id)).includes('unlimited_workspace')); -}); - -test('should switch workspace feature with overrides', async t => { - const { model, ws } = t.context; - - await model.add(ws.id, 'unlimited_workspace', 'test'); - await model.add(ws.id, 'team_plan_v1', 'test', { memberLimit: 100 }); - const f2 = await model.get(ws.id, 'team_plan_v1'); - - t.is(f2!.configs.memberLimit, 100); -}); diff --git a/packages/backend/server/src/__tests__/models/feature.spec.ts b/packages/backend/server/src/__tests__/models/feature.spec.ts index 10a8e61401..1389d3fa80 100644 --- a/packages/backend/server/src/__tests__/models/feature.spec.ts +++ b/packages/backend/server/src/__tests__/models/feature.spec.ts @@ -1,9 +1,8 @@ -import { TestingModule } from '@nestjs/testing'; -import { PrismaClient } from '@prisma/client'; import ava, { TestFn } from 'ava'; +import { FeatureType } from '../../models'; import { FeatureModel } from '../../models/feature'; -import { createTestingModule, initTestingDB } from '../utils'; +import { createTestingModule, type TestingModule } from '../utils'; interface Context { module: TestingModule; @@ -20,7 +19,7 @@ test.before(async t => { }); test.beforeEach(async t => { - await initTestingDB(t.context.module.get(PrismaClient)); + await t.context.module.initTestingDB(); }); test.after(async t => { @@ -91,7 +90,12 @@ test('should get feature if extra fields exist in feature config', async t => { test('should create feature', async t => { const { feature } = t.context; - const newFeature = await feature.upsert('new_feature' as any, {}); + const newFeature = await feature.upsert( + 'new_feature' as any, + {}, + FeatureType.Feature, + 1 + ); t.deepEqual(newFeature.configs, {}); }); @@ -100,10 +104,15 @@ test('should update feature', async t => { const { feature } = t.context; const freePlanFeature = await feature.get('free_plan_v1'); - const newFreePlanFeature = await feature.upsert('free_plan_v1', { - ...freePlanFeature.configs, - memberLimit: 10, - }); + const newFreePlanFeature = await feature.upsert( + 'free_plan_v1', + { + ...freePlanFeature.configs, + memberLimit: 10, + }, + FeatureType.Quota, + 1 + ); t.deepEqual(newFreePlanFeature.configs, { ...freePlanFeature.configs, @@ -113,7 +122,10 @@ test('should update feature', async t => { test('should throw if feature config is invalid when updating', async t => { const { feature } = t.context; - await t.throwsAsync(feature.upsert('free_plan_v1', {} as any), { - message: 'Invalid feature config for free_plan_v1', - }); + await t.throwsAsync( + feature.upsert('free_plan_v1', {} as any, FeatureType.Quota, 1), + { + message: 'Invalid feature config for free_plan_v1', + } + ); }); diff --git a/packages/backend/server/src/__tests__/models/page.spec.ts b/packages/backend/server/src/__tests__/models/page.spec.ts index 5974c86fac..d8bfd2bb4f 100644 --- a/packages/backend/server/src/__tests__/models/page.spec.ts +++ b/packages/backend/server/src/__tests__/models/page.spec.ts @@ -1,4 +1,3 @@ -import { TestingModule } from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import ava, { TestFn } from 'ava'; @@ -8,7 +7,7 @@ import { PublicPageMode } from '../../models/common'; import { PageModel } from '../../models/page'; import { type User, UserModel } from '../../models/user'; import { type Workspace, WorkspaceModel } from '../../models/workspace'; -import { createTestingModule, initTestingDB } from '../utils'; +import { createTestingModule, type TestingModule } from '../utils'; interface Context { config: Config; @@ -36,7 +35,7 @@ let user: User; let workspace: Workspace; test.beforeEach(async t => { - await initTestingDB(t.context.db); + await t.context.module.initTestingDB(); user = await t.context.user.create({ email: 'test@affine.pro', }); diff --git a/packages/backend/server/src/__tests__/models/session.spec.ts b/packages/backend/server/src/__tests__/models/session.spec.ts index bc91a58b07..83332aad13 100644 --- a/packages/backend/server/src/__tests__/models/session.spec.ts +++ b/packages/backend/server/src/__tests__/models/session.spec.ts @@ -1,11 +1,10 @@ -import { TestingModule } from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import ava, { TestFn } from 'ava'; import { Config } from '../../base/config'; import { SessionModel } from '../../models/session'; import { UserModel } from '../../models/user'; -import { createTestingModule, initTestingDB } from '../utils'; +import { createTestingModule, type TestingModule } from '../utils'; interface Context { config: Config; @@ -28,7 +27,7 @@ test.before(async t => { }); test.beforeEach(async t => { - await initTestingDB(t.context.db); + await t.context.module.initTestingDB(); }); test.after(async t => { diff --git a/packages/backend/server/src/__tests__/models/user.spec.ts b/packages/backend/server/src/__tests__/models/user.spec.ts index 0947ce4041..d48d629c7d 100644 --- a/packages/backend/server/src/__tests__/models/user.spec.ts +++ b/packages/backend/server/src/__tests__/models/user.spec.ts @@ -1,4 +1,3 @@ -import { TestingModule } from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import ava, { TestFn } from 'ava'; import Sinon from 'sinon'; @@ -7,7 +6,7 @@ import { EmailAlreadyUsed, EventBus } from '../../base'; import { WorkspaceRole } from '../../core/permission'; import { UserModel } from '../../models/user'; import { WorkspaceMemberStatus } from '../../models/workspace'; -import { createTestingModule, initTestingDB } from '../utils'; +import { createTestingModule, type TestingModule } from '../utils'; interface Context { module: TestingModule; @@ -24,7 +23,7 @@ test.before(async t => { }); test.beforeEach(async t => { - await initTestingDB(t.context.module.get(PrismaClient)); + await t.context.module.initTestingDB(); }); test.after(async t => { diff --git a/packages/backend/server/src/__tests__/models/verification-token.spec.ts b/packages/backend/server/src/__tests__/models/verification-token.spec.ts index 30814dc06d..7e65244d1f 100644 --- a/packages/backend/server/src/__tests__/models/verification-token.spec.ts +++ b/packages/backend/server/src/__tests__/models/verification-token.spec.ts @@ -1,4 +1,3 @@ -import { TestingModule } from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import ava, { TestFn } from 'ava'; @@ -6,7 +5,7 @@ import { TokenType, VerificationTokenModel, } from '../../models/verification-token'; -import { createTestingModule, initTestingDB } from '../utils'; +import { createTestingModule, type TestingModule } from '../utils'; interface Context { module: TestingModule; @@ -25,7 +24,7 @@ test.before(async t => { }); test.beforeEach(async t => { - await initTestingDB(t.context.db); + await t.context.module.initTestingDB(); }); test.after(async t => { diff --git a/packages/backend/server/src/__tests__/models/workspace.spec.ts b/packages/backend/server/src/__tests__/models/workspace.spec.ts index 32ae94637f..435a277555 100644 --- a/packages/backend/server/src/__tests__/models/workspace.spec.ts +++ b/packages/backend/server/src/__tests__/models/workspace.spec.ts @@ -1,4 +1,3 @@ -import { TestingModule } from '@nestjs/testing'; import { PrismaClient, WorkspaceMemberStatus } from '@prisma/client'; import ava, { TestFn } from 'ava'; import Sinon from 'sinon'; @@ -7,7 +6,7 @@ import { Config, EventBus } from '../../base'; import { WorkspaceRole } from '../../core/permission'; import { UserModel } from '../../models/user'; import { WorkspaceModel } from '../../models/workspace'; -import { createTestingModule, initTestingDB } from '../utils'; +import { createTestingModule, type TestingModule } from '../utils'; interface Context { config: Config; @@ -29,7 +28,7 @@ test.before(async t => { }); test.beforeEach(async t => { - await initTestingDB(t.context.db); + await t.context.module.initTestingDB(); }); test.after(async t => { diff --git a/packages/backend/server/src/__tests__/nestjs/throttler.spec.ts b/packages/backend/server/src/__tests__/nestjs/throttler.spec.ts index bf1ed49644..c139200067 100644 --- a/packages/backend/server/src/__tests__/nestjs/throttler.spec.ts +++ b/packages/backend/server/src/__tests__/nestjs/throttler.spec.ts @@ -1,13 +1,6 @@ import '../../plugins/config'; -import { - Controller, - Get, - HttpStatus, - INestApplication, - UseGuards, -} from '@nestjs/common'; -import { PrismaClient } from '@prisma/client'; +import { Controller, Get, HttpStatus, UseGuards } from '@nestjs/common'; import ava, { TestFn } from 'ava'; import Sinon from 'sinon'; import request, { type Response } from 'supertest'; @@ -21,12 +14,12 @@ import { ThrottlerStorage, } from '../../base/throttler'; import { AuthService, Public } from '../../core/auth'; -import { createTestingApp, initTestingDB, internalSignIn } from '../utils'; +import { createTestingApp, internalSignIn, TestingApp } from '../utils'; const test = ava as TestFn<{ storage: ThrottlerStorage; cookie: string; - app: INestApplication; + app: TestingApp; }>; @UseGuards(CloudThrottlerGuard) @@ -115,7 +108,7 @@ test.before(async t => { }); test.beforeEach(async t => { - await initTestingDB(t.context.app.get(PrismaClient)); + await t.context.app.initTestingDB(); const { app } = t.context; const auth = app.get(AuthService); const u1 = await auth.signUp('u1@affine.pro', 'test'); diff --git a/packages/backend/server/src/__tests__/oauth/controller.spec.ts b/packages/backend/server/src/__tests__/oauth/controller.spec.ts index b7c363d039..5fa7668bbb 100644 --- a/packages/backend/server/src/__tests__/oauth/controller.spec.ts +++ b/packages/backend/server/src/__tests__/oauth/controller.spec.ts @@ -1,6 +1,6 @@ import '../../plugins/config'; -import { HttpStatus, INestApplication } from '@nestjs/common'; +import { HttpStatus } from '@nestjs/common'; import { PrismaClient } from '@prisma/client'; import ava, { TestFn } from 'ava'; import Sinon from 'sinon'; @@ -15,7 +15,7 @@ import { Models } from '../../models'; import { OAuthProviderName } from '../../plugins/oauth/config'; import { GoogleOAuthProvider } from '../../plugins/oauth/providers/google'; import { OAuthService } from '../../plugins/oauth/service'; -import { createTestingApp, getSession, initTestingDB } from '../utils'; +import { createTestingApp, getSession, TestingApp } from '../utils'; const test = ava as TestFn<{ auth: AuthService; @@ -23,7 +23,7 @@ const test = ava as TestFn<{ models: Models; u1: CurrentUser; db: PrismaClient; - app: INestApplication; + app: TestingApp; }>; test.before(async t => { @@ -54,7 +54,7 @@ test.before(async t => { test.beforeEach(async t => { Sinon.restore(); - await initTestingDB(t.context.db); + await t.context.app.initTestingDB(); t.context.u1 = await t.context.auth.signUp('u1@affine.pro', '1'); }); @@ -247,7 +247,7 @@ test('should throw if provider is invalid in callback uri', async t => { t.pass(); }); -function mockOAuthProvider(app: INestApplication, email: string) { +function mockOAuthProvider(app: TestingApp, email: string) { const provider = app.get(GoogleOAuthProvider); const oauth = app.get(OAuthService); diff --git a/packages/backend/server/src/__tests__/payment/service.spec.ts b/packages/backend/server/src/__tests__/payment/service.spec.ts index 1e4ae8417e..202060a891 100644 --- a/packages/backend/server/src/__tests__/payment/service.spec.ts +++ b/packages/backend/server/src/__tests__/payment/service.spec.ts @@ -1,6 +1,5 @@ import '../../plugins/payment'; -import { INestApplication } from '@nestjs/common'; import { PrismaClient } from '@prisma/client'; import ava, { TestFn } from 'ava'; import Sinon from 'sinon'; @@ -11,7 +10,7 @@ import { EventBus, Runtime } from '../../base'; import { ConfigModule } from '../../base/config'; import { CurrentUser } from '../../core/auth'; import { AuthService } from '../../core/auth/service'; -import { EarlyAccessType, FeatureManagementService } from '../../core/features'; +import { EarlyAccessType, FeatureService } from '../../core/features'; import { SubscriptionService } from '../../plugins/payment/service'; import { CouponType, @@ -21,7 +20,7 @@ import { SubscriptionStatus, SubscriptionVariant, } from '../../plugins/payment/types'; -import { createTestingApp, initTestingDB } from '../utils'; +import { createTestingApp, type TestingApp } from '../utils'; const PRO_MONTHLY = `${SubscriptionPlan.Pro}_${SubscriptionRecurring.Monthly}`; const PRO_YEARLY = `${SubscriptionPlan.Pro}_${SubscriptionRecurring.Yearly}`; @@ -156,10 +155,10 @@ const sub: Stripe.Subscription = { const test = ava as TestFn<{ u1: CurrentUser; db: PrismaClient; - app: INestApplication; + app: TestingApp; service: SubscriptionService; event: Sinon.SinonStubbedInstance; - feature: Sinon.SinonStubbedInstance; + feature: Sinon.SinonStubbedInstance; runtime: Sinon.SinonStubbedInstance; stripe: { customers: Sinon.SinonStubbedInstance; @@ -200,8 +199,8 @@ test.before(async t => { AppModule, ], tapModule: m => { - m.overrideProvider(FeatureManagementService).useValue( - Sinon.createStubInstance(FeatureManagementService) + m.overrideProvider(FeatureService).useValue( + Sinon.createStubInstance(FeatureService) ); m.overrideProvider(EventBus).useValue(Sinon.createStubInstance(EventBus)); m.overrideProvider(Runtime).useValue(Sinon.createStubInstance(Runtime)); @@ -210,7 +209,7 @@ test.before(async t => { t.context.event = app.get(EventBus); t.context.service = app.get(SubscriptionService); - t.context.feature = app.get(FeatureManagementService); + t.context.feature = app.get(FeatureService); t.context.runtime = app.get(Runtime); t.context.db = app.get(PrismaClient); t.context.app = app; @@ -232,7 +231,7 @@ test.before(async t => { test.beforeEach(async t => { const { db, app, stripe } = t.context; - await initTestingDB(db); + await t.context.app.initTestingDB(); t.context.u1 = await app.get(AuthService).signUp('u1@affine.pro', '1'); await db.workspace.create({ diff --git a/packages/backend/server/src/__tests__/quota.spec.ts b/packages/backend/server/src/__tests__/quota.spec.ts deleted file mode 100644 index 6686f3e594..0000000000 --- a/packages/backend/server/src/__tests__/quota.spec.ts +++ /dev/null @@ -1,212 +0,0 @@ -/// - -import { TestingModule } from '@nestjs/testing'; -import type { TestFn } from 'ava'; -import ava from 'ava'; - -import { AuthService } from '../core/auth'; -import { - QuotaManagementService, - QuotaModule, - QuotaService, - QuotaType, -} from '../core/quota'; -import { OneGB, OneMB } from '../core/quota/constant'; -import { FreePlan, ProPlan } from '../core/quota/schema'; -import { StorageModule, WorkspaceBlobStorage } from '../core/storage'; -import { WorkspaceResolver } from '../core/workspaces/resolvers'; -import { createTestingModule } from './utils'; -import { WorkspaceResolverMock } from './utils/feature'; - -const test = ava as TestFn<{ - auth: AuthService; - quota: QuotaService; - quotaManager: QuotaManagementService; - workspace: WorkspaceResolver; - workspaceBlob: WorkspaceBlobStorage; - module: TestingModule; -}>; - -test.beforeEach(async t => { - const module = await createTestingModule({ - imports: [StorageModule, QuotaModule], - providers: [WorkspaceResolver], - tapModule: module => { - module - .overrideProvider(WorkspaceResolver) - .useClass(WorkspaceResolverMock); - }, - }); - - t.context.module = module; - 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 => { - await t.context.module.close(); -}); - -test('should be able to set quota', async t => { - const { auth, quota } = t.context; - - const u1 = await auth.signUp('test@affine.pro', '123456'); - - const q1 = await quota.getUserQuota(u1.id); - t.truthy(q1, 'should have quota'); - t.is(q1?.feature.name, QuotaType.FreePlanV1, 'should be free plan'); - t.is(q1?.feature.version, 4, 'should be version 4'); - - await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1); - - const q2 = await quota.getUserQuota(u1.id); - t.is(q2?.feature.name, QuotaType.ProPlanV1, 'should be pro plan'); - - const fail = quota.switchUserQuota(u1.id, 'not_exists_plan_v1' as QuotaType); - await t.throwsAsync(fail, { instanceOf: Error }, 'should throw error'); -}); - -test('should be able to check storage quota', async t => { - const { auth, quota, quotaManager } = t.context; - const u1 = await auth.signUp('test@affine.pro', '123456'); - const freePlan = FreePlan.configs; - const proPlan = ProPlan.configs; - - const q1 = await quotaManager.getUserQuota(u1.id); - t.is(q1?.blobLimit, freePlan.blobLimit, 'should be free plan'); - t.is(q1?.storageQuota, freePlan.storageQuota, 'should be free plan'); - - await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1); - const q2 = await quotaManager.getUserQuota(u1.id); - t.is(q2?.blobLimit, proPlan.blobLimit, 'should be pro plan'); - t.is(q2?.storageQuota, proPlan.storageQuota, 'should be pro plan'); -}); - -test('should be able revert quota', async t => { - const { auth, quota, quotaManager } = t.context; - const u1 = await auth.signUp('test@affine.pro', '123456'); - const freePlan = FreePlan.configs; - const proPlan = ProPlan.configs; - - const q1 = await quotaManager.getUserQuota(u1.id); - - t.is(q1?.blobLimit, freePlan.blobLimit, 'should be free plan'); - t.is(q1?.storageQuota, freePlan.storageQuota, 'should be free plan'); - - await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1); - const q2 = await quotaManager.getUserQuota(u1.id); - t.is(q2?.blobLimit, proPlan.blobLimit, 'should be pro plan'); - t.is(q2?.storageQuota, proPlan.storageQuota, 'should be pro plan'); - t.is( - q2?.copilotActionLimit, - proPlan.copilotActionLimit!, - 'should be pro plan' - ); - - await quota.switchUserQuota(u1.id, QuotaType.FreePlanV1); - const q3 = await quotaManager.getUserQuota(u1.id); - t.is(q3?.blobLimit, freePlan.blobLimit, 'should be free plan'); - - const quotas = await quota.getUserQuotas(u1.id); - t.is(quotas.length, 3, 'should have 3 quotas'); - t.is(quotas[0].feature.name, QuotaType.FreePlanV1, 'should be free plan'); - t.is(quotas[1].feature.name, QuotaType.ProPlanV1, 'should be pro plan'); - t.is(quotas[2].feature.name, QuotaType.FreePlanV1, 'should be free plan'); - t.is(quotas[0].activated, false, 'should be activated'); - t.is(quotas[1].activated, false, 'should be activated'); - t.is(quotas[2].activated, true, 'should be activated'); -}); - -test('should be able to check quota', async t => { - const { auth, quotaManager } = t.context; - const u1 = await auth.signUp('test@affine.pro', '123456'); - const freePlan = FreePlan.configs; - - const q1 = await quotaManager.getUserQuota(u1.id); - t.assert(q1, 'should have quota'); - t.is(q1.blobLimit, freePlan.blobLimit, 'should be free plan'); - t.is(q1.storageQuota, freePlan.storageQuota, 'should be free plan'); - t.is(q1.historyPeriod, freePlan.historyPeriod, 'should be free plan'); - t.is(q1.memberLimit, freePlan.memberLimit, 'should be free plan'); - t.is( - q1.copilotActionLimit!, - freePlan.copilotActionLimit!, - 'should be free plan' - ); -}); - -test('should be able to override quota', async t => { - 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, 10 * OneMB, 'should be 10MB'); - t.is(wq1.businessBlobLimit, 100 * OneMB, 'should be 100MB'); - t.is(wq1.memberLimit, 3, 'should be 3'); - - await quotaManager.addTeamWorkspace(w1.id, 'test'); - const wq2 = await quotaManager.getWorkspaceUsage(w1.id); - 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'); -}); - -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); - const w3 = await workspace.createWorkspace(u1, null); - await quotaManager.addTeamWorkspace(w3.id, 'test'); - - { - const wq1 = await quotaManager.getWorkspaceUsage(w1.id); - t.is(wq1.usedSize, 0, 'should be 0'); - const wq2 = await quotaManager.getWorkspaceUsage(w2.id); - t.is(wq2.usedSize, 0, 'should be 0'); - const wq3 = await quotaManager.getWorkspaceUsage(w3.id); - t.is(wq3.usedSize, 0, 'should be 0'); - } - - { - await workspaceBlob.put(w1.id, 'test', Buffer.from([0, 0])); - await workspaceBlob.put(w2.id, 'test', Buffer.from([0, 0])); - - // normal workspace - const wq1 = await quotaManager.getWorkspaceUsage(w1.id); - t.is(wq1.usedSize, 4, 'should share usage with w2'); - const wq2 = await quotaManager.getWorkspaceUsage(w2.id); - t.is(wq2.usedSize, 4, 'should share usage with w1'); - - // workspace with quota - const wq3 = await quotaManager.getWorkspaceUsage(w3.id); - t.is(wq3.usedSize, 0, 'should not share usage with w1 and w2'); - } - - { - await workspaceBlob.put(w3.id, 'test', Buffer.from([0, 0, 0])); - - // normal workspace - const wq1 = await quotaManager.getWorkspaceUsage(w1.id); - t.is(wq1.usedSize, 4, 'should not share usage with w3'); - const wq2 = await quotaManager.getWorkspaceUsage(w2.id); - t.is(wq2.usedSize, 4, 'should not share usage with w3'); - - // workspace with quota - const wq3 = await quotaManager.getWorkspaceUsage(w3.id); - t.is(wq3.usedSize, 3, 'should not share usage with w1 and w2'); - } -}); diff --git a/packages/backend/server/src/__tests__/team.e2e.ts b/packages/backend/server/src/__tests__/team.e2e.ts index e48f3d50e1..3d9e5998bc 100644 --- a/packages/backend/server/src/__tests__/team.e2e.ts +++ b/packages/backend/server/src/__tests__/team.e2e.ts @@ -3,7 +3,6 @@ import { randomUUID } from 'node:crypto'; import { getCurrentMailMessageCount } from '@affine-test/kit/utils/cloud'; -import { INestApplication } from '@nestjs/common'; import { WorkspaceMemberStatus } from '@prisma/client'; import type { TestFn } from 'ava'; import ava from 'ava'; @@ -14,8 +13,8 @@ import { EventBus } from '../base'; import { AuthService } from '../core/auth'; import { DocContentService } from '../core/doc-renderer'; import { PermissionService, WorkspaceRole } from '../core/permission'; -import { QuotaManagementService, QuotaService, QuotaType } from '../core/quota'; import { WorkspaceType } from '../core/workspaces'; +import { Models } from '../models'; import { acceptInviteById, approveMember, @@ -34,19 +33,19 @@ import { revokeUser, signUp, sleep, + TestingApp, UserAuthedType, } from './utils'; const test = ava as TestFn<{ - app: INestApplication; + app: TestingApp; auth: AuthService; event: Sinon.SinonStubbedInstance; - quota: QuotaService; - quotaManager: QuotaManagementService; + models: Models; permissions: PermissionService; }>; -test.beforeEach(async t => { +test.before(async t => { const { app } = await createTestingApp({ imports: [AppModule], tapModule: module => { @@ -67,17 +66,20 @@ test.beforeEach(async t => { t.context.app = app; t.context.auth = app.get(AuthService); t.context.event = app.get(EventBus); - t.context.quota = app.get(QuotaService); - t.context.quotaManager = app.get(QuotaManagementService); + t.context.models = app.get(Models); t.context.permissions = app.get(PermissionService); }); -test.afterEach.always(async t => { +test.beforeEach(async t => { + await t.context.app.initTestingDB(); +}); + +test.after.always(async t => { await t.context.app.close(); }); const init = async ( - app: INestApplication, + app: TestingApp, memberLimit = 10, prefix = randomUUID() ) => { @@ -87,17 +89,15 @@ const init = async ( `${prefix}owner@affine.pro`, '123456' ); + const models = app.get(Models); { - const quota = app.get(QuotaService); - await quota.switchUserQuota(owner.id, QuotaType.ProPlanV1); + await models.userFeature.add(owner.id, 'pro_plan_v1', 'test'); } const workspace = await createWorkspace(app, owner.token.token); const teamWorkspace = await createWorkspace(app, owner.token.token); { - const quota = app.get(QuotaManagementService); - await quota.addTeamWorkspace(teamWorkspace.id, 'test'); - await quota.updateWorkspaceConfig(teamWorkspace.id, QuotaType.TeamPlanV1, { + models.workspaceFeature.add(teamWorkspace.id, 'team_plan_v1', 'test', { memberLimit, }); } @@ -264,7 +264,7 @@ test('should be able to invite multiple users', async t => { }); test('should be able to check seat limit', async t => { - const { app, permissions, quotaManager } = t.context; + const { app, permissions, models } = t.context; const { invite, inviteBatch, teamWorkspace: ws } = await init(app, 4); { @@ -274,7 +274,7 @@ test('should be able to check seat limit', async t => { { message: 'You have exceeded your workspace member quota.' }, 'should throw error if exceed member limit' ); - await quotaManager.updateWorkspaceConfig(ws.id, QuotaType.TeamPlanV1, { + models.workspaceFeature.add(ws.id, 'team_plan_v1', 'test', { memberLimit: 5, }); await t.notThrowsAsync( @@ -323,7 +323,7 @@ test('should be able to check seat limit', async t => { test('should be able to grant team member permission', async t => { const { app, permissions } = t.context; - const { owner, teamWorkspace: ws, admin, write, read } = await init(app); + const { owner, teamWorkspace: ws, write, read } = await init(app); await t.throwsAsync( grantMember( @@ -350,13 +350,13 @@ test('should be able to grant team member permission', async t => { await t.throwsAsync( grantMember( app, - admin.token.token, + write.token.token, ws.id, read.id, WorkspaceRole.Collaborator ), { instanceOf: Error }, - 'should throw error if not owner' + 'should throw error if not admin' ); { @@ -571,7 +571,7 @@ test('should be able to approve team member', async t => { }); test('should be able to invite by link', async t => { - const { app, permissions, quotaManager } = t.context; + const { app, permissions, models } = t.context; const { createInviteLink, owner, @@ -631,7 +631,7 @@ test('should be able to invite by link', async t => { 'should not change status' ); - await quotaManager.updateWorkspaceConfig(tws.id, QuotaType.TeamPlanV1, { + models.workspaceFeature.add(tws.id, 'team_plan_v1', 'test', { memberLimit: 5, }); await permissions.refreshSeatStatus(tws.id, 5); @@ -646,7 +646,7 @@ test('should be able to invite by link', async t => { 'should not change status' ); - await quotaManager.updateWorkspaceConfig(tws.id, QuotaType.TeamPlanV1, { + models.workspaceFeature.add(tws.id, 'team_plan_v1', 'test', { memberLimit: 6, }); await permissions.refreshSeatStatus(tws.id, 6); diff --git a/packages/backend/server/src/__tests__/user.e2e.ts b/packages/backend/server/src/__tests__/user.e2e.ts index ef8c4fc4b5..3b41c0ef17 100644 --- a/packages/backend/server/src/__tests__/user.e2e.ts +++ b/packages/backend/server/src/__tests__/user.e2e.ts @@ -1,20 +1,23 @@ -import type { INestApplication } from '@nestjs/common'; import test from 'ava'; import request from 'supertest'; import { AppModule } from '../app.module'; -import { createTestingApp, currentUser, signUp } from './utils'; +import { createTestingApp, currentUser, signUp, TestingApp } from './utils'; -let app: INestApplication; +let app: TestingApp; -test.beforeEach(async () => { +test.before(async () => { const { app: testApp } = await createTestingApp({ imports: [AppModule], }); app = testApp; }); -test.afterEach.always(async () => { +test.beforeEach(async () => { + await app.initTestingDB(); +}); + +test.after.always(async () => { await app.close(); }); diff --git a/packages/backend/server/src/__tests__/utils/utils.ts b/packages/backend/server/src/__tests__/utils/utils.ts index cb507b4c2e..43774f8e9e 100644 --- a/packages/backend/server/src/__tests__/utils/utils.ts +++ b/packages/backend/server/src/__tests__/utils/utils.ts @@ -3,9 +3,13 @@ import { INestApplication, ModuleMetadata, } from '@nestjs/common'; -import { APP_GUARD } from '@nestjs/core'; +import { APP_GUARD, ModuleRef } from '@nestjs/core'; import { Query, Resolver } from '@nestjs/graphql'; -import { Test, TestingModuleBuilder } from '@nestjs/testing'; +import { + Test, + TestingModule as BaseTestingModule, + TestingModuleBuilder, +} from '@nestjs/testing'; import { PrismaClient } from '@prisma/client'; import cookieParser from 'cookie-parser'; import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs'; @@ -16,7 +20,7 @@ import { AppModule, FunctionalityModules } from '../../app.module'; import { GlobalExceptionFilter, Runtime } from '../../base'; import { GqlModule } from '../../base/graphql'; import { AuthGuard, AuthModule } from '../../core/auth'; -import { UserFeaturesInit1698652531198 } from '../../data/migrations/1698652531198-user-features-init'; +import { RefreshFeatures0001 } from '../../data/migrations/0001-refresh-features'; import { ModelsModule } from '../../models'; async function flushDB(client: PrismaClient) { @@ -35,20 +39,25 @@ async function flushDB(client: PrismaClient) { ); } -async function initFeatureConfigs(db: PrismaClient) { - await UserFeaturesInit1698652531198.up(db); -} - -export async function initTestingDB(db: PrismaClient) { - await flushDB(db); - await initFeatureConfigs(db); -} - interface TestingModuleMeatdata extends ModuleMetadata { tapModule?(m: TestingModuleBuilder): void; tapApp?(app: INestApplication): void; } +const initTestingDB = async (ref: ModuleRef) => { + const db = ref.get(PrismaClient, { strict: false }); + await flushDB(db); + await RefreshFeatures0001.up(db, ref); +}; + +export type TestingModule = BaseTestingModule & { + initTestingDB(): Promise; +}; + +export type TestingApp = INestApplication & { + initTestingDB(): Promise; +}; + function dedupeModules(modules: NonNullable) { const map = new Map(); @@ -73,7 +82,7 @@ class MockResolver { export async function createTestingModule( moduleDef: TestingModuleMeatdata = {}, - init = true + autoInitialize = true ) { // setting up let imports = moduleDef.imports ?? []; @@ -107,13 +116,9 @@ export async function createTestingModule( const m = await builder.compile(); - const prisma = m.get(PrismaClient); - if (prisma instanceof PrismaClient) { - await initTestingDB(prisma); - } - - if (init) { - await m.init(); + const testingModule = m as TestingModule; + testingModule.initTestingDB = async () => { + await initTestingDB(m.get(ModuleRef)); // we got a lot smoking tests try to break nestjs // can't tolerate the noisy logs // @ts-expect-error private @@ -123,9 +128,14 @@ export async function createTestingModule( const runtime = m.get(Runtime); // by pass password min length validation await runtime.set('auth/password.min', 1); + }; + + if (autoInitialize) { + await testingModule.initTestingDB(); + await testingModule.init(); } - return m; + return testingModule; } export async function createTestingApp(moduleDef: TestingModuleMeatdata = {}) { @@ -135,7 +145,7 @@ export async function createTestingApp(moduleDef: TestingModuleMeatdata = {}) { cors: true, bodyParser: true, rawBody: true, - }); + }) as TestingApp; const logger = new ConsoleLogger(); logger.setLogLevels(['fatal']); @@ -155,15 +165,14 @@ export async function createTestingApp(moduleDef: TestingModuleMeatdata = {}) { moduleDef.tapApp(app); } + await m.initTestingDB(); await app.init(); - const runtime = app.get(Runtime); - // by pass password min length validation - await runtime.set('auth/password.min', 1); + app.initTestingDB = m.initTestingDB.bind(m); return { module: m, - app, + app: app, }; } diff --git a/packages/backend/server/src/__tests__/workspace-invite.e2e.ts b/packages/backend/server/src/__tests__/workspace-invite.e2e.ts index 78f7920618..b7765c6f57 100644 --- a/packages/backend/server/src/__tests__/workspace-invite.e2e.ts +++ b/packages/backend/server/src/__tests__/workspace-invite.e2e.ts @@ -2,7 +2,6 @@ import { getCurrentMailMessageCount, getLatestMailMessage, } from '@affine-test/kit/utils/cloud'; -import type { INestApplication } from '@nestjs/common'; import { PrismaClient } from '@prisma/client'; import type { TestFn } from 'ava'; import ava from 'ava'; @@ -20,17 +19,18 @@ import { leaveWorkspace, revokeUser, signUp, + TestingApp, } from './utils'; const test = ava as TestFn<{ - app: INestApplication; + app: TestingApp; client: PrismaClient; auth: AuthService; mail: MailService; models: Models; }>; -test.beforeEach(async t => { +test.before(async t => { const { app } = await createTestingApp({ imports: [AppModule], }); @@ -41,7 +41,11 @@ test.beforeEach(async t => { t.context.models = app.get(Models); }); -test.afterEach.always(async t => { +test.beforeEach(async t => { + await t.context.app.initTestingDB(); +}); + +test.after.always(async t => { await t.context.app.close(); }); @@ -227,15 +231,12 @@ test('should support pagination for member', async t => { test('should limit member count correctly', async t => { const { app } = t.context; const u1 = await signUp(app, 'u1', 'u1@affine.pro', '1'); - for (let i = 0; i < 10; i++) { - const workspace = await createWorkspace(app, u1.token.token); - await Promise.allSettled( - Array.from({ length: 10 }).map(async (_, i) => - inviteUser(app, u1.token.token, workspace.id, `u${i}@affine.pro`) - ) - ); - - const ws = await getWorkspace(app, u1.token.token, workspace.id); - t.assert(ws.members.length <= 3, 'failed to check member list'); - } + const workspace = await createWorkspace(app, u1.token.token); + await Promise.allSettled( + Array.from({ length: 10 }).map(async (_, i) => + inviteUser(app, u1.token.token, workspace.id, `u${i}@affine.pro`) + ) + ); + const ws = await getWorkspace(app, u1.token.token, workspace.id); + t.assert(ws.members.length <= 3, 'failed to check member list'); }); diff --git a/packages/backend/server/src/__tests__/workspace.e2e.ts b/packages/backend/server/src/__tests__/workspace.e2e.ts index 2d269f6b95..c32ed40b57 100644 --- a/packages/backend/server/src/__tests__/workspace.e2e.ts +++ b/packages/backend/server/src/__tests__/workspace.e2e.ts @@ -1,30 +1,28 @@ -import type { INestApplication } from '@nestjs/common'; import { PrismaClient } from '@prisma/client'; import type { TestFn } from 'ava'; import ava from 'ava'; import request from 'supertest'; import { AppModule } from '../app.module'; -import { WorkspaceRole } from '../core/permission/types'; import { acceptInviteById, createTestingApp, createWorkspace, getWorkspacePublicPages, - grantMember, inviteUser, publishPage, revokePublicPage, signUp, + TestingApp, updateWorkspace, } from './utils'; const test = ava as TestFn<{ - app: INestApplication; + app: TestingApp; client: PrismaClient; }>; -test.beforeEach(async t => { +test.before(async t => { const { app } = await createTestingApp({ imports: [AppModule], }); @@ -33,7 +31,11 @@ test.beforeEach(async t => { t.context.app = app; }); -test.afterEach.always(async t => { +test.beforeEach(async t => { + await t.context.app.initTestingDB(); +}); + +test.after.always(async t => { await t.context.app.close(); }); @@ -132,30 +134,6 @@ test('should share a page', async t => { 'You do not have permission to access doc page2 under Space not_exists_ws.', 'unauthorized user can share page' ); - - await acceptInviteById( - app, - workspace.id, - await inviteUser(app, u1.token.token, workspace.id, u2.email) - ); - const msg3 = await publishPage(app, u2.token.token, workspace.id, 'page2'); - t.is( - msg3, - `You do not have permission to access doc page2 under Space ${workspace.id}.`, - 'WorkspaceRole and PageRole is lower than required' - ); - - await grantMember( - app, - u1.token.token, - workspace.id, - u2.id, - WorkspaceRole.Admin - ); - - const invited = await publishPage(app, u2.token.token, workspace.id, 'page2'); - t.is(invited.id, 'page2', 'failed to share page'); - const revoke = await revokePublicPage( app, u1.token.token, @@ -168,9 +146,7 @@ test('should share a page', async t => { u1.token.token, workspace.id ); - t.is(pages2.length, 1, 'failed to get shared pages'); - t.is(pages2[0].id, 'page2', 'failed to get shared page: page2'); - + t.is(pages2.length, 0, 'failed to get shared pages'); const msg4 = await revokePublicPage( app, u1.token.token, @@ -179,19 +155,12 @@ test('should share a page', async t => { ); t.is(msg4, 'Page is not public'); - const revoked = await revokePublicPage( - app, - u1.token.token, - workspace.id, - 'page2' - ); - t.false(revoked.public, 'failed to revoke page'); - const page3 = await getWorkspacePublicPages( + const pages3 = await getWorkspacePublicPages( app, u1.token.token, workspace.id ); - t.is(page3.length, 0, 'failed to get shared pages'); + t.is(pages3.length, 0, 'failed to get shared pages'); }); test('should be able to get workspace doc', async t => { diff --git a/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts b/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts index 873c8c46ca..62beaabca8 100644 --- a/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts +++ b/packages/backend/server/src/__tests__/workspace/blobs.e2e.ts @@ -1,10 +1,8 @@ -import type { INestApplication } from '@nestjs/common'; import test from 'ava'; import request from 'supertest'; import { AppModule } from '../../app.module'; -import { FeatureManagementService, FeatureType } from '../../core/features'; -import { QuotaService, QuotaType } from '../../core/quota'; +import { WorkspaceFeatureModel } from '../../models'; import { collectAllBlobSizes, createTestingApp, @@ -13,25 +11,35 @@ import { listBlobs, setBlob, signUp, + TestingApp, } from '../utils'; const OneMB = 1024 * 1024; +const RESTRICTED_QUOTA = { + seatQuota: 0, + blobLimit: OneMB, + storageQuota: 2 * OneMB - 1, + historyPeriod: 1, + memberLimit: 1, +}; -let app: INestApplication; -let quota: QuotaService; -let feature: FeatureManagementService; +let app: TestingApp; +let model: WorkspaceFeatureModel; -test.beforeEach(async () => { +test.before(async () => { const { app: testApp } = await createTestingApp({ imports: [AppModule], }); app = testApp; - quota = app.get(QuotaService); - feature = app.get(FeatureManagementService); + model = app.get(WorkspaceFeatureModel); }); -test.afterEach.always(async () => { +test.beforeEach(async () => { + await app.initTestingDB(); +}); + +test.after.always(async () => { await app.close(); }); @@ -119,32 +127,23 @@ test('should reject blob exceeded limit', async t => { const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1'); const workspace1 = await createWorkspace(app, u1.token.token); - await quota.switchUserQuota(u1.id, QuotaType.RestrictedPlanV1); + await model.add(workspace1.id, 'team_plan_v1', 'test', RESTRICTED_QUOTA); - const buffer1 = Buffer.from(Array.from({ length: OneMB + 1 }, () => 0)); + const buffer1 = Buffer.from( + Array.from({ length: RESTRICTED_QUOTA.blobLimit + 1 }, () => 0) + ); await t.throwsAsync(setBlob(app, u1.token.token, workspace1.id, buffer1)); - - await quota.switchUserQuota(u1.id, QuotaType.FreePlanV1); - - const buffer2 = Buffer.from(Array.from({ length: OneMB + 1 }, () => 0)); - await t.notThrowsAsync(setBlob(app, u1.token.token, workspace1.id, buffer2)); - - const buffer3 = Buffer.from(Array.from({ length: 100 * OneMB + 1 }, () => 0)); - await t.throwsAsync(setBlob(app, u1.token.token, workspace1.id, buffer3)); }); test('should reject blob exceeded quota', async t => { const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1'); const workspace = await createWorkspace(app, u1.token.token); - await quota.switchUserQuota(u1.id, QuotaType.RestrictedPlanV1); + await model.add(workspace.id, 'team_plan_v1', 'test', RESTRICTED_QUOTA); const buffer = Buffer.from(Array.from({ length: OneMB }, () => 0)); - for (let i = 0; i < 10; i++) { - await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer)); - } - + await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer)); await t.throwsAsync(setBlob(app, u1.token.token, workspace.id, buffer)); }); @@ -152,14 +151,10 @@ test('should accept blob even storage out of quota if workspace has unlimited fe const u1 = await signUp(app, 'darksky', 'darksky@affine.pro', '1'); const workspace = await createWorkspace(app, u1.token.token); - await quota.switchUserQuota(u1.id, QuotaType.RestrictedPlanV1); - feature.addWorkspaceFeatures(workspace.id, FeatureType.UnlimitedWorkspace); + await model.add(workspace.id, 'team_plan_v1', 'test', RESTRICTED_QUOTA); + await model.add(workspace.id, 'unlimited_workspace', 'test'); const buffer = Buffer.from(Array.from({ length: OneMB }, () => 0)); - - for (let i = 0; i < 10; i++) { - await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer)); - } - + await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer)); await t.notThrowsAsync(setBlob(app, u1.token.token, workspace.id, buffer)); }); diff --git a/packages/backend/server/src/base/error/def.ts b/packages/backend/server/src/base/error/def.ts index 2c1b85e3ea..578dcbb296 100644 --- a/packages/backend/server/src/base/error/def.ts +++ b/packages/backend/server/src/base/error/def.ts @@ -478,6 +478,10 @@ export const USER_FRIENDLY_ERRORS = { type: 'internal_server_error', message: 'Failed to store doc snapshot.', }, + action_forbidden_on_non_team_workspace: { + type: 'action_forbidden', + message: 'A Team workspace is required to perform this action.', + }, // Subscription Errors unsupported_subscription_plan: { diff --git a/packages/backend/server/src/base/error/errors.gen.ts b/packages/backend/server/src/base/error/errors.gen.ts index e578cd9fc8..3d12ed1638 100644 --- a/packages/backend/server/src/base/error/errors.gen.ts +++ b/packages/backend/server/src/base/error/errors.gen.ts @@ -399,6 +399,12 @@ export class FailedToUpsertSnapshot extends UserFriendlyError { super('internal_server_error', 'failed_to_upsert_snapshot', message); } } + +export class ActionForbiddenOnNonTeamWorkspace extends UserFriendlyError { + constructor(message?: string) { + super('action_forbidden', 'action_forbidden_on_non_team_workspace', message); + } +} @ObjectType() class UnsupportedSubscriptionPlanDataType { @Field() plan!: string @@ -754,6 +760,7 @@ export enum ErrorNames { PAGE_IS_NOT_PUBLIC, FAILED_TO_SAVE_UPDATES, FAILED_TO_UPSERT_SNAPSHOT, + ACTION_FORBIDDEN_ON_NON_TEAM_WORKSPACE, UNSUPPORTED_SUBSCRIPTION_PLAN, FAILED_TO_CHECKOUT, INVALID_CHECKOUT_PARAMETERS, diff --git a/packages/backend/server/src/base/index.ts b/packages/backend/server/src/base/index.ts index d6ad06de8a..599a19a754 100644 --- a/packages/backend/server/src/base/index.ts +++ b/packages/backend/server/src/base/index.ts @@ -35,10 +35,4 @@ export { Runtime } from './runtime'; export * from './storage'; export { type StorageProvider, StorageProviderFactory } from './storage'; export { CloudThrottlerGuard, SkipThrottle, Throttle } from './throttler'; -export { - getRequestFromHost, - getRequestResponseFromContext, - getRequestResponseFromHost, - parseCookies, -} from './utils/request'; -export * from './utils/types'; +export * from './utils'; diff --git a/packages/backend/server/src/base/runtime/service.ts b/packages/backend/server/src/base/runtime/service.ts index 68275f2980..9f401f21c4 100644 --- a/packages/backend/server/src/base/runtime/service.ts +++ b/packages/backend/server/src/base/runtime/service.ts @@ -126,14 +126,19 @@ export class Runtime implements OnModuleInit { V = FlattenedAppRuntimeConfig[K], >(key: K, value: V) { validateConfigType(key, value); - const config = await this.db.runtimeConfig.update({ + const config = await this.db.runtimeConfig.upsert({ where: { id: key, deletedAt: null, }, - data: { + create: { + ...defaultRuntimeConfig[key], value: value as any, }, + update: { + value: value as any, + deletedAt: null, + }, }); await this.setCache(key, config.value as FlattenedAppRuntimeConfig[K]); diff --git a/packages/backend/server/src/base/utils/index.ts b/packages/backend/server/src/base/utils/index.ts new file mode 100644 index 0000000000..09cc714ed4 --- /dev/null +++ b/packages/backend/server/src/base/utils/index.ts @@ -0,0 +1,4 @@ +export * from './promise'; +export * from './request'; +export * from './types'; +export * from './unit'; diff --git a/packages/backend/server/src/core/quota/constant.ts b/packages/backend/server/src/base/utils/unit.ts similarity index 64% rename from packages/backend/server/src/core/quota/constant.ts rename to packages/backend/server/src/base/utils/unit.ts index e6fb0d25a0..f0798822b6 100644 --- a/packages/backend/server/src/core/quota/constant.ts +++ b/packages/backend/server/src/base/utils/unit.ts @@ -2,4 +2,3 @@ export const OneKB = 1024; export const OneMB = OneKB * OneKB; export const OneGB = OneKB * OneMB; export const OneDay = 1000 * 60 * 60 * 24; -export const ByteUnit = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB']; diff --git a/packages/backend/server/src/core/auth/service.ts b/packages/backend/server/src/core/auth/service.ts index 59ca9197c0..1526e32ffb 100644 --- a/packages/backend/server/src/core/auth/service.ts +++ b/packages/backend/server/src/core/auth/service.ts @@ -4,9 +4,7 @@ import { assign, pick } from 'lodash-es'; import { Config, MailService, SignUpForbidden } from '../../base'; import { Models, type User, type UserSession } from '../../models'; -import { FeatureManagementService } from '../features/management'; -import { QuotaService } from '../quota/service'; -import { QuotaType } from '../quota/types'; +import { FeatureService } from '../features'; import type { CurrentUser } from './session'; export function sessionUser( @@ -45,8 +43,7 @@ export class AuthService implements OnApplicationBootstrap { private readonly config: Config, private readonly models: Models, private readonly mailer: MailService, - private readonly feature: FeatureManagementService, - private readonly quota: QuotaService + private readonly feature: FeatureService ) {} async onApplicationBootstrap() { @@ -61,17 +58,24 @@ export class AuthService implements OnApplicationBootstrap { password, }); } - await this.quota.switchUserQuota(devUser.id, QuotaType.ProPlanV1); - await this.feature.addAdmin(devUser.id); - await this.feature.addCopilot(devUser.id); + await this.models.userFeature.add( + devUser.id, + 'administrator', + 'dev user' + ); + await this.models.userFeature.add( + devUser.id, + 'unlimited_copilot', + 'dev user' + ); } catch { // ignore } } } - canSignIn(email: string) { - return this.feature.canEarlyAccess(email); + async canSignIn(email: string) { + return await this.feature.canEarlyAccess(email); } /** diff --git a/packages/backend/server/src/core/common/admin-guard.ts b/packages/backend/server/src/core/common/admin-guard.ts index 062716f48d..f964199f10 100644 --- a/packages/backend/server/src/core/common/admin-guard.ts +++ b/packages/backend/server/src/core/common/admin-guard.ts @@ -7,16 +7,16 @@ import { Injectable, UseGuards } from '@nestjs/common'; import { ModuleRef } from '@nestjs/core'; import { ActionForbidden, getRequestResponseFromContext } from '../../base'; -import { FeatureManagementService } from '../features/management'; +import { FeatureService } from '../features/service'; @Injectable() export class AdminGuard implements CanActivate, OnModuleInit { - private feature!: FeatureManagementService; + private feature!: FeatureService; constructor(private readonly ref: ModuleRef) {} onModuleInit() { - this.feature = this.ref.get(FeatureManagementService, { strict: false }); + this.feature = this.ref.get(FeatureService, { strict: false }); } async canActivate(context: ExecutionContext) { diff --git a/packages/backend/server/src/core/config/resolver.ts b/packages/backend/server/src/core/config/resolver.ts index ef64eae24c..ed6d6c3ba1 100644 --- a/packages/backend/server/src/core/config/resolver.ts +++ b/packages/backend/server/src/core/config/resolver.ts @@ -13,10 +13,10 @@ import { RuntimeConfig, RuntimeConfigType } from '@prisma/client'; import { GraphQLJSON, GraphQLJSONObject } from 'graphql-scalars'; import { Config, Runtime, URLHelper } from '../../base'; +import { Feature } from '../../models'; import { Public } from '../auth'; import { Admin } from '../common'; -import { FeatureType } from '../features'; -import { AvailableUserFeatureConfig } from '../features/resolver'; +import { AvailableUserFeatureConfig } from '../features'; import { ServerFlags } from './config'; import { ENABLED_FEATURES } from './server-feature'; import { ServerService } from './service'; @@ -139,11 +139,7 @@ export class ServerConfigResolver { @Resolver(() => ServerConfigType) export class ServerFeatureConfigResolver extends AvailableUserFeatureConfig { - constructor(config: Config) { - super(config); - } - - @ResolveField(() => [FeatureType], { + @ResolveField(() => [Feature], { description: 'Features for user that can be configured', }) override availableUserFeatures() { diff --git a/packages/backend/server/src/core/doc/options.ts b/packages/backend/server/src/core/doc/options.ts index 9f9dd1c2b3..807887aa7b 100644 --- a/packages/backend/server/src/core/doc/options.ts +++ b/packages/backend/server/src/core/doc/options.ts @@ -89,7 +89,7 @@ export class DocStorageOptions implements IDocStorageOptions { historyMaxAge = async (spaceId: string) => { const owner = await this.permission.getWorkspaceOwner(spaceId); const quota = await this.quota.getUserQuota(owner.id); - return quota.feature.historyPeriod; + return quota.historyPeriod; }; historyMinInterval = (_spaceId: string) => { diff --git a/packages/backend/server/src/core/features/feature.ts b/packages/backend/server/src/core/features/feature.ts deleted file mode 100644 index bc5a554b56..0000000000 --- a/packages/backend/server/src/core/features/feature.ts +++ /dev/null @@ -1,51 +0,0 @@ -import { PrismaTransaction } from '../../base'; -import { Feature, FeatureSchema, FeatureType } from './types'; - -class FeatureConfig { - readonly config: Feature & { feature: T }; - - constructor(data: any) { - const config = FeatureSchema.safeParse(data); - - if (config.success) { - // @ts-expect-error allow - this.config = config.data; - } else { - throw new Error(`Invalid quota config: ${config.error.message}`); - } - } - - /// feature name of quota - get name() { - return this.config.feature; - } -} - -export type FeatureConfigType = FeatureConfig; - -const FeatureCache = new Map>(); - -export async function getFeature(prisma: PrismaTransaction, featureId: number) { - const cachedFeature = FeatureCache.get(featureId); - - if (cachedFeature) { - return cachedFeature; - } - - const feature = await prisma.feature.findFirst({ - where: { - id: featureId, - }, - }); - if (!feature) { - // this should unreachable - throw new Error(`Quota config ${featureId} not found`); - } - - const config = new FeatureConfig(feature); - // we always edit quota config as a new quota config - // so we can cache it by featureId - FeatureCache.set(featureId, config); - - return config; -} diff --git a/packages/backend/server/src/core/features/index.ts b/packages/backend/server/src/core/features/index.ts index 0d953e3235..a36aa08e00 100644 --- a/packages/backend/server/src/core/features/index.ts +++ b/packages/backend/server/src/core/features/index.ts @@ -1,38 +1,20 @@ import { Module } from '@nestjs/common'; -import { UserModule } from '../user'; -import { EarlyAccessType, FeatureManagementService } from './management'; import { AdminFeatureManagementResolver, - FeatureManagementResolver, + UserFeatureResolver, } from './resolver'; -import { FeatureService } from './service'; +import { EarlyAccessType, FeatureService } from './service'; -/** - * Feature module provider pre-user feature flag management. - * includes: - * - feature query/update/permit - * - feature statistics - */ @Module({ - imports: [UserModule], providers: [ - FeatureService, - FeatureManagementService, - FeatureManagementResolver, + UserFeatureResolver, AdminFeatureManagementResolver, + FeatureService, ], - exports: [FeatureService, FeatureManagementService], + exports: [FeatureService], }) export class FeatureModule {} -export type { FeatureConfigType } from './feature'; -export { - type CommonFeature, - commonFeatureSchema, - type FeatureConfig, - FeatureKind, - Features, - FeatureType, -} from './types'; -export { EarlyAccessType, FeatureManagementService, FeatureService }; +export { EarlyAccessType, FeatureService }; +export { AvailableUserFeatureConfig } from './types'; diff --git a/packages/backend/server/src/core/features/management.ts b/packages/backend/server/src/core/features/management.ts deleted file mode 100644 index 898c139e58..0000000000 --- a/packages/backend/server/src/core/features/management.ts +++ /dev/null @@ -1,170 +0,0 @@ -import { Injectable, Logger } from '@nestjs/common'; - -import { Runtime } from '../../base'; -import { Models } from '../../models'; -import { FeatureService } from './service'; -import { FeatureType } from './types'; - -const STAFF = ['@toeverything.info', '@affine.pro']; - -export enum EarlyAccessType { - App = 'app', - AI = 'ai', -} - -@Injectable() -export class FeatureManagementService { - protected logger = new Logger(FeatureManagementService.name); - - constructor( - private readonly feature: FeatureService, - private readonly models: Models, - private readonly runtime: Runtime - ) {} - - // ======== Admin ======== - - isStaff(email: string) { - for (const domain of STAFF) { - if (email.endsWith(domain)) { - return true; - } - } - - return false; - } - - isAdmin(userId: string) { - return this.feature.hasUserFeature(userId, FeatureType.Admin); - } - - addAdmin(userId: string) { - return this.feature.addUserFeature(userId, FeatureType.Admin, 'Admin user'); - } - - // ======== Early Access ======== - async addEarlyAccess( - userId: string, - type: EarlyAccessType = EarlyAccessType.App - ) { - return this.feature.addUserFeature( - userId, - type === EarlyAccessType.App - ? FeatureType.EarlyAccess - : FeatureType.AIEarlyAccess, - 'Early access user' - ); - } - - async removeEarlyAccess( - userId: string, - type: EarlyAccessType = EarlyAccessType.App - ) { - return this.feature.removeUserFeature( - userId, - type === EarlyAccessType.App - ? FeatureType.EarlyAccess - : FeatureType.AIEarlyAccess - ); - } - - async listEarlyAccess(type: EarlyAccessType = EarlyAccessType.App) { - return this.feature.listUsersByFeature( - type === EarlyAccessType.App - ? FeatureType.EarlyAccess - : FeatureType.AIEarlyAccess - ); - } - - async isEarlyAccessUser( - userId: string, - type: EarlyAccessType = EarlyAccessType.App - ) { - return await this.feature - .hasUserFeature( - userId, - type === EarlyAccessType.App - ? FeatureType.EarlyAccess - : FeatureType.AIEarlyAccess - ) - .catch(() => false); - } - - /// check early access by email - async canEarlyAccess( - email: string, - type: EarlyAccessType = EarlyAccessType.App - ) { - const earlyAccessControlEnabled = await this.runtime.fetch( - 'flags/earlyAccessControl' - ); - - if (earlyAccessControlEnabled && !this.isStaff(email)) { - const user = await this.models.user.getUserByEmail(email); - if (!user) { - return false; - } - return this.isEarlyAccessUser(user.id, type); - } else { - return true; - } - } - - // ======== CopilotFeature ======== - async addCopilot(userId: string, reason = 'Copilot plan user') { - return this.feature.addUserFeature( - userId, - FeatureType.UnlimitedCopilot, - reason - ); - } - - async removeCopilot(userId: string) { - return this.feature.removeUserFeature(userId, FeatureType.UnlimitedCopilot); - } - - async isCopilotUser(userId: string) { - return await this.feature.hasUserFeature( - userId, - FeatureType.UnlimitedCopilot - ); - } - - // ======== User Feature ======== - async getActivatedUserFeatures(userId: string): Promise { - const features = await this.feature.getUserActivatedFeatures(userId); - return features.map(f => f.feature.name); - } - - // ======== Workspace Feature ======== - async addWorkspaceFeatures( - workspaceId: string, - feature: FeatureType, - reason?: string - ) { - return this.feature.addWorkspaceFeature( - workspaceId, - feature, - reason || 'add feature by api' - ); - } - - async getWorkspaceFeatures(workspaceId: string) { - const features = await this.feature.getWorkspaceFeatures(workspaceId); - return features.filter(f => f.activated).map(f => f.feature.name); - } - - async hasWorkspaceFeature(workspaceId: string, feature: FeatureType) { - return this.feature.hasWorkspaceFeature(workspaceId, feature); - } - - async removeWorkspaceFeature(workspaceId: string, feature: FeatureType) { - return this.feature - .removeWorkspaceFeature(workspaceId, feature) - .then(c => c > 0); - } - - async listFeatureWorkspaces(feature: FeatureType) { - return this.feature.listWorkspacesByFeature(feature); - } -} diff --git a/packages/backend/server/src/core/features/resolver.ts b/packages/backend/server/src/core/features/resolver.ts index cc672c27c0..90d8973080 100644 --- a/packages/backend/server/src/core/features/resolver.ts +++ b/packages/backend/server/src/core/features/resolver.ts @@ -8,70 +8,91 @@ import { } from '@nestjs/graphql'; import { difference } from 'lodash-es'; -import { Config } from '../../base'; +import { + Feature, + Models, + type UserFeatureName, + type WorkspaceFeatureName, +} from '../../models'; import { Admin } from '../common'; import { UserType } from '../user/types'; -import { EarlyAccessType, FeatureManagementService } from './management'; -import { FeatureService } from './service'; -import { FeatureType } from './types'; +import { AvailableUserFeatureConfig } from './types'; -registerEnumType(EarlyAccessType, { - name: 'EarlyAccessType', +registerEnumType(Feature, { + name: 'FeatureType', }); @Resolver(() => UserType) -export class FeatureManagementResolver { - constructor(private readonly feature: FeatureManagementService) {} +export class UserFeatureResolver extends AvailableUserFeatureConfig { + constructor(private readonly models: Models) { + super(); + } - @ResolveField(() => [FeatureType], { + @ResolveField(() => [Feature], { name: 'features', description: 'Enabled features of a user', }) async userFeatures(@Parent() user: UserType) { - return this.feature.getActivatedUserFeatures(user.id); - } -} - -export class AvailableUserFeatureConfig { - constructor(private readonly config: Config) {} - - async availableUserFeatures() { - return this.config.isSelfhosted - ? [FeatureType.Admin, FeatureType.UnlimitedCopilot] - : [FeatureType.EarlyAccess, FeatureType.AIEarlyAccess, FeatureType.Admin]; + const features = await this.models.userFeature.list(user.id); + const availableUserFeatures = this.availableUserFeatures(); + return features.filter(feature => availableUserFeatures.has(feature)); } } @Admin() @Resolver(() => Boolean) export class AdminFeatureManagementResolver extends AvailableUserFeatureConfig { - constructor( - config: Config, - private readonly feature: FeatureService - ) { - super(config); + constructor(private readonly models: Models) { + super(); } - @Mutation(() => [FeatureType], { + @Mutation(() => [Feature], { description: 'update user enabled feature', }) async updateUserFeatures( @Args('id') id: string, - @Args({ name: 'features', type: () => [FeatureType] }) - features: FeatureType[] + @Args({ name: 'features', type: () => [Feature] }) + features: UserFeatureName[] ) { - const configurableFeatures = await this.availableUserFeatures(); + const configurableUserFeatures = this.configurableUserFeatures(); + const removed = difference(Array.from(configurableUserFeatures), features); - const removed = difference(configurableFeatures, features); await Promise.all( - features.map(feature => - this.feature.addUserFeature(id, feature, 'admin panel') - ) + features.map(async feature => { + if (configurableUserFeatures.has(feature)) { + return this.models.userFeature.add(id, feature, 'admin panel'); + } else { + return; + } + }) ); + await Promise.all( - removed.map(feature => this.feature.removeUserFeature(id, feature)) + removed.map(feature => this.models.userFeature.remove(id, feature)) ); return features; } + + @Mutation(() => Boolean) + async addWorkspaceFeature( + @Args('workspaceId') workspaceId: string, + @Args('feature', { type: () => Feature }) feature: WorkspaceFeatureName + ) { + await this.models.workspaceFeature.add( + workspaceId, + feature, + 'by administrator' + ); + return true; + } + + @Mutation(() => Boolean) + async removeWorkspaceFeature( + @Args('workspaceId') workspaceId: string, + @Args('feature', { type: () => Feature }) feature: WorkspaceFeatureName + ) { + await this.models.workspaceFeature.remove(workspaceId, feature); + return true; + } } diff --git a/packages/backend/server/src/core/features/service.ts b/packages/backend/server/src/core/features/service.ts index a5a9fec2ef..deb9879395 100644 --- a/packages/backend/server/src/core/features/service.ts +++ b/packages/backend/server/src/core/features/service.ts @@ -1,355 +1,90 @@ -import { Injectable } from '@nestjs/common'; -import { PrismaClient } from '@prisma/client'; +import { Injectable, Logger } from '@nestjs/common'; -import { CannotDeleteAllAdminAccount } from '../../base'; -import { WorkspaceFeatureType } from '../workspaces/types'; -import { FeatureConfigType, getFeature } from './feature'; -import { FeatureKind, FeatureType } from './types'; +import { Runtime } from '../../base'; +import { Models } from '../../models'; + +const STAFF = ['@toeverything.info', '@affine.pro']; + +export enum EarlyAccessType { + App = 'app', + AI = 'ai', +} @Injectable() export class FeatureService { - constructor(private readonly prisma: PrismaClient) {} + protected logger = new Logger(FeatureService.name); - async getFeature(feature: F) { - const data = await this.prisma.feature.findFirst({ - where: { feature, type: FeatureKind.Feature }, - select: { id: true }, - orderBy: { version: 'desc' }, - }); + constructor( + private readonly models: Models, + private readonly runtime: Runtime + ) {} - if (data) { - return getFeature(this.prisma, data.id) as Promise>; + // ======== Admin ======== + isStaff(email: string) { + for (const domain of STAFF) { + if (email.endsWith(domain)) { + return true; + } } - return; + return false; } - // ======== User Features ======== + isAdmin(userId: string) { + return this.models.userFeature.has(userId, 'administrator'); + } - async addUserFeature( + addAdmin(userId: string) { + return this.models.userFeature.add(userId, 'administrator', 'Admin user'); + } + + // ======== Early Access ======== + async addEarlyAccess( userId: string, - feature: FeatureType, - reason: string, - expiredAt?: Date | string + type: EarlyAccessType = EarlyAccessType.App ) { - return this.prisma.$transaction(async tx => { - const latestFlag = await tx.userFeature.findFirst({ - where: { - userId, - feature: { - feature, - type: FeatureKind.Feature, - }, - activated: true, - }, - orderBy: { - createdAt: 'desc', - }, - }); - - if (latestFlag) { - return latestFlag.id; - } else { - const featureId = await tx.feature - .findFirst({ - where: { feature, type: FeatureKind.Feature }, - orderBy: { version: 'desc' }, - select: { id: true }, - }) - .then(r => r?.id); - - if (!featureId) { - throw new Error(`Feature ${feature} not found`); - } - - return tx.userFeature - .create({ - data: { - reason, - expiredAt, - activated: true, - userId, - featureId, - }, - }) - .then(r => r.id); - } - }); - } - - async removeUserFeature(userId: string, feature: FeatureType) { - if (feature === FeatureType.Admin) { - await this.ensureNotLastAdmin(userId); - } - return this.prisma.userFeature - .updateMany({ - where: { - userId, - feature: { - feature, - type: FeatureKind.Feature, - }, - activated: true, - }, - data: { - activated: false, - }, - }) - .then(r => r.count); - } - - async ensureNotLastAdmin(userId: string) { - const count = await this.prisma.userFeature.count({ - where: { - userId: { not: userId }, - feature: { feature: FeatureType.Admin, type: FeatureKind.Feature }, - activated: true, - }, - }); - - if (count === 0) { - throw new CannotDeleteAllAdminAccount(); - } - } - - /** - * get user's features, will included inactivated features - * @param userId user id - * @returns list of features - */ - async getUserFeatures(userId: string) { - const features = await this.prisma.userFeature.findMany({ - where: { - userId, - feature: { type: FeatureKind.Feature }, - }, - select: { - activated: true, - reason: true, - createdAt: true, - expiredAt: true, - featureId: true, - }, - }); - - const configs = await Promise.all( - features.map(async feature => ({ - ...feature, - feature: await getFeature(this.prisma, feature.featureId), - })) + return this.models.userFeature.add( + userId, + type === EarlyAccessType.App ? 'early_access' : 'ai_early_access', + 'Early access user' ); - - return configs.filter(feature => !!feature.feature); } - async getUserActivatedFeatures(userId: string) { - const features = await this.prisma.userFeature.findMany({ - where: { - userId, - feature: { type: FeatureKind.Feature }, - activated: true, - OR: [{ expiredAt: null }, { expiredAt: { gt: new Date() } }], - }, - select: { - activated: true, - reason: true, - createdAt: true, - expiredAt: true, - featureId: true, - }, - }); - - const configs = await Promise.all( - features.map(async feature => ({ - ...feature, - feature: await getFeature(this.prisma, feature.featureId), - })) - ); - - return configs.filter(feature => !!feature.feature); - } - - async listUsersByFeature(feature: FeatureType) { - return this.prisma.userFeature - .findMany({ - where: { - activated: true, - feature: { - feature: feature, - type: FeatureKind.Feature, - }, - }, - select: { - user: { - select: { - id: true, - name: true, - avatarUrl: true, - email: true, - emailVerifiedAt: true, - createdAt: true, - }, - }, - }, - }) - .then(users => users.map(user => user.user)); - } - - async hasUserFeature(userId: string, feature: FeatureType) { - return this.prisma.userFeature - .count({ - where: { - userId, - activated: true, - feature: { - feature, - type: FeatureKind.Feature, - }, - OR: [{ expiredAt: null }, { expiredAt: { gt: new Date() } }], - }, - }) - .then(count => count > 0); - } - - // ======== Workspace Features ======== - - async addWorkspaceFeature( - workspaceId: string, - feature: FeatureType, - reason: string, - expiredAt?: Date | string + async removeEarlyAccess( + userId: string, + type: EarlyAccessType = EarlyAccessType.App ) { - return this.prisma.$transaction(async tx => { - const latestFlag = await tx.workspaceFeature.findFirst({ - where: { - workspaceId, - feature: { - feature, - type: FeatureKind.Feature, - }, - activated: true, - }, - orderBy: { - createdAt: 'desc', - }, - }); - if (latestFlag) { - return latestFlag.id; - } else { - // use latest version of feature - const featureId = await tx.feature - .findFirst({ - where: { feature, type: FeatureKind.Feature }, - select: { id: true }, - orderBy: { version: 'desc' }, - }) - .then(r => r?.id); - - if (!featureId) { - throw new Error(`Feature ${feature} not found`); - } - - return tx.workspaceFeature - .create({ - data: { - reason, - expiredAt, - activated: true, - workspaceId, - featureId, - }, - }) - .then(r => r.id); - } - }); + return this.models.userFeature.remove( + userId, + type === EarlyAccessType.App ? 'early_access' : 'ai_early_access' + ); } - async removeWorkspaceFeature(workspaceId: string, feature: FeatureType) { - return this.prisma.workspaceFeature - .updateMany({ - where: { - workspaceId, - feature: { - feature, - type: FeatureKind.Feature, - }, - activated: true, - }, - data: { - activated: false, - }, - }) - .then(r => r.count); + async isEarlyAccessUser( + userId: string, + type: EarlyAccessType = EarlyAccessType.App + ) { + return await this.models.userFeature.has( + userId, + type === EarlyAccessType.App ? 'early_access' : 'ai_early_access' + ); } - /** - * get workspace's features, will included inactivated features - * @param workspaceId workspace id - * @returns list of features - */ - async getWorkspaceFeatures(workspaceId: string) { - const features = await this.prisma.workspaceFeature.findMany({ - where: { - workspace: { id: workspaceId }, - feature: { - type: FeatureKind.Feature, - }, - }, - select: { - activated: true, - reason: true, - createdAt: true, - expiredAt: true, - featureId: true, - }, - }); - - const configs = await Promise.all( - features.map(async feature => ({ - ...feature, - feature: await getFeature(this.prisma, feature.featureId), - })) + async canEarlyAccess( + email: string, + type: EarlyAccessType = EarlyAccessType.App + ) { + const earlyAccessControlEnabled = await this.runtime.fetch( + 'flags/earlyAccessControl' ); - return configs.filter(feature => !!feature.feature); - } - - async listWorkspacesByFeature( - feature: FeatureType - ): Promise { - return this.prisma.workspaceFeature - .findMany({ - where: { - activated: true, - feature: { - feature: feature, - type: FeatureKind.Feature, - }, - }, - select: { - workspace: { - select: { - id: true, - public: true, - createdAt: true, - }, - }, - }, - }) - .then(wss => wss.map(ws => ws.workspace)); - } - - async hasWorkspaceFeature(workspaceId: string, feature: FeatureType) { - return this.prisma.workspaceFeature - .count({ - where: { - workspaceId, - activated: true, - feature: { - feature, - type: FeatureKind.Feature, - }, - }, - }) - .then(count => count > 0); + if (earlyAccessControlEnabled && !this.isStaff(email)) { + const user = await this.models.user.getUserByEmail(email); + if (!user) { + return false; + } + return this.isEarlyAccessUser(user.id, type); + } else { + return true; + } } } diff --git a/packages/backend/server/src/core/features/types.ts b/packages/backend/server/src/core/features/types.ts new file mode 100644 index 0000000000..08db922999 --- /dev/null +++ b/packages/backend/server/src/core/features/types.ts @@ -0,0 +1,31 @@ +import { Inject, Injectable } from '@nestjs/common'; + +import { Config } from '../../base'; +import { Feature, UserFeatureName } from '../../models'; + +@Injectable() +export class AvailableUserFeatureConfig { + @Inject(Config) private readonly config!: Config; + + availableUserFeatures(): Set { + return new Set([ + Feature.Admin, + Feature.UnlimitedCopilot, + Feature.EarlyAccess, + Feature.AIEarlyAccess, + ]); + } + + configurableUserFeatures(): Set { + return new Set( + this.config.isSelfhosted + ? [Feature.Admin, Feature.UnlimitedCopilot] + : [ + Feature.EarlyAccess, + Feature.AIEarlyAccess, + Feature.Admin, + Feature.UnlimitedCopilot, + ] + ); + } +} diff --git a/packages/backend/server/src/core/features/types/admin.ts b/packages/backend/server/src/core/features/types/admin.ts deleted file mode 100644 index 4896415184..0000000000 --- a/packages/backend/server/src/core/features/types/admin.ts +++ /dev/null @@ -1,8 +0,0 @@ -import { z } from 'zod'; - -import { FeatureType } from './common'; - -export const featureAdministrator = z.object({ - feature: z.literal(FeatureType.Admin), - configs: z.object({}), -}); diff --git a/packages/backend/server/src/core/features/types/common.ts b/packages/backend/server/src/core/features/types/common.ts deleted file mode 100644 index 4cdfa05394..0000000000 --- a/packages/backend/server/src/core/features/types/common.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { registerEnumType } from '@nestjs/graphql'; - -export enum FeatureType { - // user feature - Admin = 'administrator', - EarlyAccess = 'early_access', - AIEarlyAccess = 'ai_early_access', - UnlimitedCopilot = 'unlimited_copilot', - // workspace feature - Copilot = 'copilot', - UnlimitedWorkspace = 'unlimited_workspace', -} - -registerEnumType(FeatureType, { - name: 'FeatureType', - description: 'The type of workspace feature', -}); diff --git a/packages/backend/server/src/core/features/types/copilot.ts b/packages/backend/server/src/core/features/types/copilot.ts deleted file mode 100644 index 0a58096354..0000000000 --- a/packages/backend/server/src/core/features/types/copilot.ts +++ /dev/null @@ -1,8 +0,0 @@ -import { z } from 'zod'; - -import { FeatureType } from './common'; - -export const featureCopilot = z.object({ - feature: z.literal(FeatureType.Copilot), - configs: z.object({}), -}); diff --git a/packages/backend/server/src/core/features/types/early-access.ts b/packages/backend/server/src/core/features/types/early-access.ts deleted file mode 100644 index bad8a9ea84..0000000000 --- a/packages/backend/server/src/core/features/types/early-access.ts +++ /dev/null @@ -1,16 +0,0 @@ -import { z } from 'zod'; - -import { FeatureType } from './common'; - -export const featureEarlyAccess = z.object({ - feature: z.literal(FeatureType.EarlyAccess), - configs: z.object({ - // field polyfill, make it optional in the future - whitelist: z.string().array(), - }), -}); - -export const featureAIEarlyAccess = z.object({ - feature: z.literal(FeatureType.AIEarlyAccess), - configs: z.object({}), -}); diff --git a/packages/backend/server/src/core/features/types/index.ts b/packages/backend/server/src/core/features/types/index.ts deleted file mode 100644 index 4623deb986..0000000000 --- a/packages/backend/server/src/core/features/types/index.ts +++ /dev/null @@ -1,100 +0,0 @@ -import { z } from 'zod'; - -import { featureAdministrator } from './admin'; -import { FeatureType } from './common'; -import { featureCopilot } from './copilot'; -import { featureAIEarlyAccess, featureEarlyAccess } from './early-access'; -import { featureUnlimitedCopilot } from './unlimited-copilot'; -import { featureUnlimitedWorkspace } from './unlimited-workspace'; - -/// ======== common schema ======== - -export enum FeatureKind { - Feature, - Quota, -} - -export const commonFeatureSchema = z.object({ - feature: z.string(), - type: z.nativeEnum(FeatureKind), - version: z.number(), - configs: z.unknown(), -}); - -export type CommonFeature = z.infer; - -/// ======== feature define ======== - -export const Features: Feature[] = [ - { - feature: FeatureType.Copilot, - type: FeatureKind.Feature, - version: 1, - configs: {}, - }, - { - feature: FeatureType.EarlyAccess, - type: FeatureKind.Feature, - version: 1, - configs: { - whitelist: ['@toeverything.info'], - }, - }, - { - feature: FeatureType.EarlyAccess, - type: FeatureKind.Feature, - version: 2, - configs: { - whitelist: [], - }, - }, - { - feature: FeatureType.UnlimitedWorkspace, - type: FeatureKind.Feature, - version: 1, - configs: {}, - }, - { - feature: FeatureType.UnlimitedCopilot, - type: FeatureKind.Feature, - version: 1, - configs: {}, - }, - { - feature: FeatureType.AIEarlyAccess, - type: FeatureKind.Feature, - version: 1, - configs: {}, - }, - { - feature: FeatureType.Admin, - type: FeatureKind.Feature, - version: 1, - configs: {}, - }, -]; - -/// ======== schema infer ======== - -export const FeatureConfigSchema = z.discriminatedUnion('feature', [ - featureCopilot, - featureEarlyAccess, - featureAIEarlyAccess, - featureUnlimitedWorkspace, - featureUnlimitedCopilot, - featureAdministrator, -]); - -export const FeatureSchema = commonFeatureSchema - .extend({ - type: z.literal(FeatureKind.Feature), - }) - .and(FeatureConfigSchema); - -export type FeatureConfig = (z.infer< - typeof FeatureConfigSchema -> & { feature: F })['configs']; - -export type Feature = z.infer; - -export { FeatureType }; diff --git a/packages/backend/server/src/core/features/types/unlimited-copilot.ts b/packages/backend/server/src/core/features/types/unlimited-copilot.ts deleted file mode 100644 index fd69e791a6..0000000000 --- a/packages/backend/server/src/core/features/types/unlimited-copilot.ts +++ /dev/null @@ -1,8 +0,0 @@ -import { z } from 'zod'; - -import { FeatureType } from './common'; - -export const featureUnlimitedCopilot = z.object({ - feature: z.literal(FeatureType.UnlimitedCopilot), - configs: z.object({}), -}); diff --git a/packages/backend/server/src/core/features/types/unlimited-workspace.ts b/packages/backend/server/src/core/features/types/unlimited-workspace.ts deleted file mode 100644 index b9b471e9e1..0000000000 --- a/packages/backend/server/src/core/features/types/unlimited-workspace.ts +++ /dev/null @@ -1,8 +0,0 @@ -import { z } from 'zod'; - -import { FeatureType } from './common'; - -export const featureUnlimitedWorkspace = z.object({ - feature: z.literal(FeatureType.UnlimitedWorkspace), - configs: z.object({}), -}); diff --git a/packages/backend/server/src/core/quota/index.ts b/packages/backend/server/src/core/quota/index.ts index 6564566121..1fc7791f85 100644 --- a/packages/backend/server/src/core/quota/index.ts +++ b/packages/backend/server/src/core/quota/index.ts @@ -1,11 +1,9 @@ import { Module } from '@nestjs/common'; -import { FeatureModule } from '../features'; import { PermissionModule } from '../permission'; import { StorageModule } from '../storage'; -import { QuotaManagementResolver } from './resolver'; +import { QuotaResolver } from './resolver'; import { QuotaService } from './service'; -import { QuotaManagementService } from './storage'; /** * Quota module provider pre-user quota management. @@ -14,18 +12,11 @@ import { QuotaManagementService } from './storage'; * - quota statistics */ @Module({ - imports: [FeatureModule, StorageModule, PermissionModule], - providers: [QuotaService, QuotaManagementResolver, QuotaManagementService], - exports: [QuotaService, QuotaManagementService], + imports: [StorageModule, PermissionModule], + providers: [QuotaService, QuotaResolver], + exports: [QuotaService], }) export class QuotaModule {} -export { QuotaManagementService, QuotaService }; -export { Quota_FreePlanV1_1, Quota_ProPlanV1 } from './schema'; -export { - formatDate, - formatSize, - type QuotaBusinessType, - QuotaQueryType, - QuotaType, -} from './types'; +export { QuotaService }; +export { WorkspaceQuotaHumanReadableType, WorkspaceQuotaType } from './types'; diff --git a/packages/backend/server/src/core/quota/quota.ts b/packages/backend/server/src/core/quota/quota.ts deleted file mode 100644 index 1930f0c136..0000000000 --- a/packages/backend/server/src/core/quota/quota.ts +++ /dev/null @@ -1,146 +0,0 @@ -import { pick } from 'lodash-es'; - -import { PrismaTransaction } from '../../base'; -import { formatDate, formatSize, Quota, QuotaSchema } from './types'; - -const QuotaCache = new Map(); - -export class QuotaConfig { - readonly config: Quota; - readonly override?: Partial; - - static async get(tx: PrismaTransaction, featureId: number) { - const cachedQuota = QuotaCache.get(featureId); - - if (cachedQuota) { - return cachedQuota; - } - - const quota = await tx.feature.findFirst({ - where: { - id: featureId, - }, - }); - - if (!quota) { - throw new Error(`Quota config ${featureId} not found`); - } - - const config = new QuotaConfig(quota); - // we always edit quota config as a new quota config - // so we can cache it by featureId - QuotaCache.set(featureId, config); - - return config; - } - - private constructor(data: any, override?: any) { - const config = QuotaSchema.safeParse(data); - if (config.success) { - this.config = config.data; - } else { - throw new Error( - `Invalid quota config: ${config.error.message}, ${JSON.stringify( - data - )})}` - ); - } - if (override) { - const overrideConfig = QuotaSchema.safeParse({ - ...config.data, - configs: Object.assign({}, config.data.configs, override), - }); - if (overrideConfig.success) { - this.override = pick( - overrideConfig.data.configs, - Object.keys(override) - ); - } else { - throw new Error( - `Invalid quota override config: ${override.error.message}, ${JSON.stringify( - data - )})}` - ); - } - } - } - - withOverride(override: any) { - if (override) { - return new QuotaConfig( - this.config, - Object.assign({}, this.override, override) - ); - } - return this; - } - - checkOverride(override: any) { - return QuotaSchema.safeParse({ - ...this.config, - configs: Object.assign({}, this.config.configs, override), - }); - } - - get version() { - return this.config.version; - } - - /// feature name of quota - get name() { - return this.config.feature; - } - - get blobLimit() { - return this.override?.blobLimit || this.config.configs.blobLimit; - } - - get businessBlobLimit() { - return ( - this.override?.businessBlobLimit || - this.config.configs.businessBlobLimit || - this.override?.blobLimit || - this.config.configs.blobLimit - ); - } - - private get additionalQuota() { - const seatQuota = - this.override?.seatQuota || this.config.configs.seatQuota || 0; - return this.memberLimit * seatQuota; - } - - get storageQuota() { - const baseQuota = - this.override?.storageQuota || this.config.configs.storageQuota; - return baseQuota + this.additionalQuota; - } - - get historyPeriod() { - return this.override?.historyPeriod || this.config.configs.historyPeriod; - } - - get memberLimit() { - return this.override?.memberLimit || this.config.configs.memberLimit; - } - - get copilotActionLimit() { - if ('copilotActionLimit' in this.config.configs) { - return this.config.configs.copilotActionLimit || undefined; - } - return undefined; - } - - get humanReadable() { - return { - name: this.config.configs.name, - blobLimit: formatSize(this.blobLimit), - storageQuota: formatSize(this.storageQuota), - historyPeriod: formatDate(this.historyPeriod), - memberLimit: this.memberLimit.toString(), - copilotActionLimit: this.copilotActionLimit - ? `${this.copilotActionLimit} times` - : 'Unlimited', - }; - } -} diff --git a/packages/backend/server/src/core/quota/resolver.ts b/packages/backend/server/src/core/quota/resolver.ts index 309f43833a..f287afd37d 100644 --- a/packages/backend/server/src/core/quota/resolver.ts +++ b/packages/backend/server/src/core/quota/resolver.ts @@ -1,86 +1,29 @@ -import { - Field, - ObjectType, - registerEnumType, - ResolveField, - Resolver, -} from '@nestjs/graphql'; -import { SafeIntResolver } from 'graphql-scalars'; +import { ResolveField, Resolver } from '@nestjs/graphql'; import { CurrentUser } from '../auth/session'; -import { EarlyAccessType } from '../features'; import { UserType } from '../user'; import { QuotaService } from './service'; -import { QuotaManagementService } from './storage'; - -registerEnumType(EarlyAccessType, { - name: 'EarlyAccessType', -}); - -@ObjectType('UserQuotaHumanReadable') -class UserQuotaHumanReadableType { - @Field({ name: 'name' }) - name!: string; - - @Field({ name: 'blobLimit' }) - blobLimit!: string; - - @Field({ name: 'storageQuota' }) - storageQuota!: string; - - @Field({ name: 'historyPeriod' }) - historyPeriod!: string; - - @Field({ name: 'memberLimit' }) - memberLimit!: string; -} - -@ObjectType('UserQuota') -class UserQuotaType { - @Field({ name: 'name' }) - name!: string; - - @Field(() => SafeIntResolver, { name: 'blobLimit' }) - blobLimit!: number; - - @Field(() => SafeIntResolver, { name: 'storageQuota' }) - storageQuota!: number; - - @Field(() => SafeIntResolver, { name: 'historyPeriod' }) - historyPeriod!: number; - - @Field({ name: 'memberLimit' }) - memberLimit!: number; - - @Field({ name: 'humanReadable' }) - humanReadable!: UserQuotaHumanReadableType; -} - -@ObjectType('UserQuotaUsage') -class UserQuotaUsageType { - @Field(() => SafeIntResolver, { name: 'storageQuota' }) - storageQuota!: number; -} +import { UserQuotaType, UserQuotaUsageType } from './types'; @Resolver(() => UserType) -export class QuotaManagementResolver { - constructor( - private readonly quota: QuotaService, - private readonly management: QuotaManagementService - ) {} +export class QuotaResolver { + constructor(private readonly quota: QuotaService) {} - @ResolveField(() => UserQuotaType, { name: 'quota', nullable: true }) - async getQuota(@CurrentUser() me: UserType) { - const quota = await this.quota.getUserQuota(me.id); + @ResolveField(() => UserQuotaType, { name: 'quota' }) + async getQuota(@CurrentUser() me: UserType): Promise { + const quota = await this.quota.getUserQuotaWithUsage(me.id); - return quota.feature; + return { + ...quota, + humanReadable: this.quota.formatUserQuota(quota), + }; } @ResolveField(() => UserQuotaUsageType, { name: 'quotaUsage' }) async getQuotaUsage( @CurrentUser() me: UserType ): Promise { - const usage = await this.management.getUserStorageUsage(me.id); + const usage = await this.quota.getUserStorageUsage(me.id); return { storageQuota: usage, diff --git a/packages/backend/server/src/core/quota/schema.ts b/packages/backend/server/src/core/quota/schema.ts deleted file mode 100644 index 535f43e854..0000000000 --- a/packages/backend/server/src/core/quota/schema.ts +++ /dev/null @@ -1,216 +0,0 @@ -import { FeatureKind } from '../features/types'; -import { OneDay, OneGB, OneMB } from './constant'; -import { Quota, QuotaType } from './types'; - -export const Quotas: Quota[] = [ - { - feature: QuotaType.FreePlanV1, - type: FeatureKind.Quota, - version: 1, - configs: { - // quota name - name: 'Free', - // single blob limit 10MB - blobLimit: 10 * OneMB, - // total blob limit 10GB - storageQuota: 10 * OneGB, - // history period of validity 7 days - historyPeriod: 7 * OneDay, - // member limit 3 - memberLimit: 3, - }, - }, - { - feature: QuotaType.ProPlanV1, - type: FeatureKind.Quota, - version: 1, - configs: { - // quota name - name: 'Pro', - // single blob limit 100MB - blobLimit: 100 * OneMB, - // total blob limit 100GB - storageQuota: 100 * OneGB, - // history period of validity 30 days - historyPeriod: 30 * OneDay, - // member limit 10 - memberLimit: 10, - }, - }, - { - feature: QuotaType.RestrictedPlanV1, - type: FeatureKind.Quota, - version: 1, - configs: { - // quota name - name: 'Restricted', - // single blob limit 10MB - blobLimit: OneMB, - // total blob limit 1GB - storageQuota: 10 * OneMB, - // history period of validity 30 days - historyPeriod: 30 * OneDay, - // member limit 10 - memberLimit: 10, - }, - }, - { - feature: QuotaType.FreePlanV1, - type: FeatureKind.Quota, - version: 2, - configs: { - // quota name - name: 'Free', - // single blob limit 10MB - blobLimit: 100 * OneMB, - // total blob limit 10GB - storageQuota: 10 * OneGB, - // history period of validity 7 days - historyPeriod: 7 * OneDay, - // member limit 3 - memberLimit: 3, - }, - }, - { - feature: QuotaType.FreePlanV1, - type: FeatureKind.Quota, - version: 3, - configs: { - // quota name - name: 'Free', - // single blob limit 10MB - blobLimit: 10 * OneMB, - // server limit will larger then client to handle a edge case: - // when a user downgrades from pro to free, he can still continue - // to upload previously added files that exceed the free limit - // NOTE: this is a product decision, may change in future - businessBlobLimit: 100 * OneMB, - // total blob limit 10GB - storageQuota: 10 * OneGB, - // history period of validity 7 days - historyPeriod: 7 * OneDay, - // member limit 3 - memberLimit: 3, - }, - }, - { - feature: QuotaType.FreePlanV1, - type: FeatureKind.Quota, - version: 4, - configs: { - // quota name - name: 'Free', - // single blob limit 10MB - blobLimit: 10 * OneMB, - // server limit will larger then client to handle a edge case: - // when a user downgrades from pro to free, he can still continue - // to upload previously added files that exceed the free limit - // NOTE: this is a product decision, may change in future - businessBlobLimit: 100 * OneMB, - // total blob limit 10GB - storageQuota: 10 * OneGB, - // history period of validity 7 days - historyPeriod: 7 * OneDay, - // member limit 3 - memberLimit: 3, - // copilot action limit 10 - copilotActionLimit: 10, - }, - }, - { - feature: QuotaType.ProPlanV1, - type: FeatureKind.Quota, - version: 2, - configs: { - // quota name - name: 'Pro', - // single blob limit 100MB - blobLimit: 100 * OneMB, - // total blob limit 100GB - storageQuota: 100 * OneGB, - // history period of validity 30 days - historyPeriod: 30 * OneDay, - // member limit 10 - memberLimit: 10, - // copilot action limit 10 - copilotActionLimit: 10, - }, - }, - { - feature: QuotaType.RestrictedPlanV1, - type: FeatureKind.Quota, - version: 2, - configs: { - // quota name - name: 'Restricted', - // single blob limit 1MB - blobLimit: OneMB, - // total blob limit 10MB - storageQuota: 10 * OneMB, - // history period of validity 30 days - historyPeriod: 30 * OneDay, - // member limit 10 - memberLimit: 10, - // copilot action limit 10 - copilotActionLimit: 10, - }, - }, - { - feature: QuotaType.LifetimeProPlanV1, - type: FeatureKind.Quota, - version: 1, - configs: { - // quota name - name: 'Lifetime Pro', - // single blob limit 100MB - blobLimit: 100 * OneMB, - // total blob limit 1TB - storageQuota: 1024 * OneGB, - // history period of validity 30 days - historyPeriod: 30 * OneDay, - // member limit 10 - memberLimit: 10, - // copilot action limit 10 - copilotActionLimit: 10, - }, - }, - { - feature: QuotaType.TeamPlanV1, - type: FeatureKind.Quota, - version: 1, - configs: { - // quota name - name: 'Team Workspace', - // single blob limit 100MB - blobLimit: 500 * OneMB, - // total blob limit 100GB - storageQuota: 100 * OneGB, - // seat quota 20GB per seat - seatQuota: 20 * OneGB, - // history period of validity 30 days - historyPeriod: 30 * OneDay, - // member limit 1, override by workspace config - memberLimit: 1, - }, - }, -]; - -export function getLatestQuota(type: Q): Quota { - const quota = Quotas.filter(f => f.feature === type); - quota.sort((a, b) => b.version - a.version); - return quota[0] as Quota; -} - -export const FreePlan = getLatestQuota(QuotaType.FreePlanV1); -export const ProPlan = getLatestQuota(QuotaType.ProPlanV1); -export const LifetimeProPlan = getLatestQuota(QuotaType.LifetimeProPlanV1); - -export const Quota_FreePlanV1_1 = { - feature: Quotas[5].feature, - version: Quotas[5].version, -}; - -export const Quota_ProPlanV1 = { - feature: Quotas[6].feature, - version: Quotas[6].version, -}; diff --git a/packages/backend/server/src/core/quota/service.ts b/packages/backend/server/src/core/quota/service.ts index e782603de4..a7e47c3888 100644 --- a/packages/backend/server/src/core/quota/service.ts +++ b/packages/backend/server/src/core/quota/service.ts @@ -1,329 +1,269 @@ -import { Injectable } from '@nestjs/common'; -import { PrismaClient } from '@prisma/client'; +import { Injectable, Logger } from '@nestjs/common'; -import { PrismaTransaction } from '../../base'; -import { FeatureKind } from '../features/types'; -import { QuotaConfig } from './quota'; -import { QuotaType } from './types'; +import { InternalServerError, MemberQuotaExceeded, OnEvent } from '../../base'; +import { + Models, + type UserQuota, + WorkspaceQuota as BaseWorkspaceQuota, +} from '../../models'; +import { PermissionService } from '../permission'; +import { WorkspaceBlobStorage } from '../storage'; +import { + UserQuotaHumanReadableType, + UserQuotaType, + WorkspaceQuotaHumanReadableType, + WorkspaceQuotaType, +} from './types'; +import { formatDate, formatSize } from './utils'; + +type UserQuotaWithUsage = Omit; +type WorkspaceQuota = Omit & { + ownerQuota?: string; +}; +type WorkspaceQuotaWithUsage = Omit; @Injectable() export class QuotaService { - constructor(private readonly prisma: PrismaClient) {} + protected logger = new Logger(QuotaService.name); - async getQuota( - quota: Q, - tx?: PrismaTransaction - ): Promise { - const executor = tx ?? this.prisma; + constructor( + private readonly models: Models, + private readonly permissions: PermissionService, + private readonly storage: WorkspaceBlobStorage + ) {} - const data = await executor.feature.findFirst({ - where: { feature: quota, type: FeatureKind.Quota }, - select: { id: true }, - orderBy: { version: 'desc' }, - }); - - if (data) { - return QuotaConfig.get(this.prisma, data.id); - } - return undefined; + @OnEvent('user.postCreated') + async onUserCreated({ id }: Events['user.postCreated']) { + await this.setupUserBaseQuota(id); } - // ======== User Quota ======== - - // get activated user quota - async getUserQuota(userId: string) { - const quota = await this.prisma.userFeature.findFirst({ - where: { - userId, - feature: { type: FeatureKind.Quota }, - activated: true, - }, - select: { - reason: true, - createdAt: true, - expiredAt: true, - featureId: true, - }, - }); + async getUserQuota(userId: string): Promise { + let quota = await this.models.userFeature.getQuota(userId); + // not possible, but just in case, we do a little fix for user to avoid system dump if (!quota) { - // this should unreachable - throw new Error(`User ${userId} has no quota`); + await this.setupUserBaseQuota(userId); + quota = await this.models.userFeature.getQuota(userId); } - const feature = await QuotaConfig.get(this.prisma, quota.featureId); - return { ...quota, feature }; - } - - // get user all quota records - async getUserQuotas(userId: string) { - const quotas = await this.prisma.userFeature.findMany({ - where: { - userId, - feature: { type: FeatureKind.Quota }, - }, - select: { - activated: true, - reason: true, - createdAt: true, - expiredAt: true, - featureId: true, - }, - orderBy: { id: 'asc' }, - }); - const configs = await Promise.all( - quotas.map(async quota => { - try { - return { - ...quota, - feature: await QuotaConfig.get(this.prisma, quota.featureId), - }; - } catch {} - return null as unknown as typeof quota & { - feature: QuotaConfig; - }; - }) + const unlimitedCopilot = await this.models.userFeature.has( + userId, + 'unlimited_copilot' ); - return configs.filter(quota => !!quota); - } - - // switch user to a new quota - // currently each user can only have one quota - async switchUserQuota( - userId: string, - quota: QuotaType, - reason?: string, - expiredAt?: Date - ) { - await this.prisma.$transaction(async tx => { - const hasSameActivatedQuota = await this.hasUserQuota(userId, quota, tx); - 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' }, - }) - .then(f => f?.id); - - if (!featureId) { - throw new Error(`Quota ${quota} not found`); - } - - // we will deactivate all exists quota for this user - await tx.userFeature.updateMany({ - where: { - id: undefined, - userId, - feature: { - type: FeatureKind.Quota, - }, - }, - data: { - activated: false, - }, - }); - - await tx.userFeature.create({ - data: { - userId, - featureId, - reason: reason ?? 'switch quota', - activated: true, - expiredAt, - }, - }); - }); - } - - async hasUserQuota(userId: string, quota: QuotaType, tx?: PrismaTransaction) { - const executor = tx ?? this.prisma; - - return executor.userFeature - .count({ - where: { - userId, - feature: { - feature: quota, - type: FeatureKind.Quota, - }, - activated: true, - }, - }) - .then(count => count > 0); - } - - // ======== Workspace Quota ======== - - // get activated workspace quota - async getWorkspaceQuota(workspaceId: string) { - const quota = await this.prisma.workspaceFeature.findFirst({ - where: { - workspaceId, - feature: { type: FeatureKind.Quota }, - activated: true, - }, - select: { - configs: true, - reason: true, - createdAt: true, - expiredAt: true, - featureId: true, - }, - }); - - if (quota) { - const feature = await QuotaConfig.get(this.prisma, quota.featureId); - const { configs, ...rest } = quota; - return { ...rest, feature: feature.withOverride(configs) }; - } - return null; - } - - // switch user to a new quota - // currently each user can only have one quota - async switchWorkspaceQuota( - workspaceId: string, - quota: QuotaType, - reason?: string, - expiredAt?: Date - ) { - await this.prisma.$transaction(async tx => { - const hasSameActivatedQuota = await this.hasWorkspaceQuota( - workspaceId, - quota, - tx - ); - 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' }, - }) - .then(f => f?.id); - - if (!featureId) { - throw new Error(`Quota ${quota} not found`); - } - - // we will deactivate all exists quota for this workspace - await this.deactivateWorkspaceQuota(workspaceId, undefined, tx); - - await tx.workspaceFeature.create({ - data: { - workspaceId, - featureId, - reason: reason ?? 'switch quota', - activated: true, - expiredAt, - }, - }); - }); - } - - async deactivateWorkspaceQuota( - workspaceId: string, - quota?: QuotaType, - tx?: PrismaTransaction - ) { - const executor = tx ?? this.prisma; - - await executor.workspaceFeature.updateMany({ - where: { - id: undefined, - workspaceId, - feature: quota - ? { feature: quota, type: FeatureKind.Quota } - : { type: FeatureKind.Quota }, - }, - data: { activated: false }, - }); - } - - async hasWorkspaceQuota( - workspaceId: string, - quota: QuotaType, - tx?: PrismaTransaction - ) { - const executor = tx ?? this.prisma; - - return executor.workspaceFeature - .count({ - where: { - workspaceId, - feature: { - feature: quota, - type: FeatureKind.Quota, - }, - activated: true, - }, - }) - .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 - ): Promise { - const quota = await this.getQuota(type); - if (quota) { - const configs = await this.prisma.workspaceFeature - .findFirst({ - where: { - workspaceId, - feature: { feature: type, type: FeatureKind.Quota }, - activated: true, - }, - select: { configs: true }, - }) - .then(q => q?.configs); - return quota.withOverride(configs); - } - return undefined; - } - - async updateWorkspaceConfig( - workspaceId: string, - quota: QuotaType, - configs: any - ) { - const current = await this.getWorkspaceConfig(workspaceId, quota); - - const ret = current?.checkOverride(configs); - if (!ret || !ret.success) { - throw new Error( - `Invalid quota config: ${ret?.error.message || 'quota not defined'}` + if (!quota) { + throw new InternalServerError( + 'User quota not found and can not be created.' ); } - const r = await this.prisma.workspaceFeature.updateMany({ - where: { - workspaceId, - feature: { feature: quota, type: FeatureKind.Quota }, - activated: true, - }, - data: { configs }, - }); - return r.count; + + return { + ...quota.configs, + copilotActionLimit: unlimitedCopilot + ? undefined + : quota.configs.copilotActionLimit, + } as UserQuotaWithUsage; + } + + async getUserQuotaWithUsage(userId: string): Promise { + const quota = await this.getUserQuota(userId); + const usedStorageQuota = await this.getUserStorageUsage(userId); + + return { ...quota, usedStorageQuota }; + } + + async getUserStorageUsage(userId: string) { + const workspaces = await this.permissions.getOwnedWorkspaces(userId); + const workspacesWithQuota = + await this.models.workspaceFeature.batchHasQuota(workspaces); + + const sizes = await Promise.allSettled( + workspaces + .filter(w => !workspacesWithQuota.includes(w)) + .map(workspace => this.storage.totalSize(workspace)) + ); + + return sizes.reduce((total, size) => { + if (size.status === 'fulfilled') { + // 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}`); + } + } else { + this.logger.error(`Failed to get workspace size: ${size.reason}`); + } + return total; + }, 0); + } + + async getWorkspaceStorageUsage(workspaceId: string) { + const totalSize = await 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; + } + + async getWorkspaceQuota(workspaceId: string): Promise { + const quota = await this.models.workspaceFeature.getQuota(workspaceId); + + if (!quota) { + // get and convert to workspace quota from owner's quota + // TODO(@forehalo): replace it with `WorkspaceRoleModel` when it's ready + const owner = await this.permissions.getWorkspaceOwner(workspaceId); + const ownerQuota = await this.getUserQuota(owner.id); + + return { + ...ownerQuota, + ownerQuota: owner.id, + }; + } + + return quota.configs; + } + + async getWorkspaceQuotaWithUsage( + workspaceId: string + ): Promise { + const quota = await this.getWorkspaceQuota(workspaceId); + const usedStorageQuota = quota.ownerQuota + ? await this.getUserStorageUsage(quota.ownerQuota) + : await this.getWorkspaceStorageUsage(workspaceId); + const memberCount = + await this.permissions.getWorkspaceMemberCount(workspaceId); + + return { + ...quota, + usedStorageQuota, + memberCount, + usedSize: usedStorageQuota, + }; + } + + formatUserQuota( + quota: Omit + ): UserQuotaHumanReadableType { + return { + name: quota.name, + blobLimit: formatSize(quota.blobLimit), + storageQuota: formatSize(quota.storageQuota), + usedStorageQuota: formatSize(quota.usedStorageQuota), + historyPeriod: formatDate(quota.historyPeriod), + memberLimit: quota.memberLimit.toString(), + copilotActionLimit: quota.copilotActionLimit + ? `${quota.copilotActionLimit} times` + : 'Unlimited', + }; + } + + async getWorkspaceSeatQuota(workspaceId: string) { + const quota = await this.getWorkspaceQuota(workspaceId); + const memberCount = + await this.permissions.getWorkspaceMemberCount(workspaceId); + + return { + memberCount, + memberLimit: quota.memberLimit, + }; + } + + async tryCheckSeat(workspaceId: string, excludeSelf = false) { + const quota = await this.getWorkspaceSeatQuota(workspaceId); + + return quota.memberCount - (excludeSelf ? 1 : 0) < quota.memberLimit; + } + + async checkSeat(workspaceId: string, excludeSelf = false) { + const available = await this.tryCheckSeat(workspaceId, excludeSelf); + + if (!available) { + throw new MemberQuotaExceeded(); + } + } + + formatWorkspaceQuota( + quota: Omit + ): WorkspaceQuotaHumanReadableType { + return { + name: quota.name, + blobLimit: formatSize(quota.blobLimit), + storageQuota: formatSize(quota.storageQuota), + storageQuotaUsed: formatSize(quota.usedStorageQuota), + historyPeriod: formatDate(quota.historyPeriod), + memberLimit: quota.memberLimit.toString(), + memberCount: quota.memberCount.toString(), + }; + } + + async getUserQuotaCalculator(userId: string) { + const quota = await this.getUserQuota(userId); + const usedSize = await this.getUserStorageUsage(userId); + + return this.generateQuotaCalculator( + quota.storageQuota, + quota.blobLimit, + usedSize + ); + } + + async getWorkspaceQuotaCalculator(workspaceId: string) { + const quota = await this.getWorkspaceQuota(workspaceId); + const unlimited = await this.models.workspaceFeature.has( + workspaceId, + 'unlimited_workspace' + ); + + // quota check will be disabled for unlimited workspace + // we save a complicated db read for used size + if (unlimited) { + return this.generateQuotaCalculator(0, quota.blobLimit, 0, true); + } + + const usedSize = quota.ownerQuota + ? await this.getUserStorageUsage(quota.ownerQuota) + : await this.getWorkspaceStorageUsage(workspaceId); + + return this.generateQuotaCalculator( + quota.storageQuota, + quota.blobLimit, + usedSize + ); + } + + private async setupUserBaseQuota(userId: string) { + await this.models.userFeature.add(userId, 'free_plan_v1', 'sign up'); + } + + private generateQuotaCalculator( + storageQuota: number, + blobLimit: number, + usedQuota: number, + unlimited = false + ) { + const checkExceeded = (recvSize: number) => { + const currentSize = usedQuota + recvSize; + // only skip total storage check if workspace has unlimited feature + if (currentSize > storageQuota && !unlimited) { + this.logger.warn( + `storage size limit exceeded: ${currentSize} > ${storageQuota}` + ); + return true; + } else if (recvSize > blobLimit) { + this.logger.warn( + `blob size limit exceeded: ${recvSize} > ${blobLimit}` + ); + return true; + } else { + return false; + } + }; + return checkExceeded; } } diff --git a/packages/backend/server/src/core/quota/storage.ts b/packages/backend/server/src/core/quota/storage.ts deleted file mode 100644 index 0f3a84bb0d..0000000000 --- a/packages/backend/server/src/core/quota/storage.ts +++ /dev/null @@ -1,264 +0,0 @@ -import { Injectable, Logger } from '@nestjs/common'; - -import { MemberQuotaExceeded } from '../../base'; -import { FeatureService, FeatureType } from '../features'; -import { PermissionService } from '../permission'; -import { WorkspaceBlobStorage } from '../storage'; -import { OneGB } from './constant'; -import { QuotaConfig } from './quota'; -import { QuotaService } from './service'; -import { formatSize, Quota, type QuotaBusinessType, QuotaType } from './types'; - -@Injectable() -export class QuotaManagementService { - protected logger = new Logger(QuotaManagementService.name); - - constructor( - private readonly feature: FeatureService, - private readonly quota: QuotaService, - private readonly permissions: PermissionService, - private readonly storage: WorkspaceBlobStorage - ) {} - - async getUserQuota(userId: string) { - const quota = await this.quota.getUserQuota(userId); - - return { - name: quota.feature.name, - reason: quota.reason, - createAt: quota.createdAt, - expiredAt: quota.expiredAt, - blobLimit: quota.feature.blobLimit, - businessBlobLimit: quota.feature.businessBlobLimit, - storageQuota: quota.feature.storageQuota, - historyPeriod: quota.feature.historyPeriod, - memberLimit: quota.feature.memberLimit, - copilotActionLimit: quota.feature.copilotActionLimit, - }; - } - - async getWorkspaceConfig( - workspaceId: string, - quota: Q - ): Promise { - return this.quota.getWorkspaceConfig(workspaceId, quota); - } - - async updateWorkspaceConfig( - workspaceId: string, - quota: Q, - configs: Partial['configs']> - ) { - const orig = await this.getWorkspaceConfig(workspaceId, quota); - return await this.quota.updateWorkspaceConfig( - workspaceId, - quota, - Object.assign({}, orig?.override, configs) - ); - } - - // ======== Team Workspace ======== - async addTeamWorkspace(workspaceId: string, reason: string) { - return this.quota.switchWorkspaceQuota( - workspaceId, - QuotaType.TeamPlanV1, - reason - ); - } - - async removeTeamWorkspace(workspaceId: string) { - return this.quota.deactivateWorkspaceQuota( - workspaceId, - QuotaType.TeamPlanV1 - ); - } - - async isTeamWorkspace(workspaceId: string) { - return this.quota.hasWorkspaceQuota(workspaceId, QuotaType.TeamPlanV1); - } - - async getUserStorageUsage(userId: string) { - const workspaces = await this.permissions.getOwnedWorkspaces(userId); - const workspacesWithQuota = await this.quota.hasWorkspacesQuota(workspaces); - - const sizes = await Promise.allSettled( - workspaces - .filter(w => !workspacesWithQuota.includes(w)) - .map(workspace => this.storage.totalSize(workspace)) - ); - - return sizes.reduce((total, size) => { - if (size.status === 'fulfilled') { - // 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}`); - } - } else { - this.logger.error(`Failed to get workspace size: ${size.reason}`); - } - return total; - }, 0); - } - - async getWorkspaceStorageUsage(workspaceId: string) { - const totalSize = await 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, - usedSize: number, - unlimited = false - ) { - const checkExceeded = (recvSize: number) => { - const total = usedSize + recvSize; - // only skip total storage check if workspace has unlimited feature - if (total > quota && !unlimited) { - this.logger.warn(`storage size limit exceeded: ${total} > ${quota}`); - return true; - } else if (recvSize > blobLimit) { - this.logger.warn( - `blob size limit exceeded: ${recvSize} > ${blobLimit}` - ); - return true; - } else { - return false; - } - }; - return checkExceeded; - } - - async getQuotaCalculator(userId: string) { - const quota = await this.getUserQuota(userId); - const { storageQuota, businessBlobLimit } = quota; - const usedSize = await this.getUserStorageUsage(userId); - - return this.generateQuotaCalculator( - storageQuota, - businessBlobLimit, - usedSize - ); - } - - async getQuotaCalculatorByWorkspace(workspaceId: string) { - const { storageQuota, usedSize, businessBlobLimit, unlimited } = - await this.getWorkspaceUsage(workspaceId); - - return this.generateQuotaCalculator( - storageQuota, - businessBlobLimit, - usedSize, - unlimited - ); - } - - 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 { - quota: workspaceQuota.withOverride({ - // override user quota with workspace quota - copilotActionLimit: userQuota.copilotActionLimit, - }), - fromUser: false, - }; - } - return { quota: userQuota, fromUser: true }; - } - - async checkWorkspaceSeat(workspaceId: string, excludeSelf = false) { - const quota = await this.getWorkspaceUsage(workspaceId); - if (quota.memberCount - (excludeSelf ? 1 : 0) >= quota.memberLimit) { - throw new MemberQuotaExceeded(); - } - } - - // get workspace's owner quota and total size of used - // quota was apply to owner's account - async getWorkspaceUsage(workspaceId: string): Promise { - const owner = await this.permissions.getWorkspaceOwner(workspaceId); - const memberCount = - await this.permissions.getWorkspaceMemberCount(workspaceId); - const { - quota: { - name, - blobLimit, - businessBlobLimit, - historyPeriod, - memberLimit, - storageQuota, - copilotActionLimit, - humanReadable, - }, - fromUser, - } = await this.getWorkspaceQuota(owner.id, workspaceId); - - 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( - workspaceId, - FeatureType.UnlimitedWorkspace - ); - - const quota: QuotaBusinessType = { - name, - blobLimit, - businessBlobLimit, - historyPeriod, - memberLimit, - storageQuota, - copilotActionLimit, - humanReadable, - usedSize, - unlimited, - memberCount, - }; - - if (quota.unlimited) { - return this.mergeUnlimitedQuota(quota); - } - - return quota; - } - - private mergeUnlimitedQuota(orig: QuotaBusinessType): QuotaBusinessType { - return { - ...orig, - storageQuota: 1000 * OneGB, - memberLimit: 1000, - humanReadable: { - ...orig.humanReadable, - name: 'Unlimited', - storageQuota: formatSize(1000 * OneGB), - memberLimit: '1000', - }, - }; - } - - async checkBlobQuota(workspaceId: string, size: number) { - const { storageQuota, usedSize } = - await this.getWorkspaceUsage(workspaceId); - - return storageQuota - (size + usedSize); - } -} diff --git a/packages/backend/server/src/core/quota/types.ts b/packages/backend/server/src/core/quota/types.ts index 122be3bb33..fbe6f41ddc 100644 --- a/packages/backend/server/src/core/quota/types.ts +++ b/packages/backend/server/src/core/quota/types.ts @@ -1,142 +1,123 @@ import { Field, ObjectType } from '@nestjs/graphql'; import { SafeIntResolver } from 'graphql-scalars'; -import { z } from 'zod'; -import { commonFeatureSchema, FeatureKind } from '../features/types'; -import { ByteUnit, OneDay, OneKB } from './constant'; - -/// ======== quota define ======== - -/** - * naming rule: - * we append Vx to the end of the feature name to indicate the version of the feature - * x is a number, start from 1, this number will be change only at the time we change the schema of config - * for example, we change the value of `blobLimit` from 10MB to 100MB, then we will only change `version` field from 1 to 2 - * but if we remove the `blobLimit` field or rename it, then we will change the Vx to Vx+1 - */ -export enum QuotaType { - FreePlanV1 = 'free_plan_v1', - ProPlanV1 = 'pro_plan_v1', - TeamPlanV1 = 'team_plan_v1', - LifetimeProPlanV1 = 'lifetime_pro_plan_v1', - // only for test, smaller quota - RestrictedPlanV1 = 'restricted_plan_v1', -} - -const basicQuota = z.object({ - name: z.string(), - blobLimit: z.number().positive().int(), - storageQuota: z.number().positive().int(), - seatQuota: z.number().positive().int().nullish(), - historyPeriod: z.number().positive().int(), - memberLimit: z.number().positive().int(), - businessBlobLimit: z.number().positive().int().nullish(), -}); - -const userQuota = basicQuota.extend({ - copilotActionLimit: z.number().positive().int().nullish(), -}); - -const userQuotaPlan = z.object({ - feature: z.enum([ - QuotaType.FreePlanV1, - QuotaType.ProPlanV1, - QuotaType.LifetimeProPlanV1, - QuotaType.RestrictedPlanV1, - ]), - configs: userQuota, -}); - -const workspaceQuotaPlan = z.object({ - feature: z.enum([QuotaType.TeamPlanV1]), - configs: basicQuota, -}); - -/// ======== schema infer ======== - -export const QuotaSchema = commonFeatureSchema - .extend({ - type: z.literal(FeatureKind.Quota), - }) - .and(z.discriminatedUnion('feature', [userQuotaPlan, workspaceQuotaPlan])); - -export type Quota = z.infer< - typeof QuotaSchema -> & { feature: Q }; -export type QuotaConfigType = Quota['configs']; - -/// ======== query types ======== +import { UserQuota, WorkspaceQuota } from '../../models'; @ObjectType() -export class HumanReadableQuotaType { - @Field(() => String) +export class UserQuotaHumanReadableType { + @Field() name!: string; - @Field(() => String) + @Field() blobLimit!: string; - @Field(() => String) + @Field() storageQuota!: string; - @Field(() => String) + @Field() + usedStorageQuota!: string; + + @Field() historyPeriod!: string; - @Field(() => String) + @Field() memberLimit!: string; - @Field(() => String, { nullable: true }) - copilotActionLimit?: string; + @Field() + copilotActionLimit!: string; } @ObjectType() -export class QuotaQueryType { - @Field(() => String) +export class UserQuotaType implements UserQuota { + @Field() name!: string; @Field(() => SafeIntResolver) blobLimit!: number; + @Field(() => SafeIntResolver) + storageQuota!: number; + + @Field(() => SafeIntResolver) + usedStorageQuota!: number; + @Field(() => SafeIntResolver) historyPeriod!: number; - @Field(() => SafeIntResolver) + @Field() memberLimit!: number; + @Field(() => Number, { nullable: true }) + copilotActionLimit?: number; + + @Field(() => UserQuotaHumanReadableType) + humanReadable!: UserQuotaHumanReadableType; +} + +@ObjectType() +export class UserQuotaUsageType { + @Field(() => SafeIntResolver, { + name: 'storageQuota', + deprecationReason: "use `UserQuotaType['usedStorageQuota']` instead", + }) + storageQuota!: number; +} + +@ObjectType() +export class WorkspaceQuotaHumanReadableType { + @Field() + name!: string; + + @Field() + blobLimit!: string; + + @Field() + storageQuota!: string; + + @Field() + storageQuotaUsed!: string; + + @Field() + historyPeriod!: string; + + @Field() + memberLimit!: string; + + @Field() + memberCount!: string; +} + +@ObjectType() +export class WorkspaceQuotaType implements Partial { + @Field() + name!: string; + @Field(() => SafeIntResolver) - memberCount!: number; + blobLimit!: number; @Field(() => SafeIntResolver) storageQuota!: number; - @Field(() => SafeIntResolver, { nullable: true }) - copilotActionLimit?: number; - - @Field(() => HumanReadableQuotaType) - humanReadable!: HumanReadableQuotaType; + @Field(() => SafeIntResolver) + usedStorageQuota!: number; @Field(() => SafeIntResolver) + historyPeriod!: number; + + @Field() + memberLimit!: number; + + @Field() + memberCount!: number; + + @Field() + humanReadable!: WorkspaceQuotaHumanReadableType; + + /** + * @deprecated + */ + @Field(() => SafeIntResolver, { + deprecationReason: 'use `usedStorageQuota` instead', + }) usedSize!: number; } - -/// ======== utils ======== - -export function formatSize(bytes: number, decimals: number = 2): string { - if (bytes === 0) return '0 B'; - - const dm = decimals < 0 ? 0 : decimals; - - const i = Math.floor(Math.log(bytes) / Math.log(OneKB)); - - return ( - parseFloat((bytes / Math.pow(OneKB, i)).toFixed(dm)) + ' ' + ByteUnit[i] - ); -} - -export function formatDate(ms: number): string { - return `${(ms / OneDay).toFixed(0)} days`; -} - -export type QuotaBusinessType = QuotaQueryType & { - businessBlobLimit: number; - unlimited: boolean; -}; diff --git a/packages/backend/server/src/core/quota/utils.ts b/packages/backend/server/src/core/quota/utils.ts new file mode 100644 index 0000000000..b4688799d1 --- /dev/null +++ b/packages/backend/server/src/core/quota/utils.ts @@ -0,0 +1,19 @@ +import { OneDay, OneKB } from '../../base'; + +export const ByteUnit = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB']; + +export function formatSize(bytes: number, decimals: number = 2): string { + if (bytes === 0) return '0 B'; + + const dm = decimals < 0 ? 0 : decimals; + + const i = Math.floor(Math.log(bytes) / Math.log(OneKB)); + + return ( + parseFloat((bytes / Math.pow(OneKB, i)).toFixed(dm)) + ' ' + ByteUnit[i] + ); +} + +export function formatDate(ms: number): string { + return `${(ms / OneDay).toFixed(0)} days`; +} diff --git a/packages/backend/server/src/core/storage/wrappers/blob.ts b/packages/backend/server/src/core/storage/wrappers/blob.ts index caf8ca972b..5c3959fa55 100644 --- a/packages/backend/server/src/core/storage/wrappers/blob.ts +++ b/packages/backend/server/src/core/storage/wrappers/blob.ts @@ -112,8 +112,17 @@ export class WorkspaceBlobStorage { } async totalSize(workspaceId: string) { - const blobs = await this.list(workspaceId); - return blobs.reduce((acc, item) => acc + item.size, 0); + const sum = await this.db.blob.aggregate({ + where: { + workspaceId, + deletedAt: null, + }, + _sum: { + size: true, + }, + }); + + return sum._sum.size ?? 0; } private trySyncBlobsMeta(workspaceId: string, blobs: ListObjectsMetadata[]) { diff --git a/packages/backend/server/src/core/workspaces/index.ts b/packages/backend/server/src/core/workspaces/index.ts index 2aaa12a974..f284de71db 100644 --- a/packages/backend/server/src/core/workspaces/index.ts +++ b/packages/backend/server/src/core/workspaces/index.ts @@ -8,7 +8,6 @@ import { QuotaModule } from '../quota'; import { StorageModule } from '../storage'; import { UserModule } from '../user'; import { WorkspacesController } from './controller'; -import { WorkspaceManagementResolver } from './management'; import { DocHistoryResolver, PagePermissionResolver, @@ -32,7 +31,6 @@ import { providers: [ WorkspaceResolver, TeamWorkspaceResolver, - WorkspaceManagementResolver, PagePermissionResolver, DocHistoryResolver, WorkspaceBlobResolver, diff --git a/packages/backend/server/src/core/workspaces/management.ts b/packages/backend/server/src/core/workspaces/management.ts deleted file mode 100644 index 95209c87b9..0000000000 --- a/packages/backend/server/src/core/workspaces/management.ts +++ /dev/null @@ -1,98 +0,0 @@ -import { - Args, - Int, - Mutation, - Parent, - Query, - ResolveField, - Resolver, -} from '@nestjs/graphql'; - -import { ActionForbidden } from '../../base'; -import { CurrentUser } from '../auth'; -import { Admin } from '../common'; -import { FeatureManagementService, FeatureType } from '../features'; -import { PermissionService } from '../permission'; -import { WorkspaceFeatureType, WorkspaceType } from './types'; - -@Resolver(() => WorkspaceType) -export class WorkspaceManagementResolver { - constructor( - private readonly feature: FeatureManagementService, - private readonly permission: PermissionService - ) {} - - @Admin() - @Mutation(() => Int) - async addWorkspaceFeature( - @Args('workspaceId') workspaceId: string, - @Args('feature', { type: () => FeatureType }) feature: FeatureType - ): Promise { - return this.feature.addWorkspaceFeatures(workspaceId, feature); - } - - @Admin() - @Mutation(() => Int) - async removeWorkspaceFeature( - @Args('workspaceId') workspaceId: string, - @Args('feature', { type: () => FeatureType }) feature: FeatureType - ): Promise { - return this.feature.removeWorkspaceFeature(workspaceId, feature); - } - - @Admin() - @Query(() => [WorkspaceFeatureType]) - async listWorkspaceFeatures( - @Args('feature', { type: () => FeatureType }) feature: FeatureType - ): Promise { - return this.feature.listFeatureWorkspaces(feature); - } - - @Mutation(() => Boolean) - async setWorkspaceExperimentalFeature( - @CurrentUser() user: CurrentUser, - @Args('workspaceId') workspaceId: string, - @Args('feature', { type: () => FeatureType }) feature: FeatureType, - @Args('enable') enable: boolean - ): Promise { - if (!(await this.feature.canEarlyAccess(user.email))) { - throw new ActionForbidden(); - } - - const owner = await this.permission.getWorkspaceOwner(workspaceId); - const availableFeatures = await this.availableFeatures(user); - if (owner.id !== user.id || !availableFeatures.includes(feature)) { - throw new ActionForbidden(); - } - - if (enable) { - return await this.feature - .addWorkspaceFeatures( - workspaceId, - feature, - 'add by experimental feature api' - ) - .then(id => id > 0); - } else { - return await this.feature.removeWorkspaceFeature(workspaceId, feature); - } - } - - @ResolveField(() => [FeatureType], { - description: 'Available features of workspace', - complexity: 2, - }) - async availableFeatures( - @CurrentUser() user: CurrentUser - ): Promise { - return await this.feature.getActivatedUserFeatures(user.id); - } - - @ResolveField(() => [FeatureType], { - description: 'Enabled features of workspace', - complexity: 2, - }) - async features(@Parent() workspace: WorkspaceType): Promise { - return this.feature.getWorkspaceFeatures(workspace.id); - } -} diff --git a/packages/backend/server/src/core/workspaces/resolvers/blob.ts b/packages/backend/server/src/core/workspaces/resolvers/blob.ts index dac1d747de..a9633ce8fc 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/blob.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/blob.ts @@ -16,7 +16,7 @@ import type { FileUpload } from '../../../base'; import { BlobQuotaExceeded, CloudThrottlerGuard } from '../../../base'; import { CurrentUser } from '../../auth'; import { PermissionService, WorkspaceRole } from '../../permission'; -import { QuotaManagementService } from '../../quota'; +import { QuotaService } from '../../quota'; import { WorkspaceBlobStorage } from '../../storage'; import { WorkspaceBlobSizes, WorkspaceType } from '../types'; @@ -41,7 +41,7 @@ export class WorkspaceBlobResolver { logger = new Logger(WorkspaceBlobResolver.name); constructor( private readonly permissions: PermissionService, - private readonly quota: QuotaManagementService, + private readonly quota: QuotaService, private readonly storage: WorkspaceBlobStorage ) {} @@ -106,7 +106,7 @@ export class WorkspaceBlobResolver { ); const checkExceeded = - await this.quota.getQuotaCalculatorByWorkspace(workspaceId); + await this.quota.getWorkspaceQuotaCalculator(workspaceId); // TODO(@darksky): need a proper way to separate `BlobQuotaExceeded` and `BlobSizeTooLarge` if (checkExceeded(0)) { diff --git a/packages/backend/server/src/core/workspaces/resolvers/service.ts b/packages/backend/server/src/core/workspaces/resolvers/service.ts index 043039fe57..4955d39c82 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/service.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/service.ts @@ -148,6 +148,9 @@ export class WorkspaceService { } // ================ Team ================ + async isTeamWorkspace(workspaceId: string) { + return this.models.workspaceFeature.has(workspaceId, 'team_plan_v1'); + } async sendTeamWorkspaceUpgradedEmail(workspaceId: string) { const workspace = await this.getWorkspaceInfo(workspaceId); diff --git a/packages/backend/server/src/core/workspaces/resolvers/team.ts b/packages/backend/server/src/core/workspaces/resolvers/team.ts index 51dc10e600..a89e33de0c 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/team.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/team.ts @@ -10,6 +10,7 @@ import { PrismaClient, WorkspaceMemberStatus } from '@prisma/client'; import { nanoid } from 'nanoid'; import { + ActionForbiddenOnNonTeamWorkspace, Cache, EventBus, MemberNotFoundInSpace, @@ -22,7 +23,7 @@ import { import { Models } from '../../../models'; import { CurrentUser } from '../../auth'; import { PermissionService, WorkspaceRole } from '../../permission'; -import { QuotaManagementService } from '../../quota'; +import { QuotaService } from '../../quota'; import { InviteLink, InviteResult, @@ -47,7 +48,7 @@ export class TeamWorkspaceResolver { private readonly prisma: PrismaClient, private readonly permissions: PermissionService, private readonly models: Models, - private readonly quota: QuotaManagementService, + private readonly quota: QuotaService, private readonly mutex: RequestMutex, private readonly workspaceService: WorkspaceService ) {} @@ -58,7 +59,7 @@ export class TeamWorkspaceResolver { complexity: 2, }) team(@Parent() workspace: WorkspaceType) { - return this.quota.isTeamWorkspace(workspace.id); + return this.workspaceService.isTeamWorkspace(workspace.id); } @Mutation(() => [InviteResult]) @@ -85,7 +86,7 @@ export class TeamWorkspaceResolver { return new TooManyRequest(); } - const quota = await this.quota.getWorkspaceUsage(workspaceId); + const quota = await this.quota.getWorkspaceSeatQuota(workspaceId); const results = []; for (const [idx, email] of emails.entries()) { @@ -285,10 +286,20 @@ export class TeamWorkspaceResolver { @Args('userId') userId: string, @Args('permission', { type: () => WorkspaceRole }) permission: WorkspaceRole ) { + // non-team workspace can only transfer ownership, but no detailed permission control + if (permission !== WorkspaceRole.Owner) { + const isTeam = await this.workspaceService.isTeamWorkspace(workspaceId); + if (!isTeam) { + throw new ActionForbiddenOnNonTeamWorkspace(); + } + } + await this.permissions.checkWorkspace( workspaceId, user.id, - WorkspaceRole.Owner + permission >= WorkspaceRole.Admin + ? WorkspaceRole.Owner + : WorkspaceRole.Admin ); try { diff --git a/packages/backend/server/src/core/workspaces/resolvers/workspace.ts b/packages/backend/server/src/core/workspaces/resolvers/workspace.ts index f0d472ef4d..d735000fca 100644 --- a/packages/backend/server/src/core/workspaces/resolvers/workspace.ts +++ b/packages/backend/server/src/core/workspaces/resolvers/workspace.ts @@ -38,7 +38,7 @@ import { mapWorkspaceRoleToWorkspaceActions, WorkspacePermissionsList, } from '../../permission/types'; -import { QuotaManagementService, QuotaQueryType } from '../../quota'; +import { QuotaService, WorkspaceQuotaType } from '../../quota'; import { UserType } from '../../user'; import { InvitationType, @@ -122,7 +122,7 @@ export class WorkspaceResolver { private readonly cache: Cache, private readonly prisma: PrismaClient, private readonly permissions: PermissionService, - private readonly quota: QuotaManagementService, + private readonly quota: QuotaService, private readonly models: Models, private readonly event: EventBus, private readonly mutex: RequestMutex, @@ -260,13 +260,19 @@ export class WorkspaceResolver { }; } - @ResolveField(() => QuotaQueryType, { + @ResolveField(() => WorkspaceQuotaType, { name: 'quota', description: 'quota of workspace', complexity: 2, }) - workspaceQuota(@Parent() workspace: WorkspaceType) { - return this.quota.getWorkspaceUsage(workspace.id); + async workspaceQuota( + @Parent() workspace: WorkspaceType + ): Promise { + const quota = await this.quota.getWorkspaceQuotaWithUsage(workspace.id); + return { + ...quota, + humanReadable: this.quota.formatWorkspaceQuota(quota), + }; } @Query(() => Boolean, { @@ -421,12 +427,7 @@ export class WorkspaceResolver { @Args({ name: 'input', type: () => UpdateWorkspaceInput }) { id, ...updates }: UpdateWorkspaceInput ) { - const isTeam = await this.quota.isTeamWorkspace(id); - await this.permissions.checkWorkspace( - id, - user.id, - isTeam ? WorkspaceRole.Owner : WorkspaceRole.Admin - ); + await this.permissions.checkWorkspace(id, user.id, WorkspaceRole.Admin); return this.prisma.workspace.update({ where: { @@ -483,7 +484,7 @@ export class WorkspaceResolver { } // member limit check - await this.quota.checkWorkspaceSeat(workspaceId); + await this.quota.checkSeat(workspaceId); let target = await this.models.user.getUserByEmail(email); if (target) { @@ -569,14 +570,14 @@ export class WorkspaceResolver { @Args('workspaceId') workspaceId: string, @Args('userId') userId: string ) { - const isTeam = await this.quota.isTeamWorkspace(workspaceId); const isAdmin = await this.permissions.tryCheckWorkspaceIs( workspaceId, userId, WorkspaceRole.Admin ); - if (isTeam && isAdmin) { - // only owner can revoke team workspace admin + + if (isAdmin) { + // only owner can revoke workspace admin await this.permissions.checkWorkspaceIs( workspaceId, user.id, @@ -607,7 +608,6 @@ export class WorkspaceResolver { throw new TooManyRequest(); } - const isTeam = await this.quota.isTeamWorkspace(workspaceId); if (user) { const status = await this.permissions.getWorkspaceMemberStatus( workspaceId, @@ -622,8 +622,9 @@ export class WorkspaceResolver { `workspace:inviteLink:${workspaceId}` ); if (invite?.inviteId === inviteId) { - const quota = await this.quota.getWorkspaceUsage(workspaceId); - if (quota.memberCount >= quota.memberLimit) { + const isTeam = await this.workspaceService.isTeamWorkspace(workspaceId); + const seatAvailable = await this.quota.tryCheckSeat(workspaceId); + if (!seatAvailable) { // only team workspace allow over limit if (isTeam) { await this.permissions.grant( @@ -661,10 +662,6 @@ export class WorkspaceResolver { } } - // we added seats when sending invitation emails, but the payment may fail - // so we need to check seat again here - await this.quota.checkWorkspaceSeat(workspaceId, true); - if (sendAcceptMail) { const success = await this.workspaceService.sendAcceptedEmail(inviteId); if (!success) throw new UserNotFound(); diff --git a/packages/backend/server/src/data/commands/run.ts b/packages/backend/server/src/data/commands/run.ts index 33b4482b02..50f5016020 100644 --- a/packages/backend/server/src/data/commands/run.ts +++ b/packages/backend/server/src/data/commands/run.ts @@ -5,16 +5,18 @@ import { fileURLToPath, pathToFileURL } from 'node:url'; import { Logger } from '@nestjs/common'; import { ModuleRef } from '@nestjs/core'; import { PrismaClient } from '@prisma/client'; +import { once } from 'lodash-es'; import { Command, CommandRunner } from 'nest-commander'; interface Migration { file: string; name: string; + always?: boolean; up: (db: PrismaClient, injector: ModuleRef) => Promise; down: (db: PrismaClient, injector: ModuleRef) => Promise; } -export async function collectMigrations(): Promise { +export const collectMigrations = once(async () => { const folder = join(fileURLToPath(import.meta.url), '../../migrations'); const migrationFiles = readdirSync(folder) @@ -33,6 +35,7 @@ export async function collectMigrations(): Promise { return { file, name: migration.name, + always: migration.always, up: migration.up, down: migration.down, }; @@ -41,7 +44,8 @@ export async function collectMigrations(): Promise { ); return migrations; -} +}); + @Command({ name: 'run', description: 'Run all pending data migrations', @@ -65,7 +69,7 @@ export class RunCommand extends CommandRunner { }, }); - if (exists) { + if (exists && !migration.always) { continue; } @@ -100,8 +104,14 @@ export class RunCommand extends CommandRunner { private async runMigration(migration: Migration) { this.logger.log(`Running ${migration.name}...`); - const record = await this.db.dataMigration.create({ - data: { + const record = await this.db.dataMigration.upsert({ + where: { + name: migration.name, + }, + update: { + startedAt: new Date(), + }, + create: { name: migration.name, startedAt: new Date(), }, diff --git a/packages/backend/server/src/data/migrations/0001-refresh-features.ts b/packages/backend/server/src/data/migrations/0001-refresh-features.ts new file mode 100644 index 0000000000..42e16193ab --- /dev/null +++ b/packages/backend/server/src/data/migrations/0001-refresh-features.ts @@ -0,0 +1,16 @@ +import { ModuleRef } from '@nestjs/core'; +import { PrismaClient } from '@prisma/client'; + +import { FeatureModel } from '../../models'; + +export class RefreshFeatures0001 { + static always = true; + + // do the migration + static async up(_db: PrismaClient, ref: ModuleRef) { + await ref.get(FeatureModel, { strict: false }).refreshFeatures(); + } + + // revert the migration + static async down(_db: PrismaClient) {} +} diff --git a/packages/backend/server/src/data/migrations/1698652531198-user-features-init.ts b/packages/backend/server/src/data/migrations/1698652531198-user-features-init.ts deleted file mode 100644 index d15a00864e..0000000000 --- a/packages/backend/server/src/data/migrations/1698652531198-user-features-init.ts +++ /dev/null @@ -1,24 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { Features } from '../../core/features'; -import { Quotas } from '../../core/quota/schema'; -import { upsertFeature } from './utils/user-features'; - -export class UserFeaturesInit1698652531198 { - // do the migration - static async up(db: PrismaClient) { - // upgrade features from lower version to higher version - for (const feature of Features) { - await upsertFeature(db, feature); - } - - for (const quota of Quotas) { - await upsertFeature(db, quota); - } - } - - // revert the migration - static async down(_db: PrismaClient) { - // noop - } -} diff --git a/packages/backend/server/src/data/migrations/1702620653283-old-user-feature.ts b/packages/backend/server/src/data/migrations/1702620653283-old-user-feature.ts deleted file mode 100644 index 8fe978336d..0000000000 --- a/packages/backend/server/src/data/migrations/1702620653283-old-user-feature.ts +++ /dev/null @@ -1,36 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { QuotaType } from '../../core/quota/types'; -export class OldUserFeature1702620653283 { - // do the migration - static async up(db: PrismaClient) { - await db.$transaction(async tx => { - const latestFreePlan = await tx.feature.findFirstOrThrow({ - where: { feature: QuotaType.FreePlanV1 }, - orderBy: { version: 'desc' }, - select: { id: true }, - }); - - // find all users that don't have any features - const userIds = await db.user.findMany({ - where: { NOT: { features: { some: { NOT: { id: { gt: 0 } } } } } }, - select: { id: true }, - }); - - await tx.userFeature.createMany({ - data: userIds.map(({ id: userId }) => ({ - userId, - featureId: latestFreePlan.id, - reason: 'old user feature migration', - activated: true, - })), - }); - }); - } - - // revert the migration - // WARN: this will drop all user features - static async down(db: PrismaClient) { - await db.userFeature.deleteMany({}); - } -} diff --git a/packages/backend/server/src/data/migrations/1704352562369-refresh-user-features.ts b/packages/backend/server/src/data/migrations/1704352562369-refresh-user-features.ts deleted file mode 100644 index 69555563c1..0000000000 --- a/packages/backend/server/src/data/migrations/1704352562369-refresh-user-features.ts +++ /dev/null @@ -1,17 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { Features } from '../../core/features'; -import { upsertFeature } from './utils/user-features'; - -export class RefreshUserFeatures1704352562369 { - // do the migration - static async up(db: PrismaClient) { - // add early access v2 & copilot feature - for (const feature of Features) { - await upsertFeature(db, feature); - } - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1705395933447-new-free-plan.ts b/packages/backend/server/src/data/migrations/1705395933447-new-free-plan.ts deleted file mode 100644 index 51b869e9c7..0000000000 --- a/packages/backend/server/src/data/migrations/1705395933447-new-free-plan.ts +++ /dev/null @@ -1,16 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { Quotas } from '../../core/quota/schema'; -import { upgradeQuotaVersion } from './utils/user-quotas'; - -export class NewFreePlan1705395933447 { - // do the migration - static async up(db: PrismaClient) { - // free plan 1.0 - const quota = Quotas[3]; - await upgradeQuotaVersion(db, quota, 'free plan 1.0 migration'); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1706513866287-business-blob-limit.ts b/packages/backend/server/src/data/migrations/1706513866287-business-blob-limit.ts deleted file mode 100644 index 4c61590057..0000000000 --- a/packages/backend/server/src/data/migrations/1706513866287-business-blob-limit.ts +++ /dev/null @@ -1,16 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { Quotas } from '../../core/quota/schema'; -import { upgradeQuotaVersion } from './utils/user-quotas'; - -export class BusinessBlobLimit1706513866287 { - // do the migration - static async up(db: PrismaClient) { - // free plan 1.1 - const quota = Quotas[4]; - await upgradeQuotaVersion(db, quota, 'free plan 1.1 migration'); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1708321519830-refresh-unlimited-workspace-feature.ts b/packages/backend/server/src/data/migrations/1708321519830-refresh-unlimited-workspace-feature.ts deleted file mode 100644 index 347f90ded9..0000000000 --- a/packages/backend/server/src/data/migrations/1708321519830-refresh-unlimited-workspace-feature.ts +++ /dev/null @@ -1,15 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { FeatureType } from '../../core/features'; -import { upsertLatestFeatureVersion } from './utils/user-features'; - -export class RefreshUnlimitedWorkspaceFeature1708321519830 { - // do the migration - static async up(db: PrismaClient) { - // add unlimited workspace feature - await upsertLatestFeatureVersion(db, FeatureType.UnlimitedWorkspace); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1712224382221-refresh-free-plan.ts b/packages/backend/server/src/data/migrations/1712224382221-refresh-free-plan.ts deleted file mode 100644 index 5db27509a8..0000000000 --- a/packages/backend/server/src/data/migrations/1712224382221-refresh-free-plan.ts +++ /dev/null @@ -1,19 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { QuotaType } from '../../core/quota/types'; -import { upgradeLatestQuotaVersion } from './utils/user-quotas'; - -export class RefreshFreePlan1712224382221 { - // do the migration - static async up(db: PrismaClient) { - // free plan 1.1 - await upgradeLatestQuotaVersion( - db, - QuotaType.FreePlanV1, - 'free plan 1.1 migration' - ); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1713164714634-copilot-feature.ts b/packages/backend/server/src/data/migrations/1713164714634-copilot-feature.ts deleted file mode 100644 index 9b6e2033b3..0000000000 --- a/packages/backend/server/src/data/migrations/1713164714634-copilot-feature.ts +++ /dev/null @@ -1,23 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { QuotaType } from '../../core/quota/types'; -import { upgradeLatestQuotaVersion } from './utils/user-quotas'; - -export class CopilotFeature1713164714634 { - // do the migration - static async up(db: PrismaClient) { - await upgradeLatestQuotaVersion( - db, - QuotaType.ProPlanV1, - 'pro plan 1.1 migration' - ); - await upgradeLatestQuotaVersion( - db, - QuotaType.RestrictedPlanV1, - 'restricted plan 1.1 migration' - ); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1713176777814-ai-early-access.ts b/packages/backend/server/src/data/migrations/1713176777814-ai-early-access.ts deleted file mode 100644 index 058c0cceef..0000000000 --- a/packages/backend/server/src/data/migrations/1713176777814-ai-early-access.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 AiEarlyAccess1713176777814 { - // do the migration - static async up(db: PrismaClient) { - await upsertLatestFeatureVersion(db, FeatureType.AIEarlyAccess); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1713285638427-unlimited-copilot.ts b/packages/backend/server/src/data/migrations/1713285638427-unlimited-copilot.ts deleted file mode 100644 index c2521302c5..0000000000 --- a/packages/backend/server/src/data/migrations/1713285638427-unlimited-copilot.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 UnlimitedCopilot1713285638427 { - // do the migration - static async up(db: PrismaClient) { - await upsertLatestFeatureVersion(db, FeatureType.UnlimitedCopilot); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1716195522794-administrator-feature.ts b/packages/backend/server/src/data/migrations/1716195522794-administrator-feature.ts deleted file mode 100644 index 904363a397..0000000000 --- a/packages/backend/server/src/data/migrations/1716195522794-administrator-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 AdministratorFeature1716195522794 { - // do the migration - static async up(db: PrismaClient) { - await upsertLatestFeatureVersion(db, FeatureType.Admin); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1719917815802-lifetime-pro-quota.ts b/packages/backend/server/src/data/migrations/1719917815802-lifetime-pro-quota.ts deleted file mode 100644 index a673b7f466..0000000000 --- a/packages/backend/server/src/data/migrations/1719917815802-lifetime-pro-quota.ts +++ /dev/null @@ -1,14 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { QuotaType } from '../../core/quota'; -import { upsertLatestQuotaVersion } from './utils/user-quotas'; - -export class LifetimeProQuota1719917815802 { - // do the migration - static async up(db: PrismaClient) { - await upsertLatestQuotaVersion(db, QuotaType.LifetimeProPlanV1); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1733804966417-team-quota.ts b/packages/backend/server/src/data/migrations/1733804966417-team-quota.ts deleted file mode 100644 index 6fa3ca105d..0000000000 --- a/packages/backend/server/src/data/migrations/1733804966417-team-quota.ts +++ /dev/null @@ -1,14 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { QuotaType } from '../../core/quota'; -import { upsertLatestQuotaVersion } from './utils/user-quotas'; - -export class TeamQuota1733804966417 { - // do the migration - static async up(db: PrismaClient) { - await upsertLatestQuotaVersion(db, QuotaType.TeamPlanV1); - } - - // revert the migration - static async down(_db: PrismaClient) {} -} diff --git a/packages/backend/server/src/data/migrations/1738590347632-feature-redundant.ts b/packages/backend/server/src/data/migrations/1738590347632-feature-redundant.ts new file mode 100644 index 0000000000..cf6c01f7f7 --- /dev/null +++ b/packages/backend/server/src/data/migrations/1738590347632-feature-redundant.ts @@ -0,0 +1,57 @@ +import { PrismaClient } from '@prisma/client'; + +import { FeatureConfigs, FeatureName, FeatureType } from '../../models'; + +export class FeatureRedundant1738590347632 { + // do the migration + static async up(db: PrismaClient) { + const features = await db.feature.findMany(); + const validFeatures = new Map< + number, + { + name: string; + type: FeatureType; + } + >(); + + for (const feature of features) { + const def = FeatureConfigs[feature.name as FeatureName]; + if (!def || def.deprecatedVersion !== feature.deprecatedVersion) { + await db.feature.delete({ + where: { id: feature.id }, + }); + } else { + validFeatures.set(feature.id, { + name: feature.name, + type: def.type, + }); + } + } + + for (const [id, def] of validFeatures.entries()) { + await db.userFeature.updateMany({ + where: { + featureId: id, + }, + data: { + name: def.name, + type: def.type, + }, + }); + await db.workspaceFeature.updateMany({ + where: { + featureId: id, + }, + data: { + name: def.name, + type: def.type, + }, + }); + } + } + + // revert the migration + static async down(_db: PrismaClient) { + // noop + } +} diff --git a/packages/backend/server/src/data/migrations/99999-self-host-admin.ts b/packages/backend/server/src/data/migrations/99999-self-host-admin.ts deleted file mode 100644 index f21b022488..0000000000 --- a/packages/backend/server/src/data/migrations/99999-self-host-admin.ts +++ /dev/null @@ -1,29 +0,0 @@ -import { ModuleRef } from '@nestjs/core'; -import { PrismaClient } from '@prisma/client'; - -import { Config } from '../../base'; -import { FeatureManagementService } from '../../core/features'; - -export class SelfHostAdmin1 { - // do the migration - static async up(db: PrismaClient, ref: ModuleRef) { - const config = ref.get(Config, { strict: false }); - if (config.isSelfhosted) { - const feature = ref.get(FeatureManagementService, { strict: false }); - - const firstUser = await db.user.findFirst({ - orderBy: { - createdAt: 'asc', - }, - }); - if (firstUser) { - await feature.addAdmin(firstUser.id); - } - } - } - - // revert the migration - static async down() { - // - } -} diff --git a/packages/backend/server/src/data/migrations/utils/user-features.ts b/packages/backend/server/src/data/migrations/utils/user-features.ts deleted file mode 100644 index 0f30cb6118..0000000000 --- a/packages/backend/server/src/data/migrations/utils/user-features.ts +++ /dev/null @@ -1,40 +0,0 @@ -import { Prisma, PrismaClient } from '@prisma/client'; - -import { CommonFeature, Features, FeatureType } from '../../../core/features'; - -// upgrade features from lower version to higher version -export async function upsertFeature( - db: PrismaClient, - feature: CommonFeature -): Promise { - const hasEqualOrGreaterVersion = - (await db.feature.count({ - where: { - feature: feature.feature, - version: { - gte: feature.version, - }, - }, - })) > 0; - // will not update exists version - if (!hasEqualOrGreaterVersion) { - await db.feature.create({ - data: { - feature: feature.feature, - type: feature.type, - version: feature.version, - configs: feature.configs as Prisma.InputJsonValue, - }, - }); - } -} - -export async function upsertLatestFeatureVersion( - db: PrismaClient, - type: FeatureType -) { - const feature = Features.filter(f => f.feature === type); - feature.sort((a, b) => b.version - a.version); - const latestFeature = feature[0]; - await upsertFeature(db, latestFeature); -} diff --git a/packages/backend/server/src/data/migrations/utils/user-quotas.ts b/packages/backend/server/src/data/migrations/utils/user-quotas.ts deleted file mode 100644 index 771ab46f0e..0000000000 --- a/packages/backend/server/src/data/migrations/utils/user-quotas.ts +++ /dev/null @@ -1,89 +0,0 @@ -import { PrismaClient } from '@prisma/client'; - -import { FeatureKind } from '../../../core/features'; -import { getLatestQuota } from '../../../core/quota/schema'; -import { Quota, QuotaType } from '../../../core/quota/types'; -import { upsertFeature } from './user-features'; - -export async function upgradeQuotaVersion( - db: PrismaClient, - quota: Quota, - reason: string -) { - // add new quota - await upsertFeature(db, quota); - // migrate all users that using old quota to new quota - await db.$transaction( - async tx => { - const latestQuotaVersion = await tx.feature.findFirstOrThrow({ - where: { feature: quota.feature }, - orderBy: { version: 'desc' }, - select: { id: true }, - }); - - // find all users that have old free plan - const userIds = await tx.user.findMany({ - where: { - features: { - some: { - feature: { - type: FeatureKind.Quota, - feature: quota.feature, - version: { lt: quota.version }, - }, - activated: true, - }, - }, - }, - select: { id: true }, - }); - - // deactivate all old quota for the user - await tx.userFeature.updateMany({ - where: { - id: undefined, - userId: { - in: userIds.map(({ id }) => id), - }, - feature: { - type: FeatureKind.Quota, - }, - activated: true, - }, - data: { - activated: false, - }, - }); - - await tx.userFeature.createMany({ - data: userIds.map(({ id: userId }) => ({ - userId, - featureId: latestQuotaVersion.id, - reason, - activated: true, - })), - }); - }, - { - maxWait: 10000, - timeout: 20000, - } - ); -} - -export async function upsertLatestQuotaVersion( - db: PrismaClient, - type: QuotaType -) { - const latestQuota = getLatestQuota(type); - await upsertFeature(db, latestQuota); -} - -export async function upgradeLatestQuotaVersion( - db: PrismaClient, - type: QuotaType, - reason: string -) { - const latestQuota = getLatestQuota(type); - await upgradeQuotaVersion(db, latestQuota, reason); -} diff --git a/packages/backend/server/src/models/common/feature.ts b/packages/backend/server/src/models/common/feature.ts index 055c1959cd..1fc5117fa9 100644 --- a/packages/backend/server/src/models/common/feature.ts +++ b/packages/backend/server/src/models/common/feature.ts @@ -1,63 +1,75 @@ import { z } from 'zod'; -export enum FeatureType { - Feature = 0, - Quota = 1, -} +import { OneDay, OneGB, OneMB } from '../../base'; -// TODO(@forehalo): quota is a useless extra concept, merge it with feature -export const UserPlanQuotaConfig = z.object({ +const UserPlanQuotaConfig = z.object({ // quota name name: z.string(), // single blob limit blobLimit: z.number(), + // server limit will larger then client to handle a edge case: + // when a user downgrades from pro to free, he can still continue + // to upload previously added files that exceed the free limit + // NOTE: this is a product decision, may change in future + businessBlobLimit: z.number().optional(), // total blob limit storageQuota: z.number(), // history period of validity historyPeriod: z.number(), // member limit memberLimit: z.number(), - // copilot action limit 10 - copilotActionLimit: z.number(), + // copilot action limit + copilotActionLimit: z.number().optional(), }); -export const WorkspaceQuotaConfig = UserPlanQuotaConfig.extend({ +export type UserQuota = z.infer; + +const WorkspaceQuotaConfig = UserPlanQuotaConfig.extend({ // seat quota seatQuota: z.number(), }).omit({ copilotActionLimit: true, }); -function feature(configs: z.ZodObject) { - return z.object({ - type: z.literal(FeatureType.Feature), - configs: configs, - }); +export type WorkspaceQuota = z.infer; + +const EMPTY_CONFIG = z.object({}); + +export enum FeatureType { + Feature, + Quota, } -function quota(configs: z.ZodObject) { - return z.object({ - type: z.literal(FeatureType.Quota), - configs: configs, - }); +export enum Feature { + // user + Admin = 'administrator', + EarlyAccess = 'early_access', + AIEarlyAccess = 'ai_early_access', + UnlimitedCopilot = 'unlimited_copilot', + FreePlan = 'free_plan_v1', + ProPlan = 'pro_plan_v1', + LifetimeProPlan = 'lifetime_pro_plan_v1', + + // workspace + UnlimitedWorkspace = 'unlimited_workspace', + TeamPlan = 'team_plan_v1', } -export const Features = { - copilot: feature(z.object({})), - early_access: feature(z.object({ whitelist: z.array(z.string()) })), - unlimited_workspace: feature(z.object({})), - unlimited_copilot: feature(z.object({})), - ai_early_access: feature(z.object({})), - administrator: feature(z.object({})), - free_plan_v1: quota(UserPlanQuotaConfig), - pro_plan_v1: quota(UserPlanQuotaConfig), - lifetime_pro_plan_v1: quota(UserPlanQuotaConfig), - restricted_plan_v1: quota(UserPlanQuotaConfig), - team_plan_v1: quota(WorkspaceQuotaConfig), -}; +// TODO(@forehalo): may merge `FeatureShapes` and `FeatureConfigs`? +export const FeaturesShapes = { + early_access: z.object({ whitelist: z.array(z.string()) }), + unlimited_workspace: EMPTY_CONFIG, + unlimited_copilot: EMPTY_CONFIG, + ai_early_access: EMPTY_CONFIG, + administrator: EMPTY_CONFIG, + free_plan_v1: UserPlanQuotaConfig, + pro_plan_v1: UserPlanQuotaConfig, + lifetime_pro_plan_v1: UserPlanQuotaConfig, + team_plan_v1: WorkspaceQuotaConfig, +} satisfies Record>; export type UserFeatureName = keyof Pick< - typeof Features, + typeof FeaturesShapes, | 'early_access' | 'ai_early_access' | 'unlimited_copilot' @@ -65,14 +77,97 @@ export type UserFeatureName = keyof Pick< | 'free_plan_v1' | 'pro_plan_v1' | 'lifetime_pro_plan_v1' - | 'restricted_plan_v1' >; export type WorkspaceFeatureName = keyof Pick< - typeof Features, + typeof FeaturesShapes, 'unlimited_workspace' | 'team_plan_v1' >; export type FeatureName = UserFeatureName | WorkspaceFeatureName; -export type FeatureConfigs = z.infer< - (typeof Features)[T]['shape']['configs'] +export type FeatureConfig = z.infer< + (typeof FeaturesShapes)[T] >; + +export const FeatureConfigs: { + [K in FeatureName]: { + type: FeatureType; + configs: FeatureConfig; + deprecatedVersion: number; + }; +} = { + free_plan_v1: { + type: FeatureType.Quota, + deprecatedVersion: 4, + configs: { + // quota name + name: 'Free', + blobLimit: 10 * OneMB, + businessBlobLimit: 100 * OneMB, + storageQuota: 10 * OneGB, + historyPeriod: 7 * OneDay, + memberLimit: 3, + copilotActionLimit: 10, + }, + }, + pro_plan_v1: { + type: FeatureType.Quota, + deprecatedVersion: 2, + configs: { + name: 'Pro', + blobLimit: 100 * OneMB, + storageQuota: 100 * OneGB, + historyPeriod: 30 * OneDay, + memberLimit: 10, + copilotActionLimit: 10, + }, + }, + lifetime_pro_plan_v1: { + type: FeatureType.Quota, + deprecatedVersion: 1, + configs: { + name: 'Lifetime Pro', + blobLimit: 100 * OneMB, + storageQuota: 1024 * OneGB, + historyPeriod: 30 * OneDay, + memberLimit: 10, + copilotActionLimit: 10, + }, + }, + team_plan_v1: { + type: FeatureType.Quota, + deprecatedVersion: 1, + configs: { + name: 'Team Workspace', + blobLimit: 500 * OneMB, + storageQuota: 100 * OneGB, + seatQuota: 20 * OneGB, + historyPeriod: 30 * OneDay, + memberLimit: 1, + }, + }, + early_access: { + type: FeatureType.Feature, + deprecatedVersion: 2, + configs: { whitelist: [] }, + }, + unlimited_workspace: { + type: FeatureType.Feature, + deprecatedVersion: 1, + configs: {}, + }, + unlimited_copilot: { + type: FeatureType.Feature, + deprecatedVersion: 1, + configs: {}, + }, + ai_early_access: { + type: FeatureType.Feature, + deprecatedVersion: 1, + configs: {}, + }, + administrator: { + type: FeatureType.Feature, + deprecatedVersion: 1, + configs: {}, + }, +}; diff --git a/packages/backend/server/src/models/feature.ts b/packages/backend/server/src/models/feature.ts index 6f7803b8a2..6e87405e4a 100644 --- a/packages/backend/server/src/models/feature.ts +++ b/packages/backend/server/src/models/feature.ts @@ -5,9 +5,10 @@ import { z } from 'zod'; import { BaseModel } from './base'; import { - type FeatureConfigs, + type FeatureConfig, + FeatureConfigs, type FeatureName, - Features, + FeaturesShapes, FeatureType, } from './common'; @@ -29,7 +30,12 @@ export class FeatureModel extends BaseModel { } @Transactional() - async upsert(name: T, configs: FeatureConfigs) { + async upsert( + name: T, + configs: FeatureConfig, + deprecatedType: FeatureType, + deprecatedVersion: number + ) { const parsedConfigs = this.check(name, configs); // TODO(@forehalo): @@ -41,8 +47,9 @@ export class FeatureModel extends BaseModel { if (!latest) { feature = await this.db.feature.create({ data: { - type: FeatureType.Feature, - feature: name, + name, + deprecatedType, + deprecatedVersion, configs: parsedConfigs, }, }); @@ -57,7 +64,7 @@ export class FeatureModel extends BaseModel { this.logger.verbose(`Feature ${name} upserted`); - return feature as Feature & { configs: FeatureConfigs }; + return feature as Feature & { configs: FeatureConfig }; } /** @@ -67,8 +74,7 @@ export class FeatureModel extends BaseModel { */ async try_get_unchecked(name: T) { const feature = await this.db.feature.findFirst({ - where: { feature: name }, - orderBy: { version: 'desc' }, + where: { name }, }); return feature as Omit & { @@ -104,10 +110,22 @@ export class FeatureModel extends BaseModel { }); } - return parseResult.data as FeatureConfigs; + return parseResult.data as FeatureConfig; } getConfigShape(name: FeatureName): z.ZodObject { - return Features[name]?.shape.configs ?? z.object({}); + return FeaturesShapes[name] ?? z.object({}); + } + + getFeatureType(name: FeatureName): FeatureType { + return FeatureConfigs[name].type; + } + + async refreshFeatures() { + for (const key in FeatureConfigs) { + const name = key as FeatureName; + const def = FeatureConfigs[name]; + await this.upsert(name, def.configs, def.type, def.deprecatedVersion); + } } } diff --git a/packages/backend/server/src/models/index.ts b/packages/backend/server/src/models/index.ts index 7e7ccb38b4..65522a0f71 100644 --- a/packages/backend/server/src/models/index.ts +++ b/packages/backend/server/src/models/index.ts @@ -77,6 +77,7 @@ const ModelsSymbolProvider: ExistingProvider = { }) export class ModelsModule {} +export * from './common'; export * from './feature'; export * from './page'; export * from './session'; diff --git a/packages/backend/server/src/models/user-feature.ts b/packages/backend/server/src/models/user-feature.ts index 356bad49c1..d8bc5a6ecb 100644 --- a/packages/backend/server/src/models/user-feature.ts +++ b/packages/backend/server/src/models/user-feature.ts @@ -1,47 +1,49 @@ import { Injectable } from '@nestjs/common'; import { Transactional } from '@nestjs-cls/transactional'; +import { Prisma } from '@prisma/client'; import { BaseModel } from './base'; -import type { UserFeatureName } from './common'; +import { FeatureType, type UserFeatureName } from './common'; @Injectable() export class UserFeatureModel extends BaseModel { async get(userId: string, name: T) { - // TODO(@forehalo): - // all feature query-and-use queries should be simplified like the below when `feature(name)` becomes a unique index - // - // this.db.userFeature.findFirst({ - // include: { - // feature: true - // }, - // where: { - // userId, - // activated: true, - // feature: { - // feature: name, - // } - // } - // }) - const feature = await this.models.feature.get(name); - const count = await this.db.userFeature.count({ where: { userId, - featureId: feature.id, + name, activated: true, }, }); - return count > 0 ? feature : null; + if (count === 0) { + return null; + } + + return await this.models.feature.get(name); + } + + async getQuota(userId: string) { + const quota = await this.db.userFeature.findFirst({ + where: { + userId, + type: FeatureType.Quota, + activated: true, + }, + }); + + if (!quota) { + return null; + } + + return await this.models.feature.get<'free_plan_v1'>(quota.name as any); } async has(userId: string, name: UserFeatureName) { - const feature = await this.models.feature.get_unchecked(name); - const count = await this.db.userFeature.count({ where: { userId, - featureId: feature.id, + name, activated: true, }, }); @@ -49,29 +51,37 @@ export class UserFeatureModel extends BaseModel { return count > 0; } - async list(userId: string) { + async list(userId: string, type?: FeatureType) { + const filter: Prisma.UserFeatureWhereInput = + type === undefined + ? { + userId, + activated: true, + } + : { + userId, + activated: true, + type, + }; + const userFeatures = await this.db.userFeature.findMany({ - include: { - feature: true, - }, - where: { - userId, - activated: true, + where: filter, + select: { + name: true, }, }); return userFeatures.map( - userFeature => userFeature.feature.feature + userFeature => userFeature.name ) as UserFeatureName[]; } - async add(userId: string, featureName: UserFeatureName, reason: string) { - const feature = await this.models.feature.get_unchecked(featureName); - + async add(userId: string, name: UserFeatureName, reason: string) { + const feature = await this.models.feature.get_unchecked(name); const existing = await this.db.userFeature.findFirst({ where: { userId, - featureId: feature.id, + name: name, activated: true, }, }); @@ -84,37 +94,56 @@ export class UserFeatureModel extends BaseModel { data: { userId, featureId: feature.id, + name, + type: this.models.feature.getFeatureType(name), activated: true, reason, }, }); - this.logger.verbose(`Feature ${featureName} added to user ${userId}`); + this.logger.verbose(`Feature ${name} added to user ${userId}`); return userFeature; } async remove(userId: string, featureName: UserFeatureName) { - const feature = await this.models.feature.get_unchecked(featureName); - - await this.db.userFeature.deleteMany({ + await this.db.userFeature.updateMany({ where: { userId, - featureId: feature.id, + name: featureName, + }, + data: { + activated: false, }, }); - this.logger.verbose(`Feature ${featureName} removed from user ${userId}`); + this.logger.verbose( + `Feature ${featureName} deactivated for user ${userId}` + ); } @Transactional() - async switch( - userId: string, - from: UserFeatureName, - to: UserFeatureName, - reason: string - ) { - await this.remove(userId, from); + async switchQuota(userId: string, to: UserFeatureName, reason: string) { + const quotas = await this.list(userId, FeatureType.Quota); + + // deactivate the previous quota + if (quotas.length) { + // db state error + if (quotas.length > 1) { + this.logger.error( + `User ${userId} has multiple quotas, please check the database state.` + ); + } + + const from = quotas.at(-1) as UserFeatureName; + + if (from === to) { + return; + } + + await this.remove(userId, from); + } + await this.add(userId, to, reason); } } diff --git a/packages/backend/server/src/models/user.ts b/packages/backend/server/src/models/user.ts index 021b77958c..b122b3ab84 100644 --- a/packages/backend/server/src/models/user.ts +++ b/packages/backend/server/src/models/user.ts @@ -9,7 +9,6 @@ import { WrongSignInCredentials, WrongSignInMethod, } from '../base'; -import { Quota_FreePlanV1_1 } from '../core/quota/schema'; import { BaseModel } from './base'; import type { Workspace } from './workspace'; @@ -22,23 +21,6 @@ const publicUserSelect = { type CreateUserInput = Omit & { name?: string }; type UpdateUserInput = Omit, 'id'>; -const defaultUserCreatingData = { - name: 'Unnamed', - // TODO(@forehalo): it's actually a external dependency for user - // how could we avoid user model's knowledge of feature? - features: { - create: { - reason: 'sign up', - activated: true, - feature: { - connect: { - feature_version: Quota_FreePlanV1_1, - }, - }, - }, - }, -}; - declare global { interface Events { 'user.created': User; @@ -48,6 +30,7 @@ declare global { // dealing of owned workspaces of deleted users to workspace model ownedWorkspaces: Workspace['id'][]; }; + 'user.postCreated': User; } } @@ -134,17 +117,16 @@ export class UserModel extends BaseModel { data.password = await this.crypto.encryptPassword(data.password); } - if (!data.name) { - data.name = data.email.split('@')[0]; - } - user = await this.db.user.create({ data: { - ...defaultUserCreatingData, ...data, + name: data.name ?? data.email.split('@')[0], }, }); + // delegate the responsibility of finish user creating setup to the corresponding models + await this.event.emitAsync('user.postCreated', user); + this.logger.debug(`User [${user.id}] created with email [${user.email}]`); this.event.emit('user.created', user); diff --git a/packages/backend/server/src/models/workspace-feature.ts b/packages/backend/server/src/models/workspace-feature.ts index 34f89f7433..90d82b6bee 100644 --- a/packages/backend/server/src/models/workspace-feature.ts +++ b/packages/backend/server/src/models/workspace-feature.ts @@ -1,18 +1,21 @@ import { Injectable } from '@nestjs/common'; import { Transactional } from '@nestjs-cls/transactional'; +import { Prisma } from '@prisma/client'; import { BaseModel } from './base'; -import type { FeatureConfigs, WorkspaceFeatureName } from './common'; +import { + type FeatureConfig, + FeatureType, + type WorkspaceFeatureName, +} from './common'; @Injectable() export class WorkspaceFeatureModel extends BaseModel { async get(workspaceId: string, name: T) { - const feature = await this.models.feature.get_unchecked(name); - const workspaceFeature = await this.db.workspaceFeature.findFirst({ where: { workspaceId, - featureId: feature.id, + name, activated: true, }, }); @@ -21,6 +24,8 @@ export class WorkspaceFeatureModel extends BaseModel { return null; } + const feature = await this.models.feature.get_unchecked(name); + return { ...feature, configs: this.models.feature.check(name, { @@ -30,13 +35,47 @@ export class WorkspaceFeatureModel extends BaseModel { }; } - async has(workspaceId: string, name: WorkspaceFeatureName) { - const feature = await this.models.feature.get_unchecked(name); + async getQuota(workspaceId: string) { + const quota = await this.db.workspaceFeature.findFirst({ + where: { + workspaceId, + type: FeatureType.Quota, + activated: true, + }, + orderBy: { + createdAt: 'desc', + }, + }); + if (!quota) { + return null; + } + + const rawFeature = await this.models.feature.get_unchecked( + quota.name as WorkspaceFeatureName + ); + + const feature = { + ...rawFeature, + configs: this.models.feature.check(quota.name as 'team_plan_v1', { + ...rawFeature.configs, + ...(quota?.configs as {}), + }), + }; + + // workspace's storage quota is the sum of base quota and seats * quota per seat + feature.configs.storageQuota = + feature.configs.seatQuota * feature.configs.memberLimit + + feature.configs.storageQuota; + + return feature; + } + + async has(workspaceId: string, name: WorkspaceFeatureName) { const count = await this.db.workspaceFeature.count({ where: { workspaceId, - featureId: feature.id, + name, activated: true, }, }); @@ -44,35 +83,62 @@ export class WorkspaceFeatureModel extends BaseModel { return count > 0; } - async list(workspaceId: string) { + /** + * helper function to check if a list of workspaces have a standalone quota feature when calculating owner's quota usage + */ + async batchHasQuota(workspaceIds: string[]) { const workspaceFeatures = await this.db.workspaceFeature.findMany({ - include: { - feature: true, + select: { + workspaceId: true, }, where: { - workspaceId, + workspaceId: { in: workspaceIds }, + type: FeatureType.Quota, activated: true, }, }); + return workspaceFeatures.map(feature => feature.workspaceId); + } + + async list(workspaceId: string, type?: FeatureType) { + const filter: Prisma.WorkspaceFeatureWhereInput = + type === undefined + ? { + workspaceId, + activated: true, + } + : { + workspaceId, + activated: true, + type, + }; + + const workspaceFeatures = await this.db.workspaceFeature.findMany({ + select: { + name: true, + }, + where: filter, + }); + return workspaceFeatures.map( - workspaceFeature => workspaceFeature.feature.feature + workspaceFeature => workspaceFeature.name ) as WorkspaceFeatureName[]; } @Transactional() async add( workspaceId: string, - featureName: T, + name: T, reason: string, - overrides?: Partial> + overrides?: Partial> ) { - const feature = await this.models.feature.get_unchecked(featureName); + const feature = await this.models.feature.get_unchecked(name); const existing = await this.db.workspaceFeature.findFirst({ where: { workspaceId, - featureId: feature.id, + name: name, activated: true, }, }); @@ -87,12 +153,12 @@ export class WorkspaceFeatureModel extends BaseModel { }; const parseResult = this.models.feature - .getConfigShape(featureName) + .getConfigShape(name) .partial() .safeParse(configs); if (!parseResult.success) { - throw new Error(`Invalid feature config for ${featureName}`, { + throw new Error(`Invalid feature config for ${name}`, { cause: parseResult.error, }); } @@ -113,6 +179,8 @@ export class WorkspaceFeatureModel extends BaseModel { data: { workspaceId, featureId: feature.id, + name, + type: this.models.feature.getFeatureType(name), activated: true, reason, configs: parseResult.data, @@ -120,20 +188,16 @@ export class WorkspaceFeatureModel extends BaseModel { }); } - this.logger.verbose( - `Feature ${featureName} added to workspace ${workspaceId}` - ); + this.logger.verbose(`Feature ${name} added to workspace ${workspaceId}`); return workspaceFeature; } async remove(workspaceId: string, featureName: WorkspaceFeatureName) { - const feature = await this.models.feature.get_unchecked(featureName); - await this.db.workspaceFeature.deleteMany({ where: { workspaceId, - featureId: feature.id, + name: featureName, }, }); @@ -141,16 +205,4 @@ export class WorkspaceFeatureModel extends BaseModel { `Feature ${featureName} removed from workspace ${workspaceId}` ); } - - @Transactional() - async switch( - workspaceId: string, - from: WorkspaceFeatureName, - to: T, - reason: string, - overrides?: Partial> - ) { - await this.remove(workspaceId, from); - return await this.add(workspaceId, to, reason, overrides); - } } diff --git a/packages/backend/server/src/plugins/copilot/session.ts b/packages/backend/server/src/plugins/copilot/session.ts index bb126bc747..ffa8c1a9e2 100644 --- a/packages/backend/server/src/plugins/copilot/session.ts +++ b/packages/backend/server/src/plugins/copilot/session.ts @@ -12,8 +12,8 @@ import { CopilotSessionNotFound, PrismaTransaction, } from '../../base'; -import { FeatureManagementService } from '../../core/features'; import { QuotaService } from '../../core/quota'; +import { Models } from '../../models'; import { ChatMessageCache } from './message'; import { PromptService } from './prompt'; import { @@ -195,10 +195,10 @@ export class ChatSessionService { constructor( private readonly db: PrismaClient, - private readonly feature: FeatureManagementService, private readonly quota: QuotaService, private readonly messageCache: ChatMessageCache, - private readonly prompt: PromptService + private readonly prompt: PromptService, + private readonly models: Models ) {} private async haveSession( @@ -545,12 +545,15 @@ export class ChatSessionService { } async getQuota(userId: string) { - const isCopilotUser = await this.feature.isCopilotUser(userId); + const isCopilotUser = await this.models.userFeature.has( + userId, + 'unlimited_copilot' + ); let limit: number | undefined; if (!isCopilotUser) { const quota = await this.quota.getUserQuota(userId); - limit = quota.feature.copilotActionLimit; + limit = quota.copilotActionLimit; } const used = await this.countUserMessages(userId); diff --git a/packages/backend/server/src/plugins/copilot/storage.ts b/packages/backend/server/src/plugins/copilot/storage.ts index e592274321..c004b911c7 100644 --- a/packages/backend/server/src/plugins/copilot/storage.ts +++ b/packages/backend/server/src/plugins/copilot/storage.ts @@ -12,7 +12,7 @@ import { StorageProviderFactory, URLHelper, } from '../../base'; -import { QuotaManagementService } from '../../core/quota'; +import { QuotaService } from '../../core/quota'; @Injectable() export class CopilotStorage { @@ -22,7 +22,7 @@ export class CopilotStorage { private readonly config: Config, private readonly url: URLHelper, private readonly storageFactory: StorageProviderFactory, - private readonly quota: QuotaManagementService + private readonly quota: QuotaService ) { this.provider = this.storageFactory.create( this.config.plugins.copilot.storage @@ -57,7 +57,7 @@ export class CopilotStorage { @CallMetric('ai', 'blob_upload') async handleUpload(userId: string, blob: FileUpload) { - const checkExceeded = await this.quota.getQuotaCalculator(userId); + const checkExceeded = await this.quota.getUserQuotaCalculator(userId); if (checkExceeded(0)) { throw new BlobQuotaExceeded(); diff --git a/packages/backend/server/src/plugins/license/service.ts b/packages/backend/server/src/plugins/license/service.ts index e31992af78..93f6a553f7 100644 --- a/packages/backend/server/src/plugins/license/service.ts +++ b/packages/backend/server/src/plugins/license/service.ts @@ -12,7 +12,7 @@ import { WorkspaceLicenseAlreadyExists, } from '../../base'; import { PermissionService } from '../../core/permission'; -import { QuotaManagementService, QuotaType } from '../../core/quota'; +import { Models } from '../../models'; import { SubscriptionPlan, SubscriptionRecurring } from '../payment/types'; interface License { @@ -29,9 +29,9 @@ export class LicenseService { constructor( private readonly config: Config, private readonly db: PrismaClient, - private readonly quota: QuotaManagementService, private readonly event: EventBus, - private readonly permission: PermissionService + private readonly permission: PermissionService, + private readonly models: Models ) {} async getLicense(workspaceId: string) { @@ -316,14 +316,13 @@ export class LicenseService { }: Events['workspace.subscription.activated']) { switch (plan) { case SubscriptionPlan.SelfHostedTeam: - await this.quota.addTeamWorkspace( + await this.models.workspaceFeature.add( workspaceId, - `${recurring} team subscription activated` - ); - await this.quota.updateWorkspaceConfig( - workspaceId, - QuotaType.TeamPlanV1, - { memberLimit: quantity } + 'team_plan_v1', + `${recurring} team subscription activated`, + { + memberLimit: quantity, + } ); await this.permission.refreshSeatStatus(workspaceId, quantity); break; @@ -339,7 +338,7 @@ export class LicenseService { }: Events['workspace.subscription.canceled']) { switch (plan) { case SubscriptionPlan.SelfHostedTeam: - await this.quota.removeTeamWorkspace(workspaceId); + await this.models.workspaceFeature.remove(workspaceId, 'team_plan_v1'); break; default: break; diff --git a/packages/backend/server/src/plugins/payment/manager/user.ts b/packages/backend/server/src/plugins/payment/manager/user.ts index 5f54974033..437610d721 100644 --- a/packages/backend/server/src/plugins/payment/manager/user.ts +++ b/packages/backend/server/src/plugins/payment/manager/user.ts @@ -15,10 +15,7 @@ import { TooManyRequest, URLHelper, } from '../../../base'; -import { - EarlyAccessType, - FeatureManagementService, -} from '../../../core/features'; +import { EarlyAccessType, FeatureService } from '../../../core/features'; import { CouponType, KnownStripeInvoice, @@ -59,7 +56,7 @@ export class UserSubscriptionManager extends SubscriptionManager { stripe: Stripe, db: PrismaClient, private readonly runtime: Runtime, - private readonly feature: FeatureManagementService, + private readonly feature: FeatureService, private readonly event: EventBus, private readonly url: URLHelper, private readonly mutex: Mutex diff --git a/packages/backend/server/src/plugins/payment/quota.ts b/packages/backend/server/src/plugins/payment/quota.ts index a7296d82c4..0294139fe1 100644 --- a/packages/backend/server/src/plugins/payment/quota.ts +++ b/packages/backend/server/src/plugins/payment/quota.ts @@ -1,25 +1,17 @@ import { Injectable } from '@nestjs/common'; import { OnEvent } from '@nestjs/event-emitter'; -import { FeatureManagementService } from '../../core/features'; import { PermissionService } from '../../core/permission'; -import { - QuotaManagementService, - QuotaService, - QuotaType, -} from '../../core/quota'; import { WorkspaceService } from '../../core/workspaces/resolvers'; +import { Models } from '../../models'; import { SubscriptionPlan } from './types'; @Injectable() export class QuotaOverride { constructor( - private readonly quota: QuotaService, - private readonly manager: QuotaManagementService, private readonly permission: PermissionService, private readonly workspace: WorkspaceService, - private readonly feature: FeatureManagementService, - private readonly quotaService: QuotaService + private readonly models: Models ) {} @OnEvent('workspace.subscription.activated') @@ -31,21 +23,17 @@ export class QuotaOverride { }: Events['workspace.subscription.activated']) { switch (plan) { case 'team': { - const hasTeamWorkspace = await this.quota.hasWorkspaceQuota( + const isTeam = await this.workspace.isTeamWorkspace(workspaceId); + await this.models.workspaceFeature.add( workspaceId, - QuotaType.TeamPlanV1 - ); - await this.manager.addTeamWorkspace( - workspaceId, - `${recurring} team subscription activated` - ); - await this.quota.updateWorkspaceConfig( - workspaceId, - QuotaType.TeamPlanV1, - { memberLimit: quantity } + 'team_plan_v1', + `${recurring} team subscription activated`, + { + memberLimit: quantity, + } ); await this.permission.refreshSeatStatus(workspaceId, quantity); - if (!hasTeamWorkspace) { + if (!isTeam) { // this event will triggered when subscription is activated or changed // we only send emails when the team workspace is activated await this.workspace.sendTeamWorkspaceUpgradedEmail(workspaceId); @@ -64,7 +52,7 @@ export class QuotaOverride { }: Events['workspace.subscription.canceled']) { switch (plan) { case SubscriptionPlan.Team: - await this.manager.removeTeamWorkspace(workspaceId); + await this.models.workspaceFeature.remove(workspaceId, 'team_plan_v1'); break; default: break; @@ -79,14 +67,16 @@ export class QuotaOverride { }: Events['user.subscription.activated']) { switch (plan) { case SubscriptionPlan.AI: - await this.feature.addCopilot(userId, 'subscription activated'); + await this.models.userFeature.add( + userId, + 'unlimited_copilot', + 'subscription activated' + ); break; case SubscriptionPlan.Pro: - await this.quotaService.switchUserQuota( + await this.models.userFeature.add( userId, - recurring === 'lifetime' - ? QuotaType.LifetimeProPlanV1 - : QuotaType.ProPlanV1, + recurring === 'lifetime' ? 'lifetime_pro_plan_v1' : 'pro_plan_v1', 'subscription activated' ); break; @@ -102,16 +92,20 @@ export class QuotaOverride { }: Events['user.subscription.canceled']) { switch (plan) { case SubscriptionPlan.AI: - await this.feature.removeCopilot(userId); + await this.models.userFeature.remove(userId, 'unlimited_copilot'); break; case SubscriptionPlan.Pro: { // edge case: when user switch from recurring Pro plan to `Lifetime` plan, // a subscription canceled event will be triggered because `Lifetime` plan is not subscription based - const quota = await this.quotaService.getUserQuota(userId); - if (quota.feature.name !== QuotaType.LifetimeProPlanV1) { - await this.quotaService.switchUserQuota( + const isLifetimeUser = await this.models.userFeature.has( + userId, + 'lifetime_pro_plan_v1' + ); + + if (!isLifetimeUser) { + await this.models.userFeature.switchQuota( userId, - QuotaType.FreePlanV1, + 'free_plan_v1', 'subscription canceled' ); } diff --git a/packages/backend/server/src/plugins/payment/service.ts b/packages/backend/server/src/plugins/payment/service.ts index dbbcf727b3..eb08825a9c 100644 --- a/packages/backend/server/src/plugins/payment/service.ts +++ b/packages/backend/server/src/plugins/payment/service.ts @@ -26,7 +26,7 @@ import { UserNotFound, } from '../../base'; import { CurrentUser } from '../../core/auth'; -import { FeatureManagementService } from '../../core/features'; +import { FeatureService } from '../../core/features'; import { Models } from '../../models'; import { CheckoutParams, @@ -83,7 +83,7 @@ export class SubscriptionService implements OnApplicationBootstrap { private readonly config: Config, private readonly stripe: Stripe, private readonly db: PrismaClient, - private readonly feature: FeatureManagementService, + private readonly feature: FeatureService, private readonly models: Models, private readonly userManager: UserSubscriptionManager, private readonly workspaceManager: WorkspaceSubscriptionManager, diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index a8878d0524..8fda75ec32 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -230,6 +230,7 @@ union ErrorDataUnion = AlreadyInSpaceDataType | BlobNotFoundDataType | CopilotMe enum ErrorNames { ACCESS_DENIED ACTION_FORBIDDEN + ACTION_FORBIDDEN_ON_NON_TEAM_WORKSPACE ALREADY_IN_SPACE AUTHENTICATION_REQUIRED BLOB_NOT_FOUND @@ -336,12 +337,14 @@ type ExpectToUpdateDocUserRoleDataType { spaceId: String! } -"""The type of workspace feature""" enum FeatureType { AIEarlyAccess Admin - Copilot EarlyAccess + FreePlan + LifetimeProPlan + ProPlan + TeamPlan UnlimitedCopilot UnlimitedWorkspace } @@ -380,15 +383,6 @@ type GrantedDocUsersConnection { totalCount: Int! } -type HumanReadableQuotaType { - blobLimit: String! - copilotActionLimit: String - historyPeriod: String! - memberLimit: String! - name: String! - storageQuota: String! -} - type InvalidEmailDataType { email: String! } @@ -565,7 +559,7 @@ type MissingOauthQueryParameterDataType { type Mutation { acceptInviteById(inviteId: String!, sendAcceptMail: Boolean, workspaceId: String!): Boolean! activateLicense(license: String!, workspaceId: String!): License! - addWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Int! + addWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Boolean! approveMember(userId: String!, workspaceId: String!): String! cancelSubscription(idempotencyKey: String @deprecated(reason: "use header `Idempotency-Key`"), plan: SubscriptionPlan = Pro, workspaceId: String): SubscriptionType! changeEmail(email: String!, token: String!): UserType! @@ -621,7 +615,7 @@ type Mutation { """Remove user avatar""" removeAvatar: RemoveAvatar! - removeWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Int! + removeWorkspaceFeature(feature: FeatureType!, workspaceId: String!): Boolean! resumeSubscription(idempotencyKey: String @deprecated(reason: "use header `Idempotency-Key`"), plan: SubscriptionPlan = Pro, workspaceId: String): SubscriptionType! revoke(userId: String!, workspaceId: String!): Boolean! revokeDocUserRoles(docId: String!, userIds: [String!]!): Boolean! @@ -634,7 +628,6 @@ type Mutation { sendVerifyChangeEmail(callbackUrl: String!, email: String!, token: String!): Boolean! sendVerifyEmail(callbackUrl: String!): Boolean! setBlob(blob: Upload!, workspaceId: String!): String! - setWorkspaceExperimentalFeature(enable: Boolean!, feature: FeatureType!, workspaceId: String!): Boolean! sharePage(pageId: String!, workspaceId: String!): Boolean! @deprecated(reason: "renamed to publishPage") """Update a copilot prompt""" @@ -733,7 +726,6 @@ type Query { """List all copilot prompts""" listCopilotPrompts: [CopilotPromptType!]! - listWorkspaceFeatures(feature: FeatureType!): [WorkspaceFeatureType!]! prices: [SubscriptionPrice!]! """server config""" @@ -782,18 +774,6 @@ type QueryTooLongDataType { max: Int! } -type QuotaQueryType { - blobLimit: SafeInt! - copilotActionLimit: SafeInt - historyPeriod: SafeInt! - humanReadable: HumanReadableQuotaType! - memberCount: SafeInt! - memberLimit: SafeInt! - name: String! - storageQuota: SafeInt! - usedSize: SafeInt! -} - type RemoveAvatar { success: Boolean! } @@ -1043,25 +1023,29 @@ scalar Upload union UserOrLimitedUser = LimitedUserType | UserType -type UserQuota { - blobLimit: SafeInt! - historyPeriod: SafeInt! - humanReadable: UserQuotaHumanReadable! - memberLimit: Int! - name: String! - storageQuota: SafeInt! -} - -type UserQuotaHumanReadable { +type UserQuotaHumanReadableType { blobLimit: String! + copilotActionLimit: String! historyPeriod: String! memberLimit: String! name: String! storageQuota: String! + usedStorageQuota: String! } -type UserQuotaUsage { +type UserQuotaType { + blobLimit: SafeInt! + copilotActionLimit: Int + historyPeriod: SafeInt! + humanReadable: UserQuotaHumanReadableType! + memberLimit: Int! + name: String! storageQuota: SafeInt! + usedStorageQuota: SafeInt! +} + +type UserQuotaUsageType { + storageQuota: SafeInt! @deprecated(reason: "use `UserQuotaType['usedStorageQuota']` instead") } type UserType { @@ -1091,8 +1075,8 @@ type UserType { """User name""" name: String! - quota: UserQuota - quotaUsage: UserQuotaUsage! + quota: UserQuotaType! + quotaUsage: UserQuotaUsageType! subscriptions: [SubscriptionType!]! token: tokenType! @deprecated(reason: "use [/api/auth/sign-in?native=true] instead") } @@ -1106,15 +1090,6 @@ type WorkspaceBlobSizes { size: SafeInt! } -type WorkspaceFeatureType { - """Workspace created date""" - createdAt: DateTime! - id: ID! - - """is Public workspace""" - public: Boolean! -} - """Workspace invite link expire time""" enum WorkspaceInviteLinkExpireTime { OneDay @@ -1170,15 +1145,34 @@ type WorkspacePermissions { Workspace_Users_Read: Boolean! } +type WorkspaceQuotaHumanReadableType { + blobLimit: String! + historyPeriod: String! + memberCount: String! + memberLimit: String! + name: String! + storageQuota: String! + storageQuotaUsed: String! +} + +type WorkspaceQuotaType { + blobLimit: SafeInt! + historyPeriod: SafeInt! + humanReadable: WorkspaceQuotaHumanReadableType! + memberCount: Int! + memberLimit: Int! + name: String! + storageQuota: SafeInt! + usedSize: SafeInt! @deprecated(reason: "use `usedStorageQuota` instead") + usedStorageQuota: SafeInt! +} + type WorkspaceRolePermissions { permissions: WorkspacePermissions! role: Permission! } type WorkspaceType { - """Available features of workspace""" - availableFeatures: [FeatureType!]! - """List blobs of workspace""" blobs: [ListedBlob!]! @@ -1193,9 +1187,6 @@ type WorkspaceType { """Enable url previous when sharing""" enableUrlPreview: Boolean! - - """Enabled features of workspace""" - features: [FeatureType!]! histories(before: DateTime, guid: String!, take: Int): [DocHistoryType!]! id: ID! @@ -1240,7 +1231,7 @@ type WorkspaceType { publicPages: [WorkspacePage!]! """quota of workspace""" - quota: QuotaQueryType! + quota: WorkspaceQuotaType! """Role of current signed in user in workspace""" role: Permission! diff --git a/packages/frontend/core/src/modules/cloud/entities/user-quota.ts b/packages/frontend/core/src/modules/cloud/entities/user-quota.ts index 01110aa062..265508d45d 100644 --- a/packages/frontend/core/src/modules/cloud/entities/user-quota.ts +++ b/packages/frontend/core/src/modules/cloud/entities/user-quota.ts @@ -19,7 +19,9 @@ import type { AuthService } from '../services/auth'; import type { UserQuotaStore } from '../stores/user-quota'; export class UserQuota extends Entity { - quota$ = new LiveData['quota']>(null); + quota$ = new LiveData['quota'] | null>( + null + ); /** Used storage in bytes */ used$ = new LiveData(null); /** Formatted used storage */ diff --git a/packages/frontend/core/src/modules/quota/entities/quota.ts b/packages/frontend/core/src/modules/quota/entities/quota.ts index 7b56d9777d..e220a5c372 100644 --- a/packages/frontend/core/src/modules/quota/entities/quota.ts +++ b/packages/frontend/core/src/modules/quota/entities/quota.ts @@ -74,7 +74,7 @@ export class WorkspaceQuota extends Entity { this.workspaceService.workspace.id, signal ); - return { quota: data, used: data.usedSize }; + return { quota: data, used: data.usedStorageQuota }; }).pipe( backoffRetry({ when: isNetworkError, diff --git a/packages/frontend/graphql/src/graphql/get-workspace-features.gql b/packages/frontend/graphql/src/graphql/get-workspace-features.gql deleted file mode 100644 index 8644712835..0000000000 --- a/packages/frontend/graphql/src/graphql/get-workspace-features.gql +++ /dev/null @@ -1,5 +0,0 @@ -query getWorkspaceFeatures($workspaceId: String!) { - workspace(id: $workspaceId) { - features - } -} diff --git a/packages/frontend/graphql/src/graphql/index.ts b/packages/frontend/graphql/src/graphql/index.ts index ce583d93c5..a95a830055 100644 --- a/packages/frontend/graphql/src/graphql/index.ts +++ b/packages/frontend/graphql/src/graphql/index.ts @@ -206,17 +206,6 @@ mutation createCopilotSession($options: CreateChatSessionInput!) { }`, }; -export const updateCopilotSessionMutation = { - id: 'updateCopilotSessionMutation' as const, - operationName: 'updateCopilotSession', - definitionName: 'updateCopilotSession', - containsFile: false, - query: ` -mutation updateCopilotSession($options: UpdateChatSessionInput!) { - updateCopilotSession(options: $options) -}`, -}; - export const createCustomerPortalMutation = { id: 'createCustomerPortalMutation' as const, operationName: 'createCustomerPortal', @@ -637,19 +626,6 @@ query getUsersCount { }`, }; -export const getWorkspaceFeaturesQuery = { - id: 'getWorkspaceFeaturesQuery' as const, - operationName: 'getWorkspaceFeatures', - definitionName: 'workspace', - containsFile: false, - query: ` -query getWorkspaceFeatures($workspaceId: String!) { - workspace(id: $workspaceId) { - features - } -}`, -}; - export const getWorkspaceInfoQuery = { id: 'getWorkspaceInfoQuery' as const, operationName: 'getWorkspaceInfo', @@ -1143,6 +1119,17 @@ mutation updateAccount($id: String!, $input: ManageUserInput!) { }`, }; +export const updateCopilotSessionMutation = { + id: 'updateCopilotSessionMutation' as const, + operationName: 'updateCopilotSession', + definitionName: 'updateCopilotSession', + containsFile: false, + query: ` +mutation updateCopilotSession($options: UpdateChatSessionInput!) { + updateCopilotSession(options: $options) +}`, +}; + export const updatePromptMutation = { id: 'updatePromptMutation' as const, operationName: 'updatePrompt', @@ -1289,89 +1276,6 @@ mutation setEnableUrlPreview($id: ID!, $enableUrlPreview: Boolean!) { }`, }; -export const enabledFeaturesQuery = { - id: 'enabledFeaturesQuery' as const, - operationName: 'enabledFeatures', - definitionName: 'workspace', - containsFile: false, - query: ` -query enabledFeatures($id: String!) { - workspace(id: $id) { - features - } -}`, -}; - -export const availableFeaturesQuery = { - id: 'availableFeaturesQuery' as const, - operationName: 'availableFeatures', - definitionName: 'workspace', - containsFile: false, - query: ` -query availableFeatures($id: String!) { - workspace(id: $id) { - availableFeatures - } -}`, -}; - -export const setWorkspaceExperimentalFeatureMutation = { - id: 'setWorkspaceExperimentalFeatureMutation' as const, - operationName: 'setWorkspaceExperimentalFeature', - definitionName: 'setWorkspaceExperimentalFeature', - containsFile: false, - query: ` -mutation setWorkspaceExperimentalFeature($workspaceId: String!, $feature: FeatureType!, $enable: Boolean!) { - setWorkspaceExperimentalFeature( - workspaceId: $workspaceId - feature: $feature - enable: $enable - ) -}`, -}; - -export const addWorkspaceFeatureMutation = { - id: 'addWorkspaceFeatureMutation' as const, - operationName: 'addWorkspaceFeature', - definitionName: 'addWorkspaceFeature', - containsFile: false, - query: ` -mutation addWorkspaceFeature($workspaceId: String!, $feature: FeatureType!) { - addWorkspaceFeature(workspaceId: $workspaceId, feature: $feature) -}`, -}; - -export const listWorkspaceFeaturesQuery = { - id: 'listWorkspaceFeaturesQuery' as const, - operationName: 'listWorkspaceFeatures', - definitionName: 'listWorkspaceFeatures', - containsFile: false, - query: ` -query listWorkspaceFeatures($feature: FeatureType!) { - listWorkspaceFeatures(feature: $feature) { - id - public - createdAt - memberCount - owner { - id - } - features - } -}`, -}; - -export const removeWorkspaceFeatureMutation = { - id: 'removeWorkspaceFeatureMutation' as const, - operationName: 'removeWorkspaceFeature', - definitionName: 'removeWorkspaceFeature', - containsFile: false, - query: ` -mutation removeWorkspaceFeature($workspaceId: String!, $feature: FeatureType!) { - removeWorkspaceFeature(workspaceId: $workspaceId, feature: $feature) -}`, -}; - export const inviteByEmailMutation = { id: 'inviteByEmailMutation' as const, operationName: 'inviteByEmail', @@ -1500,6 +1404,7 @@ query workspaceQuota($id: String!) { name blobLimit storageQuota + usedStorageQuota historyPeriod memberLimit memberCount @@ -1510,7 +1415,6 @@ query workspaceQuota($id: String!) { historyPeriod memberLimit } - usedSize } } }`, diff --git a/packages/frontend/graphql/src/graphql/workspace-enabled-features.gql b/packages/frontend/graphql/src/graphql/workspace-enabled-features.gql deleted file mode 100644 index 837dddc86c..0000000000 --- a/packages/frontend/graphql/src/graphql/workspace-enabled-features.gql +++ /dev/null @@ -1,5 +0,0 @@ -query enabledFeatures($id: String!) { - workspace(id: $id) { - features - } -} diff --git a/packages/frontend/graphql/src/graphql/workspace-experimental-feature-get.gql b/packages/frontend/graphql/src/graphql/workspace-experimental-feature-get.gql deleted file mode 100644 index 2720fa6fe3..0000000000 --- a/packages/frontend/graphql/src/graphql/workspace-experimental-feature-get.gql +++ /dev/null @@ -1,5 +0,0 @@ -query availableFeatures($id: String!) { - workspace(id: $id) { - availableFeatures - } -} diff --git a/packages/frontend/graphql/src/graphql/workspace-experimental-feature-set.gql b/packages/frontend/graphql/src/graphql/workspace-experimental-feature-set.gql deleted file mode 100644 index 1b9ce1f683..0000000000 --- a/packages/frontend/graphql/src/graphql/workspace-experimental-feature-set.gql +++ /dev/null @@ -1,11 +0,0 @@ -mutation setWorkspaceExperimentalFeature( - $workspaceId: String! - $feature: FeatureType! - $enable: Boolean! -) { - setWorkspaceExperimentalFeature( - workspaceId: $workspaceId - feature: $feature - enable: $enable - ) -} diff --git a/packages/frontend/graphql/src/graphql/workspace-feature-add.gql b/packages/frontend/graphql/src/graphql/workspace-feature-add.gql deleted file mode 100644 index d77eafbd16..0000000000 --- a/packages/frontend/graphql/src/graphql/workspace-feature-add.gql +++ /dev/null @@ -1,3 +0,0 @@ -mutation addWorkspaceFeature($workspaceId: String!, $feature: FeatureType!) { - addWorkspaceFeature(workspaceId: $workspaceId, feature: $feature) -} diff --git a/packages/frontend/graphql/src/graphql/workspace-feature-list.gql b/packages/frontend/graphql/src/graphql/workspace-feature-list.gql deleted file mode 100644 index 4bd21fbe45..0000000000 --- a/packages/frontend/graphql/src/graphql/workspace-feature-list.gql +++ /dev/null @@ -1,12 +0,0 @@ -query listWorkspaceFeatures($feature: FeatureType!) { - listWorkspaceFeatures(feature: $feature) { - id - public - createdAt - memberCount - owner { - id - } - features - } -} diff --git a/packages/frontend/graphql/src/graphql/workspace-feature-remove.gql b/packages/frontend/graphql/src/graphql/workspace-feature-remove.gql deleted file mode 100644 index e856885b56..0000000000 --- a/packages/frontend/graphql/src/graphql/workspace-feature-remove.gql +++ /dev/null @@ -1,3 +0,0 @@ -mutation removeWorkspaceFeature($workspaceId: String!, $feature: FeatureType!) { - removeWorkspaceFeature(workspaceId: $workspaceId, feature: $feature) -} diff --git a/packages/frontend/graphql/src/graphql/workspace-quota.gql b/packages/frontend/graphql/src/graphql/workspace-quota.gql index 176e1fc7b5..7d51730c15 100644 --- a/packages/frontend/graphql/src/graphql/workspace-quota.gql +++ b/packages/frontend/graphql/src/graphql/workspace-quota.gql @@ -4,6 +4,7 @@ query workspaceQuota($id: String!) { name blobLimit storageQuota + usedStorageQuota historyPeriod memberLimit memberCount @@ -14,7 +15,6 @@ query workspaceQuota($id: String!) { historyPeriod memberLimit } - usedSize } } } diff --git a/packages/frontend/graphql/src/schema.ts b/packages/frontend/graphql/src/schema.ts index b6eac30be6..1a0c69b124 100644 --- a/packages/frontend/graphql/src/schema.ts +++ b/packages/frontend/graphql/src/schema.ts @@ -1,4 +1,3 @@ -/* oxlint-disable */ export type Maybe = T | null; export type InputMaybe = T | null; export type Exact = { @@ -188,11 +187,6 @@ export interface CreateChatSessionInput { workspaceId: Scalars['String']['input']; } -export interface UpdateChatSessionInput { - sessionId: Scalars['String']['input']; - promptName: Scalars['String']['input']; -} - export interface CreateCheckoutSessionInput { args?: InputMaybe; coupon?: InputMaybe; @@ -282,6 +276,7 @@ export type ErrorDataUnion = | MemberNotFoundInSpaceDataType | MissingOauthQueryParameterDataType | NotInSpaceDataType + | QueryTooLongDataType | RuntimeConfigNotFoundDataType | SameSubscriptionRecurringDataType | SpaceAccessDeniedDataType @@ -356,6 +351,7 @@ export enum ErrorNames { OAUTH_STATE_EXPIRED = 'OAUTH_STATE_EXPIRED', PAGE_IS_NOT_PUBLIC = 'PAGE_IS_NOT_PUBLIC', PASSWORD_REQUIRED = 'PASSWORD_REQUIRED', + QUERY_TOO_LONG = 'QUERY_TOO_LONG', RUNTIME_CONFIG_NOT_FOUND = 'RUNTIME_CONFIG_NOT_FOUND', SAME_EMAIL_PROVIDED = 'SAME_EMAIL_PROVIDED', SAME_SUBSCRIPTION_RECURRING = 'SAME_SUBSCRIPTION_RECURRING', @@ -384,12 +380,14 @@ export enum ErrorNames { WRONG_SIGN_IN_METHOD = 'WRONG_SIGN_IN_METHOD', } -/** The type of workspace feature */ export enum FeatureType { AIEarlyAccess = 'AIEarlyAccess', Admin = 'Admin', - Copilot = 'Copilot', EarlyAccess = 'EarlyAccess', + FreePlan = 'FreePlan', + LifetimeProPlan = 'LifetimeProPlan', + ProPlan = 'ProPlan', + TeamPlan = 'TeamPlan', UnlimitedCopilot = 'UnlimitedCopilot', UnlimitedWorkspace = 'UnlimitedWorkspace', } @@ -402,16 +400,6 @@ export interface ForkChatSessionInput { workspaceId: Scalars['String']['input']; } -export interface HumanReadableQuotaType { - __typename?: 'HumanReadableQuotaType'; - blobLimit: Scalars['String']['output']; - copilotActionLimit: Maybe; - historyPeriod: Scalars['String']['output']; - memberLimit: Scalars['String']['output']; - name: Scalars['String']['output']; - storageQuota: Scalars['String']['output']; -} - export interface InvalidEmailDataType { __typename?: 'InvalidEmailDataType'; email: Scalars['String']['output']; @@ -533,6 +521,15 @@ export interface InvoiceType { updatedAt: Scalars['DateTime']['output']; } +export interface License { + __typename?: 'License'; + expiredAt: Maybe; + installedAt: Scalars['DateTime']['output']; + quantity: Scalars['Int']['output']; + recurring: SubscriptionRecurring; + validatedAt: Scalars['DateTime']['output']; +} + export interface LimitedUserType { __typename?: 'LimitedUserType'; /** User email */ @@ -574,7 +571,8 @@ export interface MissingOauthQueryParameterDataType { export interface Mutation { __typename?: 'Mutation'; acceptInviteById: Scalars['Boolean']['output']; - addWorkspaceFeature: Scalars['Int']['output']; + activateLicense: License; + addWorkspaceFeature: Scalars['Boolean']['output']; approveMember: Scalars['String']['output']; cancelSubscription: SubscriptionType; changeEmail: UserType; @@ -594,10 +592,12 @@ export interface Mutation { /** Create a stripe customer portal to manage payment methods */ createCustomerPortal: Scalars['String']['output']; createInviteLink: InviteLink; + createSelfhostWorkspaceCustomerPortal: Scalars['String']['output']; /** Create a new user */ createUser: UserType; /** Create a new workspace */ createWorkspace: WorkspaceType; + deactivateLicense: Scalars['Boolean']['output']; deleteAccount: DeleteAccount; deleteBlob: Scalars['Boolean']['output']; /** Delete a user account */ @@ -615,7 +615,7 @@ export interface Mutation { releaseDeletedBlobs: Scalars['Boolean']['output']; /** Remove user avatar */ removeAvatar: RemoveAvatar; - removeWorkspaceFeature: Scalars['Int']['output']; + removeWorkspaceFeature: Scalars['Boolean']['output']; resumeSubscription: SubscriptionType; revoke: Scalars['Boolean']['output']; revokeInviteLink: Scalars['Boolean']['output']; @@ -628,11 +628,12 @@ export interface Mutation { sendVerifyChangeEmail: Scalars['Boolean']['output']; sendVerifyEmail: Scalars['Boolean']['output']; setBlob: Scalars['String']['output']; - setWorkspaceExperimentalFeature: Scalars['Boolean']['output']; /** @deprecated renamed to publishPage */ sharePage: Scalars['Boolean']['output']; /** Update a copilot prompt */ updateCopilotPrompt: CopilotPromptType; + /** Update a chat session */ + updateCopilotSession: Scalars['String']['output']; updateProfile: UserType; /** update server runtime configurable setting */ updateRuntimeConfig: ServerRuntimeConfigType; @@ -656,6 +657,11 @@ export interface MutationAcceptInviteByIdArgs { workspaceId: Scalars['String']['input']; } +export interface MutationActivateLicenseArgs { + license: Scalars['String']['input']; + workspaceId: Scalars['String']['input']; +} + export interface MutationAddWorkspaceFeatureArgs { feature: FeatureType; workspaceId: Scalars['String']['input']; @@ -708,15 +714,15 @@ export interface MutationCreateCopilotSessionArgs { options: CreateChatSessionInput; } -export interface MutationUpdateCopilotSessionArgs { - options: UpdateChatSessionInput; -} - export interface MutationCreateInviteLinkArgs { expireTime: WorkspaceInviteLinkExpireTime; workspaceId: Scalars['String']['input']; } +export interface MutationCreateSelfhostWorkspaceCustomerPortalArgs { + workspaceId: Scalars['String']['input']; +} + export interface MutationCreateUserArgs { input: CreateUserInput; } @@ -725,6 +731,10 @@ export interface MutationCreateWorkspaceArgs { init?: InputMaybe; } +export interface MutationDeactivateLicenseArgs { + workspaceId: Scalars['String']['input']; +} + export interface MutationDeleteBlobArgs { hash?: InputMaybe; key?: InputMaybe; @@ -849,12 +859,6 @@ export interface MutationSetBlobArgs { workspaceId: Scalars['String']['input']; } -export interface MutationSetWorkspaceExperimentalFeatureArgs { - enable: Scalars['Boolean']['input']; - feature: FeatureType; - workspaceId: Scalars['String']['input']; -} - export interface MutationSharePageArgs { pageId: Scalars['String']['input']; workspaceId: Scalars['String']['input']; @@ -865,6 +869,10 @@ export interface MutationUpdateCopilotPromptArgs { name: Scalars['String']['input']; } +export interface MutationUpdateCopilotSessionArgs { + options: UpdateChatSessionInput; +} + export interface MutationUpdateProfileArgs { input: UpdateUserInput; } @@ -958,7 +966,6 @@ export interface Query { listBlobs: Array; /** List all copilot prompts */ listCopilotPrompts: Array; - listWorkspaceFeatures: Array; prices: Array; /** server config */ serverConfig: ServerConfigType; @@ -1001,10 +1008,6 @@ export interface QueryListBlobsArgs { workspaceId: Scalars['String']['input']; } -export interface QueryListWorkspaceFeaturesArgs { - feature: FeatureType; -} - export interface QueryUserArgs { email: Scalars['String']['input']; } @@ -1035,17 +1038,9 @@ export interface QueryChatHistoriesInput { skip?: InputMaybe; } -export interface QuotaQueryType { - __typename?: 'QuotaQueryType'; - blobLimit: Scalars['SafeInt']['output']; - copilotActionLimit: Maybe; - historyPeriod: Scalars['SafeInt']['output']; - humanReadable: HumanReadableQuotaType; - memberCount: Scalars['SafeInt']['output']; - memberLimit: Scalars['SafeInt']['output']; - name: Scalars['String']['output']; - storageQuota: Scalars['SafeInt']['output']; - usedSize: Scalars['SafeInt']['output']; +export interface QueryTooLongDataType { + __typename?: 'QueryTooLongDataType'; + max: Scalars['Int']['output']; } export interface RemoveAvatar { @@ -1241,6 +1236,12 @@ export interface UnsupportedSubscriptionPlanDataType { plan: Scalars['String']['output']; } +export interface UpdateChatSessionInput { + /** The prompt name to use for the session */ + promptName: Scalars['String']['input']; + sessionId: Scalars['String']['input']; +} + export interface UpdateUserInput { /** User name */ name?: InputMaybe; @@ -1258,27 +1259,32 @@ export interface UpdateWorkspaceInput { export type UserOrLimitedUser = LimitedUserType | UserType; -export interface UserQuota { - __typename?: 'UserQuota'; - blobLimit: Scalars['SafeInt']['output']; - historyPeriod: Scalars['SafeInt']['output']; - humanReadable: UserQuotaHumanReadable; - memberLimit: Scalars['Int']['output']; - name: Scalars['String']['output']; - storageQuota: Scalars['SafeInt']['output']; -} - -export interface UserQuotaHumanReadable { - __typename?: 'UserQuotaHumanReadable'; +export interface UserQuotaHumanReadableType { + __typename?: 'UserQuotaHumanReadableType'; blobLimit: Scalars['String']['output']; + copilotActionLimit: Scalars['String']['output']; historyPeriod: Scalars['String']['output']; memberLimit: Scalars['String']['output']; name: Scalars['String']['output']; storageQuota: Scalars['String']['output']; + usedStorageQuota: Scalars['String']['output']; } -export interface UserQuotaUsage { - __typename?: 'UserQuotaUsage'; +export interface UserQuotaType { + __typename?: 'UserQuotaType'; + blobLimit: Scalars['SafeInt']['output']; + copilotActionLimit: Maybe; + historyPeriod: Scalars['SafeInt']['output']; + humanReadable: UserQuotaHumanReadableType; + memberLimit: Scalars['Int']['output']; + name: Scalars['String']['output']; + storageQuota: Scalars['SafeInt']['output']; + usedStorageQuota: Scalars['SafeInt']['output']; +} + +export interface UserQuotaUsageType { + __typename?: 'UserQuotaUsageType'; + /** @deprecated use `UserQuotaType['usedStorageQuota']` instead */ storageQuota: Scalars['SafeInt']['output']; } @@ -1306,8 +1312,8 @@ export interface UserType { invoices: Array; /** User name */ name: Scalars['String']['output']; - quota: Maybe; - quotaUsage: UserQuotaUsage; + quota: UserQuotaType; + quotaUsage: UserQuotaUsageType; subscriptions: Array; /** @deprecated use [/api/auth/sign-in?native=true] instead */ token: TokenType; @@ -1371,10 +1377,33 @@ export interface WorkspacePageMeta { updatedBy: Maybe; } +export interface WorkspaceQuotaHumanReadableType { + __typename?: 'WorkspaceQuotaHumanReadableType'; + blobLimit: Scalars['String']['output']; + historyPeriod: Scalars['String']['output']; + memberCount: Scalars['String']['output']; + memberLimit: Scalars['String']['output']; + name: Scalars['String']['output']; + storageQuota: Scalars['String']['output']; + storageQuotaUsed: Scalars['String']['output']; +} + +export interface WorkspaceQuotaType { + __typename?: 'WorkspaceQuotaType'; + blobLimit: Scalars['SafeInt']['output']; + historyPeriod: Scalars['SafeInt']['output']; + humanReadable: WorkspaceQuotaHumanReadableType; + memberCount: Scalars['Int']['output']; + memberLimit: Scalars['Int']['output']; + name: Scalars['String']['output']; + storageQuota: Scalars['SafeInt']['output']; + /** @deprecated use `usedStorageQuota` instead */ + usedSize: Scalars['SafeInt']['output']; + usedStorageQuota: Scalars['SafeInt']['output']; +} + export interface WorkspaceType { __typename?: 'WorkspaceType'; - /** Available features of workspace */ - availableFeatures: Array; /** List blobs of workspace */ blobs: Array; /** Blobs size of workspace */ @@ -1385,8 +1414,6 @@ export interface WorkspaceType { enableAi: Scalars['Boolean']['output']; /** Enable url previous when sharing */ enableUrlPreview: Scalars['Boolean']['output']; - /** Enabled features of workspace */ - features: Array; histories: Array; id: Scalars['ID']['output']; /** is current workspace initialized */ @@ -1396,6 +1423,8 @@ export interface WorkspaceType { /** Get user invoice count */ invoiceCount: Scalars['Int']['output']; invoices: Array; + /** The selfhost license of the workspace */ + license: Maybe; /** member count of workspace */ memberCount: Scalars['Int']['output']; /** Members of workspace */ @@ -1413,7 +1442,7 @@ export interface WorkspaceType { /** Public pages of a workspace */ publicPages: Array; /** quota of workspace */ - quota: QuotaQueryType; + quota: WorkspaceQuotaType; /** * Shared pages of workspace * @deprecated use WorkspaceType.publicPages @@ -1437,6 +1466,7 @@ export interface WorkspaceTypeInvoicesArgs { } export interface WorkspaceTypeMembersArgs { + query?: InputMaybe; skip?: InputMaybe; take?: InputMaybe; } @@ -1630,15 +1660,6 @@ export type CreateCopilotSessionMutation = { createCopilotSession: string; }; -export type UpdateCopilotSessionMutationVariables = Exact<{ - options: UpdateChatSessionInput; -}>; - -export type UpdateCopilotSessionMutation = { - __typename?: 'Mutation'; - updateCopilotSession: string; -}; - export type CreateCustomerPortalMutationVariables = Exact<{ [key: string]: never; }>; @@ -1975,16 +1996,16 @@ export type GetUserByEmailQuery = { emailVerified: boolean; avatarUrl: string | null; quota: { - __typename?: 'UserQuota'; + __typename?: 'UserQuotaType'; humanReadable: { - __typename?: 'UserQuotaHumanReadable'; + __typename?: 'UserQuotaHumanReadableType'; blobLimit: string; historyPeriod: string; memberLimit: string; name: string; storageQuota: string; }; - } | null; + }; } | null; }; @@ -2026,15 +2047,6 @@ export type GetUsersCountQueryVariables = Exact<{ [key: string]: never }>; export type GetUsersCountQuery = { __typename?: 'Query'; usersCount: number }; -export type GetWorkspaceFeaturesQueryVariables = Exact<{ - workspaceId: Scalars['String']['input']; -}>; - -export type GetWorkspaceFeaturesQuery = { - __typename?: 'Query'; - workspace: { __typename?: 'WorkspaceType'; features: Array }; -}; - export type GetWorkspaceInfoQueryVariables = Exact<{ workspaceId: Scalars['String']['input']; }>; @@ -2281,22 +2293,22 @@ export type QuotaQuery = { __typename?: 'UserType'; id: string; quota: { - __typename?: 'UserQuota'; + __typename?: 'UserQuotaType'; name: string; blobLimit: number; storageQuota: number; historyPeriod: number; memberLimit: number; humanReadable: { - __typename?: 'UserQuotaHumanReadable'; + __typename?: 'UserQuotaHumanReadableType'; name: string; blobLimit: string; storageQuota: string; historyPeriod: string; memberLimit: string; }; - } | null; - quotaUsage: { __typename?: 'UserQuotaUsage'; storageQuota: number }; + }; + quotaUsage: { __typename?: 'UserQuotaUsageType'; storageQuota: number }; } | null; }; @@ -2487,6 +2499,15 @@ export type UpdateAccountMutation = { }; }; +export type UpdateCopilotSessionMutationVariables = Exact<{ + options: UpdateChatSessionInput; +}>; + +export type UpdateCopilotSessionMutation = { + __typename?: 'Mutation'; + updateCopilotSession: string; +}; + export type UpdatePromptMutationVariables = Exact<{ name: Scalars['String']['input']; messages: Array | CopilotPromptMessageInput; @@ -2617,75 +2638,6 @@ export type SetEnableUrlPreviewMutation = { updateWorkspace: { __typename?: 'WorkspaceType'; id: string }; }; -export type EnabledFeaturesQueryVariables = Exact<{ - id: Scalars['String']['input']; -}>; - -export type EnabledFeaturesQuery = { - __typename?: 'Query'; - workspace: { __typename?: 'WorkspaceType'; features: Array }; -}; - -export type AvailableFeaturesQueryVariables = Exact<{ - id: Scalars['String']['input']; -}>; - -export type AvailableFeaturesQuery = { - __typename?: 'Query'; - workspace: { - __typename?: 'WorkspaceType'; - availableFeatures: Array; - }; -}; - -export type SetWorkspaceExperimentalFeatureMutationVariables = Exact<{ - workspaceId: Scalars['String']['input']; - feature: FeatureType; - enable: Scalars['Boolean']['input']; -}>; - -export type SetWorkspaceExperimentalFeatureMutation = { - __typename?: 'Mutation'; - setWorkspaceExperimentalFeature: boolean; -}; - -export type AddWorkspaceFeatureMutationVariables = Exact<{ - workspaceId: Scalars['String']['input']; - feature: FeatureType; -}>; - -export type AddWorkspaceFeatureMutation = { - __typename?: 'Mutation'; - addWorkspaceFeature: number; -}; - -export type ListWorkspaceFeaturesQueryVariables = Exact<{ - feature: FeatureType; -}>; - -export type ListWorkspaceFeaturesQuery = { - __typename?: 'Query'; - listWorkspaceFeatures: Array<{ - __typename?: 'WorkspaceType'; - id: string; - public: boolean; - createdAt: string; - memberCount: number; - features: Array; - owner: { __typename?: 'UserType'; id: string }; - }>; -}; - -export type RemoveWorkspaceFeatureMutationVariables = Exact<{ - workspaceId: Scalars['String']['input']; - feature: FeatureType; -}>; - -export type RemoveWorkspaceFeatureMutation = { - __typename?: 'Mutation'; - removeWorkspaceFeature: number; -}; - export type InviteByEmailMutationVariables = Exact<{ workspaceId: Scalars['String']['input']; email: Scalars['String']['input']; @@ -2794,16 +2746,16 @@ export type WorkspaceQuotaQuery = { workspace: { __typename?: 'WorkspaceType'; quota: { - __typename?: 'QuotaQueryType'; + __typename?: 'WorkspaceQuotaType'; name: string; blobLimit: number; storageQuota: number; + usedStorageQuota: number; historyPeriod: number; memberLimit: number; memberCount: number; - usedSize: number; humanReadable: { - __typename?: 'HumanReadableQuotaType'; + __typename?: 'WorkspaceQuotaHumanReadableType'; name: string; blobLimit: string; storageQuota: string; @@ -2941,11 +2893,6 @@ export type Queries = variables: GetUsersCountQueryVariables; response: GetUsersCountQuery; } - | { - name: 'getWorkspaceFeaturesQuery'; - variables: GetWorkspaceFeaturesQueryVariables; - response: GetWorkspaceFeaturesQuery; - } | { name: 'getWorkspaceInfoQuery'; variables: GetWorkspaceInfoQueryVariables; @@ -3031,21 +2978,6 @@ export type Queries = variables: GetWorkspaceConfigQueryVariables; response: GetWorkspaceConfigQuery; } - | { - name: 'enabledFeaturesQuery'; - variables: EnabledFeaturesQueryVariables; - response: EnabledFeaturesQuery; - } - | { - name: 'availableFeaturesQuery'; - variables: AvailableFeaturesQueryVariables; - response: AvailableFeaturesQuery; - } - | { - name: 'listWorkspaceFeaturesQuery'; - variables: ListWorkspaceFeaturesQueryVariables; - response: ListWorkspaceFeaturesQuery; - } | { name: 'workspaceInvoicesQuery'; variables: WorkspaceInvoicesQueryVariables; @@ -3113,11 +3045,6 @@ export type Mutations = variables: CreateCopilotSessionMutationVariables; response: CreateCopilotSessionMutation; } - | { - name: 'updateCopilotSessionMutation'; - variables: UpdateCopilotSessionMutationVariables; - response: UpdateCopilotSessionMutation; - } | { name: 'createCustomerPortalMutation'; variables: CreateCustomerPortalMutationVariables; @@ -3228,6 +3155,11 @@ export type Mutations = variables: UpdateAccountMutationVariables; response: UpdateAccountMutation; } + | { + name: 'updateCopilotSessionMutation'; + variables: UpdateCopilotSessionMutationVariables; + response: UpdateCopilotSessionMutation; + } | { name: 'updatePromptMutation'; variables: UpdatePromptMutationVariables; @@ -3268,21 +3200,6 @@ export type Mutations = variables: SetEnableUrlPreviewMutationVariables; response: SetEnableUrlPreviewMutation; } - | { - name: 'setWorkspaceExperimentalFeatureMutation'; - variables: SetWorkspaceExperimentalFeatureMutationVariables; - response: SetWorkspaceExperimentalFeatureMutation; - } - | { - name: 'addWorkspaceFeatureMutation'; - variables: AddWorkspaceFeatureMutationVariables; - response: AddWorkspaceFeatureMutation; - } - | { - name: 'removeWorkspaceFeatureMutation'; - variables: RemoveWorkspaceFeatureMutationVariables; - response: RemoveWorkspaceFeatureMutation; - } | { name: 'inviteByEmailMutation'; variables: InviteByEmailMutationVariables; diff --git a/tests/affine-cloud/e2e/collaboration.spec.ts b/tests/affine-cloud/e2e/collaboration.spec.ts index f29a5f6df4..e84d7d7364 100644 --- a/tests/affine-cloud/e2e/collaboration.spec.ts +++ b/tests/affine-cloud/e2e/collaboration.spec.ts @@ -22,11 +22,8 @@ let user: { password: string; }; -test.beforeEach(async () => { - user = await createRandomUser(); -}); - test.beforeEach(async ({ page }) => { + user = await createRandomUser(); await loginUser(page, user); }); diff --git a/tests/kit/src/utils/cloud.ts b/tests/kit/src/utils/cloud.ts index fc4bcbbcd9..1e4658ef50 100644 --- a/tests/kit/src/utils/cloud.ts +++ b/tests/kit/src/utils/cloud.ts @@ -119,9 +119,8 @@ export async function createRandomUser(): Promise<{ const result = await runPrisma(async client => { const featureId = await client.feature .findFirst({ - where: { feature: 'free_plan_v1' }, + where: { name: 'free_plan_v1' }, select: { id: true }, - orderBy: { version: 'desc' }, }) .then(f => f!.id); @@ -135,6 +134,8 @@ export async function createRandomUser(): Promise<{ reason: 'created by test case', activated: true, featureId, + name: 'free_plan_v1', + type: 1, }, }, }, @@ -169,16 +170,14 @@ export async function createRandomAIUser(): Promise<{ const result = await runPrisma(async client => { const freeFeatureId = await client.feature .findFirst({ - where: { feature: 'free_plan_v1' }, + where: { name: 'free_plan_v1' }, select: { id: true }, - orderBy: { version: 'desc' }, }) .then(f => f!.id); const aiFeatureId = await client.feature .findFirst({ - where: { feature: 'unlimited_copilot' }, + where: { name: 'unlimited_copilot' }, select: { id: true }, - orderBy: { version: 'desc' }, }) .then(f => f!.id); @@ -193,11 +192,15 @@ export async function createRandomAIUser(): Promise<{ reason: 'created by test case', activated: true, featureId: freeFeatureId, + name: 'free_plan_v1', + type: 1, }, { reason: 'created by test case', activated: true, featureId: aiFeatureId, + name: 'unlimited_copilot', + type: 0, }, ], },