PB-845: Calculate folder size and expose it via shared & Drive

migrations/20240112041245_create-calculate-folder-size-function.js

@@ -0,0 +1,34 @@
+'use strict';
+
+const functionName = 'calculate_folder_size';
+
+/** @type {import('sequelize-cli').Migration} */
+module.exports = {
+  up: async (queryInterface) => {
+    await queryInterface.sequelize.query(`
+        CREATE OR REPLACE FUNCTION ${functionName}(folder_uuid_param UUID) RETURNS BIGINT AS $$
+        DECLARE
+            total_size BIGINT := 0;
+            subfolder_uuid UUID;
+        BEGIN
+        total_size := total_size + COALESCE(
+            (SELECT SUM(f.size) FROM files f WHERE f.folder_uuid = folder_uuid_param),
+            0
+        );
+        
+        FOR subfolder_uuid IN (SELECT uuid FROM folders WHERE parent_uuid = folder_uuid_param) LOOP
+            total_size := total_size + ${functionName}(subfolder_uuid);
+        END LOOP;
+        
+        RETURN total_size;
+        END;
+        $$ LANGUAGE plpgsql;    
+    `);
+  },
+
+  down: async (queryInterface) => {
+    await queryInterface.sequelize.query(
+      `DROP FUNCTION IF EXISTS ${functionName}(UUID);`,
+    );
+  },
+};

src/common/base-http.exception.ts

@@ -1,25 +1,11 @@
 import { HttpStatus } from '@nestjs/common';
 
 export class BaseHttpException extends Error {
-  private readonly _statusCode: number;
-  private readonly _code: string;
-
   constructor(
-    message: string,
-    statusCode = HttpStatus.INTERNAL_SERVER_ERROR,
-    code?: string,
+    public readonly message: string,
+    public readonly statusCode = HttpStatus.INTERNAL_SERVER_ERROR,
+    public readonly code?: string,
   ) {
     super(message);
-    this.message = message;
-    this._statusCode = statusCode;
-    this._code = code;
-  }
-
-  get statusCode(): number {
-    return this._statusCode;
-  }
-
-  get code(): string {
-    return this._code;
   }
 }

src/modules/folder/exception/calculate-folder-size-timeout.exception.ts

@@ -0,0 +1,12 @@
+import { HttpStatus } from '@nestjs/common';
+import { BaseHttpException } from '../../../common/base-http.exception';
+
+export class CalculateFolderSizeTimeoutException extends BaseHttpException {
+  constructor(
+    message = 'Calculate folder size timeout',
+    code = 'CALCULATE_FOLDER_SIZE_TIMEOUT',
+    statusCode = HttpStatus.UNPROCESSABLE_ENTITY,
+  ) {
+    super(message, statusCode, code);
+  }
+}

src/modules/folder/folder.controller.spec.ts

@@ -0,0 +1,49 @@
+import { createMock } from '@golevelup/ts-jest';
+import { Test, TestingModule } from '@nestjs/testing';
+import { FileUseCases } from '../file/file.usecase';
+import { FolderController } from './folder.controller';
+import { FolderUseCases } from './folder.usecase';
+import { CalculateFolderSizeTimeoutException } from './exception/calculate-folder-size-timeout.exception';
+
+describe('FolderController', () => {
+  let folderController: FolderController;
+  let folderUseCases: FolderUseCases;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      controllers: [FolderController],
+      providers: [
+        { provide: FolderUseCases, useValue: createMock() },
+        { provide: FileUseCases, useValue: createMock() },
+      ],
+    }).compile();
+
+    folderController = module.get<FolderController>(FolderController);
+    folderUseCases = module.get<FolderUseCases>(FolderUseCases);
+  });
+
+  describe('getFolderSize', () => {
+    it('When getFolderSizeController is called, Then return folder size', async () => {
+      const folderUuid = 'valid-uuid';
+      const expectedSize = 100;
+      jest
+        .spyOn(folderUseCases, 'getFolderSizeByUuid')
+        .mockResolvedValue(expectedSize);
+
+      const result = await folderController.getFolderSize(folderUuid);
+      expect(result).toEqual({ size: expectedSize });
+    });
+
+    it('When getFolderSizeController throws timeout exception, Then throws CalculateFolderSizeTimeoutException instance', async () => {
+      const folderUuid = 'valid-uuid';
+
+      jest
+        .spyOn(folderUseCases, 'getFolderSizeByUuid')
+        .mockRejectedValue(new CalculateFolderSizeTimeoutException());
+
+      await expect(folderController.getFolderSize(folderUuid)).rejects.toThrow(
+        CalculateFolderSizeTimeoutException,
+      );
+    });
+  });
+});

src/modules/folder/folder.controller.ts

