test-assignment-datawrapper/src/screenshots/screenshots.controller.ts

import { InjectQueue } from '@nestjs/bull';
import {
  Body,
  Controller,
  Get,
  NotFoundException,
  Param,
  Post,
  Res,
  StreamableFile,
} from '@nestjs/common';
import { ApiResponse } from '@nestjs/swagger';
import { Response } from 'express';
import {
  CreateJobRequestDto,
  GetJobResponseDto,
  JobStatusDto,
} from './screenshots.dto';
import { QUEUE_NAME, ScreenshotsQueue } from './shared';

@Controller('screenshots')
export class ScreenshotsController {
  constructor(
    @InjectQueue(QUEUE_NAME)
    private readonly screenshotsQueue: ScreenshotsQueue,
  ) {}

  @Post()
  async createScreenshotJob(
    @Body() createScreenshotJobDto: CreateJobRequestDto,
  ): Promise<{ jobId: string }> {
    const result = await this.screenshotsQueue.add({
      pageUrl: new URL(createScreenshotJobDto.pageUrl),
      imageType: createScreenshotJobDto.imageType,
    });
    return {
      jobId: result.id.toString(),
    };
  }

  @Get(':id')
  @ApiResponse({ status: 404 })
  async getJob(@Param('id') id: string): Promise<GetJobResponseDto> {
    const jobInfo = await this.screenshotsQueue.getJob(id);

    if (!jobInfo) {
      throw new NotFoundException();
    }

    switch (await jobInfo.getState()) {
      case 'completed':
        return { status: JobStatusDto.Completed };
      case 'failed':
        return { status: JobStatusDto.Failed };
      default:
        return { status: JobStatusDto.Queued };
    }
  }

  @Get(':id/result')
  @ApiResponse({ status: 404 })
  async getScreenshot(
    @Param('id') id: string,
    @Res({ passthrough: true }) res: Response,
  ): Promise<StreamableFile> {
    const jobInfo = await this.screenshotsQueue.getJob(id);

    if (!jobInfo) {
      throw new NotFoundException();
    }

    if (!(await jobInfo.isCompleted())) {
      throw new NotFoundException();
    }

    switch (jobInfo.data.imageType) {
      case 'jpeg':
        res.setHeader('Content-Type', 'image/jpeg');
        break;
      case 'png':
        res.setHeader('Content-Type', 'image/png');
        break;
    }

    const buffer = Buffer.from((await jobInfo.returnvalue) as string, 'base64');
    return new StreamableFile(buffer);
  }
}