feat: struct type feature config (#5142)

This commit is contained in:
DarkSky
2023-12-14 09:50:51 +00:00
parent 2b7f6f8b74
commit 0c2d2f8d16
33 changed files with 1223 additions and 1015 deletions

View File

@@ -82,7 +82,8 @@
"socket.io": "^4.7.2", "socket.io": "^4.7.2",
"stripe": "^14.5.0", "stripe": "^14.5.0",
"ws": "^8.14.2", "ws": "^8.14.2",
"yjs": "^13.6.10" "yjs": "^13.6.10",
"zod": "^3.22.4"
}, },
"devDependencies": { "devDependencies": {
"@affine-test/kit": "workspace:*", "@affine-test/kit": "workspace:*",

View File

@@ -1,10 +1,12 @@
import { Prisma } from '@prisma/client';
import { import {
CommonFeature, CommonFeature,
FeatureKind, FeatureKind,
Features, Features,
FeatureType, FeatureType,
} from '../../modules/features'; } from '../../modules/features';
import { Quotas } from '../../modules/quota/types'; import { Quotas } from '../../modules/quota/schema';
import { PrismaService } from '../../prisma'; import { PrismaService } from '../../prisma';
export class UserFeaturesInit1698652531198 { export class UserFeaturesInit1698652531198 {
@@ -48,7 +50,7 @@ async function upsertFeature(
feature: feature.feature, feature: feature.feature,
type: feature.type, type: feature.type,
version: feature.version, version: feature.version,
configs: feature.configs, configs: feature.configs as Prisma.InputJsonValue,
}, },
}); });
} }

View File

@@ -240,8 +240,7 @@ export class DocHistoryManager {
} }
const quota = await this.quota.getUserQuota(permission.userId); const quota = await this.quota.getUserQuota(permission.userId);
return quota.feature.historyPeriodFromNow;
return new Date(Date.now() + quota.feature.configs.historyPeriod);
} }
@Cron(CronExpression.EVERY_DAY_AT_MIDNIGHT /* everyday at 12am */) @Cron(CronExpression.EVERY_DAY_AT_MIDNIGHT /* everyday at 12am */)

View File

@@ -1,80 +1,78 @@
import { Injectable, Logger } from '@nestjs/common';
import { Config } from '../../config';
import { PrismaService } from '../../prisma'; import { PrismaService } from '../../prisma';
import { FeatureService } from './configure'; import { Feature, FeatureSchema, FeatureType } from './types';
import { FeatureType } from './types';
enum NewFeaturesKind { class FeatureConfig {
EarlyAccess, readonly config: Feature;
}
@Injectable() constructor(data: any) {
export class FeatureManagementService { const config = FeatureSchema.safeParse(data);
protected logger = new Logger(FeatureManagementService.name); if (config.success) {
constructor( this.config = config.data;
private readonly feature: FeatureService,
private readonly prisma: PrismaService,
private readonly config: Config
) {}
isStaff(email: string) {
return email.endsWith('@toeverything.info');
}
async addEarlyAccess(userId: string) {
return this.feature.addUserFeature(
userId,
FeatureType.EarlyAccess,
1,
'Early access user'
);
}
async removeEarlyAccess(userId: string) {
return this.feature.removeUserFeature(userId, FeatureType.EarlyAccess);
}
async listEarlyAccess() {
return this.feature.listFeatureUsers(FeatureType.EarlyAccess);
}
/// check early access by email
async canEarlyAccess(email: string) {
if (
this.config.featureFlags.earlyAccessPreview &&
!email.endsWith('@toeverything.info')
) {
const user = await this.prisma.user.findFirst({
where: {
email,
},
});
if (user) {
const canEarlyAccess = await this.feature
.hasFeature(user.id, FeatureType.EarlyAccess)
.catch(() => false);
if (canEarlyAccess) {
return true;
}
// TODO: Outdated, switch to feature gates
const oldCanEarlyAccess = await this.prisma.newFeaturesWaitingList
.findUnique({
where: { email, type: NewFeaturesKind.EarlyAccess },
})
.then(x => !!x)
.catch(() => false);
if (oldCanEarlyAccess) {
this.logger.warn(
`User ${email} has early access in old table but not in new table`
);
}
return oldCanEarlyAccess;
}
return false;
} else { } else {
return true; throw new Error(`Invalid quota config: ${config.error.message}`);
} }
} }
/// feature name of quota
get name() {
return this.config.feature;
}
}
export class EarlyAccessFeatureConfig extends FeatureConfig {
constructor(data: any) {
super(data);
if (this.config.feature !== FeatureType.EarlyAccess) {
throw new Error('Invalid feature config: type is not EarlyAccess');
}
}
checkWhiteList(email: string) {
for (const domain in this.config.configs.whitelist) {
if (email.endsWith(domain)) {
return true;
}
}
return false;
}
}
const FeatureConfigMap = {
[FeatureType.EarlyAccess]: EarlyAccessFeatureConfig,
};
const FeatureCache = new Map<
number,
InstanceType<(typeof FeatureConfigMap)[FeatureType]>
>();
export async function getFeature(prisma: PrismaService, featureId: number) {
const cachedQuota = FeatureCache.get(featureId);
if (cachedQuota) {
return cachedQuota;
}
const feature = await prisma.features.findFirst({
where: {
id: featureId,
},
});
if (!feature) {
// this should unreachable
throw new Error(`Quota config ${featureId} not found`);
}
const ConfigClass = FeatureConfigMap[feature.feature as FeatureType];
if (!ConfigClass) {
throw new Error(`Feature config ${featureId} not found`);
}
const config = new ConfigClass(feature);
// we always edit quota config as a new quota config
// so we can cache it by featureId
FeatureCache.set(featureId, config);
return config;
} }

View File