@@ -9,6 +9,7 @@ import {
   NotImplementedException,
   Param,
   Query,
+  UseFilters,
 } from '@nestjs/common';
 import { ApiTags } from '@nestjs/swagger';
 import { FolderUseCases } from './folder.usecase';
@@ -19,6 +20,7 @@ import { Folder, SortableFolderAttributes } from './folder.domain';
 import { FileStatus, SortableFileAttributes } from '../file/file.domain';
 import logger from '../../externals/logger';
 import { validate } from 'uuid';
+import { HttpExceptionFilter } from '../../lib/http/http-exception.filter';
 
 const foldersStatuses = ['ALL', 'EXISTS', 'TRASHED', 'DELETED'] as const;
 
@@ -255,7 +257,7 @@ export class FolderController {
     @UserDecorator() user: User,
     @Query('limit') limit: number,
     @Query('offset') offset: number,
-    @Query('status') status: typeof foldersStatuses[number],
+    @Query('status') status: (typeof foldersStatuses)[number],
     @Query('updatedAt') updatedAt?: string,
   ) {
     if (!status) {
@@ -418,4 +420,12 @@ export class FolderController {
       });
     }
   }
+
+  @UseFilters(new HttpExceptionFilter())
+  @Get(':uuid/size')
+  async getFolderSize(@Param('uuid') folderUuid: Folder['uuid']) {
+    const size = await this.folderUseCases.getFolderSizeByUuid(folderUuid);
+
+    return { size };
+  }
 }

src/modules/folder/folder.repository.spec.ts

@@ -0,0 +1,67 @@
+import { createMock } from '@golevelup/ts-jest';
+import { CalculateFolderSizeTimeoutException } from './exception/calculate-folder-size-timeout.exception';
+import { SequelizeFolderRepository } from './folder.repository';
+import { FolderModel } from './folder.model';
+
+jest.mock('./folder.model', () => ({
+  FolderModel: {
+    sequelize: {
+      query: jest.fn(() => Promise.resolve([[{ totalsize: 100 }]])),
+    },
+  },
+}));
+
+describe('SequelizeFolderRepository', () => {
+  let repository: SequelizeFolderRepository;
+  let folderModel: typeof FolderModel;
+
+  beforeEach(async () => {
+    folderModel = createMock<typeof FolderModel>();
+
+    repository = new SequelizeFolderRepository(folderModel);
+  });
+
+  describe('calculateFolderSize', () => {
+    const folderUuid = 'your-folder-uuid';
+    const timeoutSeconds = 15;
+
+    it('When calculateFolderSize is called, Then call the query correctly', async () => {
+      const calculateSizeQuery = `
+          SET LOCAL statement_timeout = :timeoutSeconds;
+          SELECT calculate_folder_size(:folderUuid) AS totalSize;
+      `;
+
+      jest
+        .spyOn(FolderModel.sequelize, 'query')
+        .mockResolvedValue([[{ totalsize: 100 }]] as any);
+
+      const size = await repository.calculateFolderSize(
+        folderUuid,
+        timeoutSeconds,
+      );
+
+      expect(size).toBeGreaterThanOrEqual(0);
+      expect(FolderModel.sequelize.query).toHaveBeenCalledWith(
+        calculateSizeQuery,
+        {
+          replacements: {
+            folderUuid,
+            timeoutSeconds: timeoutSeconds * 1000,
+          },
+        },
+      );
+    });
+
+    it('When calculateFolderSize query throws a timeout error, Then throw CalculateFolderSizeTimeoutException', async () => {
+      jest.spyOn(FolderModel.sequelize, 'query').mockRejectedValue({
+        original: {
+          code: '57014',
+        },
+      });
+
+      await expect(
+        repository.calculateFolderSize(folderUuid, timeoutSeconds),
+      ).rejects.toThrow(CalculateFolderSizeTimeoutException);
+    });
+  });
+});

src/modules/folder/folder.repository.ts

@@ -6,12 +6,12 @@ import { v4 } from 'uuid';
 import { Folder } from './folder.domain';
 import { FolderAttributes } from './folder.attributes';
 
-import { UserModel } from '../user/user.model';
 import { User } from '../user/user.domain';
 import { UserAttributes } from '../user/user.attributes';
 import { Pagination } from '../../lib/pagination';
 import { FolderModel } from './folder.model';
 import { SharingModel } from '../sharing/models';
