feat: add query quota of workspace (#5603)

This commit is contained in:
DarkSky
2024-01-16 09:45:55 +00:00
parent 4f4d057aad
commit ee2520ec18
9 changed files with 98 additions and 14 deletions

View File

@@ -21,4 +21,4 @@ export class QuotaModule {}
export { QuotaManagementService, QuotaService };
export { Quota_FreePlanV1, Quota_ProPlanV1, Quotas } from './schema';
export { QuotaType } from './types';
export { QuotaQueryType, QuotaType } from './types';

View File

@@ -3,6 +3,7 @@ import { Injectable, NotFoundException } from '@nestjs/common';
import { WorkspaceBlobStorage } from '../storage';
import { PermissionService } from '../workspaces/permission';
import { QuotaService } from './service';
import { QuotaQueryType } from './types';
@Injectable()
export class QuotaManagementService {
@@ -40,21 +41,21 @@ export class QuotaManagementService {
// get workspace's owner quota and total size of used
// quota was apply to owner's account
async getWorkspaceUsage(workspaceId: string) {
async getWorkspaceUsage(workspaceId: string): Promise<QuotaQueryType> {
const { user: owner } =
await this.permissions.getWorkspaceOwner(workspaceId);
if (!owner) throw new NotFoundException('Workspace owner not found');
const { storageQuota, blobLimit } = await this.getUserQuota(owner.id);
// get all workspaces size of owner used
const usageSize = await this.getUserUsage(owner.id);
const usedSize = await this.getUserUsage(owner.id);
return { quota: storageQuota, size: usageSize, limit: blobLimit };
return { storageQuota, usedSize, blobLimit };
}
async checkBlobQuota(workspaceId: string, size: number) {
const { quota, size: usageSize } =
const { storageQuota, usedSize } =
await this.getWorkspaceUsage(workspaceId);
return quota - (size + usageSize);
return storageQuota - (size + usedSize);
}
}

View File

@@ -1,3 +1,4 @@
import { Field, Int, ObjectType } from '@nestjs/graphql';
import { z } from 'zod';
import { commonFeatureSchema, FeatureKind } from '../features';
@@ -37,6 +38,20 @@ export const QuotaSchema = commonFeatureSchema
export type Quota = z.infer<typeof QuotaSchema>;
/// ======== query types ========
@ObjectType()
export class QuotaQueryType {
@Field(() => Int)
storageQuota!: number;
@Field(() => Int)
usedSize!: number;
@Field(() => Int)
blobLimit!: number;
}
/// ======== utils ========
export function formatSize(bytes: number, decimals: number = 2): string {

View File

@@ -128,7 +128,7 @@ export class WorkspaceBlobResolver {
Permission.Write
);
const { quota, size, limit } =
const { storageQuota, usedSize, blobLimit } =
await this.quota.getWorkspaceUsage(workspaceId);
const unlimited = await this.feature.hasWorkspaceFeature(
@@ -137,7 +137,7 @@ export class WorkspaceBlobResolver {
);
const checkExceeded = (recvSize: number) => {
if (!quota) {
if (!storageQuota) {
throw new GraphQLError('cannot find user quota', {
extensions: {
status: HttpStatus[HttpStatus.FORBIDDEN],
@@ -145,13 +145,15 @@ export class WorkspaceBlobResolver {
},
});
}
const total = size + recvSize;
const total = usedSize + recvSize;
// only skip total storage check if workspace has unlimited feature
if (total > quota && !unlimited) {
this.logger.log(`storage size limit exceeded: ${total} > ${quota}`);
if (total > storageQuota && !unlimited) {
this.logger.log(
`storage size limit exceeded: ${total} > ${storageQuota}`
);
return true;
} else if (recvSize > limit) {
this.logger.log(`blob size limit exceeded: ${recvSize} > ${limit}`);
} else if (recvSize > blobLimit) {
this.logger.log(`blob size limit exceeded: ${recvSize} > ${blobLimit}`);
return true;
} else {
return false;

View File

@@ -31,7 +31,7 @@ import {
import { Auth, CurrentUser, Public } from '../../auth';
import { AuthService } from '../../auth/service';
import { FeatureManagementService, FeatureType } from '../../features';
import { QuotaManagementService } from '../../quota';
import { QuotaManagementService, QuotaQueryType } from '../../quota';
import { WorkspaceBlobStorage } from '../../storage';
import { UsersService, UserType } from '../../users';
import { PermissionService } from '../permission';
@@ -149,6 +149,15 @@ export class WorkspaceResolver {
}));
}
@ResolveField(() => QuotaQueryType, {
name: 'quota',
description: 'quota of workspace',
complexity: 2,
})
workspaceQuota(@Parent() workspace: WorkspaceType) {
return this.quota.getWorkspaceUsage(workspace.id);
}
@Query(() => Boolean, {
description: 'Get is owner of workspace',
complexity: 2,

View File

@@ -131,6 +131,9 @@ type WorkspaceType {
"""Owner of workspace"""
owner: UserType!
"""quota of workspace"""
quota: QuotaQueryType!
"""Available features of workspace"""
availableFeatures: [FeatureType!]!
@@ -183,6 +186,12 @@ type InvitationType {
invitee: UserType!
}
type QuotaQueryType {
storageQuota: Int!
usedSize: Int!
blobLimit: Int!
}
type TokenType {
token: String!
refresh: String!