refactor(server): use feature model (#9932)

This commit is contained in:
forehalo
2025-02-05 10:27:26 +00:00
parent 0ff8d3af6f
commit 7826e2b7c8
121 changed files with 1723 additions and 3826 deletions

View File

@@ -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<typeof UserPlanQuotaConfig>;
const WorkspaceQuotaConfig = UserPlanQuotaConfig.extend({
// seat quota
seatQuota: z.number(),
}).omit({
copilotActionLimit: true,
});
function feature<T extends z.ZodRawShape>(configs: z.ZodObject<T>) {
return z.object({
type: z.literal(FeatureType.Feature),
configs: configs,
});
export type WorkspaceQuota = z.infer<typeof WorkspaceQuotaConfig>;
const EMPTY_CONFIG = z.object({});
export enum FeatureType {
Feature,
Quota,
}
function quota<T extends z.ZodRawShape>(configs: z.ZodObject<T>) {
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<Feature, z.ZodObject<any>>;
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<T extends FeatureName> = z.infer<
(typeof Features)[T]['shape']['configs']
export type FeatureConfig<T extends FeatureName> = z.infer<
(typeof FeaturesShapes)[T]
>;
export const FeatureConfigs: {
[K in FeatureName]: {
type: FeatureType;
configs: FeatureConfig<K>;
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: {},
},
};

View File

@@ -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<T extends FeatureName>(name: T, configs: FeatureConfigs<T>) {
async upsert<T extends FeatureName>(
name: T,
configs: FeatureConfig<T>,
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<T> };
return feature as Feature & { configs: FeatureConfig<T> };
}
/**
@@ -67,8 +74,7 @@ export class FeatureModel extends BaseModel {
*/
async try_get_unchecked<T extends FeatureName>(name: T) {
const feature = await this.db.feature.findFirst({
where: { feature: name },
orderBy: { version: 'desc' },
where: { name },
});
return feature as Omit<Feature, 'configs'> & {
@@ -104,10 +110,22 @@ export class FeatureModel extends BaseModel {
});
}
return parseResult.data as FeatureConfigs<T>;
return parseResult.data as FeatureConfig<T>;
}
getConfigShape(name: FeatureName): z.ZodObject<any> {
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);
}
}
}

View File

@@ -77,6 +77,7 @@ const ModelsSymbolProvider: ExistingProvider = {
})
export class ModelsModule {}
export * from './common';
export * from './feature';
export * from './page';
export * from './session';

View File

@@ -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<T extends UserFeatureName>(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);
}
}

View File

@@ -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<Prisma.UserCreateInput, 'name'> & { name?: string };
type UpdateUserInput = Omit<Partial<Prisma.UserCreateInput>, '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);

View File

@@ -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<T extends WorkspaceFeatureName>(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<T extends WorkspaceFeatureName>(
workspaceId: string,
featureName: T,
name: T,
reason: string,
overrides?: Partial<FeatureConfigs<T>>
overrides?: Partial<FeatureConfig<T>>
) {
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<T extends WorkspaceFeatureName>(
workspaceId: string,
from: WorkspaceFeatureName,
to: T,
reason: string,
overrides?: Partial<FeatureConfigs<T>>
) {
await this.remove(workspaceId, from);
return await this.add(workspaceId, to, reason, overrides);
}
}