@@ -1,8 +1,8 @@
import { Module } from '@nestjs/common'; import { Module } from '@nestjs/common';
import { PrismaService } from '../../prisma'; import { PrismaService } from '../../prisma';
import { FeatureService } from './configure'; import { FeatureManagementService } from './management';
import { FeatureManagementService } from './feature'; import { FeatureService } from './service';
/** /**
* Feature module provider pre-user feature flag management. * Feature module provider pre-user feature flag management.
@@ -16,6 +16,6 @@ import { FeatureManagementService } from './feature';
}) })
export class FeatureModule {} export class FeatureModule {}
export type { CommonFeature, Feature } from './types'; export { type CommonFeature, commonFeatureSchema } from './types';
export { FeatureKind, Features, FeatureType } from './types'; export { FeatureKind, Features, FeatureType } from './types';
export { FeatureManagementService, FeatureService, PrismaService }; export { FeatureManagementService, FeatureService, PrismaService };

View File

@@ -0,0 +1,89 @@
import { Injectable, Logger, OnModuleInit } from '@nestjs/common';
import { Config } from '../../config';
import { PrismaService } from '../../prisma';
import { EarlyAccessFeatureConfig } from './feature';
import { FeatureService } from './service';
import { FeatureType } from './types';
enum NewFeaturesKind {
EarlyAccess,
}
@Injectable()
export class FeatureManagementService implements OnModuleInit {
protected logger = new Logger(FeatureManagementService.name);
private earlyAccessFeature?: EarlyAccessFeatureConfig;
constructor(
private readonly feature: FeatureService,
private readonly prisma: PrismaService,
private readonly config: Config
) {}
async onModuleInit() {
this.earlyAccessFeature = await this.feature.getFeature(
FeatureType.EarlyAccess
);
}
// ======== Admin ========
// todo(@darkskygit): replace this with abac
isStaff(email: string) {
return this.earlyAccessFeature?.checkWhiteList(email) ?? false;
}
// ======== Early Access ========
async addEarlyAccess(userId: string) {
return this.feature.addUserFeature(
userId,
FeatureType.EarlyAccess,
1,
'Early access user'
);
}
async removeEarlyAccess(userId: string) {
return this.feature.removeUserFeature(userId, FeatureType.EarlyAccess);
}
async listEarlyAccess() {
return this.feature.listFeatureUsers(FeatureType.EarlyAccess);
}
/// check early access by email
async canEarlyAccess(email: string) {
if (this.config.featureFlags.earlyAccessPreview && !this.isStaff(email)) {
const user = await this.prisma.user.findFirst({
where: {
email,
},
});
if (user) {
const canEarlyAccess = await this.feature
.hasFeature(user.id, FeatureType.EarlyAccess)
.catch(() => false);
if (canEarlyAccess) {
return true;
}
// TODO: Outdated, switch to feature gates
const oldCanEarlyAccess = await this.prisma.newFeaturesWaitingList
.findUnique({
where: { email, type: NewFeaturesKind.EarlyAccess },
})
.then(x => !!x)
.catch(() => false);
if (oldCanEarlyAccess) {
this.logger.warn(
`User ${email} has early access in old table but not in new table`
);
}
return oldCanEarlyAccess;
}
return false;
} else {
return true;
}
}
}

View File

@@ -1,7 +1,9 @@
import { Injectable } from '@nestjs/common'; import { Injectable } from '@nestjs/common';
import { PrismaService } from '../../prisma'; import { PrismaService } from '../../prisma';
import { Feature, FeatureKind, FeatureType } from './types'; import { UserType } from '../users/types';
import { getFeature } from './feature';
import { FeatureKind, FeatureType } from './types';
@Injectable() @Injectable()
export class FeatureService { export class FeatureService {
@@ -27,15 +29,20 @@ export class FeatureService {
} }
async getFeature(feature: FeatureType) { async getFeature(feature: FeatureType) {
return this.prisma.features.findFirst({ const data = await this.prisma.features.findFirst({
where: { where: {
feature, feature,
type: FeatureKind.Feature, type: FeatureKind.Feature,
}, },
select: { id: true },
orderBy: { orderBy: {
version: 'desc', version: 'desc',
}, },
}); });
if (data) {
return getFeature(this.prisma, data.id);
}
return undefined;
} }
async addUserFeature( async addUserFeature(
@@ -120,21 +127,21 @@ export class FeatureService {
reason: true, reason: true,
createdAt: true, createdAt: true,
expiredAt: true, expiredAt: true,
feature: { featureId: true,
select: {
feature: true,
configs: true,
},
},
}, },
}); });
return features as typeof features &
{ const configs = await Promise.all(
feature: Pick<Feature, 'feature' | 'configs'>; features.map(async feature => ({
}[]; ...feature,
feature: await getFeature(this.prisma, feature.featureId),
}))
);
return configs.filter(feature => !!feature.feature);
} }
async listFeatureUsers(feature: FeatureType) { async listFeatureUsers(feature: FeatureType): Promise<UserType[]> {
return this.prisma.userFeatures return this.prisma.userFeatures
.findMany({ .findMany({
where: { where: {

View File

@@ -1,26 +1,48 @@
import type { Prisma } from '@prisma/client'; import { URL } from 'node:url';
import { z } from 'zod';
/// ======== common schema ========
export enum FeatureKind { export enum FeatureKind {
Feature, Feature,
Quota, Quota,
} }
export type CommonFeature = { export const commonFeatureSchema = z.object({
feature: string; feature: z.string(),
type: FeatureKind; type: z.nativeEnum(FeatureKind),
version: number; version: z.number(),
configs: Prisma.InputJsonValue; configs: z.unknown(),
}; });
export type Feature = CommonFeature & { export type CommonFeature = z.infer<typeof commonFeatureSchema>;
type: FeatureKind.Feature;
feature: FeatureType; /// ======== feature define ========
};
export enum FeatureType { export enum FeatureType {
EarlyAccess = 'early_access', EarlyAccess = 'early_access',
} }
function checkHostname(host: string) {
try {
return new URL(`https://${host}`).hostname === host;
} catch (_) {
return false;
}
}
const featureEarlyAccess = z.object({
feature: z.literal(FeatureType.EarlyAccess),
configs: z.object({
whitelist: z
.string()
.startsWith('@')
.refine(domain => checkHostname(domain.slice(1)))
.array(),
}),
});
export const Features: Feature[] = [ export const Features: Feature[] = [
{ {
feature: FeatureType.EarlyAccess, feature: FeatureType.EarlyAccess,
@@ -31,3 +53,13 @@ export const Features: Feature[] = [
}, },
}, },
]; ];
/// ======== schema infer ========
export const FeatureSchema = commonFeatureSchema
.extend({
type: z.literal(FeatureKind.Feature),
})
.and(z.discriminatedUnion('feature', [featureEarlyAccess]));
export type Feature = z.infer<typeof FeatureSchema>;

View File

@@ -480,9 +480,9 @@ export class SubscriptionService {
private getPlanQuota(plan: SubscriptionPlan) { private getPlanQuota(plan: SubscriptionPlan) {
if (plan === SubscriptionPlan.Free) { if (plan === SubscriptionPlan.Free) {
return QuotaType.Quota_FreePlanV1; return QuotaType.FreePlanV1;
} else if (plan === SubscriptionPlan.Pro) { } else if (plan === SubscriptionPlan.Pro) {
return QuotaType.Quota_ProPlanV1; return QuotaType.ProPlanV1;
} else { } else {
throw new Error(`Unknown plan: ${plan}`); throw new Error(`Unknown plan: ${plan}`);
} }
@@ -520,7 +520,7 @@ export class SubscriptionService {
} }
} else { } else {
// switch to free plan if subscription is canceled // switch to free plan if subscription is canceled
await this.quota.switchUserQuota(user.id, QuotaType.Quota_FreePlanV1); await this.quota.switchUserQuota(user.id, QuotaType.FreePlanV1);
} }
const commonData = { const commonData = {

View File

@@ -0,0 +1,5 @@
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'];

View File

@@ -1,7 +1,7 @@
import { Module } from '@nestjs/common'; import { Module } from '@nestjs/common';
import { PermissionService } from '../workspaces/permission'; import { PermissionService } from '../workspaces/permission';
import { QuotaService } from './quota'; import { QuotaService } from './service';
import { QuotaManagementService } from './storage'; import { QuotaManagementService } from './storage';
/** /**
@@ -17,4 +17,5 @@ import { QuotaManagementService } from './storage';
export class QuotaModule {} export class QuotaModule {}
export { QuotaManagementService, QuotaService }; export { QuotaManagementService, QuotaService };
export { Quota_FreePlanV1, Quota_ProPlanV1, Quotas, QuotaType } from './types'; export { Quota_FreePlanV1, Quota_ProPlanV1, Quotas } from './schema';
export { QuotaType } from './types';

View File

@@ -1,157 +1,81 @@
import { Injectable } from '@nestjs/common';
import { PrismaService } from '../../prisma'; import { PrismaService } from '../../prisma';
import { FeatureKind } from '../features'; import { formatDate, formatSize, Quota, QuotaSchema } from './types';
import {
formatDate,
formatSize,
getQuotaName,
Quota,
QuotaType,
} from './types';
@Injectable() const QuotaCache = new Map<number, QuotaConfig>();
export class QuotaService {
constructor(private readonly prisma: PrismaService) {}
// get activated user quota export class QuotaConfig {
async getUserQuota(userId: string) { readonly config: Quota;
const quota = await this.prisma.userFeatures.findFirst({
static async get(prisma: PrismaService, featureId: number) {
const cachedQuota = QuotaCache.get(featureId);
if (cachedQuota) {
return cachedQuota;
}
const quota = await prisma.features.findFirst({
where: { where: {
user: { id: featureId,
id: userId,
},
feature: {
type: FeatureKind.Quota,
},
activated: true,
},
select: {
reason: true,
createdAt: true,
expiredAt: true,
feature: {
select: {
feature: true,
configs: true,
},
},
}, },
}); });
console.error(userId, quota);
return quota as typeof quota & { if (!quota) {
feature: Pick<Quota, 'feature' | 'configs'>; 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;
} }
getHumanReadableQuota(feature: QuotaType, configs: Quota['configs']) { private constructor(data: 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
)})}`
);
}
}
/// feature name of quota
get name() {
return this.config.feature;
}
get blobLimit() {
return this.config.configs.blobLimit;
}
get storageQuota() {
return this.config.configs.storageQuota;
}
get historyPeriod() {
return this.config.configs.historyPeriod;
}
get historyPeriodFromNow() {
return new Date(Date.now() + this.historyPeriod);
}
get memberLimit() {
return this.config.configs.memberLimit;
}
get humanReadable() {
return { return {
name: getQuotaName(feature), name: this.config.configs.name,
blobLimit: formatSize(configs.blobLimit), blobLimit: formatSize(this.blobLimit),
storageQuota: formatSize(configs.storageQuota), storageQuota: formatSize(this.storageQuota),
historyPeriod: formatDate(configs.historyPeriod), historyPeriod: formatDate(this.historyPeriod),
memberLimit: configs.memberLimit.toString(), memberLimit: this.memberLimit.toString(),
}; };
} }
// get all user quota records
async getUserQuotas(userId: string) {
const quotas = await this.prisma.userFeatures.findMany({
where: {
user: {
id: userId,
},
feature: {
type: FeatureKind.Quota,
},
},
select: {
activated: true,
reason: true,
createdAt: true,
expiredAt: true,
feature: {
select: {
feature: true,
configs: true,
},
},
},
});
return quotas as typeof quotas &
{
feature: Pick<Quota, 'feature' | 'configs'>;
}[];
}
// 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 latestFreePlan = await tx.features.aggregate({
where: {
feature: QuotaType.Quota_FreePlanV1,
},
_max: {
version: true,
},
});
// we will deactivate all exists quota for this user
await tx.userFeatures.updateMany({
where: {
id: undefined,
userId,
feature: {
type: FeatureKind.Quota,
},
},
data: {
activated: false,
},
});
await tx.userFeatures.create({
data: {
user: {
connect: {
id: userId,
},
},
feature: {
connect: {
feature_version: {
feature: quota,
version: latestFreePlan._max.version || 1,
},
type: FeatureKind.Quota,
},
},
reason: reason ?? 'switch quota',
activated: true,
expiredAt,
},
});
});
}
async hasQuota(userId: string, quota: QuotaType) {
return this.prisma.userFeatures
.count({
where: {
userId,
feature: {
feature: quota,
type: FeatureKind.Quota,
},
activated: true,
},
})
.then(count => count > 0);
}
} }

View File

@@ -0,0 +1,50 @@
import { FeatureKind } from '../features';
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,
},
},
];
export const Quota_FreePlanV1 = {
feature: Quotas[0].feature,
version: Quotas[0].version,
};
export const Quota_ProPlanV1 = {
feature: Quotas[1].feature,
version: Quotas[1].version,
};

View File

@@ -0,0 +1,147 @@
import { Injectable } from '@nestjs/common';
import { PrismaService } from '../../prisma';
import { FeatureKind } from '../features';
import { QuotaConfig } from './quota';
import { QuotaType } from './types';
@Injectable()
export class QuotaService {
constructor(private readonly prisma: PrismaService) {}
// get activated user quota
async getUserQuota(userId: string) {
const quota = await this.prisma.userFeatures.findFirst({
where: {
user: {
id: userId,
},
feature: {
type: FeatureKind.Quota,
},
activated: true,
},
select: {
reason: true,
createdAt: true,
expiredAt: true,
featureId: true,
},
});
if (!quota) {
// this should unreachable
throw new Error(`User ${userId} has no quota`);
}
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.userFeatures.findMany({
where: {
user: {
id: userId,
},
feature: {
type: FeatureKind.Quota,
},
},
select: {
activated: true,
reason: true,
createdAt: true,
expiredAt: true,
featureId: true,
},
});
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;
};
})
);
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 latestFreePlan = await tx.features.aggregate({
where: {
feature: QuotaType.FreePlanV1,
},
_max: {
version: true,
},
});
// we will deactivate all exists quota for this user
await tx.userFeatures.updateMany({
where: {
id: undefined,
userId,
feature: {
type: FeatureKind.Quota,
},
},
data: {
activated: false,
},
});
await tx.userFeatures.create({
data: {
user: {
connect: {
id: userId,
},
},
feature: {
connect: {
feature_version: {
feature: quota,
version: latestFreePlan._max.version || 1,
},
type: FeatureKind.Quota,
},
},
reason: reason ?? 'switch quota',
activated: true,
expiredAt,
},
});
});
}
async hasQuota(userId: string, quota: QuotaType) {
return this.prisma.userFeatures
.count({
where: {
userId,
feature: {
feature: quota,
type: FeatureKind.Quota,
},
activated: true,
},
})
.then(count => count > 0);
}
}

View File

@@ -1,14 +1,9 @@
import type { Storage } from '@affine/storage'; import type { Storage } from '@affine/storage';
import { import { Inject, Injectable, NotFoundException } from '@nestjs/common';
ForbiddenException,
Inject,
Injectable,
NotFoundException,
} from '@nestjs/common';
import { StorageProvide } from '../../storage'; import { StorageProvide } from '../../storage';
import { PermissionService } from '../workspaces/permission'; import { PermissionService } from '../workspaces/permission';
import { QuotaService } from './quota'; import { QuotaService } from './service';
@Injectable() @Injectable()
export class QuotaManagementService { export class QuotaManagementService {
@@ -20,17 +15,15 @@ export class QuotaManagementService {
async getUserQuota(userId: string) { async getUserQuota(userId: string) {
const quota = await this.quota.getUserQuota(userId); const quota = await this.quota.getUserQuota(userId);
if (quota) {
return { return {
name: quota.feature.feature, name: quota.feature.name,
reason: quota.reason, reason: quota.reason,
createAt: quota.createdAt, createAt: quota.createdAt,
expiredAt: quota.expiredAt, expiredAt: quota.expiredAt,
blobLimit: quota.feature.configs.blobLimit, blobLimit: quota.feature.blobLimit,
storageQuota: quota.feature.configs.storageQuota, storageQuota: quota.feature.storageQuota,
}; };
}
return null;
} }
// TODO: lazy calc, need to be optimized with cache // TODO: lazy calc, need to be optimized with cache
@@ -45,7 +38,7 @@ export class QuotaManagementService {
const { user: owner } = const { user: owner } =
await this.permissions.getWorkspaceOwner(workspaceId); await this.permissions.getWorkspaceOwner(workspaceId);
if (!owner) throw new NotFoundException('Workspace owner not found'); if (!owner) throw new NotFoundException('Workspace owner not found');
const { storageQuota } = (await this.getUserQuota(owner.id)) || {}; const { storageQuota } = await this.getUserQuota(owner.id);
// get all workspaces size of owner used // get all workspaces size of owner used
const usageSize = await this.getUserUsage(owner.id); const usageSize = await this.getUserUsage(owner.id);
@@ -55,9 +48,6 @@ export class QuotaManagementService {
async checkBlobQuota(workspaceId: string, size: number) { async checkBlobQuota(workspaceId: string, size: number) {
const { quota, size: usageSize } = const { quota, size: usageSize } =
await this.getWorkspaceUsage(workspaceId); await this.getWorkspaceUsage(workspaceId);
if (typeof quota !== 'number') {
throw new ForbiddenException(`user's quota not exists`);
}
return quota - (size + usageSize); return quota - (size + usageSize);
} }

View File

@@ -1,31 +1,37 @@
import { CommonFeature, FeatureKind } from '../features'; import { z } from 'zod';
import { commonFeatureSchema, FeatureKind } from '../features';
import { ByteUnit, OneDay, OneKB } from './constant';
/// ======== quota define ========
export enum QuotaType { export enum QuotaType {
Quota_FreePlanV1 = 'free_plan_v1', FreePlanV1 = 'free_plan_v1',
Quota_ProPlanV1 = 'pro_plan_v1', ProPlanV1 = 'pro_plan_v1',
} }
export enum QuotaName { const quotaPlan = z.object({
free_plan_v1 = 'Free Plan', feature: z.enum([QuotaType.FreePlanV1, QuotaType.ProPlanV1]),
pro_plan_v1 = 'Pro Plan', configs: z.object({
} name: z.string(),
blobLimit: z.number().positive().int(),
storageQuota: z.number().positive().int(),
historyPeriod: z.number().positive().int(),
memberLimit: z.number().positive().int(),
}),
});
export type Quota = CommonFeature & { /// ======== schema infer ========
type: FeatureKind.Quota;
feature: QuotaType;
configs: {
blobLimit: number;
storageQuota: number;
historyPeriod: number;
memberLimit: number;
};
};
const OneKB = 1024; export const QuotaSchema = commonFeatureSchema
const OneMB = OneKB * OneKB; .extend({
const OneGB = OneKB * OneMB; type: z.literal(FeatureKind.Quota),
})
.and(z.discriminatedUnion('feature', [quotaPlan]));
const sizes = ['B', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB', 'ZB', 'YB']; export type Quota = z.infer<typeof QuotaSchema>;
/// ======== utils ========
export function formatSize(bytes: number, decimals: number = 2): string { export function formatSize(bytes: number, decimals: number = 2): string {
if (bytes === 0) return '0 B'; if (bytes === 0) return '0 B';
@@ -34,60 +40,11 @@ export function formatSize(bytes: number, decimals: number = 2): string {
const i = Math.floor(Math.log(bytes) / Math.log(OneKB)); const i = Math.floor(Math.log(bytes) / Math.log(OneKB));
return parseFloat((bytes / Math.pow(OneKB, i)).toFixed(dm)) + ' ' + sizes[i]; return (
parseFloat((bytes / Math.pow(OneKB, i)).toFixed(dm)) + ' ' + ByteUnit[i]
);
} }
const OneDay = 1000 * 60 * 60 * 24;
export function formatDate(ms: number): string { export function formatDate(ms: number): string {
return `${(ms / OneDay).toFixed(0)} days`; return `${(ms / OneDay).toFixed(0)} days`;
} }
export function getQuotaName(quota: QuotaType): string {
return QuotaName[quota];
}
export const Quotas: Quota[] = [
{
feature: QuotaType.Quota_FreePlanV1,
type: FeatureKind.Quota,
version: 1,
configs: {
// 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.Quota_ProPlanV1,
type: FeatureKind.Quota,
version: 1,
configs: {
// 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,
},
},
];
// ======== payload ========
export const Quota_FreePlanV1 = {
feature: Quotas[0].feature,
version: Quotas[0].version,
};
export const Quota_ProPlanV1 = {
feature: Quotas[1].feature,
version: Quotas[1].version,
};

View File

@@ -114,18 +114,8 @@ export class UserResolver {
@ResolveField(() => UserQuotaType, { name: 'quota', nullable: true }) @ResolveField(() => UserQuotaType, { name: 'quota', nullable: true })
async getQuota(@CurrentUser() me: User) { async getQuota(@CurrentUser() me: User) {
const quota = await this.quota.getUserQuota(me.id); const quota = await this.quota.getUserQuota(me.id);
const configs = quota.feature.configs;
return Object.assign( return quota.feature;
{
name: quota.feature.feature,
humanReadable: this.quota.getHumanReadableQuota(
quota.feature.feature,
configs
),
},
configs
);
} }
@Throttle({ default: { limit: 10, ttl: 60 } }) @Throttle({ default: { limit: 10, ttl: 60 } })

View File

@@ -12,6 +12,7 @@ import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs';
import request from 'supertest'; import request from 'supertest';
import { AppModule } from '../src/app'; import { AppModule } from '../src/app';
import { FeatureManagementService } from '../src/modules/features';
import { PrismaService } from '../src/prisma/service'; import { PrismaService } from '../src/prisma/service';
const gql = '/graphql'; const gql = '/graphql';
@@ -60,6 +61,8 @@ test.beforeEach(async t => {
}) })
.overrideProvider(PrismaService) .overrideProvider(PrismaService)
.useClass(FakePrisma) .useClass(FakePrisma)
.overrideProvider(FeatureManagementService)
.useValue({ canEarlyAccess: () => true })
.compile(); .compile();
t.context.app = module.createNestApplication({ t.context.app = module.createNestApplication({
cors: true, cors: true,

View File

@@ -6,6 +6,7 @@ import request from 'supertest';
import { AppModule } from '../src/app'; import { AppModule } from '../src/app';
import { ExceptionLogger } from '../src/middleware/exception-logger'; import { ExceptionLogger } from '../src/middleware/exception-logger';
import { FeatureManagementService } from '../src/modules/features';
import { PrismaService } from '../src/prisma'; import { PrismaService } from '../src/prisma';
const gql = '/graphql'; const gql = '/graphql';
@@ -38,6 +39,8 @@ test.beforeEach(async () => {
}) })
.overrideProvider(PrismaService) .overrideProvider(PrismaService)
.useClass(FakePrisma) .useClass(FakePrisma)
.overrideProvider(FeatureManagementService)
.useValue({})
.compile(); .compile();
app = module.createNestApplication({ app = module.createNestApplication({
cors: true, cors: true,

View File

@@ -81,12 +81,8 @@ test('should be able to set feature', async t => {
await feature.addUserFeature(u1.id, FeatureType.EarlyAccess, 1, 'test'); await feature.addUserFeature(u1.id, FeatureType.EarlyAccess, 1, 'test');
const f2 = await feature.getUserFeatures(u1.id); const f2 = await feature.getUserFeatures(u1.id);
t.is(f2.length, 1, 'should have one feature'); t.is(f2.length, 1, 'should have 1 feature');
t.is( t.is(f2[0].feature.name, FeatureType.EarlyAccess, 'should be early access');
f2[0].feature.feature,
FeatureType.EarlyAccess,
'should be early access'
);
}); });
test('should be able to check early access', async t => { test('should be able to check early access', async t => {
@@ -101,7 +97,7 @@ test('should be able to check early access', async t => {
t.true(f2, 'should have early access'); t.true(f2, 'should have early access');
const f3 = await feature.listFeatureUsers(FeatureType.EarlyAccess); const f3 = await feature.listFeatureUsers(FeatureType.EarlyAccess);
t.is(f3.length, 1, 'should have one user'); t.is(f3.length, 1, 'should have 1 user');
t.is(f3[0].id, u1.id, 'should be the same user'); t.is(f3[0].id, u1.id, 'should be the same user');
}); });
@@ -116,7 +112,7 @@ test('should be able revert quota', async t => {
const f2 = await early_access.canEarlyAccess(u1.email); const f2 = await early_access.canEarlyAccess(u1.email);
t.true(f2, 'should have early access'); t.true(f2, 'should have early access');
const q1 = await early_access.listEarlyAccess(); const q1 = await early_access.listEarlyAccess();
t.is(q1.length, 1, 'should have one user'); t.is(q1.length, 1, 'should have 1 user');
t.is(q1[0].id, u1.id, 'should be the same user'); t.is(q1[0].id, u1.id, 'should be the same user');
await early_access.removeEarlyAccess(u1.id); await early_access.removeEarlyAccess(u1.id);
@@ -127,10 +123,21 @@ test('should be able revert quota', async t => {
const q3 = await feature.getUserFeatures(u1.id); const q3 = await feature.getUserFeatures(u1.id);
t.is(q3.length, 1, 'should have 1 feature'); t.is(q3.length, 1, 'should have 1 feature');
t.is( t.is(q3[0].feature.name, FeatureType.EarlyAccess, 'should be early access');
q3[0].feature.feature,
FeatureType.EarlyAccess,
'should be early access'
);
t.is(q3[0].activated, false, 'should be deactivated'); t.is(q3[0].activated, false, 'should be deactivated');
}); });
test('should be same instance after reset the feature', async t => {
const { auth, feature, early_access } = t.context;
const u1 = await auth.signUp('DarkSky', 'darksky@example.org', '123456');
await early_access.addEarlyAccess(u1.id);
const f1 = (await feature.getUserFeatures(u1.id))[0];
await early_access.removeEarlyAccess(u1.id);
await early_access.addEarlyAccess(u1.id);
const f2 = (await feature.getUserFeatures(u1.id))[1];
t.is(f1.feature, f2.feature, 'should be same instance');
});

View File

@@ -9,6 +9,7 @@ import graphqlUploadExpress from 'graphql-upload/graphqlUploadExpress.mjs';
import { AppModule } from '../src/app'; import { AppModule } from '../src/app';
import { MailService } from '../src/modules/auth/mailer'; import { MailService } from '../src/modules/auth/mailer';
import { FeatureManagementService } from '../src/modules/features';
import { PrismaService } from '../src/prisma'; import { PrismaService } from '../src/prisma';
import { createWorkspace, getInviteInfo, inviteUser, signUp } from './utils'; import { createWorkspace, getInviteInfo, inviteUser, signUp } from './utils';
@@ -100,6 +101,8 @@ test.beforeEach(async t => {
}) })
.overrideProvider(PrismaService) .overrideProvider(PrismaService)
.useValue(FakePrisma) .useValue(FakePrisma)
.overrideProvider(FeatureManagementService)
.useValue({})
.compile(); .compile();
const app = module.createNestApplication(); const app = module.createNestApplication();
app.use( app.use(

View File

@@ -80,12 +80,12 @@ test('should be able to set quota', async t => {
const q1 = await quota.getUserQuota(u1.id); const q1 = await quota.getUserQuota(u1.id);
t.truthy(q1, 'should have quota'); t.truthy(q1, 'should have quota');
t.is(q1?.feature.feature, QuotaType.Quota_FreePlanV1, 'should be free plan'); t.is(q1?.feature.name, QuotaType.FreePlanV1, 'should be free plan');
await quota.switchUserQuota(u1.id, QuotaType.Quota_ProPlanV1); await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1);
const q2 = await quota.getUserQuota(u1.id); const q2 = await quota.getUserQuota(u1.id);
t.is(q2?.feature.feature, QuotaType.Quota_ProPlanV1, 'should be pro plan'); t.is(q2?.feature.name, QuotaType.ProPlanV1, 'should be pro plan');
const fail = quota.switchUserQuota(u1.id, 'not_exists_plan_v1' as QuotaType); const fail = quota.switchUserQuota(u1.id, 'not_exists_plan_v1' as QuotaType);
await t.throwsAsync(fail, { instanceOf: Error }, 'should throw error'); await t.throwsAsync(fail, { instanceOf: Error }, 'should throw error');
@@ -99,7 +99,7 @@ test('should be able to check storage quota', async t => {
t.is(q1?.blobLimit, Quotas[0].configs.blobLimit, 'should be free plan'); t.is(q1?.blobLimit, Quotas[0].configs.blobLimit, 'should be free plan');
t.is(q1?.storageQuota, Quotas[0].configs.storageQuota, 'should be free plan'); t.is(q1?.storageQuota, Quotas[0].configs.storageQuota, 'should be free plan');
await quota.switchUserQuota(u1.id, QuotaType.Quota_ProPlanV1); await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1);
const q2 = await storageQuota.getUserQuota(u1.id); const q2 = await storageQuota.getUserQuota(u1.id);
t.is(q2?.blobLimit, Quotas[1].configs.blobLimit, 'should be pro plan'); t.is(q2?.blobLimit, Quotas[1].configs.blobLimit, 'should be pro plan');
t.is(q2?.storageQuota, Quotas[1].configs.storageQuota, 'should be pro plan'); t.is(q2?.storageQuota, Quotas[1].configs.storageQuota, 'should be pro plan');
@@ -113,32 +113,20 @@ test('should be able revert quota', async t => {
t.is(q1?.blobLimit, Quotas[0].configs.blobLimit, 'should be free plan'); t.is(q1?.blobLimit, Quotas[0].configs.blobLimit, 'should be free plan');
t.is(q1?.storageQuota, Quotas[0].configs.storageQuota, 'should be free plan'); t.is(q1?.storageQuota, Quotas[0].configs.storageQuota, 'should be free plan');
await quota.switchUserQuota(u1.id, QuotaType.Quota_ProPlanV1); await quota.switchUserQuota(u1.id, QuotaType.ProPlanV1);
const q2 = await storageQuota.getUserQuota(u1.id); const q2 = await storageQuota.getUserQuota(u1.id);
t.is(q2?.blobLimit, Quotas[1].configs.blobLimit, 'should be pro plan'); t.is(q2?.blobLimit, Quotas[1].configs.blobLimit, 'should be pro plan');
t.is(q2?.storageQuota, Quotas[1].configs.storageQuota, 'should be pro plan'); t.is(q2?.storageQuota, Quotas[1].configs.storageQuota, 'should be pro plan');
await quota.switchUserQuota(u1.id, QuotaType.Quota_FreePlanV1); await quota.switchUserQuota(u1.id, QuotaType.FreePlanV1);
const q3 = await storageQuota.getUserQuota(u1.id); const q3 = await storageQuota.getUserQuota(u1.id);
t.is(q3?.blobLimit, Quotas[0].configs.blobLimit, 'should be free plan'); t.is(q3?.blobLimit, Quotas[0].configs.blobLimit, 'should be free plan');
const quotas = await quota.getUserQuotas(u1.id); const quotas = await quota.getUserQuotas(u1.id);
t.is(quotas.length, 3, 'should have 3 quotas'); t.is(quotas.length, 3, 'should have 3 quotas');
t.is( t.is(quotas[0].feature.name, QuotaType.FreePlanV1, 'should be free plan');
quotas[0].feature.feature, t.is(quotas[1].feature.name, QuotaType.ProPlanV1, 'should be pro plan');
QuotaType.Quota_FreePlanV1, t.is(quotas[2].feature.name, QuotaType.FreePlanV1, 'should be free plan');
'should be free plan'
);
t.is(
quotas[1].feature.feature,
QuotaType.Quota_ProPlanV1,
'should be pro plan'
);
t.is(
quotas[2].feature.feature,
QuotaType.Quota_FreePlanV1,
'should be free plan'
);
t.is(quotas[0].activated, false, 'should be activated'); t.is(quotas[0].activated, false, 'should be activated');
t.is(quotas[1].activated, false, 'should be activated'); t.is(quotas[1].activated, false, 'should be activated');
t.is(quotas[2].activated, true, 'should be activated'); t.is(quotas[2].activated, true, 'should be activated');

View File

@@ -1,619 +0,0 @@
import { randomUUID } from 'node:crypto';
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';
import request from 'supertest';
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';
async function signUp(
app: INestApplication,
name: string,
email: string,
password: string
): Promise<UserType & { token: TokenType }> {
const res = await request(app.getHttpServer())
.post(gql)
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
signUp(name: "${name}", email: "${email}", password: "${password}") {
id, name, email, token { token }
}
}
`,
})
.expect(200);
return res.body.data.signUp;
}
async function currentUser(app: INestApplication, token: string) {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
currentUser {
id, name, email, emailVerified, avatarUrl, createdAt, hasPassword,
token { token }
}
}
`,
})
.expect(200);
return res.body.data.currentUser;
}
async function createWorkspace(
app: INestApplication,
token: string
): Promise<WorkspaceType> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.field(
'operations',
JSON.stringify({
name: 'createWorkspace',
query: `mutation createWorkspace($init: Upload!) {
createWorkspace(init: $init) {
id
}
}`,
variables: { init: null },
})
)
.field('map', JSON.stringify({ '0': ['variables.init'] }))
.attach('0', Buffer.from([0, 0]), 'init.data')
.expect(200);
return res.body.data.createWorkspace;
}
export async function getWorkspacePublicPages(
app: INestApplication,
token: string,
workspaceId: string
) {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
workspace(id: "${workspaceId}") {
publicPages {
id
mode
}
}
}
`,
})
.expect(200);
return res.body.data.workspace.publicPages;
}
async function getWorkspace(
app: INestApplication,
token: string,
workspaceId: string,
skip = 0,
take = 8
): Promise<WorkspaceType> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
workspace(id: "${workspaceId}") {
id, members(skip: ${skip}, take: ${take}) { id, name, email, permission, inviteId }
}
}
`,
})
.expect(200);
return res.body.data.workspace;
}
async function getPublicWorkspace(
app: INestApplication,
workspaceId: string
): Promise<WorkspaceType> {
const res = await request(app.getHttpServer())
.post(gql)
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
publicWorkspace(id: "${workspaceId}") {
id
}
}
`,
})
.expect(200);
return res.body.data.publicWorkspace;
}
async function updateWorkspace(
app: INestApplication,
token: string,
workspaceId: string,
isPublic: boolean
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
updateWorkspace(input: { id: "${workspaceId}", public: ${isPublic} }) {
public
}
}
`,
})
.expect(200);
return res.body.data.updateWorkspace.public;
}
async function inviteUser(
app: INestApplication,
token: string,
workspaceId: string,
email: string,
permission: string,
sendInviteMail = false
): Promise<string> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
invite(workspaceId: "${workspaceId}", email: "${email}", permission: ${permission}, sendInviteMail: ${sendInviteMail})
}
`,
})
.expect(200);
return res.body.data.invite;
}
async function acceptInviteById(
app: INestApplication,
workspaceId: string,
inviteId: string,
sendAcceptMail = false
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
acceptInviteById(workspaceId: "${workspaceId}", inviteId: "${inviteId}", sendAcceptMail: ${sendAcceptMail})
}
`,
})
.expect(200);
return res.body.data.acceptInviteById;
}
async function leaveWorkspace(
app: INestApplication,
token: string,
workspaceId: string,
sendLeaveMail = false
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
leaveWorkspace(workspaceId: "${workspaceId}", workspaceName: "test workspace", sendLeaveMail: ${sendLeaveMail})
}
`,
})
.expect(200);
return res.body.data.leaveWorkspace;
}
async function revokeUser(
app: INestApplication,
token: string,
workspaceId: string,
userId: string
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
revoke(workspaceId: "${workspaceId}", userId: "${userId}")
}
`,
})
.expect(200);
return res.body.data.revoke;
}
async function publishPage(
app: INestApplication,
token: string,
workspaceId: string,
pageId: string
) {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
publishPage(workspaceId: "${workspaceId}", pageId: "${pageId}") {
id
mode
}
}
`,
})
.expect(200);
return res.body.errors?.[0]?.message || res.body.data?.publishPage;
}
async function revokePublicPage(
app: INestApplication,
token: string,
workspaceId: string,
pageId: string
) {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
revokePublicPage(workspaceId: "${workspaceId}", pageId: "${pageId}") {
id
mode
public
}
}
`,
})
.expect(200);
return res.body.errors?.[0]?.message || res.body.data?.revokePublicPage;
}
async function listBlobs(
app: INestApplication,
token: string,
workspaceId: string
): Promise<string[]> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
listBlobs(workspaceId: "${workspaceId}")
}
`,
})
.expect(200);
return res.body.data.listBlobs;
}
async function getWorkspaceBlobsSize(
app: INestApplication,
token: string,
workspaceId: string
): Promise<number> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.send({
query: `
query {
workspace(id: "${workspaceId}") {
blobsSize
}
}
`,
})
.expect(200);
return res.body.data.workspace.blobsSize;
}
async function collectAllBlobSizes(
app: INestApplication,
token: string
): Promise<number> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.send({
query: `
query {
collectAllBlobSizes {
size
}
}
`,
})
.expect(200);
return res.body.data.collectAllBlobSizes.size;
}
async function checkBlobSize(
app: INestApplication,
token: string,
workspaceId: string,
size: number
): Promise<number> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.send({
query: `query checkBlobSize($workspaceId: String!, $size: Float!) {
checkBlobSize(workspaceId: $workspaceId, size: $size) {
size
}
}`,
variables: { workspaceId, size },
})
.expect(200);
return res.body.data.checkBlobSize.size;
}
async function setBlob(
app: INestApplication,
token: string,
workspaceId: string,
buffer: Buffer
): Promise<string> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.field(
'operations',
JSON.stringify({
name: 'setBlob',
query: `mutation setBlob($blob: Upload!) {
setBlob(workspaceId: "${workspaceId}", blob: $blob)
}`,
variables: { blob: null },
})
)
.field('map', JSON.stringify({ '0': ['variables.blob'] }))
.attach('0', buffer, 'blob.data')
.expect(200);
return res.body.data.setBlob;
}
async function flushDB() {
const client = new PrismaClient();
await client.$connect();
const result: { tablename: string }[] =
await client.$queryRaw`SELECT tablename
FROM pg_catalog.pg_tables
WHERE schemaname != 'pg_catalog'
AND schemaname != 'information_schema'`;
// remove all table data
await client.$executeRawUnsafe(
`TRUNCATE TABLE ${result
.map(({ tablename }) => tablename)
.filter(name => !name.includes('migrations'))
.join(', ')}`
);
await client.$disconnect();
}
async function getInviteInfo(
app: INestApplication,
token: string,
inviteId: string
): Promise<InvitationType> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
getInviteInfo(inviteId: "${inviteId}") {
workspace {
id
name
avatar
}
user {
id
name
avatarUrl
}
}
}
`,
})
.expect(200);
return res.body.data.getInviteInfo;
}
async function sendChangeEmail(
app: INestApplication,
userToken: string,
email: string,
callbackUrl: string
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(userToken, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
sendChangeEmail(email: "${email}", callbackUrl: "${callbackUrl}")
}
`,
})
.expect(200);
return res.body.data.sendChangeEmail;
}
async function sendVerifyChangeEmail(
app: INestApplication,
userToken: string,
token: string,
email: string,
callbackUrl: string
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(userToken, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
sendVerifyChangeEmail(token:"${token}", email: "${email}", callbackUrl: "${callbackUrl}")
}
`,
})
.expect(200);
return res.body.data.sendVerifyChangeEmail;
}
async function changeEmail(
app: INestApplication,
userToken: string,
token: string
): Promise<UserType & { token: TokenType }> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(userToken, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
changeEmail(token: "${token}") {
id
name
avatarUrl
email
}
}
`,
})
.expect(200);
return res.body.data.changeEmail;
}
export class FakePrisma {
fakeUser: User = {
id: randomUUID(),
name: 'Alex Yang',
avatarUrl: '',
email: 'alex.yang@example.org',
password: hashSync('123456'),
emailVerified: new Date(),
createdAt: new Date(),
};
get user() {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const prisma = this;
return {
async findFirst() {
return prisma.fakeUser;
},
async findUnique() {
return this.findFirst();
},
async update() {
return this.findFirst();
},
};
}
}
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);
await Promise.allSettled([revert.run(['UserFeaturesInit1698652531198'])]);
await run.runOne('UserFeaturesInit1698652531198');
}
export {
acceptInviteById,
changeEmail,
checkBlobSize,
collectAllBlobSizes,
createWorkspace,
currentUser,
flushDB,
getInviteInfo,
getPublicWorkspace,
getWorkspace,
getWorkspaceBlobsSize,
inviteUser,
leaveWorkspace,
listBlobs,
publishPage,
revokePublicPage,
revokeUser,
sendChangeEmail,
sendVerifyChangeEmail,
setBlob,
signUp,
updateWorkspace,
};

View File

@@ -0,0 +1,112 @@
import type { INestApplication } from '@nestjs/common';
import request from 'supertest';
import { gql } from './common';
export async function listBlobs(
app: INestApplication,
token: string,
workspaceId: string
): Promise<string[]> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
listBlobs(workspaceId: "${workspaceId}")
}
`,
})
.expect(200);
return res.body.data.listBlobs;
}
export async function getWorkspaceBlobsSize(
app: INestApplication,
token: string,
workspaceId: string
): Promise<number> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.send({
query: `
query {
workspace(id: "${workspaceId}") {
blobsSize
}
}
`,
})
.expect(200);
return res.body.data.workspace.blobsSize;
}
export async function collectAllBlobSizes(
app: INestApplication,
token: string
): Promise<number> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.send({
query: `
query {
collectAllBlobSizes {
size
}
}
`,
})
.expect(200);
return res.body.data.collectAllBlobSizes.size;
}
export async function checkBlobSize(
app: INestApplication,
token: string,
workspaceId: string,
size: number
): Promise<number> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.send({
query: `query checkBlobSize($workspaceId: String!, $size: Float!) {
checkBlobSize(workspaceId: $workspaceId, size: $size) {
size
}
}`,
variables: { workspaceId, size },
})
.expect(200);
return res.body.data.checkBlobSize.size;
}
export async function setBlob(
app: INestApplication,
token: string,
workspaceId: string,
buffer: Buffer
): Promise<string> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.field(
'operations',
JSON.stringify({
name: 'setBlob',
query: `mutation setBlob($blob: Upload!) {
setBlob(workspaceId: "${workspaceId}", blob: $blob)
}`,
variables: { blob: null },
})
)
.field('map', JSON.stringify({ '0': ['variables.blob'] }))
.attach('0', buffer, 'blob.data')
.expect(200);
return res.body.data.setBlob;
}

View File

@@ -0,0 +1 @@
export const gql = '/graphql';

View File

@@ -0,0 +1,5 @@
export * from './blobs';
export * from './invite';
export * from './user';
export * from './utils';
export * from './workspace';

View File

@@ -0,0 +1,121 @@
import type { INestApplication } from '@nestjs/common';
import request from 'supertest';
import type { InvitationType } from '../../src/modules/workspaces';
import { gql } from './common';
export async function inviteUser(
app: INestApplication,
token: string,
workspaceId: string,
email: string,
permission: string,
sendInviteMail = false
): Promise<string> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
invite(workspaceId: "${workspaceId}", email: "${email}", permission: ${permission}, sendInviteMail: ${sendInviteMail})
}
`,
})
.expect(200);
return res.body.data.invite;
}
export async function acceptInviteById(
app: INestApplication,
workspaceId: string,
inviteId: string,
sendAcceptMail = false
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
acceptInviteById(workspaceId: "${workspaceId}", inviteId: "${inviteId}", sendAcceptMail: ${sendAcceptMail})
}
`,
})
.expect(200);
return res.body.data.acceptInviteById;
}
export async function leaveWorkspace(
app: INestApplication,
token: string,
workspaceId: string,
sendLeaveMail = false
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
leaveWorkspace(workspaceId: "${workspaceId}", workspaceName: "test workspace", sendLeaveMail: ${sendLeaveMail})
}
`,
})
.expect(200);
return res.body.data.leaveWorkspace;
}
export async function revokeUser(
app: INestApplication,
token: string,
workspaceId: string,
userId: string
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
revoke(workspaceId: "${workspaceId}", userId: "${userId}")
}
`,
})
.expect(200);
return res.body.data.revoke;
}
export async function getInviteInfo(
app: INestApplication,
token: string,
inviteId: string
): Promise<InvitationType> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
getInviteInfo(inviteId: "${inviteId}") {
workspace {
id
name
avatar
}
user {
id
name
avatarUrl
}
}
}
`,
})
.expect(200);
return res.body.data.getInviteInfo;
}

View File

@@ -0,0 +1,117 @@
import type { INestApplication } from '@nestjs/common';
import request from 'supertest';
import type { TokenType } from '../../src/modules/auth';
import type { UserType } from '../../src/modules/users';
import { gql } from './common';
export async function signUp(
app: INestApplication,
name: string,
email: string,
password: string
): Promise<UserType & { token: TokenType }> {
const res = await request(app.getHttpServer())
.post(gql)
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
signUp(name: "${name}", email: "${email}", password: "${password}") {
id, name, email, token { token }
}
}
`,
})
.expect(200);
return res.body.data.signUp;
}
export async function currentUser(app: INestApplication, token: string) {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
currentUser {
id, name, email, emailVerified, avatarUrl, createdAt, hasPassword,
token { token }
}
}
`,
})
.expect(200);
return res.body.data.currentUser;
}
export async function sendChangeEmail(
app: INestApplication,
userToken: string,
email: string,
callbackUrl: string
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(userToken, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
sendChangeEmail(email: "${email}", callbackUrl: "${callbackUrl}")
}
`,
})
.expect(200);
return res.body.data.sendChangeEmail;
}
export async function sendVerifyChangeEmail(
app: INestApplication,
userToken: string,
token: string,
email: string,
callbackUrl: string
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(userToken, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
sendVerifyChangeEmail(token:"${token}", email: "${email}", callbackUrl: "${callbackUrl}")
}
`,
})
.expect(200);
return res.body.data.sendVerifyChangeEmail;
}
export async function changeEmail(
app: INestApplication,
userToken: string,
token: string
): Promise<UserType & { token: TokenType }> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(userToken, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
changeEmail(token: "${token}") {
id
name
avatarUrl
email
}
}
`,
})
.expect(200);
return res.body.data.changeEmail;
}

View File

@@ -0,0 +1,82 @@
import { randomUUID } from 'node:crypto';
import type { DynamicModule, FactoryProvider } from '@nestjs/common';
import { TestingModule } from '@nestjs/testing';
import { hashSync } from '@node-rs/argon2';
import { PrismaClient, type User } from '@prisma/client';
import { RevertCommand, RunCommand } from '../../src/data/commands/run';
import { StorageProvide } from '../../src/storage';
export async function flushDB() {
const client = new PrismaClient();
await client.$connect();
const result: { tablename: string }[] =
await client.$queryRaw`SELECT tablename
FROM pg_catalog.pg_tables
WHERE schemaname != 'pg_catalog'
AND schemaname != 'information_schema'`;
// remove all table data
await client.$executeRawUnsafe(
`TRUNCATE TABLE ${result
.map(({ tablename }) => tablename)
.filter(name => !name.includes('migrations'))
.join(', ')}`
);
await client.$disconnect();
}
export class FakePrisma {
fakeUser: User = {
id: randomUUID(),
name: 'Alex Yang',
avatarUrl: '',
email: 'alex.yang@example.org',
password: hashSync('123456'),
emailVerified: new Date(),
createdAt: new Date(),
};
get user() {
// eslint-disable-next-line @typescript-eslint/no-this-alias
const prisma = this;
return {
async findFirst() {
return prisma.fakeUser;
},
async findUnique() {
return this.findFirst();
},
async update() {
return this.findFirst();
},
};
}
}
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);
await Promise.allSettled([revert.run(['UserFeaturesInit1698652531198'])]);
await run.runOne('UserFeaturesInit1698652531198');
}

View File

@@ -0,0 +1,172 @@
import type { INestApplication } from '@nestjs/common';
import request from 'supertest';
import type { WorkspaceType } from '../../src/modules/workspaces';
import { gql } from './common';
export async function createWorkspace(
app: INestApplication,
token: string
): Promise<WorkspaceType> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.field(
'operations',
JSON.stringify({
name: 'createWorkspace',
query: `mutation createWorkspace($init: Upload!) {
createWorkspace(init: $init) {
id
}
}`,
variables: { init: null },
})
)
.field('map', JSON.stringify({ '0': ['variables.init'] }))
.attach('0', Buffer.from([0, 0]), 'init.data')
.expect(200);
return res.body.data.createWorkspace;
}
export async function getWorkspacePublicPages(
app: INestApplication,
token: string,
workspaceId: string
) {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
workspace(id: "${workspaceId}") {
publicPages {
id
mode
}
}
}
`,
})
.expect(200);
return res.body.data.workspace.publicPages;
}
export async function getWorkspace(
app: INestApplication,
token: string,
workspaceId: string,
skip = 0,
take = 8
): Promise<WorkspaceType> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
workspace(id: "${workspaceId}") {
id, members(skip: ${skip}, take: ${take}) { id, name, email, permission, inviteId }
}
}
`,
})
.expect(200);
return res.body.data.workspace;
}
export async function getPublicWorkspace(
app: INestApplication,
workspaceId: string
): Promise<WorkspaceType> {
const res = await request(app.getHttpServer())
.post(gql)
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
query {
publicWorkspace(id: "${workspaceId}") {
id
}
}
`,
})
.expect(200);
return res.body.data.publicWorkspace;
}
export async function updateWorkspace(
app: INestApplication,
token: string,
workspaceId: string,
isPublic: boolean
): Promise<boolean> {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
updateWorkspace(input: { id: "${workspaceId}", public: ${isPublic} }) {
public
}
}
`,
})
.expect(200);
return res.body.data.updateWorkspace.public;
}
export async function publishPage(
app: INestApplication,
token: string,
workspaceId: string,
pageId: string
) {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
publishPage(workspaceId: "${workspaceId}", pageId: "${pageId}") {
id
mode
}
}
`,
})
.expect(200);
return res.body.errors?.[0]?.message || res.body.data?.publishPage;
}
export async function revokePublicPage(
app: INestApplication,
token: string,
workspaceId: string,
pageId: string
) {
const res = await request(app.getHttpServer())
.post(gql)
.auth(token, { type: 'bearer' })
.set({ 'x-request-id': 'test', 'x-operation-name': 'test' })
.send({
query: `
mutation {
revokePublicPage(workspaceId: "${workspaceId}", pageId: "${pageId}") {
id
mode
public
}
}
`,
})
.expect(200);
return res.body.errors?.[0]?.message || res.body.data?.revokePublicPage;
}

View File

@@ -179,7 +179,7 @@ test('should be able calc quota after switch plan', async t => {
); );
t.is(size1, 0, 'failed to check free plan blob size'); t.is(size1, 0, 'failed to check free plan blob size');
quota.switchUserQuota(u1.id, QuotaType.Quota_ProPlanV1); quota.switchUserQuota(u1.id, QuotaType.ProPlanV1);
const size2 = await checkBlobSize( const size2 = await checkBlobSize(
app, app,

View File

@@ -4,6 +4,7 @@ import ava, { type TestFn } from 'ava';
import { stub } from 'sinon'; import { stub } from 'sinon';
import { AppModule } from '../src/app'; import { AppModule } from '../src/app';
import { FeatureManagementService } from '../src/modules/features';
import { Quotas } from '../src/modules/quota'; import { Quotas } from '../src/modules/quota';
import { UsersService } from '../src/modules/users'; import { UsersService } from '../src/modules/users';
import { PermissionService } from '../src/modules/workspaces/permission'; import { PermissionService } from '../src/modules/workspaces/permission';
@@ -59,6 +60,23 @@ test.beforeEach(async t => {
}; };
}, },
}, },
features: {
async findFirst() {
return {
id: 0,
feature: 'free_plan_v1',
version: 1,
type: 1,
configs: {
name: 'Free',
blobLimit: 1,
storageQuota: 1,
historyPeriod: 1,
memberLimit: 3,
},
};
},
},
}) })
.overrideProvider(PermissionService) .overrideProvider(PermissionService)
.useClass(FakePermission) .useClass(FakePermission)
@@ -70,6 +88,8 @@ test.beforeEach(async t => {
return 1024 * 10; return 1024 * 10;
}, },
}) })
.overrideProvider(FeatureManagementService)
.useValue({})
.compile(); .compile();
t.context.app = module.createNestApplication(); t.context.app = module.createNestApplication();
t.context.resolver = t.context.app.get(WorkspaceResolver); t.context.resolver = t.context.app.get(WorkspaceResolver);

View File

@@ -812,6 +812,7 @@ __metadata:
typescript: "npm:^5.3.2" typescript: "npm:^5.3.2"
ws: "npm:^8.14.2" ws: "npm:^8.14.2"
yjs: "npm:^13.6.10" yjs: "npm:^13.6.10"
zod: "npm:^3.22.4"
bin: bin:
run-test: ./scripts/run-test.ts run-test: ./scripts/run-test.ts
languageName: unknown languageName: unknown