+import { CalculateFolderSizeTimeoutException } from './exception/calculate-folder-size-timeout.exception';
 
 function mapSnakeCaseToCamelCase(data) {
   const camelCasedObject = {};
@@ -59,15 +59,17 @@ export interface FolderRepository {
   ): Promise<void>;
   deleteById(folderId: FolderAttributes['id']): Promise<void>;
   clearOrphansFolders(userId: FolderAttributes['userId']): Promise<number>;
+  calculateFolderSize(
+    folderUuid: string,
+    timeoutSeconds: number,
+  ): Promise<number>;
 }
 
 @Injectable()
 export class SequelizeFolderRepository implements FolderRepository {
   constructor(
     @InjectModel(FolderModel)
     private folderModel: typeof FolderModel,
-    @InjectModel(UserModel)
-    private userModel: typeof UserModel,
   ) {}
 
   async findAllCursor(
@@ -447,6 +449,36 @@ export class SequelizeFolderRepository implements FolderRepository {
     return folders.map((folder) => this.toDomain(folder));
   }
 
+  async calculateFolderSize(
+    folderUuid: string,
+    timeoutSeconds = 15,
+  ): Promise<number> {
+    try {
+      const calculateSizeQuery = `
+          SET LOCAL statement_timeout = :timeoutSeconds;
+          SELECT calculate_folder_size(:folderUuid) AS totalSize;
+      `;
+
+      const [[{ totalsize }]]: any = await FolderModel.sequelize.query(
+        calculateSizeQuery,
+        {
+          replacements: {
+            folderUuid,
+            timeoutSeconds: timeoutSeconds * 1000,
+          },
+        },
+      );
+
+      return +totalsize;
+    } catch (error) {
+      if (error.original?.code === '57014') {
+        throw new CalculateFolderSizeTimeoutException();
+      }
+
+      throw error;
+    }
+  }
+
   private toDomain(model: FolderModel): Folder {
     return Folder.build({
       ...model.toJSON(),
@@ -459,4 +491,6 @@ export class SequelizeFolderRepository implements FolderRepository {
     return domain.toJSON();
   }
 }
+
+// QUESTION: Why is it exporting from here?
 export { FolderModel };

src/modules/folder/folder.usecase.spec.ts

@@ -26,6 +26,7 @@ import { FolderModel } from './folder.model';
 import { FileModel } from '../file/file.model';
 import { SequelizeThumbnailRepository } from '../thumbnail/thumbnail.repository';
 import { ThumbnailModel } from '../thumbnail/thumbnail.model';
+import { CalculateFolderSizeTimeoutException } from './exception/calculate-folder-size-timeout.exception';
 
 const folderId = 4;
 const userId = 1;
@@ -492,4 +493,32 @@ describe('FolderUseCases', () => {
       }
     });
   });
+
+  describe('get folder size', () => {
+    it('When getFolderSizeByUuid is called, Then return the folder size', async () => {
+      const mockSize = 123456789;
+
+      jest
+        .spyOn(folderRepository, 'calculateFolderSize')
+        .mockResolvedValueOnce(mockSize);
+
+      const result = await service.getFolderSizeByUuid('an-uuid');
+
+      expect(result).toBe(mockSize);
+      expect(folderRepository.calculateFolderSize).toHaveBeenCalledTimes(1);
+      expect(folderRepository.calculateFolderSize).toHaveBeenCalledWith(
+        'an-uuid',
+      );
+    });
+
+    it('When getFolderSizeByUuid times out, Then throw CalculateFolderSizeTimeoutException', async () => {
+      jest
+        .spyOn(folderRepository, 'calculateFolderSize')
+        .mockRejectedValueOnce(new CalculateFolderSizeTimeoutException());
+
+      await expect(service.getFolderSizeByUuid('an-uuid')).rejects.toThrow(
+        CalculateFolderSizeTimeoutException,
+      );
+    });
+  });
 });

src/modules/folder/folder.usecase.ts

@@ -16,6 +16,7 @@ import {
 } from './folder.domain';
 import { FolderAttributes } from './folder.attributes';
 import { SequelizeFolderRepository } from './folder.repository';
+import { SequelizeFileRepository } from '../file/file.repository';
 
 const invalidName = /[\\/]|^\s*$/;
 
@@ -26,6 +27,7 @@ export class FolderUseCases {
   constructor(
     private folderRepository: SequelizeFolderRepository,
     private userRepository: SequelizeUserRepository,
+    private readonly fileRepository: SequelizeFileRepository,
     private readonly cryptoService: CryptoService,
   ) {}
 
@@ -498,9 +500,8 @@ export class FolderUseCases {
   }
 
   async deleteOrphansFolders(userId: UserAttributes['id']): Promise<number> {
-    let remainingFolders = await this.folderRepository.clearOrphansFolders(
-      userId,
-    );
+    let remainingFolders =
+      await this.folderRepository.clearOrphansFolders(userId);
 
     if (remainingFolders > 0) {
       remainingFolders += await this.deleteOrphansFolders(userId);
@@ -536,4 +537,8 @@ export class FolderUseCases {
   async deleteByUser(user: User, folders: Folder[]): Promise<void> {
     await this.folderRepository.deleteByUser(user, folders);
   }
+
+  getFolderSizeByUuid(folderUuid: Folder['uuid']): Promise<number> {
+    return this.folderRepository.calculateFolderSize(folderUuid);
+  }
 }