feat: integrate user usage into apis (#5075)

This commit is contained in:
DarkSky
2023-12-14 09:50:36 +00:00
parent 63de73a815
commit ad23ead5e4
36 changed files with 984 additions and 282 deletions

View File

@@ -188,11 +188,6 @@ export interface AFFiNEConfig {
fs: {
path: string;
};
/**
* default storage quota
* @default 10 * 1024 * 1024 * 1024 (10GB)
*/
quota: number;
};
/**

View File

@@ -58,7 +58,6 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => {
AFFINE_SERVER_HOST: 'host',
AFFINE_SERVER_SUB_PATH: 'path',
AFFINE_ENV: 'affineEnv',
AFFINE_FREE_USER_QUOTA: 'objectStorage.quota',
DATABASE_URL: 'db.url',
ENABLE_R2_OBJECT_STORAGE: ['objectStorage.r2.enabled', 'boolean'],
R2_OBJECT_STORAGE_ACCOUNT_ID: 'objectStorage.r2.accountId',
@@ -192,8 +191,6 @@ export const getDefaultAFFiNEConfig: () => AFFiNEConfig = () => {
fs: {
path: join(homedir(), '.affine-storage'),
},
// 10GB
quota: 10 * 1024 * 1024 * 1024,
},
rateLimiter: {
ttl: 60,

View File

@@ -1,12 +1,39 @@
import {
CommonFeature,
FeatureKind,
Features,
FeatureType,
upsertFeature,
} from '../../modules/features';
import { Quotas } from '../../modules/quota';
} from '../../modules/features/types';
import { Quotas } from '../../modules/quota/types';
import { PrismaService } from '../../prisma';
// upgrade features from lower version to higher version
async function upsertFeature(
db: PrismaService,
feature: CommonFeature
): Promise<void> {
const hasEqualOrGreaterVersion =
(await db.features.count({
where: {
feature: feature.feature,
version: {
gte: feature.version,
},
},
})) > 0;
// will not update exists version
if (!hasEqualOrGreaterVersion) {
await db.features.create({
data: {
feature: feature.feature,
type: feature.type,
version: feature.version,
configs: feature.configs,
},
});
}
}
export class UserFeaturesInit1698652531198 {
// do the migration
static async up(db: PrismaService) {

View File

@@ -11,8 +11,8 @@ import Google from 'next-auth/providers/google';
import { Config } from '../../config';
import { PrismaService } from '../../prisma';
import { SessionService } from '../../session';
import { NewFeaturesKind } from '../users/types';
import { isStaff } from '../users/utils';
import { FeatureType } from '../features';
import { Quota_FreePlanV1 } from '../quota';
import { MailService } from './mailer';
import {
decode,
@@ -44,6 +44,17 @@ export const NextAuthOptionsProvider: FactoryProvider<NextAuthOptions> = {
email: data.email,
avatarUrl: '',
emailVerified: data.emailVerified,
features: {
create: {
reason: 'created by email sign up',
activated: true,
feature: {
connect: {
feature_version: Quota_FreePlanV1,
},
},
},
},
};
if (data.email && !data.name) {
userData.name = data.email.split('@')[0];
@@ -223,18 +234,23 @@ export const NextAuthOptionsProvider: FactoryProvider<NextAuthOptions> = {
}
const email = profile?.email ?? user.email;
if (email) {
if (isStaff(email)) {
return true;
}
return prisma.newFeaturesWaitingList
.findUnique({
// FIXME: cannot inject FeatureManagementService here
// it will cause prisma.account to be undefined
// then prismaAdapter.getUserByAccount will throw error
if (email.endsWith('@toeverything.info')) return true;
return prisma.userFeatures
.count({
where: {
email,
type: NewFeaturesKind.EarlyAccess,
user: {
email,
},
feature: {
feature: FeatureType.EarlyAccess,
},
activated: true,
},
})
.then(user => !!user)
.catch(() => false);
.then(count => count > 0);
}
return false;
},

View File

@@ -14,6 +14,7 @@ import { nanoid } from 'nanoid';
import { Config } from '../../config';
import { PrismaService } from '../../prisma';
import { verifyChallengeResponse } from '../../storage';
import { Quota_FreePlanV1 } from '../quota';
import { MailService } from './mailer';
export type UserClaim = Pick<
@@ -190,6 +191,17 @@ export class AuthService {
name,
email,
password: hashedPassword,
features: {
create: {
reason: 'created by api sign up',
activated: true,
feature: {
connect: {
feature_version: Quota_FreePlanV1,
},
},
},
},
},
});
}
@@ -209,6 +221,17 @@ export class AuthService {
data: {
name: 'Unnamed',
email,
features: {
create: {
reason: 'created by invite sign up',
activated: true,
feature: {
connect: {
feature_version: Quota_FreePlanV1,
},
},
},
},
},
});
}
@@ -258,6 +281,7 @@ export class AuthService {
},
});
}
async changeEmail(id: string, newEmail: string): Promise<User> {
const user = await this.prisma.user.findUnique({
where: {

View File

@@ -3,34 +3,6 @@ import { Module } from '@nestjs/common';
import { PrismaService } from '../../prisma';
import { FeatureService } from './configure';
import { FeatureManagementService } from './feature';
import type { CommonFeature } from './types';
// upgrade features from lower version to higher version
async function upsertFeature(
db: PrismaService,
feature: CommonFeature
): Promise<void> {
const hasEqualOrGreaterVersion =
(await db.features.count({
where: {
feature: feature.feature,
version: {
gte: feature.version,
},
},
})) > 0;
// will not update exists version
if (!hasEqualOrGreaterVersion) {
await db.features.create({
data: {
feature: feature.feature,
type: feature.type,
version: feature.version,
configs: feature.configs,
},
});
}
}
/**
* Feature module provider pre-user feature flag management.
@@ -46,9 +18,4 @@ export class FeatureModule {}
export type { CommonFeature, Feature } from './types';
export { FeatureKind, Features, FeatureType } from './types';
export {
FeatureManagementService,
FeatureService,
PrismaService,
upsertFeature,
};
export { FeatureManagementService, FeatureService, PrismaService };

View File

@@ -6,6 +6,7 @@ import { GqlModule } from '../graphql.module';
import { ServerConfigModule } from './config';
import { DocModule } from './doc';
import { PaymentModule } from './payment';
import { QuotaModule } from './quota';
import { SelfHostedModule } from './self-hosted';
import { SyncModule } from './sync';
import { UsersModule } from './users';
@@ -37,7 +38,8 @@ switch (SERVER_FLAVOR) {
WorkspaceModule,
UsersModule,
DocModule,
PaymentModule
PaymentModule,
QuotaModule
);
break;
case 'allinone':
@@ -48,6 +50,7 @@ switch (SERVER_FLAVOR) {
GqlModule,
WorkspaceModule,
UsersModule,
QuotaModule,
SyncModule,
DocModule,
PaymentModule

View File

@@ -1,6 +1,6 @@
import { Module } from '@nestjs/common';
import { UsersModule } from '../users';
import { FeatureModule } from '../features';
import { SubscriptionResolver, UserSubscriptionResolver } from './resolver';
import { ScheduleManager } from './schedule';
import { SubscriptionService } from './service';
@@ -8,7 +8,7 @@ import { StripeProvider } from './stripe';
import { StripeWebhook } from './webhook';
@Module({
imports: [UsersModule],
imports: [FeatureModule],
providers: [
ScheduleManager,
StripeProvider,

View File

@@ -11,7 +11,7 @@ import Stripe from 'stripe';
import { Config } from '../../config';
import { PrismaService } from '../../prisma';
import { UsersService } from '../users';
import { FeatureManagementService } from '../features';
import { ScheduleManager } from './schedule';
const OnEvent = (
@@ -82,8 +82,8 @@ export class SubscriptionService {
config: Config,
private readonly stripe: Stripe,
private readonly db: PrismaService,
private readonly user: UsersService,
private readonly scheduleManager: ScheduleManager
private readonly scheduleManager: ScheduleManager,
private readonly features: FeatureManagementService
) {
this.paymentConfig = config.payment;
@@ -658,7 +658,7 @@ export class SubscriptionService {
user: User,
couponType: CouponType
): Promise<string | null> {
const earlyAccess = await this.user.isEarlyAccessUser(user.email);
const earlyAccess = await this.features.canEarlyAccess(user.email);
if (earlyAccess) {
try {
const coupon = await this.stripe.coupons.retrieve(couponType);

View File

@@ -1,42 +0,0 @@
type FeatureEarlyAccessPreview = {
whitelist: RegExp[];
};
type FeatureStorageLimit = {
storageQuota: number;
};
type UserFeatureGate = {
earlyAccessPreview: FeatureEarlyAccessPreview;
freeUser: FeatureStorageLimit;
proUser: FeatureStorageLimit;
};
const UserLevel = {
freeUser: {
storageQuota: 10 * 1024 * 1024 * 1024,
},
proUser: {
storageQuota: 100 * 1024 * 1024 * 1024,
},
} satisfies Pick<UserFeatureGate, 'freeUser' | 'proUser'>;
export function getStorageQuota(features: string[]) {
for (const feature of features) {
if (feature in UserLevel) {
return UserLevel[feature as keyof typeof UserLevel].storageQuota;
}
}
return null;
}
const UserType = {
earlyAccessPreview: {
whitelist: [/@toeverything\.info$/],
},
} satisfies Pick<UserFeatureGate, 'earlyAccessPreview'>;
export const FeatureGates = {
...UserType,
...UserLevel,
} satisfies UserFeatureGate;

View File

@@ -1,11 +1,12 @@
import { Module } from '@nestjs/common';
import { FeatureModule } from '../features';
import { StorageModule } from '../storage';
import { UserResolver } from './resolver';
import { UsersService } from './users';
@Module({
imports: [StorageModule],
imports: [StorageModule, FeatureModule],
providers: [UserResolver, UsersService],
exports: [UsersService],
})

View File

@@ -12,7 +12,6 @@ import {
Mutation,
ObjectType,
Query,
registerEnumType,
ResolveField,
Resolver,
} from '@nestjs/graphql';
@@ -24,14 +23,10 @@ import { PrismaService } from '../../prisma/service';
import { CloudThrottlerGuard, Throttle } from '../../throttler';
import type { FileUpload } from '../../types';
import { Auth, CurrentUser, Public, Publicable } from '../auth/guard';
import { AuthService } from '../auth/service';
import { FeatureManagementService } from '../features';
import { StorageService } from '../storage/storage.service';
import { NewFeaturesKind } from './types';
import { UsersService } from './users';
import { isStaff } from './utils';
registerEnumType(NewFeaturesKind, {
name: 'NewFeaturesKind',
});
@ObjectType()
export class UserType implements Partial<User> {
@@ -71,14 +66,6 @@ export class RemoveAvatar {
success!: boolean;
}
@ObjectType()
export class AddToNewFeaturesWaitingList {
@Field()
email!: string;
@Field(() => NewFeaturesKind, { description: 'New features kind' })
type!: NewFeaturesKind;
}
/**
* User resolver
* All op rate limit: 10 req/m
@@ -88,9 +75,11 @@ export class AddToNewFeaturesWaitingList {
@Resolver(() => UserType)
export class UserResolver {
constructor(
private readonly auth: AuthService,
private readonly prisma: PrismaService,
private readonly storage: StorageService,
private readonly users: UsersService
private readonly users: UsersService,
private readonly feature: FeatureManagementService
) {}
@Throttle({
@@ -138,7 +127,7 @@ export class UserResolver {
})
@Public()
async user(@Args('email') email: string) {
if (!(await this.users.canEarlyAccess(email))) {
if (!(await this.feature.canEarlyAccess(email))) {
return new GraphQLError(
`You don't have early access permission\nVisit https://community.affine.pro/c/insider-general/ for more information`,
{
@@ -233,27 +222,55 @@ export class UserResolver {
ttl: 60,
},
})
@Mutation(() => AddToNewFeaturesWaitingList)
async addToNewFeaturesWaitingList(
@CurrentUser() user: UserType,
@Args('type', {
type: () => NewFeaturesKind,
})
type: NewFeaturesKind,
@Mutation(() => Int)
async addToEarlyAccess(
@CurrentUser() currentUser: UserType,
@Args('email') email: string
): Promise<AddToNewFeaturesWaitingList> {
if (!isStaff(user.email)) {
): Promise<number> {
if (!this.feature.isStaff(currentUser.email)) {
throw new ForbiddenException('You are not allowed to do this');
}
await this.prisma.newFeaturesWaitingList.create({
data: {
email,
type,
},
});
return {
email,
type,
};
const user = await this.users.findUserByEmail(email);
if (user) {
return this.feature.addEarlyAccess(user.id);
} else {
const user = await this.auth.createAnonymousUser(email);
return this.feature.addEarlyAccess(user.id);
}
}
@Throttle({
default: {
limit: 10,
ttl: 60,
},
})
@Mutation(() => Int)
async removeEarlyAccess(
@CurrentUser() currentUser: UserType,
@Args('email') email: string
): Promise<number> {
if (!this.feature.isStaff(currentUser.email)) {
throw new ForbiddenException('You are not allowed to do this');
}
const user = await this.users.findUserByEmail(email);
if (!user) {
throw new BadRequestException(`User ${email} not found`);
}
return this.feature.removeEarlyAccess(user.id);
}
@Throttle({
default: {
limit: 10,
ttl: 60,
},
})
@Query(() => [UserType])
async listEarlyAccess(@CurrentUser() user: UserType): Promise<UserType[]> {
if (!this.feature.isStaff(user.email)) {
throw new ForbiddenException('You are not allowed to do this');
}
return this.feature.listEarlyAccess();
}
}

View File

@@ -1,3 +0,0 @@
export enum NewFeaturesKind {
EarlyAccess,
}

View File

@@ -1,54 +1,10 @@
import { Injectable } from '@nestjs/common';
import { Config } from '../../config';
import { PrismaService } from '../../prisma';
import { getStorageQuota } from './gates';
import { NewFeaturesKind } from './types';
import { isStaff } from './utils';
@Injectable()
export class UsersService {
constructor(
private readonly prisma: PrismaService,
private readonly config: Config
) {}
async canEarlyAccess(email: string) {
if (this.config.featureFlags.earlyAccessPreview && !isStaff(email)) {
return this.isEarlyAccessUser(email);
} else {
return true;
}
}
async isEarlyAccessUser(email: string) {
return this.prisma.newFeaturesWaitingList
.count({
where: { email, type: NewFeaturesKind.EarlyAccess },
})
.then(count => count > 0)
.catch(() => false);
}
async getStorageQuotaById(id: string) {
const features = await this.prisma.user
.findUnique({
where: { id },
select: {
features: {
select: {
feature: true,
},
},
},
})
.then(user => user?.features.map(f => f.feature) ?? []);
return (
getStorageQuota(features.map(f => f.feature)) ||
this.config.objectStorage.quota
);
}
constructor(private readonly prisma: PrismaService) {}
async findUserByEmail(email: string) {
return this.prisma.user

View File

@@ -1,3 +0,0 @@
export function isStaff(email: string) {
return email.endsWith('@toeverything.info');
}

View File

@@ -1,6 +1,7 @@
import { Module } from '@nestjs/common';
import { DocModule } from '../doc';
import { QuotaModule } from '../quota';
import { UsersService } from '../users';
import { WorkspacesController } from './controller';
import { DocHistoryResolver } from './history.resolver';
@@ -8,7 +9,7 @@ import { PermissionService } from './permission';
import { PagePermissionResolver, WorkspaceResolver } from './resolver';
@Module({
imports: [DocModule],
imports: [DocModule, QuotaModule],
controllers: [WorkspacesController],
providers: [
WorkspaceResolver,

View File

@@ -42,6 +42,7 @@ import { DocID } from '../../utils/doc';
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 { PermissionService, PublicPageMode } from './permission';
@@ -148,6 +149,7 @@ export class WorkspaceResolver {
private readonly permissions: PermissionService,
private readonly users: UsersService,
private readonly event: EventEmitter,
private readonly quota: QuotaManagementService,
@Inject(StorageProvide) private readonly storage: Storage
) {}
@@ -233,6 +235,14 @@ export class WorkspaceResolver {
}));
}
@ResolveField(() => Int, {
description: 'Blobs size of workspace',
complexity: 2,
})
async blobsSize(@Parent() workspace: WorkspaceType) {
return this.storage.blobsSize([workspace.id]);
}
@Query(() => Boolean, {
description: 'Get is owner of workspace',
complexity: 2,
@@ -656,36 +666,9 @@ export class WorkspaceResolver {
return this.storage.listBlobs(workspaceId);
}
@Query(() => WorkspaceBlobSizes)
async collectBlobSizes(
@CurrentUser() user: UserType,
@Args('workspaceId') workspaceId: string
) {
await this.permissions.checkWorkspace(workspaceId, user.id);
return this.storage.blobsSize([workspaceId]).then(size => ({ size }));
}
@Query(() => WorkspaceBlobSizes)
async collectAllBlobSizes(@CurrentUser() user: UserType) {
const workspaces = await this.prisma.workspaceUserPermission
.findMany({
where: {
userId: user.id,
accepted: true,
type: Permission.Owner,
},
select: {
workspace: {
select: {
id: true,
},
},
},
})
.then(data => data.map(({ workspace }) => workspace.id));
const size = await this.storage.blobsSize(workspaces);
const size = await this.quota.getUserUsage(user.id);
return { size };
}
@@ -693,7 +676,7 @@ export class WorkspaceResolver {
async checkBlobSize(
@CurrentUser() user: UserType,
@Args('workspaceId') workspaceId: string,
@Args('size', { type: () => Float }) size: number
@Args('size', { type: () => Float }) blobSize: number
) {
const canWrite = await this.permissions.tryCheckWorkspace(
workspaceId,
@@ -701,13 +684,8 @@ export class WorkspaceResolver {
Permission.Write
);
if (canWrite) {
const { user } = await this.permissions.getWorkspaceOwner(workspaceId);
if (user) {
const quota = await this.users.getStorageQuotaById(user.id);
const { size: currentSize } = await this.collectAllBlobSizes(user);
return { size: quota - (size + currentSize) };
}
const size = await this.quota.checkBlobQuota(workspaceId, blobSize);
return { size };
}
return false;
}
@@ -725,14 +703,12 @@ export class WorkspaceResolver {
Permission.Write
);
// quota was apply to owner's account
const { user: owner } =
await this.permissions.getWorkspaceOwner(workspaceId);
if (!owner) return new NotFoundException('Workspace owner not found');
const quota = await this.users.getStorageQuotaById(owner.id);
const { size } = await this.collectAllBlobSizes(owner);
const { quota, size } = await this.quota.getWorkspaceUsage(workspaceId);
const checkExceeded = (recvSize: number) => {
if (!quota) {
throw new ForbiddenException('cannot find user quota');
}
if (size + recvSize > quota) {
this.logger.log(
`storage size limit exceeded: ${size + recvSize} > ${quota}`

View File

@@ -51,17 +51,6 @@ type RemoveAvatar {
success: Boolean!
}
type AddToNewFeaturesWaitingList {
email: String!
"""New features kind"""
type: NewFeaturesKind!
}
enum NewFeaturesKind {
EarlyAccess
}
type TokenType {
token: String!
refresh: String!
@@ -196,6 +185,9 @@ type WorkspaceType {
"""Owner of workspace"""
owner: UserType!
"""Blobs size of workspace"""
blobsSize: Int!
"""Shared pages of workspace"""
sharedPages: [String!]! @deprecated(reason: "use WorkspaceType.publicPages")
@@ -269,7 +261,6 @@ type Query {
"""List blobs of workspace"""
listBlobs(workspaceId: String!): [String!]!
collectBlobSizes(workspaceId: String!): WorkspaceBlobSizes!
collectAllBlobSizes: WorkspaceBlobSizes!
checkBlobSize(workspaceId: String!, size: Float!): WorkspaceBlobSizes!
@@ -278,6 +269,7 @@ type Query {
"""Get user by email"""
user(email: String!): UserType
listEarlyAccess: [UserType!]!
prices: [SubscriptionPrice!]!
}
@@ -315,7 +307,8 @@ type Mutation {
"""Remove user avatar"""
removeAvatar: RemoveAvatar!
deleteAccount: DeleteAccount!
addToNewFeaturesWaitingList(type: NewFeaturesKind!, email: String!): AddToNewFeaturesWaitingList!
addToEarlyAccess(email: String!): Int!
removeEarlyAccess(email: String!): Int!
"""Create a subscription checkout link of stripe"""
checkout(recurring: SubscriptionRecurring!, idempotencyKey: String!): String!