Enable AWS S3 image storage : Issue #208

.env

@@ -12,6 +12,15 @@ POSTGRES_USER=postgres
 POSTGRES_PASSWORD=postgres
 POSTGRES_DB=vrt_db_dev
 
+#If you store your images in AWS S3 bucket, then you can give this config and it will take that config. A small amout of data is needed to temporarily keep the downloaded images. You can configure the size of that below.
+USE_AWS_S3_BUCKET=false
+AWS_ACCESS_KEY_ID=
+AWS_SECRET_ACCESS_KEY=
+AWS_REGION=
+AWS_S3_BUCKET_NAME=
+# Below is the maximum allowed usage of local space to optimize/prevent multiple download from S3 bucket. Give the value in MB. E.g. enter 1024 for 1GB
+MAX_TEMP_STORAGE_FOR_S3_DOWNLOAD=
+
 # optional
 #HTTPS_KEY_PATH='./secrets/ssl.key'
 #HTTPS_CERT_PATH='./secrets/ssl.cert'

package.json

@@ -26,6 +26,7 @@
     "node": ">=18.12.0"
   },
   "dependencies": {
+    "@aws-sdk/client-s3": "^3.715.0",
     "@nestjs/cache-manager": "^2.1.0",
     "@nestjs/common": "^10.2.5",
     "@nestjs/config": "^3.1.1",

src/compare/libs/looks-same/looks-same.service.ts

@@ -32,8 +32,8 @@ export class LookSameService implements ImageComparator {
       ...NO_BASELINE_RESULT,
     };
 
-    const baseline = this.staticService.getImage(data.baseline);
-    const image = this.staticService.getImage(data.image);
+    const baseline = await this.staticService.getImage(data.baseline);
+    const image = await this.staticService.getImage(data.image);
 
     if (!baseline) {
       return NO_BASELINE_RESULT;

src/compare/libs/pixelmatch/pixelmatch.service.ts

@@ -27,8 +27,8 @@ export class PixelmatchService implements ImageComparator {
       ...NO_BASELINE_RESULT,
     };
 
-    const baseline = this.staticService.getImage(data.baseline);
-    const image = this.staticService.getImage(data.image);
+    const baseline = await this.staticService.getImage(data.baseline);
+    const image = await this.staticService.getImage(data.image);
 
     if (!baseline) {
       return NO_BASELINE_RESULT;
@@ -68,7 +68,7 @@ export class PixelmatchService implements ImageComparator {
     if (result.diffPercent > data.diffTollerancePercent) {
       // save diff
       if (data.saveDiffAsFile) {
-        result.diffName = this.staticService.saveImage('diff', PNG.sync.write(diff));
+        result.diffName = await this.staticService.saveImage('diff', PNG.sync.write(diff));
       }
       result.status = TestStatus.unresolved;
     } else {

src/shared/shared.module.ts

@@ -4,12 +4,13 @@ import { EventsGateway } from '../shared/events/events.gateway';
 import { PrismaService } from '../prisma/prisma.service';
 import { TasksService } from './tasks/tasks.service';
 import { TestVariationsModule } from '../test-variations/test-variations.module';
+import { StaticController } from './static/static.controller';
 
 @Global()
 @Module({
   providers: [StaticService, EventsGateway, PrismaService, TasksService],
   exports: [StaticService, EventsGateway, PrismaService],
   imports: [forwardRef(() => TestVariationsModule)],
-  controllers: [],
+  controllers: [StaticController],
 })
 export class SharedModule {}

src/shared/static/static.controller.ts

@@ -0,0 +1,40 @@
+import { Controller, Get, Logger, Param, Res } from '@nestjs/common';
+import { Response } from 'express';
+import { StaticService } from './static.service';
+import { ApiOkResponse, ApiTags } from '@nestjs/swagger';
+
+@ApiTags('images')
+@Controller('images')
+export class StaticController {
+  private readonly logger: Logger = new Logger(StaticController.name);
+
+  constructor(private staticService: StaticService) {}
+
+  @Get('/:fileName')
+  @ApiOkResponse()
+  async downloadPngAndRedirect(@Param('fileName') fileName: string, @Res() res: Response) {
+    try {
+      if (!fileName.endsWith('.png')) {
+        return res.status(400).send('Invalid file type. Only PNG files are allowed.');
+      }
+      if (this.staticService.doesFileExist(fileName)) {
+        res.redirect('/' + fileName);
+      } else {
+        const localFileStream = await this.staticService.saveFileToServerFromS3(fileName);
+        localFileStream.on('finish', () => {
+          this.staticService.scheduleFileDeletion(fileName);
+          this.staticService.checkDiskUsageAndClean();
+          // After saving the file from S3, just redirect to the local file.
+          res.redirect('/' + fileName);
+        });
+        localFileStream.on('error', (error) => {
+          this.logger.error('Error writing file:', error);
+          res.status(500).send('Error occurred while saving the file.');
+        });
+      }
+    } catch (error) {
+      this.logger.error('Error fetching file from S3:' + fileName, error);
+      res.status(500).send('Error occurred while downloading the file.');
+    }
+  }
+}

src/shared/static/static.service.ts

@@ -1,14 +1,69 @@
 import { Injectable, Logger } from '@nestjs/common';
 import path from 'path';
-import { writeFileSync, readFileSync, unlink, mkdirSync, existsSync } from 'fs';
+import {
+  writeFileSync,
+  readFileSync,
+  unlink,
+  mkdirSync,
+  existsSync,
+  createWriteStream,
+  unlinkSync,
+  statSync,
+  readdir,
+} from 'fs';
 import { PNG, PNGWithMetadata } from 'pngjs';
 
+import { S3Client, PutObjectCommand, DeleteObjectCommand, GetObjectCommand } from '@aws-sdk/client-s3';
+import { Readable } from 'stream';
+
 export const IMAGE_PATH = 'imageUploads/';
 
 @Injectable()
 export class StaticService {
   private readonly logger: Logger = new Logger(StaticService.name);
 
+  private readonly USE_AWS_S3_BUCKET = process.env.USE_AWS_S3_BUCKET?.trim().toLowerCase();
+  private readonly AWS_ACCESS_KEY_ID = process.env.AWS_ACCESS_KEY_ID;
+  private readonly AWS_SECRET_ACCESS_KEY = process.env.AWS_SECRET_ACCESS_KEY;
+  private readonly AWS_REGION = process.env.AWS_REGION;
+  private readonly AWS_S3_BUCKET_NAME = process.env.AWS_S3_BUCKET_NAME;
+  private readonly MAX_DISK_USAGE: number = parseInt(process.env.MAX_TEMP_STORAGE_FOR_S3_DOWNLOAD) * 1 * 1024 * 1024;
+  private s3Client: S3Client;
+
+  private readonly DELETE_INTERVAL = 60 * 60 * 1000; // Check for deletions every hour
+  private deletionQueue: { filePath: string; size: number }[] = [];
+
+  constructor() {
+    setInterval(() => this.cleanupQueuedFiles(), this.DELETE_INTERVAL);
+    this.cleanUpAllFiles();
+  }
+
+  getEnvBoolean(value: string, defaultValue: boolean = false): boolean {
+    if (value === 'true' || value === '1') return true;
+    if (value === 'false' || value === '0') return false;
+    return defaultValue;
+  }
+
+  isAWSDefined(): boolean {
+    const areAWSVariablesValid =
+      this.getEnvBoolean(this.USE_AWS_S3_BUCKET) &&
+      this.AWS_ACCESS_KEY_ID?.trim().length > 1 &&
+      this.AWS_SECRET_ACCESS_KEY?.trim().length > 1 &&
+      this.AWS_S3_BUCKET_NAME?.trim().length > 1;
+
+    if (areAWSVariablesValid && !this.s3Client) {
+      this.s3Client = new S3Client({
+        credentials: {
+          accessKeyId: this.AWS_ACCESS_KEY_ID,
+          secretAccessKey: this.AWS_SECRET_ACCESS_KEY,
+        },
+        region: this.AWS_REGION,
+      });
+      this.logger.log('Using AWS S3 bucket.');
+    }
+    return areAWSVariablesValid;
+  }
+
   generateNewImage(type: 'screenshot' | 'diff' | 'baseline'): { imageName: string; imagePath: string } {
     const imageName = `${Date.now()}.${type}.png`;
     return {
@@ -22,37 +77,148 @@ export class StaticService {
     return path.resolve(IMAGE_PATH, imageName);
   }
 
-  saveImage(type: 'screenshot' | 'diff' | 'baseline', imageBuffer: Buffer): string {
+  async saveImage(type: 'screenshot' | 'diff' | 'baseline', imageBuffer: Buffer): Promise<string> {
+    const { imageName, imagePath } = this.generateNewImage(type);
     try {
-      new PNG().parse(imageBuffer);
+      if (this.isAWSDefined()) {
+        await this.s3Client.send(
+          new PutObjectCommand({
+            Bucket: this.AWS_S3_BUCKET_NAME,
+            Key: imageName,
+            ContentType: 'image/png',
+            Body: imageBuffer,
+          })
+        );
+        return imageName;
+      } else {
+        new PNG().parse(imageBuffer);
+        writeFileSync(imagePath, imageBuffer);
+        return imageName;
+      }
     } catch (ex) {
-      throw new Error('Cannot parse image as PNG file');
+      throw new Error('Cannot parse image as PNG file: ' + ex);
     }
+  }
 
-    const { imageName, imagePath } = this.generateNewImage(type);
-    writeFileSync(imagePath, imageBuffer);
-    return imageName;
+  scheduleFileDeletion(fileName: string): void {
+    const filePath = this.getImagePath(fileName);
+    const fileSize = statSync(filePath).size;
+    this.deletionQueue.push({ filePath, size: fileSize });
+    this.logger.log(`Scheduled deletion for file: ${filePath} with size: ${fileSize} bytes`);
+  }
+
+  checkDiskUsageAndClean(): void {
+    const totalDiskUsage = this.getTotalDiskUsage();
+    if (totalDiskUsage > this.MAX_DISK_USAGE) {
+      this.logger.log(`Disk usage exceeded. Triggering cleanup.`);
+      this.cleanupQueuedFiles();
+    }
+  }
+
+  private getTotalDiskUsage(): number {
+    return this.deletionQueue.reduce((total, file) => total + file.size, 0);
   }
 
-  getImage(imageName: string): PNGWithMetadata {
-    if (!imageName) return;
+  private cleanUpAllFiles() {
+    readdir(path.resolve(IMAGE_PATH), (err, files) => {
+      if (err) throw err;
+      for (const file of files) {
+        unlink(path.join(path.resolve(IMAGE_PATH), file), (err) => {
+          if (err) throw err;
+        });
+      }
+    });
+  }
+
+  private cleanupQueuedFiles(): void {
+    const totalDiskUsage = this.getTotalDiskUsage();
+    this.logger.log(`Cleaning up files. Total disk usage: ${totalDiskUsage} bytes`);
+
+    // Delete files until the total disk usage is within the limit
+    let currentUsage = totalDiskUsage;
+
+    if (currentUsage > this.MAX_DISK_USAGE) {
+      // Sort files by the earliest (oldest) first for deletion
+      this.deletionQueue.sort((a, b) => statSync(a.filePath).birthtimeMs - statSync(b.filePath).birthtimeMs);
+
+      // Reduce the size by half
+      while (currentUsage > this.MAX_DISK_USAGE / 2 && this.deletionQueue.length > 0) {
+        const fileToDelete = this.deletionQueue.shift();
+        if (fileToDelete) {
+          try {
+            unlinkSync(fileToDelete.filePath);
+            currentUsage -= fileToDelete.size;
+            this.logger.log(`Deleted file: ${fileToDelete.filePath}`);
+          } catch (error) {
+            this.logger.log(`Failed to delete file: ${fileToDelete.filePath}. Error: ${error.message}`);
+          }
+        }
+      }
+    }
+  }
+
+  doesFileExist(fileName: string) {
+    return existsSync(this.getImagePath(fileName));
+  }
+
+  async getImage(fileName: string): Promise<PNGWithMetadata> {
+    if (!fileName) return null;
     try {
-      return PNG.sync.read(readFileSync(this.getImagePath(imageName)));
+      if (!this.doesFileExist(fileName) && this.isAWSDefined()) {
+        const localFileStream = await this.saveFileToServerFromS3(fileName);
+        await new Promise<void>((resolve, reject) => {
+          localFileStream.on('finish', () => {
+            this.scheduleFileDeletion(fileName);
+            this.checkDiskUsageAndClean();
+            resolve();
+          });
+          localFileStream.on('error', (error) => {
+            this.logger.error('Error writing file:', error);
+            reject(error);
+          });
+        });
+      }
+      return PNG.sync.read(readFileSync(this.getImagePath(fileName)));
     } catch (ex) {
-      this.logger.error(`Cannot get image: ${imageName}. ${ex}`);
+      this.logger.error(`Error from read : Cannot get image: ${fileName}. ${ex}`);
+    }
+  }
+
+  async saveFileToServerFromS3(fileName: string) {
+    if (this.isAWSDefined()) {
+      const command = new GetObjectCommand({ Bucket: this.AWS_S3_BUCKET_NAME, Key: fileName });
+      const s3Response = await this.s3Client.send(command);
+      const fileStream = s3Response.Body as Readable;
+      const localFileStream = createWriteStream(this.getImagePath(fileName));
+      fileStream.pipe(localFileStream);
+      this.logger.log(`File saved locally at ${this.getImagePath(fileName)}`);
+      return localFileStream;
+    } else {
+      throw Error('Error connecting to AWS');
     }
   }
 
   async deleteImage(imageName: string): Promise<boolean> {
-    if (!imageName) return;
-    return new Promise((resolvePromise) => {
-      unlink(this.getImagePath(imageName), (err) => {
-        if (err) {
-          this.logger.error(err);
-        }
-        resolvePromise(true);
+    if (!imageName) return false;
+    if (this.isAWSDefined()) {
+      try {
+        await this.s3Client.send(new DeleteObjectCommand({ Bucket: this.AWS_S3_BUCKET_NAME, Key: imageName }));
+        return true;
+      } catch (error) {
+        this.logger.log(`Failed to delete image ${imageName}:`, error);
+        return false; // Return `false` if an error occurs.
+      }
+    } else {
+      return new Promise((resolvePromise) => {
+        unlink(this.getImagePath(imageName), (err) => {
+          if (err) {
+            this.logger.error(err);
+            resolvePromise(false);
+          }
+          resolvePromise(true);
+        });
       });
-    });
+    }
   }
 
   private ensureDirectoryExistence(dir: string) {

src/test-runs/test-runs.module.ts

@@ -5,11 +5,12 @@ import { PrismaService } from '../prisma/prisma.service';
 import { TestRunsController } from './test-runs.controller';
 import { TestVariationsModule } from '../test-variations/test-variations.module';
 import { CompareModule } from '../compare/compare.module';
+import { StaticController } from 'src/shared/static/static.controller';
 
 @Module({
   imports: [SharedModule, forwardRef(() => TestVariationsModule), CompareModule],
   providers: [TestRunsService, PrismaService],
-  controllers: [TestRunsController],
+  controllers: [TestRunsController, StaticController],
   exports: [TestRunsService],
 })
 export class TestRunsModule {}

src/test-runs/test-runs.service.ts

@@ -103,8 +103,8 @@ export class TestRunsService {
     }
 
     // save new baseline
-    const baseline = this.staticService.getImage(testRun.imageName);
-    const baselineName = this.staticService.saveImage('baseline', PNG.sync.write(baseline));
+    const baseline = await this.staticService.getImage(testRun.imageName);
+    const baselineName = await this.staticService.saveImage('baseline', PNG.sync.write(baseline));
 
     if (testRun.baselineBranchName !== testRun.branchName && !merge && !autoApprove) {
       // replace main branch with feature branch test variation
@@ -206,7 +206,7 @@ export class TestRunsService {
     imageBuffer: Buffer;
   }): Promise<TestRun> {
     // save image
-    const imageName = this.staticService.saveImage('screenshot', imageBuffer);
+    const imageName = await this.staticService.saveImage('screenshot', imageBuffer);
 
     const testRun = await this.prismaService.testRun.create({
       data: {