feat: integrate user usage into apis (#5075)

packages/backend/server/src/modules/workspaces/index.ts

@@ -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,

packages/backend/server/src/modules/workspaces/resolver.ts

@@ -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}`