EnrollmentController.java

package edu.ucsb.cs156.courses.controllers;

import com.opencsv.bean.StatefulBeanToCsv;
import com.opencsv.exceptions.CsvDataTypeMismatchException;
import com.opencsv.exceptions.CsvRequiredFieldEmptyException;
import edu.ucsb.cs156.courses.entities.EnrollmentDataPoint;
import edu.ucsb.cs156.courses.models.EnrollmentCSV;
import edu.ucsb.cs156.courses.repositories.EnrollmentDataPointRepository;
import edu.ucsb.cs156.courses.services.EnrollmentCSVService;
import edu.ucsb.cs156.courses.utilities.CourseUtilities;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.Parameter;
import io.swagger.v3.oas.annotations.media.Content;
import io.swagger.v3.oas.annotations.media.Schema;
import io.swagger.v3.oas.annotations.responses.ApiResponse;
import io.swagger.v3.oas.annotations.tags.Tag;
import java.io.IOException;
import java.io.OutputStreamWriter;
import java.io.Writer;
import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.stream.Collectors;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.util.Streamable;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.StreamingResponseBody;

@Slf4j
@Tag(name = "API for enrollment data")
@RequestMapping("/api/enrollment")
@RestController
public class EnrollmentController extends ApiController {

  @Autowired EnrollmentDataPointRepository enrollmentDataPointRepository;
  @Autowired private EnrollmentCSVService enrollmentCSVService;

  @Operation(summary = "Get enrollment history for a course or section")
  @GetMapping(value = "/search", produces = "application/json")
  public Iterable<EnrollmentDataPoint> search(
      @Parameter(
              name = "startQtr",
              description =
                  "starting quarter in yyyyq format, e.g. 20231 for W23, 20232 for S23, etc. (1=Winter, 2=Spring, 3=Summer, 4=Fall)",
              example = "20231",
              required = true)
          @RequestParam
          String startQtr,
      @Parameter(
              name = "endQtr",
              description =
                  "ending quarter in yyyyq format, e.g. 20231 for W23, 20232 for S23, etc. (1=Winter, 2=Spring, 3=Summer, 4=Fall)",
              example = "20231",
              required = true)
          @RequestParam
          String endQtr,
      @Parameter(
              name = "subjectArea",
              description = "simplified area name, e.g. CMPSC for computer science",
              example = "CMPSC",
              required = true)
          @RequestParam
          String subjectArea,
      @Parameter(
              name = "courseNumber",
              description = "the specific course number, e.g. 130A for CS130A",
              example = "130A",
              required = true)
          @RequestParam
          String courseNumber,
      @Parameter(name = "enrollCd", description = "enroll code", example = "08268")
          @RequestParam(required = false)
          String enrollCd,
      @Parameter(name = "section", description = "section number", example = "0100")
          @RequestParam(required = false)
          String section) {
    String courseId = CourseUtilities.makeFormattedCourseId(subjectArea, courseNumber);
    return enrollmentDataPointRepository.findByQuarterRangeAndCourseIdAndOptionalFilters(
        startQtr, endQtr, courseId, enrollCd, section);
  }

  @Operation(
      summary = "Download Enrollment Data as CSV File",
      description = "Returns a CSV file as a response",
      responses = {
        @ApiResponse(
            responseCode = "200",
            description = "CSV file",
            content =
                @Content(
                    mediaType = "text/csv",
                    schema = @Schema(type = "string", format = "binary"))),
        @ApiResponse(responseCode = "500", description = "Internal Server Error")
      })
  @GetMapping(value = "/csv/quarter", produces = "text/csv")
  public ResponseEntity<StreamingResponseBody> csvForQuarter(
      @Parameter(name = "yyyyq", description = "quarter in yyyyq format", example = "20252")
          @RequestParam
          String yyyyq)
      throws IOException {

    StreamingResponseBody stream =
        (outputStream) -> {
          Iterable<EnrollmentDataPoint> iterable = enrollmentDataPointRepository.findByYyyyq(yyyyq);
          List<EnrollmentCSV> list =
              Streamable.of(iterable).toList().stream()
                  .map(enrollmentDataPoint -> EnrollmentCSV.fromEntity(enrollmentDataPoint))
                  .collect(Collectors.toList());

          try (Writer writer = new OutputStreamWriter(outputStream, StandardCharsets.UTF_8)) {
            try {
              StatefulBeanToCsv<EnrollmentCSV> beanToCsvWriter =
                  enrollmentCSVService.getStatefulBeanToCSV(writer);
              beanToCsvWriter.write(list);
            } catch (CsvDataTypeMismatchException | CsvRequiredFieldEmptyException e) {
              log.error("Error writing CSV file", e);
              throw new IOException("Error writing CSV file: " + e.getMessage());
            }
          }
        };

    return ResponseEntity.ok()
        .contentType(MediaType.parseMediaType("text/csv; charset=UTF-8"))
        .header(
            HttpHeaders.CONTENT_DISPOSITION,
            String.format("attachment;filename=enrollment_%s.csv", yyyyq))
        .header(HttpHeaders.CONTENT_TYPE, "text/csv; charset=UTF-8")
        .header(HttpHeaders.ACCESS_CONTROL_EXPOSE_HEADERS, HttpHeaders.CONTENT_DISPOSITION)
        .body(stream);
  }
}