diff --git a/packages/backend/server/scripts/init-db.ts b/packages/backend/server/scripts/init-db.ts index a43ebb4f58..8e850ea716 100644 --- a/packages/backend/server/scripts/init-db.ts +++ b/packages/backend/server/scripts/init-db.ts @@ -8,6 +8,20 @@ async function main() { data: { ...userA, password: await hash(userA.password), + features: { + create: { + reason: 'created by api sign up', + activated: true, + feature: { + connect: { + feature_version: { + feature: 'free_plan_v1', + version: 1, + }, + }, + }, + }, + }, }, }); } 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 index d6336dcfd0..0245770706 100644 --- a/packages/backend/server/src/data/migrations/1698652531198-user-features-init.ts +++ b/packages/backend/server/src/data/migrations/1698652531198-user-features-init.ts @@ -3,10 +3,30 @@ import { FeatureKind, Features, FeatureType, -} from '../../modules/features/types'; +} from '../../modules/features'; import { Quotas } from '../../modules/quota/types'; import { PrismaService } from '../../prisma'; +export class UserFeaturesInit1698652531198 { + // do the migration + static async up(db: PrismaService) { + // upgrade features from lower version to higher version + for (const feature of Features) { + await upsertFeature(db, feature); + } + await migrateNewFeatureTable(db); + + for (const quota of Quotas) { + await upsertFeature(db, quota); + } + } + + // revert the migration + static async down(_db: PrismaService) { + // TODO: revert the migration + } +} + // upgrade features from lower version to higher version async function upsertFeature( db: PrismaService, @@ -34,26 +54,6 @@ async function upsertFeature( } } -export class UserFeaturesInit1698652531198 { - // do the migration - static async up(db: PrismaService) { - // upgrade features from lower version to higher version - for (const feature of Features) { - await upsertFeature(db, feature); - } - await migrateNewFeatureTable(db); - - for (const quota of Quotas) { - await upsertFeature(db, quota); - } - } - - // revert the migration - static async down(_db: PrismaService) { - // TODO: revert the migration - } -} - async function migrateNewFeatureTable(prisma: PrismaService) { const waitingList = await prisma.newFeaturesWaitingList.findMany(); for (const oldUser of waitingList) { diff --git a/packages/backend/server/src/modules/auth/resolver.ts b/packages/backend/server/src/modules/auth/resolver.ts index 680d0a37ec..697812520f 100644 --- a/packages/backend/server/src/modules/auth/resolver.ts +++ b/packages/backend/server/src/modules/auth/resolver.ts @@ -19,7 +19,7 @@ import { nanoid } from 'nanoid'; import { Config } from '../../config'; import { SessionService } from '../../session'; import { CloudThrottlerGuard, Throttle } from '../../throttler'; -import { UserType } from '../users/resolver'; +import { UserType } from '../users'; import { Auth, CurrentUser } from './guard'; import { AuthService } from './service'; diff --git a/packages/backend/server/src/modules/doc/history.ts b/packages/backend/server/src/modules/doc/history.ts index fa52ecda97..108560bd45 100644 --- a/packages/backend/server/src/modules/doc/history.ts +++ b/packages/backend/server/src/modules/doc/history.ts @@ -7,7 +7,7 @@ import { Config } from '../../config'; import { type EventPayload, OnEvent } from '../../event'; import { metrics } from '../../metrics'; import { PrismaService } from '../../prisma'; -import { SubscriptionStatus } from '../payment/service'; +import { QuotaService } from '../quota'; import { Permission } from '../workspaces/types'; import { isEmptyBuffer } from './manager'; @@ -16,7 +16,8 @@ export class DocHistoryManager { private readonly logger = new Logger(DocHistoryManager.name); constructor( private readonly config: Config, - private readonly db: PrismaService + private readonly db: PrismaService, + private readonly quota: QuotaService ) {} @OnEvent('workspace.deleted') @@ -222,9 +223,6 @@ export class DocHistoryManager { return history.timestamp; } - /** - * @todo(@darkskygit) refactor with [Usage Control] system - */ async getExpiredDateFromNow(workspaceId: string) { const permission = await this.db.workspaceUserPermission.findFirst({ select: { @@ -241,25 +239,9 @@ export class DocHistoryManager { throw new Error('Workspace owner not found'); } - const sub = await this.db.userSubscription.findFirst({ - select: { - id: true, - }, - where: { - userId: permission.userId, - status: SubscriptionStatus.Active, - }, - }); + const quota = await this.quota.getUserQuota(permission.userId); - return new Date( - Date.now() + - 1000 * - 60 * - 60 * - 24 * - // 30 days for subscription user, 7 days for free user - (sub ? 30 : 7) - ); + return new Date(Date.now() + quota.feature.configs.historyPeriod); } @Cron(CronExpression.EVERY_DAY_AT_MIDNIGHT /* everyday at 12am */) diff --git a/packages/backend/server/src/modules/doc/index.ts b/packages/backend/server/src/modules/doc/index.ts index 52e00763a1..3630ddc943 100644 --- a/packages/backend/server/src/modules/doc/index.ts +++ b/packages/backend/server/src/modules/doc/index.ts @@ -1,9 +1,11 @@ import { Module } from '@nestjs/common'; +import { QuotaModule } from '../quota'; import { DocHistoryManager } from './history'; import { DocManager } from './manager'; @Module({ + imports: [QuotaModule], providers: [DocManager, DocHistoryManager], exports: [DocManager, DocHistoryManager], }) diff --git a/packages/backend/server/src/modules/features/feature.ts b/packages/backend/server/src/modules/features/feature.ts index 02886265e1..0e842c5946 100644 --- a/packages/backend/server/src/modules/features/feature.ts +++ b/packages/backend/server/src/modules/features/feature.ts @@ -5,7 +5,7 @@ import { PrismaService } from '../../prisma'; import { FeatureService } from './configure'; import { FeatureType } from './types'; -export enum NewFeaturesKind { +enum NewFeaturesKind { EarlyAccess, } diff --git a/packages/backend/server/src/modules/payment/index.ts b/packages/backend/server/src/modules/payment/index.ts index f7aec52c49..86cea35d4a 100644 --- a/packages/backend/server/src/modules/payment/index.ts +++ b/packages/backend/server/src/modules/payment/index.ts @@ -1,6 +1,7 @@ import { Module } from '@nestjs/common'; import { FeatureModule } from '../features'; +import { QuotaModule } from '../quota'; import { SubscriptionResolver, UserSubscriptionResolver } from './resolver'; import { ScheduleManager } from './schedule'; import { SubscriptionService } from './service'; @@ -8,7 +9,7 @@ import { StripeProvider } from './stripe'; import { StripeWebhook } from './webhook'; @Module({ - imports: [FeatureModule], + imports: [FeatureModule, QuotaModule], providers: [ ScheduleManager, StripeProvider, diff --git a/packages/backend/server/src/modules/payment/service.ts b/packages/backend/server/src/modules/payment/service.ts index aabb3c3a86..25a0fedd8f 100644 --- a/packages/backend/server/src/modules/payment/service.ts +++ b/packages/backend/server/src/modules/payment/service.ts @@ -12,6 +12,7 @@ import Stripe from 'stripe'; import { Config } from '../../config'; import { PrismaService } from '../../prisma'; import { FeatureManagementService } from '../features'; +import { QuotaService, QuotaType } from '../quota'; import { ScheduleManager } from './schedule'; const OnEvent = ( @@ -60,6 +61,11 @@ export enum SubscriptionStatus { Trialing = 'trialing', } +const SubscriptionActivated: Stripe.Subscription.Status[] = [ + SubscriptionStatus.Active, + SubscriptionStatus.Trialing, +]; + export enum InvoiceStatus { Draft = 'draft', Open = 'open', @@ -83,7 +89,8 @@ export class SubscriptionService { private readonly stripe: Stripe, private readonly db: PrismaService, private readonly scheduleManager: ScheduleManager, - private readonly features: FeatureManagementService + private readonly features: FeatureManagementService, + private readonly quota: QuotaService ) { this.paymentConfig = config.payment; @@ -471,6 +478,16 @@ export class SubscriptionService { } } + private getPlanQuota(plan: SubscriptionPlan) { + if (plan === SubscriptionPlan.Free) { + return QuotaType.Quota_FreePlanV1; + } else if (plan === SubscriptionPlan.Pro) { + return QuotaType.Quota_ProPlanV1; + } else { + throw new Error(`Unknown plan: ${plan}`); + } + } + private async saveSubscription( user: User, subscription: Stripe.Subscription, @@ -483,23 +500,28 @@ export class SubscriptionService { subscription = await this.stripe.subscriptions.retrieve(subscription.id); } - // get next bill date from upcoming invoice - // see https://stripe.com/docs/api/invoices/upcoming - let nextBillAt: Date | null = null; - if ( - (subscription.status === SubscriptionStatus.Active || - subscription.status === SubscriptionStatus.Trialing) && - !subscription.canceled_at - ) { - nextBillAt = new Date(subscription.current_period_end * 1000); - } - const price = subscription.items.data[0].price; if (!price.lookup_key) { throw new Error('Unexpected subscription with no key'); } const [plan, recurring] = decodeLookupKey(price.lookup_key); + const planActivated = SubscriptionActivated.includes(subscription.status); + + let nextBillAt: Date | null = null; + if (planActivated) { + // update user's quota if plan activated + await this.quota.switchUserQuota(user.id, this.getPlanQuota(plan)); + + // get next bill date from upcoming invoice + // see https://stripe.com/docs/api/invoices/upcoming + if (!subscription.canceled_at) { + nextBillAt = new Date(subscription.current_period_end * 1000); + } + } else { + // switch to free plan if subscription is canceled + await this.quota.switchUserQuota(user.id, QuotaType.Quota_FreePlanV1); + } const commonData = { start: new Date(subscription.current_period_start * 1000), diff --git a/packages/backend/server/src/modules/quota/quota.ts b/packages/backend/server/src/modules/quota/quota.ts index 9fbc53daf2..d22113c9c7 100644 --- a/packages/backend/server/src/modules/quota/quota.ts +++ b/packages/backend/server/src/modules/quota/quota.ts @@ -2,7 +2,13 @@ import { Injectable } from '@nestjs/common'; import { PrismaService } from '../../prisma'; import { FeatureKind } from '../features'; -import { Quota, QuotaType } from './types'; +import { + formatDate, + formatSize, + getQuotaName, + Quota, + QuotaType, +} from './types'; @Injectable() export class QuotaService { @@ -32,11 +38,22 @@ export class QuotaService { }, }, }); + console.error(userId, quota); return quota as typeof quota & { feature: Pick; }; } + getHumanReadableQuota(feature: QuotaType, configs: Quota['configs']) { + return { + name: getQuotaName(feature), + blobLimit: formatSize(configs.blobLimit), + storageQuota: formatSize(configs.storageQuota), + historyPeriod: formatDate(configs.historyPeriod), + memberLimit: configs.memberLimit.toString(), + }; + } + // get all user quota records async getUserQuotas(userId: string) { const quotas = await this.prisma.userFeatures.findMany({ diff --git a/packages/backend/server/src/modules/quota/types.ts b/packages/backend/server/src/modules/quota/types.ts index d1ac8335f2..498aba31a7 100644 --- a/packages/backend/server/src/modules/quota/types.ts +++ b/packages/backend/server/src/modules/quota/types.ts @@ -5,15 +5,48 @@ export enum QuotaType { Quota_ProPlanV1 = 'pro_plan_v1', } +export enum QuotaName { + free_plan_v1 = 'Free Plan', + pro_plan_v1 = 'Pro Plan', +} + export type Quota = CommonFeature & { type: FeatureKind.Quota; feature: QuotaType; configs: { blobLimit: number; storageQuota: number; + historyPeriod: number; + memberLimit: number; }; }; +const OneKB = 1024; +const OneMB = OneKB * OneKB; +const OneGB = OneKB * OneMB; + +const sizes = ['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)) + ' ' + sizes[i]; +} + +const OneDay = 1000 * 60 * 60 * 24; + +export function formatDate(ms: number): string { + return `${(ms / OneDay).toFixed(0)} days`; +} + +export function getQuotaName(quota: QuotaType): string { + return QuotaName[quota]; +} + export const Quotas: Quota[] = [ { feature: QuotaType.Quota_FreePlanV1, @@ -21,9 +54,13 @@ export const Quotas: Quota[] = [ version: 1, configs: { // single blob limit 10MB - blobLimit: 10 * 1024 * 1024, + blobLimit: 10 * OneMB, // total blob limit 10GB - storageQuota: 10 * 1024 * 1024 * 1024, + storageQuota: 10 * OneGB, + // history period of validity 7 days + historyPeriod: 7 * OneDay, + // member limit 3 + memberLimit: 3, }, }, { @@ -32,9 +69,13 @@ export const Quotas: Quota[] = [ version: 1, configs: { // single blob limit 100MB - blobLimit: 100 * 1024 * 1024, + blobLimit: 100 * OneMB, // total blob limit 100GB - storageQuota: 100 * 1024 * 1024 * 1024, + storageQuota: 100 * OneGB, + // history period of validity 30 days + historyPeriod: 30 * OneDay, + // member limit 10 + memberLimit: 10, }, }, ]; diff --git a/packages/backend/server/src/modules/users/index.ts b/packages/backend/server/src/modules/users/index.ts index 8b506241d8..e48f44e7b7 100644 --- a/packages/backend/server/src/modules/users/index.ts +++ b/packages/backend/server/src/modules/users/index.ts @@ -1,16 +1,17 @@ import { Module } from '@nestjs/common'; import { FeatureModule } from '../features'; +import { QuotaModule } from '../quota'; import { StorageModule } from '../storage'; import { UserResolver } from './resolver'; import { UsersService } from './users'; @Module({ - imports: [StorageModule, FeatureModule], + imports: [StorageModule, FeatureModule, QuotaModule], providers: [UserResolver, UsersService], exports: [UsersService], }) export class UsersModule {} -export { UserType } from './resolver'; +export { UserType } from './types'; export { UsersService } from './users'; diff --git a/packages/backend/server/src/modules/users/resolver.ts b/packages/backend/server/src/modules/users/resolver.ts index 75d1c19940..f002040a5d 100644 --- a/packages/backend/server/src/modules/users/resolver.ts +++ b/packages/backend/server/src/modules/users/resolver.ts @@ -7,11 +7,8 @@ import { import { Args, Context, - Field, - ID, Int, Mutation, - ObjectType, Query, ResolveField, Resolver, @@ -26,47 +23,11 @@ import type { FileUpload } from '../../types'; import { Auth, CurrentUser, Public, Publicable } from '../auth/guard'; import { AuthService } from '../auth/service'; import { FeatureManagementService } from '../features'; +import { QuotaService } from '../quota'; import { StorageService } from '../storage/storage.service'; +import { DeleteAccount, RemoveAvatar, UserQuotaType, UserType } from './types'; import { UsersService } from './users'; -@ObjectType() -export class UserType implements Partial { - @Field(() => ID) - id!: string; - - @Field({ description: 'User name' }) - name!: string; - - @Field({ description: 'User email' }) - email!: string; - - @Field(() => String, { description: 'User avatar url', nullable: true }) - avatarUrl: string | null = null; - - @Field(() => Date, { description: 'User email verified', nullable: true }) - emailVerified: Date | null = null; - - @Field({ description: 'User created date', nullable: true }) - createdAt!: Date; - - @Field(() => Boolean, { - description: 'User password has been set', - nullable: true, - }) - hasPassword?: boolean; -} - -@ObjectType() -export class DeleteAccount { - @Field() - success!: boolean; -} -@ObjectType() -export class RemoveAvatar { - @Field() - success!: boolean; -} - /** * User resolver * All op rate limit: 10 req/m @@ -80,7 +41,8 @@ export class UserResolver { private readonly prisma: PrismaService, private readonly storage: StorageService, private readonly users: UsersService, - private readonly feature: FeatureManagementService + private readonly feature: FeatureManagementService, + private readonly quota: QuotaService ) {} @Throttle({ @@ -148,6 +110,24 @@ export class UserResolver { return user; } + @Throttle({ default: { limit: 10, ttl: 60 } }) + @ResolveField(() => UserQuotaType, { name: 'quota', nullable: true }) + async getQuota(@CurrentUser() me: User) { + const quota = await this.quota.getUserQuota(me.id); + const configs = quota.feature.configs; + + return Object.assign( + { + name: quota.feature.feature, + humanReadable: this.quota.getHumanReadableQuota( + quota.feature.feature, + configs + ), + }, + configs + ); + } + @Throttle({ default: { limit: 10, ttl: 60 } }) @ResolveField(() => Int, { name: 'invoiceCount', diff --git a/packages/backend/server/src/modules/users/types.ts b/packages/backend/server/src/modules/users/types.ts new file mode 100644 index 0000000000..8e95409f09 --- /dev/null +++ b/packages/backend/server/src/modules/users/types.ts @@ -0,0 +1,79 @@ +import { Field, Float, ID, ObjectType } from '@nestjs/graphql'; +import type { User } from '@prisma/client'; + +@ObjectType('UserQuotaHumanReadable') +export 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') +export class UserQuotaType { + @Field({ name: 'name' }) + name!: string; + + @Field(() => Float, { name: 'blobLimit' }) + blobLimit!: number; + + @Field(() => Float, { name: 'storageQuota' }) + storageQuota!: number; + + @Field(() => Float, { name: 'historyPeriod' }) + historyPeriod!: number; + + @Field({ name: 'memberLimit' }) + memberLimit!: number; + + @Field({ name: 'humanReadable' }) + humanReadable!: UserQuotaHumanReadableType; +} + +@ObjectType() +export class UserType implements Partial { + @Field(() => ID) + id!: string; + + @Field({ description: 'User name' }) + name!: string; + + @Field({ description: 'User email' }) + email!: string; + + @Field(() => String, { description: 'User avatar url', nullable: true }) + avatarUrl: string | null = null; + + @Field(() => Date, { description: 'User email verified', nullable: true }) + emailVerified: Date | null = null; + + @Field({ description: 'User created date', nullable: true }) + createdAt!: Date; + + @Field(() => Boolean, { + description: 'User password has been set', + nullable: true, + }) + hasPassword?: boolean; +} + +@ObjectType() +export class DeleteAccount { + @Field() + success!: boolean; +} +@ObjectType() +export class RemoveAvatar { + @Field() + success!: boolean; +} diff --git a/packages/backend/server/src/modules/workspaces/resolver.ts b/packages/backend/server/src/modules/workspaces/resolver.ts index 7a81860ff6..f8e30cf30f 100644 --- a/packages/backend/server/src/modules/workspaces/resolver.ts +++ b/packages/backend/server/src/modules/workspaces/resolver.ts @@ -43,8 +43,7 @@ import { Auth, CurrentUser, Public } from '../auth'; import { MailService } from '../auth/mailer'; import { AuthService } from '../auth/service'; import { QuotaManagementService } from '../quota'; -import { UsersService } from '../users'; -import { UserType } from '../users/resolver'; +import { UsersService, UserType } from '../users'; import { PermissionService, PublicPageMode } from './permission'; import { Permission } from './types'; import { defaultWorkspaceAvatar } from './utils'; diff --git a/packages/backend/server/src/schema.gql b/packages/backend/server/src/schema.gql index 001dc645ed..6707d700df 100644 --- a/packages/backend/server/src/schema.gql +++ b/packages/backend/server/src/schema.gql @@ -10,6 +10,23 @@ type ServerConfigType { flavor: String! } +type UserQuotaHumanReadable { + name: String! + blobLimit: String! + storageQuota: String! + historyPeriod: String! + memberLimit: String! +} + +type UserQuota { + name: String! + blobLimit: Float! + storageQuota: Float! + historyPeriod: Float! + memberLimit: Int! + humanReadable: UserQuotaHumanReadable! +} + type UserType { id: ID! @@ -31,6 +48,7 @@ type UserType { """User password has been set""" hasPassword: Boolean token: TokenType! + quota: UserQuota """Get user invoice count""" invoiceCount: Int! diff --git a/packages/backend/server/tests/doc.spec.ts b/packages/backend/server/tests/doc.spec.ts index ac5262748a..f2a8e5ac7e 100644 --- a/packages/backend/server/tests/doc.spec.ts +++ b/packages/backend/server/tests/doc.spec.ts @@ -14,10 +14,16 @@ import { import { CacheModule } from '../src/cache'; import { Config, ConfigModule } from '../src/config'; +import { + collectMigrations, + RevertCommand, + RunCommand, +} from '../src/data/commands/run'; import { EventModule } from '../src/event'; import { DocManager, DocModule } from '../src/modules/doc'; +import { QuotaModule } from '../src/modules/quota'; import { PrismaModule, PrismaService } from '../src/prisma'; -import { flushDB } from './utils'; +import { FakeStorageModule, flushDB } from './utils'; const createModule = () => { return Test.createTestingModule({ @@ -25,8 +31,12 @@ const createModule = () => { PrismaModule, CacheModule, EventModule, + QuotaModule, + FakeStorageModule.forRoot(), ConfigModule.forRoot(), DocModule, + RevertCommand, + RunCommand, ], }).compile(); }; @@ -45,6 +55,13 @@ test.beforeEach(async () => { app = m.createNestApplication(); app.enableShutdownHooks(); await app.init(); + + // init features + const run = m.get(RunCommand); + const revert = m.get(RevertCommand); + const migrations = await collectMigrations(); + await Promise.allSettled(migrations.map(m => revert.run([m.name]))); + await run.run(); }); test.afterEach.always(async () => { diff --git a/packages/backend/server/tests/feature.spec.ts b/packages/backend/server/tests/feature.spec.ts index ac8b776a97..daaa5efed6 100644 --- a/packages/backend/server/tests/feature.spec.ts +++ b/packages/backend/server/tests/feature.spec.ts @@ -57,14 +57,10 @@ test.beforeEach(async t => { ], }).compile(); - const quota = module.get(FeatureService); - const storageQuota = module.get(FeatureManagementService); - const auth = module.get(AuthService); - t.context.app = module; - t.context.feature = quota; - t.context.early_access = storageQuota; - t.context.auth = auth; + t.context.auth = module.get(AuthService); + t.context.feature = module.get(FeatureService); + t.context.early_access = module.get(FeatureManagementService); // init features await initFeatureConfigs(module); diff --git a/packages/backend/server/tests/history.spec.ts b/packages/backend/server/tests/history.spec.ts index 0c77b1facd..d496816461 100644 --- a/packages/backend/server/tests/history.spec.ts +++ b/packages/backend/server/tests/history.spec.ts @@ -8,8 +8,9 @@ import * as Sinon from 'sinon'; import { ConfigModule } from '../src/config'; import type { EventPayload } from '../src/event'; import { DocHistoryManager } from '../src/modules/doc'; +import { QuotaModule } from '../src/modules/quota'; import { PrismaModule, PrismaService } from '../src/prisma'; -import { flushDB } from './utils'; +import { FakeStorageModule, flushDB } from './utils'; let app: INestApplication; let m: TestingModule; @@ -20,7 +21,13 @@ let db: PrismaService; test.beforeEach(async () => { await flushDB(); m = await Test.createTestingModule({ - imports: [PrismaModule, ScheduleModule.forRoot(), ConfigModule.forRoot()], + imports: [ + PrismaModule, + QuotaModule, + FakeStorageModule.forRoot(), + ScheduleModule.forRoot(), + ConfigModule.forRoot(), + ], providers: [DocHistoryManager], }).compile(); @@ -277,8 +284,8 @@ test('should be able to recover from history', async t => { t.is(history2.timestamp.getTime(), snapshot.updatedAt.getTime()); // new history data force created with snapshot state before recovered - t.deepEqual(history2?.blob, Buffer.from([1, 1])); - t.deepEqual(history2?.state, Buffer.from([1, 1])); + t.deepEqual(history2.blob, Buffer.from([1, 1])); + t.deepEqual(history2.state, Buffer.from([1, 1])); }); test('should be able to cleanup expired history', async t => { diff --git a/packages/backend/server/tests/quota.spec.ts b/packages/backend/server/tests/quota.spec.ts index fe53b2ed5e..7f56190de3 100644 --- a/packages/backend/server/tests/quota.spec.ts +++ b/packages/backend/server/tests/quota.spec.ts @@ -16,9 +16,8 @@ import { QuotaType, } from '../src/modules/quota'; import { PrismaModule } from '../src/prisma'; -import { StorageModule } from '../src/storage'; import { RateLimiterModule } from '../src/throttler'; -import { initFeatureConfigs } from './utils'; +import { FakeStorageModule, initFeatureConfigs } from './utils'; const test = ava as TestFn<{ auth: AuthService; @@ -47,10 +46,10 @@ test.beforeEach(async t => { host: 'example.org', https: true, }), - StorageModule.forRoot(), PrismaModule, AuthModule, QuotaModule, + FakeStorageModule.forRoot(), RateLimiterModule, RevertCommand, RunCommand, diff --git a/packages/backend/server/tests/user.e2e.ts b/packages/backend/server/tests/user.e2e.ts index 065925af21..7087ba3f9f 100644 --- a/packages/backend/server/tests/user.e2e.ts +++ b/packages/backend/server/tests/user.e2e.ts @@ -6,7 +6,8 @@ import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs'; import request from 'supertest'; import { AppModule } from '../src/app'; -import { currentUser, signUp } from './utils'; +import { RevertCommand, RunCommand } from '../src/data/commands/run'; +import { currentUser, initFeatureConfigs, signUp } from './utils'; let app: INestApplication; @@ -21,6 +22,7 @@ test.beforeEach(async () => { test.beforeEach(async () => { const module = await Test.createTestingModule({ imports: [AppModule], + providers: [RevertCommand, RunCommand], }).compile(); app = module.createNestApplication(); app.use( @@ -30,6 +32,9 @@ test.beforeEach(async () => { }) ); await app.init(); + + // init features + await initFeatureConfigs(module); }); test.afterEach.always(async () => { diff --git a/packages/backend/server/tests/utils.ts b/packages/backend/server/tests/utils.ts index 98741c99c1..e413de583e 100644 --- a/packages/backend/server/tests/utils.ts +++ b/packages/backend/server/tests/utils.ts @@ -1,6 +1,10 @@ import { randomUUID } from 'node:crypto'; -import type { INestApplication } from '@nestjs/common'; +import type { + DynamicModule, + FactoryProvider, + INestApplication, +} from '@nestjs/common'; import { TestingModule } from '@nestjs/testing'; import { hashSync } from '@node-rs/argon2'; import { PrismaClient, type User } from '@prisma/client'; @@ -10,6 +14,7 @@ import { RevertCommand, RunCommand } from '../src/data/commands/run'; import type { TokenType } from '../src/modules/auth'; import type { UserType } from '../src/modules/users'; import type { InvitationType, WorkspaceType } from '../src/modules/workspaces'; +import { StorageProvide } from '../src/storage'; const gql = '/graphql'; @@ -563,6 +568,24 @@ export class FakePrisma { } } +export class FakeStorageModule { + static forRoot(): DynamicModule { + const storageProvider: FactoryProvider = { + provide: StorageProvide, + useFactory: async () => { + return null; + }, + }; + + return { + global: true, + module: FakeStorageModule, + providers: [storageProvider], + exports: [storageProvider], + }; + } +} + export async function initFeatureConfigs(module: TestingModule) { const run = module.get(RunCommand); const revert = module.get(RevertCommand); diff --git a/packages/frontend/core/src/components/affine/new-workspace-setting-detail/members.tsx b/packages/frontend/core/src/components/affine/new-workspace-setting-detail/members.tsx index f1d5a0adde..b29f39ae99 100644 --- a/packages/frontend/core/src/components/affine/new-workspace-setting-detail/members.tsx +++ b/packages/frontend/core/src/components/affine/new-workspace-setting-detail/members.tsx @@ -15,7 +15,7 @@ import { Menu, MenuItem } from '@affine/component/ui/menu'; import { Tooltip } from '@affine/component/ui/tooltip'; import type { AffineOfficialWorkspace } from '@affine/env/workspace'; import { WorkspaceFlavour } from '@affine/env/workspace'; -import { Permission, SubscriptionPlan } from '@affine/graphql'; +import { Permission } from '@affine/graphql'; import { useAFFiNEI18N } from '@affine/i18n/hooks'; import { ArrowRightBigIcon, MoreVerticalIcon } from '@blocksuite/icons'; import clsx from 'clsx'; @@ -37,17 +37,11 @@ import { useInviteMember } from '../../../hooks/affine/use-invite-member'; import { useMemberCount } from '../../../hooks/affine/use-member-count'; import { type Member, useMembers } from '../../../hooks/affine/use-members'; import { useRevokeMemberPermission } from '../../../hooks/affine/use-revoke-member-permission'; -import { useUserSubscription } from '../../../hooks/use-subscription'; +import { useUserQuota } from '../../../hooks/use-quota'; import { AffineErrorBoundary } from '../affine-error-boundary'; import * as style from './style.css'; import type { WorkspaceSettingDetailProps } from './types'; -enum MemberLimitCount { - Free = '3', - Pro = '10', - Other = '?', -} - const COUNT_PER_PAGE = 8; export interface MembersPanelProps extends WorkspaceSettingDetailProps { upgradable: boolean; @@ -148,23 +142,17 @@ export const CloudWorkspaceMembersPanel = ({ }); }, [setSettingModalAtom]); - const [subscription] = useUserSubscription(); - const plan = subscription?.plan ?? SubscriptionPlan.Free; - const memberLimit = useMemo(() => { - if (plan === SubscriptionPlan.Free) { - return MemberLimitCount.Free; - } - if (plan === SubscriptionPlan.Pro) { - return MemberLimitCount.Pro; - } - return MemberLimitCount.Other; - }, [plan]); + const quota = useUserQuota(); + const desc = useMemo(() => { + if (!quota) return null; + + const humanReadable = quota.humanReadable; return ( {t['com.affine.payment.member.description']({ - planName: plan, - memberLimit, + planName: humanReadable.name, + memberLimit: humanReadable.memberLimit, })} {upgradable ? ( <> @@ -179,7 +167,7 @@ export const CloudWorkspaceMembersPanel = ({ ) : null} ); - }, [handleUpgrade, memberLimit, plan, t, upgradable]); + }, [handleUpgrade, quota, t, upgradable]); return ( <> diff --git a/packages/frontend/core/src/hooks/use-quota.ts b/packages/frontend/core/src/hooks/use-quota.ts new file mode 100644 index 0000000000..5d2d8db58a --- /dev/null +++ b/packages/frontend/core/src/hooks/use-quota.ts @@ -0,0 +1,10 @@ +import { quotaQuery } from '@affine/graphql'; +import { useQuery } from '@affine/workspace/affine/gql'; + +export const useUserQuota = () => { + const { data } = useQuery({ + query: quotaQuery, + }); + + return data.currentUser?.quota || null; +}; diff --git a/packages/frontend/graphql/src/graphql/index.ts b/packages/frontend/graphql/src/graphql/index.ts index 972fc945ca..5d2e01f89d 100644 --- a/packages/frontend/graphql/src/graphql/index.ts +++ b/packages/frontend/graphql/src/graphql/index.ts @@ -508,6 +508,32 @@ mutation publishPage($workspaceId: String!, $pageId: String!, $mode: PublicPageM }`, }; +export const quotaQuery = { + id: 'quotaQuery' as const, + operationName: 'quota', + definitionName: 'currentUser', + containsFile: false, + query: ` +query quota { + currentUser { + quota { + name + blobLimit + storageQuota + historyPeriod + memberLimit + humanReadable { + name + blobLimit + storageQuota + historyPeriod + memberLimit + } + } + } +}`, +}; + export const recoverDocMutation = { id: 'recoverDocMutation' as const, operationName: 'recoverDoc', diff --git a/packages/frontend/graphql/src/graphql/quota.gql b/packages/frontend/graphql/src/graphql/quota.gql new file mode 100644 index 0000000000..e02b1c2784 --- /dev/null +++ b/packages/frontend/graphql/src/graphql/quota.gql @@ -0,0 +1,18 @@ +query quota { + currentUser { + quota { + name + blobLimit + storageQuota + historyPeriod + memberLimit + humanReadable { + name + blobLimit + storageQuota + historyPeriod + memberLimit + } + } + } +} diff --git a/packages/frontend/graphql/src/schema.ts b/packages/frontend/graphql/src/schema.ts index 483711bfd2..cc7c293213 100644 --- a/packages/frontend/graphql/src/schema.ts +++ b/packages/frontend/graphql/src/schema.ts @@ -499,6 +499,31 @@ export type PublishPageMutation = { }; }; +export type QuotaQueryVariables = Exact<{ [key: string]: never }>; + +export type QuotaQuery = { + __typename?: 'Query'; + currentUser: { + __typename?: 'UserType'; + quota: { + __typename?: 'UserQuota'; + name: string; + blobLimit: number; + storageQuota: number; + historyPeriod: number; + memberLimit: number; + humanReadable: { + __typename?: 'UserQuotaHumanReadable'; + name: string; + blobLimit: string; + storageQuota: string; + historyPeriod: string; + memberLimit: string; + }; + } | null; + } | null; +}; + export type RecoverDocMutationVariables = Exact<{ workspaceId: Scalars['String']['input']; docId: Scalars['String']['input']; @@ -819,6 +844,11 @@ export type Queries = variables: PricesQueryVariables; response: PricesQuery; } + | { + name: 'quotaQuery'; + variables: QuotaQueryVariables; + response: QuotaQuery; + } | { name: 'serverConfigQuery'; variables: ServerConfigQueryVariables;