websocket fun!
This commit is contained in:
Anthony Stirling
@@ -0,0 +1,19 @@
|
||||
package stirling.software.common.annotations;
|
||||
|
||||
import java.lang.annotation.*;
|
||||
|
||||
import org.springframework.core.annotation.AliasFor;
|
||||
import org.springframework.web.bind.annotation.RequestMapping;
|
||||
import org.springframework.web.bind.annotation.RequestMethod;
|
||||
|
||||
@Target(ElementType.METHOD)
|
||||
@Retention(RetentionPolicy.RUNTIME)
|
||||
@Documented
|
||||
@RequestMapping(method = RequestMethod.POST)
|
||||
public @interface AutoJobPostMapping {
|
||||
@AliasFor(annotation = RequestMapping.class, attribute = "value")
|
||||
String[] value() default {};
|
||||
|
||||
@AliasFor(annotation = RequestMapping.class, attribute = "consumes")
|
||||
String[] consumes() default {"multipart/form-data"};
|
||||
}
|
||||
@@ -0,0 +1,91 @@
|
||||
package stirling.software.common.aop;
|
||||
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.io.IOException;
|
||||
|
||||
import org.aspectj.lang.ProceedingJoinPoint;
|
||||
import org.aspectj.lang.annotation.*;
|
||||
import org.springframework.stereotype.Component;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
|
||||
import jakarta.servlet.http.HttpServletRequest;
|
||||
|
||||
import lombok.RequiredArgsConstructor;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.common.annotations.AutoJobPostMapping;
|
||||
import stirling.software.common.model.api.PDFFile;
|
||||
import stirling.software.common.service.FileOrUploadService;
|
||||
import stirling.software.common.service.FileStorage;
|
||||
import stirling.software.common.service.JobExecutorService;
|
||||
|
||||
@Aspect
|
||||
@Component
|
||||
@RequiredArgsConstructor
|
||||
@Slf4j
|
||||
public class AutoJobAspect {
|
||||
|
||||
private final JobExecutorService jobExecutorService;
|
||||
private final HttpServletRequest request;
|
||||
private final FileOrUploadService fileOrUploadService;
|
||||
private final FileStorage fileStorage;
|
||||
|
||||
@Around("@annotation(autoJobPostMapping)")
|
||||
public Object wrapWithJobExecution(
|
||||
ProceedingJoinPoint joinPoint, AutoJobPostMapping autoJobPostMapping) {
|
||||
boolean async = Boolean.parseBoolean(request.getParameter("async"));
|
||||
|
||||
// Inspect and possibly mutate arguments
|
||||
Object[] args = joinPoint.getArgs();
|
||||
boolean isAsyncRequest = async;
|
||||
|
||||
for (int i = 0; i < args.length; i++) {
|
||||
Object arg = args[i];
|
||||
|
||||
if (arg instanceof PDFFile pdfFile) {
|
||||
// Case 1: fileId is provided but no fileInput
|
||||
if (pdfFile.getFileInput() == null && pdfFile.getFileId() != null) {
|
||||
try {
|
||||
log.debug("Using fileId {} to get file content", pdfFile.getFileId());
|
||||
MultipartFile file = fileStorage.retrieveFile(pdfFile.getFileId());
|
||||
pdfFile.setFileInput(file);
|
||||
} catch (Exception e) {
|
||||
throw new RuntimeException(
|
||||
"Failed to resolve file by ID: " + pdfFile.getFileId(), e);
|
||||
}
|
||||
}
|
||||
// Case 2: For async requests, we need to make a copy of the MultipartFile
|
||||
else if (isAsyncRequest && pdfFile.getFileInput() != null) {
|
||||
try {
|
||||
log.debug("Making persistent copy of uploaded file for async processing");
|
||||
MultipartFile originalFile = pdfFile.getFileInput();
|
||||
String fileId = fileStorage.storeFile(originalFile);
|
||||
|
||||
// Store the fileId for later reference
|
||||
pdfFile.setFileId(fileId);
|
||||
|
||||
// Replace the original MultipartFile with our persistent copy
|
||||
MultipartFile persistentFile = fileStorage.retrieveFile(fileId);
|
||||
pdfFile.setFileInput(persistentFile);
|
||||
|
||||
log.debug("Created persistent file copy with fileId: {}", fileId);
|
||||
} catch (IOException e) {
|
||||
throw new RuntimeException("Failed to create persistent copy of uploaded file", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Wrap job execution
|
||||
return jobExecutorService.runJobGeneric(
|
||||
async,
|
||||
() -> {
|
||||
try {
|
||||
return joinPoint.proceed(args);
|
||||
} catch (Throwable ex) {
|
||||
throw new RuntimeException(ex);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
package stirling.software.common.config;
|
||||
|
||||
import org.springframework.context.annotation.Configuration;
|
||||
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
|
||||
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
|
||||
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
|
||||
import org.springframework.web.socket.config.annotation.WebSocketMessageBrokerConfigurer;
|
||||
|
||||
@Configuration
|
||||
@EnableWebSocketMessageBroker
|
||||
public class WebSocketConfig implements WebSocketMessageBrokerConfigurer {
|
||||
|
||||
@Override
|
||||
public void configureMessageBroker(MessageBrokerRegistry config) {
|
||||
config.enableSimpleBroker("/topic");
|
||||
config.setApplicationDestinationPrefixes("/app");
|
||||
}
|
||||
|
||||
@Override
|
||||
public void registerStompEndpoints(StompEndpointRegistry registry) {
|
||||
registry.addEndpoint("/ws").setAllowedOriginPatterns("*").withSockJS();
|
||||
}
|
||||
}
|
||||
+27
@@ -0,0 +1,27 @@
|
||||
package stirling.software.common.controller;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.messaging.simp.SimpMessagingTemplate;
|
||||
import org.springframework.stereotype.Controller;
|
||||
|
||||
import lombok.NoArgsConstructor;
|
||||
|
||||
import stirling.software.common.model.job.JobProgress;
|
||||
|
||||
@Controller
|
||||
@NoArgsConstructor
|
||||
public class WebSocketProgressController {
|
||||
|
||||
private SimpMessagingTemplate messagingTemplate;
|
||||
|
||||
@Autowired(required = false)
|
||||
public void setMessagingTemplate(SimpMessagingTemplate messagingTemplate) {
|
||||
this.messagingTemplate = messagingTemplate;
|
||||
}
|
||||
|
||||
public void sendProgress(String jobId, JobProgress progress) {
|
||||
if (messagingTemplate != null) {
|
||||
messagingTemplate.convertAndSend("/topic/progress/" + jobId, progress);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -14,8 +14,12 @@ import lombok.NoArgsConstructor;
|
||||
public class PDFFile {
|
||||
@Schema(
|
||||
description = "The input PDF file",
|
||||
requiredMode = Schema.RequiredMode.REQUIRED,
|
||||
contentMediaType = "application/pdf",
|
||||
format = "binary")
|
||||
private MultipartFile fileInput;
|
||||
|
||||
@Schema(
|
||||
description = "File ID for server-side files (can be used instead of fileInput)",
|
||||
example = "a1b2c3d4-5678-90ab-cdef-ghijklmnopqr")
|
||||
private String fileId;
|
||||
}
|
||||
|
||||
@@ -0,0 +1,15 @@
|
||||
package stirling.software.common.model.job;
|
||||
|
||||
import lombok.AllArgsConstructor;
|
||||
import lombok.Data;
|
||||
import lombok.NoArgsConstructor;
|
||||
|
||||
@Data
|
||||
@NoArgsConstructor
|
||||
@AllArgsConstructor
|
||||
public class JobProgress {
|
||||
private String jobId;
|
||||
private String status;
|
||||
private int percentComplete;
|
||||
private String message;
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
package stirling.software.common.model.job;
|
||||
|
||||
import lombok.AllArgsConstructor;
|
||||
import lombok.Data;
|
||||
import lombok.NoArgsConstructor;
|
||||
|
||||
@Data
|
||||
@NoArgsConstructor
|
||||
@AllArgsConstructor
|
||||
public class JobResponse<T> {
|
||||
private boolean async;
|
||||
private String jobId;
|
||||
private T result;
|
||||
}
|
||||
@@ -0,0 +1,115 @@
|
||||
package stirling.software.common.model.job;
|
||||
|
||||
import java.time.LocalDateTime;
|
||||
|
||||
import lombok.AllArgsConstructor;
|
||||
import lombok.Builder;
|
||||
import lombok.Data;
|
||||
import lombok.NoArgsConstructor;
|
||||
|
||||
/**
|
||||
* Represents the result of a job execution.
|
||||
* Used by the TaskManager to store job results.
|
||||
*/
|
||||
@Data
|
||||
@Builder
|
||||
@NoArgsConstructor
|
||||
@AllArgsConstructor
|
||||
public class JobResult {
|
||||
|
||||
/**
|
||||
* The job ID
|
||||
*/
|
||||
private String jobId;
|
||||
|
||||
/**
|
||||
* Flag indicating if the job is complete
|
||||
*/
|
||||
private boolean complete;
|
||||
|
||||
/**
|
||||
* Error message if the job failed
|
||||
*/
|
||||
private String error;
|
||||
|
||||
/**
|
||||
* The file ID of the result file, if applicable
|
||||
*/
|
||||
private String fileId;
|
||||
|
||||
/**
|
||||
* Original file name, if applicable
|
||||
*/
|
||||
private String originalFileName;
|
||||
|
||||
/**
|
||||
* MIME type of the result, if applicable
|
||||
*/
|
||||
private String contentType;
|
||||
|
||||
/**
|
||||
* Time when the job was created
|
||||
*/
|
||||
private LocalDateTime createdAt;
|
||||
|
||||
/**
|
||||
* Time when the job was completed
|
||||
*/
|
||||
private LocalDateTime completedAt;
|
||||
|
||||
/**
|
||||
* The actual result object, if not a file
|
||||
*/
|
||||
private Object result;
|
||||
|
||||
/**
|
||||
* Create a new JobResult with the given job ID
|
||||
*
|
||||
* @param jobId The job ID
|
||||
* @return A new JobResult
|
||||
*/
|
||||
public static JobResult createNew(String jobId) {
|
||||
return JobResult.builder()
|
||||
.jobId(jobId)
|
||||
.complete(false)
|
||||
.createdAt(LocalDateTime.now())
|
||||
.build();
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark this job as complete with a file result
|
||||
*
|
||||
* @param fileId The file ID of the result
|
||||
* @param originalFileName The original file name
|
||||
* @param contentType The content type of the file
|
||||
*/
|
||||
public void completeWithFile(String fileId, String originalFileName, String contentType) {
|
||||
this.complete = true;
|
||||
this.fileId = fileId;
|
||||
this.originalFileName = originalFileName;
|
||||
this.contentType = contentType;
|
||||
this.completedAt = LocalDateTime.now();
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark this job as complete with a general result
|
||||
*
|
||||
* @param result The result object
|
||||
*/
|
||||
public void completeWithResult(Object result) {
|
||||
this.complete = true;
|
||||
this.result = result;
|
||||
this.completedAt = LocalDateTime.now();
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark this job as failed with an error message
|
||||
*
|
||||
* @param error The error message
|
||||
*/
|
||||
public void failWithError(String error) {
|
||||
this.complete = true;
|
||||
this.error = error;
|
||||
this.completedAt = LocalDateTime.now();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
package stirling.software.common.service;
|
||||
|
||||
import java.io.ByteArrayInputStream;
|
||||
import java.io.IOException;
|
||||
import java.nio.file.*;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
|
||||
import lombok.RequiredArgsConstructor;
|
||||
|
||||
@Service
|
||||
@RequiredArgsConstructor
|
||||
public class FileOrUploadService {
|
||||
|
||||
@Value("${stirling.tempDir:/tmp/stirling-files}")
|
||||
private String tempDirPath;
|
||||
|
||||
public Path resolveFilePath(String fileId) {
|
||||
return Path.of(tempDirPath).resolve(fileId);
|
||||
}
|
||||
|
||||
public MultipartFile toMockMultipartFile(String name, byte[] data) throws IOException {
|
||||
return new CustomMultipartFile(name, data);
|
||||
}
|
||||
|
||||
// Custom implementation of MultipartFile
|
||||
private static class CustomMultipartFile implements MultipartFile {
|
||||
private final String name;
|
||||
private final byte[] content;
|
||||
|
||||
public CustomMultipartFile(String name, byte[] content) {
|
||||
this.name = name;
|
||||
this.content = content;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String getName() {
|
||||
return name;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String getOriginalFilename() {
|
||||
return name;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String getContentType() {
|
||||
return "application/pdf";
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean isEmpty() {
|
||||
return content == null || content.length == 0;
|
||||
}
|
||||
|
||||
@Override
|
||||
public long getSize() {
|
||||
return content.length;
|
||||
}
|
||||
|
||||
@Override
|
||||
public byte[] getBytes() throws IOException {
|
||||
return content;
|
||||
}
|
||||
|
||||
@Override
|
||||
public java.io.InputStream getInputStream() throws IOException {
|
||||
return new ByteArrayInputStream(content);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void transferTo(java.io.File dest) throws IOException, IllegalStateException {
|
||||
Files.write(dest.toPath(), content);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,152 @@
|
||||
package stirling.software.common.service;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.nio.file.Files;
|
||||
import java.nio.file.Path;
|
||||
import java.util.UUID;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
|
||||
import lombok.RequiredArgsConstructor;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
/**
|
||||
* Service for storing and retrieving files with unique file IDs.
|
||||
* Used by the AutoJobPostMapping system to handle file references.
|
||||
*/
|
||||
@Service
|
||||
@RequiredArgsConstructor
|
||||
@Slf4j
|
||||
public class FileStorage {
|
||||
|
||||
@Value("${stirling.tempDir:/tmp/stirling-files}")
|
||||
private String tempDirPath;
|
||||
|
||||
private final FileOrUploadService fileOrUploadService;
|
||||
|
||||
/**
|
||||
* Store a file and return its unique ID
|
||||
*
|
||||
* @param file The file to store
|
||||
* @return The unique ID assigned to the file
|
||||
* @throws IOException If there is an error storing the file
|
||||
*/
|
||||
public String storeFile(MultipartFile file) throws IOException {
|
||||
String fileId = generateFileId();
|
||||
Path filePath = getFilePath(fileId);
|
||||
|
||||
// Ensure the directory exists
|
||||
Files.createDirectories(filePath.getParent());
|
||||
|
||||
// Transfer the file to the storage location
|
||||
file.transferTo(filePath.toFile());
|
||||
|
||||
log.debug("Stored file with ID: {}", fileId);
|
||||
return fileId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Store a byte array as a file and return its unique ID
|
||||
*
|
||||
* @param bytes The byte array to store
|
||||
* @param originalName The original name of the file (for extension)
|
||||
* @return The unique ID assigned to the file
|
||||
* @throws IOException If there is an error storing the file
|
||||
*/
|
||||
public String storeBytes(byte[] bytes, String originalName) throws IOException {
|
||||
String fileId = generateFileId();
|
||||
Path filePath = getFilePath(fileId);
|
||||
|
||||
// Ensure the directory exists
|
||||
Files.createDirectories(filePath.getParent());
|
||||
|
||||
// Write the bytes to the file
|
||||
Files.write(filePath, bytes);
|
||||
|
||||
log.debug("Stored byte array with ID: {}", fileId);
|
||||
return fileId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieve a file by its ID as a MultipartFile
|
||||
*
|
||||
* @param fileId The ID of the file to retrieve
|
||||
* @return The file as a MultipartFile
|
||||
* @throws IOException If the file doesn't exist or can't be read
|
||||
*/
|
||||
public MultipartFile retrieveFile(String fileId) throws IOException {
|
||||
Path filePath = getFilePath(fileId);
|
||||
|
||||
if (!Files.exists(filePath)) {
|
||||
throw new IOException("File not found with ID: " + fileId);
|
||||
}
|
||||
|
||||
byte[] fileData = Files.readAllBytes(filePath);
|
||||
return fileOrUploadService.toMockMultipartFile(fileId, fileData);
|
||||
}
|
||||
|
||||
/**
|
||||
* Retrieve a file by its ID as a byte array
|
||||
*
|
||||
* @param fileId The ID of the file to retrieve
|
||||
* @return The file as a byte array
|
||||
* @throws IOException If the file doesn't exist or can't be read
|
||||
*/
|
||||
public byte[] retrieveBytes(String fileId) throws IOException {
|
||||
Path filePath = getFilePath(fileId);
|
||||
|
||||
if (!Files.exists(filePath)) {
|
||||
throw new IOException("File not found with ID: " + fileId);
|
||||
}
|
||||
|
||||
return Files.readAllBytes(filePath);
|
||||
}
|
||||
|
||||
/**
|
||||
* Delete a file by its ID
|
||||
*
|
||||
* @param fileId The ID of the file to delete
|
||||
* @return true if the file was deleted, false otherwise
|
||||
*/
|
||||
public boolean deleteFile(String fileId) {
|
||||
try {
|
||||
Path filePath = getFilePath(fileId);
|
||||
return Files.deleteIfExists(filePath);
|
||||
} catch (IOException e) {
|
||||
log.error("Error deleting file with ID: {}", fileId, e);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a file exists by its ID
|
||||
*
|
||||
* @param fileId The ID of the file to check
|
||||
* @return true if the file exists, false otherwise
|
||||
*/
|
||||
public boolean fileExists(String fileId) {
|
||||
Path filePath = getFilePath(fileId);
|
||||
return Files.exists(filePath);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the path for a file ID
|
||||
*
|
||||
* @param fileId The ID of the file
|
||||
* @return The path to the file
|
||||
*/
|
||||
private Path getFilePath(String fileId) {
|
||||
return Path.of(tempDirPath).resolve(fileId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Generate a unique file ID
|
||||
*
|
||||
* @return A unique file ID
|
||||
*/
|
||||
private String generateFileId() {
|
||||
return UUID.randomUUID().toString();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,262 @@
|
||||
package stirling.software.common.service;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.function.Supplier;
|
||||
|
||||
import org.springframework.http.HttpHeaders;
|
||||
import org.springframework.http.MediaType;
|
||||
import org.springframework.http.ResponseEntity;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.web.multipart.MultipartFile;
|
||||
|
||||
import lombok.RequiredArgsConstructor;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.common.controller.WebSocketProgressController;
|
||||
import stirling.software.common.model.job.JobProgress;
|
||||
import stirling.software.common.model.job.JobResponse;
|
||||
|
||||
/**
|
||||
* Service for executing jobs asynchronously or synchronously
|
||||
*/
|
||||
@Service
|
||||
@RequiredArgsConstructor
|
||||
@Slf4j
|
||||
public class JobExecutorService {
|
||||
|
||||
private final TaskManager taskManager;
|
||||
private final WebSocketProgressController webSocketSender;
|
||||
private final FileStorage fileStorage;
|
||||
private final ExecutorService executor = Executors.newCachedThreadPool();
|
||||
|
||||
/**
|
||||
* Run a job either asynchronously or synchronously
|
||||
*
|
||||
* @param async Whether to run the job asynchronously
|
||||
* @param work The work to be done
|
||||
* @return The response
|
||||
*/
|
||||
public ResponseEntity<?> runJobGeneric(boolean async, Supplier<Object> work) {
|
||||
String jobId = UUID.randomUUID().toString();
|
||||
log.debug("Running job with ID: {}, async: {}", jobId, async);
|
||||
|
||||
if (async) {
|
||||
taskManager.createTask(jobId);
|
||||
webSocketSender.sendProgress(jobId, new JobProgress(jobId, "Started", 0, "Running"));
|
||||
|
||||
executor.execute(
|
||||
() -> {
|
||||
try {
|
||||
Object result = work.get();
|
||||
processJobResult(jobId, result);
|
||||
webSocketSender.sendProgress(
|
||||
jobId, new JobProgress(jobId, "Done", 100, "Complete"));
|
||||
} catch (Exception e) {
|
||||
log.error("Error executing job {}: {}", jobId, e.getMessage(), e);
|
||||
taskManager.setError(jobId, e.getMessage());
|
||||
webSocketSender.sendProgress(
|
||||
jobId, new JobProgress(jobId, "Error", 100, e.getMessage()));
|
||||
}
|
||||
});
|
||||
|
||||
return ResponseEntity.ok().body(new JobResponse<>(true, jobId, null));
|
||||
} else {
|
||||
try {
|
||||
Object result = work.get();
|
||||
|
||||
// If the result is already a ResponseEntity, return it directly
|
||||
if (result instanceof ResponseEntity) {
|
||||
return (ResponseEntity<?>) result;
|
||||
}
|
||||
|
||||
// Process different result types
|
||||
return handleResultForSyncJob(result);
|
||||
} catch (Exception e) {
|
||||
log.error("Error executing synchronous job: {}", e.getMessage(), e);
|
||||
return ResponseEntity.internalServerError().body("Job failed: " + e.getMessage());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Process the result of an asynchronous job
|
||||
*
|
||||
* @param jobId The job ID
|
||||
* @param result The result
|
||||
*/
|
||||
private void processJobResult(String jobId, Object result) {
|
||||
try {
|
||||
if (result instanceof byte[]) {
|
||||
// Store byte array as a file
|
||||
String fileId = fileStorage.storeBytes((byte[]) result, "result.pdf");
|
||||
taskManager.setFileResult(jobId, fileId, "result.pdf", "application/pdf");
|
||||
log.debug("Stored byte[] result with fileId: {}", fileId);
|
||||
} else if (result instanceof ResponseEntity) {
|
||||
ResponseEntity<?> response = (ResponseEntity<?>) result;
|
||||
Object body = response.getBody();
|
||||
|
||||
if (body instanceof byte[]) {
|
||||
// Extract filename from content-disposition header if available
|
||||
String filename = "result.pdf";
|
||||
String contentType = "application/pdf";
|
||||
|
||||
if (response.getHeaders().getContentDisposition() != null) {
|
||||
String disposition = response.getHeaders().getContentDisposition().toString();
|
||||
if (disposition.contains("filename=")) {
|
||||
filename = disposition.substring(
|
||||
disposition.indexOf("filename=") + 9,
|
||||
disposition.lastIndexOf("\""));
|
||||
}
|
||||
}
|
||||
|
||||
if (response.getHeaders().getContentType() != null) {
|
||||
contentType = response.getHeaders().getContentType().toString();
|
||||
}
|
||||
|
||||
String fileId = fileStorage.storeBytes((byte[]) body, filename);
|
||||
taskManager.setFileResult(jobId, fileId, filename, contentType);
|
||||
log.debug("Stored ResponseEntity<byte[]> result with fileId: {}", fileId);
|
||||
} else {
|
||||
// Check if the response body contains a fileId
|
||||
if (body != null && body.toString().contains("fileId")) {
|
||||
try {
|
||||
// Try to extract fileId using reflection
|
||||
java.lang.reflect.Method getFileId = body.getClass().getMethod("getFileId");
|
||||
String fileId = (String) getFileId.invoke(body);
|
||||
|
||||
if (fileId != null && !fileId.isEmpty()) {
|
||||
// Try to get filename and content type
|
||||
String filename = "result.pdf";
|
||||
String contentType = "application/pdf";
|
||||
|
||||
try {
|
||||
java.lang.reflect.Method getOriginalFileName = body.getClass().getMethod("getOriginalFilename");
|
||||
String origName = (String) getOriginalFileName.invoke(body);
|
||||
if (origName != null && !origName.isEmpty()) {
|
||||
filename = origName;
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.debug("Could not get original filename: {}", e.getMessage());
|
||||
}
|
||||
|
||||
try {
|
||||
java.lang.reflect.Method getContentType = body.getClass().getMethod("getContentType");
|
||||
String ct = (String) getContentType.invoke(body);
|
||||
if (ct != null && !ct.isEmpty()) {
|
||||
contentType = ct;
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.debug("Could not get content type: {}", e.getMessage());
|
||||
}
|
||||
|
||||
taskManager.setFileResult(jobId, fileId, filename, contentType);
|
||||
log.debug("Extracted fileId from response body: {}", fileId);
|
||||
|
||||
taskManager.setComplete(jobId);
|
||||
return;
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.debug("Failed to extract fileId from response body: {}", e.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
// Store generic result
|
||||
taskManager.setResult(jobId, body);
|
||||
}
|
||||
} else if (result instanceof MultipartFile) {
|
||||
MultipartFile file = (MultipartFile) result;
|
||||
String fileId = fileStorage.storeFile(file);
|
||||
taskManager.setFileResult(
|
||||
jobId,
|
||||
fileId,
|
||||
file.getOriginalFilename(),
|
||||
file.getContentType());
|
||||
log.debug("Stored MultipartFile result with fileId: {}", fileId);
|
||||
} else {
|
||||
// Check if result has a fileId field
|
||||
if (result != null) {
|
||||
try {
|
||||
// Try to extract fileId using reflection
|
||||
java.lang.reflect.Method getFileId = result.getClass().getMethod("getFileId");
|
||||
String fileId = (String) getFileId.invoke(result);
|
||||
|
||||
if (fileId != null && !fileId.isEmpty()) {
|
||||
// Try to get filename and content type
|
||||
String filename = "result.pdf";
|
||||
String contentType = "application/pdf";
|
||||
|
||||
try {
|
||||
java.lang.reflect.Method getOriginalFileName = result.getClass().getMethod("getOriginalFilename");
|
||||
String origName = (String) getOriginalFileName.invoke(result);
|
||||
if (origName != null && !origName.isEmpty()) {
|
||||
filename = origName;
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.debug("Could not get original filename: {}", e.getMessage());
|
||||
}
|
||||
|
||||
try {
|
||||
java.lang.reflect.Method getContentType = result.getClass().getMethod("getContentType");
|
||||
String ct = (String) getContentType.invoke(result);
|
||||
if (ct != null && !ct.isEmpty()) {
|
||||
contentType = ct;
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.debug("Could not get content type: {}", e.getMessage());
|
||||
}
|
||||
|
||||
taskManager.setFileResult(jobId, fileId, filename, contentType);
|
||||
log.debug("Extracted fileId from result object: {}", fileId);
|
||||
|
||||
taskManager.setComplete(jobId);
|
||||
return;
|
||||
}
|
||||
} catch (Exception e) {
|
||||
log.debug("Failed to extract fileId from result object: {}", e.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
// Default case: store the result as is
|
||||
taskManager.setResult(jobId, result);
|
||||
}
|
||||
|
||||
taskManager.setComplete(jobId);
|
||||
} catch (Exception e) {
|
||||
log.error("Error processing job result: {}", e.getMessage(), e);
|
||||
taskManager.setError(jobId, "Error processing result: " + e.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Handle different result types for synchronous jobs
|
||||
*
|
||||
* @param result The result object
|
||||
* @return The appropriate ResponseEntity
|
||||
* @throws IOException If there is an error processing the result
|
||||
*/
|
||||
private ResponseEntity<?> handleResultForSyncJob(Object result) throws IOException {
|
||||
if (result instanceof byte[]) {
|
||||
// Return byte array as PDF
|
||||
return ResponseEntity.ok()
|
||||
.contentType(MediaType.APPLICATION_PDF)
|
||||
.header(HttpHeaders.CONTENT_DISPOSITION,
|
||||
"form-data; name=\"attachment\"; filename=\"result.pdf\"")
|
||||
.body(result);
|
||||
} else if (result instanceof MultipartFile) {
|
||||
// Return MultipartFile content
|
||||
MultipartFile file = (MultipartFile) result;
|
||||
return ResponseEntity.ok()
|
||||
.contentType(MediaType.parseMediaType(file.getContentType()))
|
||||
.header(HttpHeaders.CONTENT_DISPOSITION,
|
||||
"form-data; name=\"attachment\"; filename=\"" +
|
||||
file.getOriginalFilename() + "\"")
|
||||
.body(file.getBytes());
|
||||
} else {
|
||||
// Default case: return as JSON
|
||||
return ResponseEntity.ok(result);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,186 @@
|
||||
package stirling.software.common.service;
|
||||
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
|
||||
import org.springframework.beans.factory.annotation.Value;
|
||||
import org.springframework.http.ResponseEntity;
|
||||
import org.springframework.stereotype.Service;
|
||||
import org.springframework.web.bind.annotation.GetMapping;
|
||||
import org.springframework.web.bind.annotation.PathVariable;
|
||||
import org.springframework.web.bind.annotation.RestController;
|
||||
|
||||
import lombok.RequiredArgsConstructor;
|
||||
import lombok.extern.slf4j.Slf4j;
|
||||
|
||||
import stirling.software.common.model.job.JobResult;
|
||||
|
||||
/**
|
||||
* Manages async tasks and their results
|
||||
*/
|
||||
@Service
|
||||
@RequiredArgsConstructor
|
||||
@Slf4j
|
||||
public class TaskManager {
|
||||
private final Map<String, JobResult> jobResults = new ConcurrentHashMap<>();
|
||||
|
||||
@Value("${stirling.jobResultExpiryMinutes:30}")
|
||||
private int jobResultExpiryMinutes = 30;
|
||||
|
||||
private final FileStorage fileStorage;
|
||||
|
||||
/**
|
||||
* Create a new task with the given job ID
|
||||
*
|
||||
* @param jobId The job ID
|
||||
*/
|
||||
public void createTask(String jobId) {
|
||||
jobResults.put(jobId, JobResult.createNew(jobId));
|
||||
log.debug("Created task with job ID: {}", jobId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the result of a task as a general object
|
||||
*
|
||||
* @param jobId The job ID
|
||||
* @param result The result object
|
||||
*/
|
||||
public void setResult(String jobId, Object result) {
|
||||
JobResult jobResult = getOrCreateJobResult(jobId);
|
||||
jobResult.completeWithResult(result);
|
||||
log.debug("Set result for job ID: {}", jobId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the result of a task as a file
|
||||
*
|
||||
* @param jobId The job ID
|
||||
* @param fileId The file ID
|
||||
* @param originalFileName The original file name
|
||||
* @param contentType The content type of the file
|
||||
*/
|
||||
public void setFileResult(String jobId, String fileId, String originalFileName, String contentType) {
|
||||
JobResult jobResult = getOrCreateJobResult(jobId);
|
||||
jobResult.completeWithFile(fileId, originalFileName, contentType);
|
||||
log.debug("Set file result for job ID: {} with file ID: {}", jobId, fileId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Set an error for a task
|
||||
*
|
||||
* @param jobId The job ID
|
||||
* @param error The error message
|
||||
*/
|
||||
public void setError(String jobId, String error) {
|
||||
JobResult jobResult = getOrCreateJobResult(jobId);
|
||||
jobResult.failWithError(error);
|
||||
log.debug("Set error for job ID: {}: {}", jobId, error);
|
||||
}
|
||||
|
||||
/**
|
||||
* Mark a task as complete
|
||||
*
|
||||
* @param jobId The job ID
|
||||
*/
|
||||
public void setComplete(String jobId) {
|
||||
JobResult jobResult = getOrCreateJobResult(jobId);
|
||||
if (jobResult.getResult() == null && jobResult.getFileId() == null && jobResult.getError() == null) {
|
||||
// If no result or error has been set, mark it as complete with an empty result
|
||||
jobResult.completeWithResult("Task completed successfully");
|
||||
}
|
||||
log.debug("Marked job ID: {} as complete", jobId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Check if a task is complete
|
||||
*
|
||||
* @param jobId The job ID
|
||||
* @return true if the task is complete, false otherwise
|
||||
*/
|
||||
public boolean isComplete(String jobId) {
|
||||
JobResult result = jobResults.get(jobId);
|
||||
return result != null && result.isComplete();
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the result of a task
|
||||
*
|
||||
* @param jobId The job ID
|
||||
* @return The result object, or null if the task doesn't exist or is not complete
|
||||
*/
|
||||
public JobResult getJobResult(String jobId) {
|
||||
return jobResults.get(jobId);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get or create a job result
|
||||
*
|
||||
* @param jobId The job ID
|
||||
* @return The job result
|
||||
*/
|
||||
private JobResult getOrCreateJobResult(String jobId) {
|
||||
return jobResults.computeIfAbsent(jobId, JobResult::createNew);
|
||||
}
|
||||
|
||||
/**
|
||||
* REST controller for job-related endpoints
|
||||
*/
|
||||
@RestController
|
||||
public class JobController {
|
||||
|
||||
/**
|
||||
* Get the status of a job
|
||||
*
|
||||
* @param jobId The job ID
|
||||
* @return The job result
|
||||
*/
|
||||
@GetMapping("/api/v1/general/job/{jobId}")
|
||||
public ResponseEntity<?> getJobStatus(@PathVariable("jobId") String jobId) {
|
||||
JobResult result = jobResults.get(jobId);
|
||||
if (result == null) {
|
||||
return ResponseEntity.notFound().build();
|
||||
}
|
||||
return ResponseEntity.ok(result);
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the result of a job
|
||||
*
|
||||
* @param jobId The job ID
|
||||
* @return The job result
|
||||
*/
|
||||
@GetMapping("/api/v1/general/job/{jobId}/result")
|
||||
public ResponseEntity<?> getJobResult(@PathVariable("jobId") String jobId) {
|
||||
JobResult result = jobResults.get(jobId);
|
||||
if (result == null) {
|
||||
return ResponseEntity.notFound().build();
|
||||
}
|
||||
|
||||
if (!result.isComplete()) {
|
||||
return ResponseEntity.badRequest().body("Job is not complete yet");
|
||||
}
|
||||
|
||||
if (result.getError() != null) {
|
||||
return ResponseEntity.badRequest().body("Job failed: " + result.getError());
|
||||
}
|
||||
|
||||
if (result.getFileId() != null) {
|
||||
try {
|
||||
byte[] fileContent = fileStorage.retrieveBytes(result.getFileId());
|
||||
return ResponseEntity.ok()
|
||||
.header("Content-Type", result.getContentType())
|
||||
.header("Content-Disposition",
|
||||
"form-data; name=\"attachment\"; filename=\"" +
|
||||
result.getOriginalFileName() + "\"")
|
||||
.body(fileContent);
|
||||
} catch (Exception e) {
|
||||
log.error("Error retrieving file for job {}: {}", jobId, e.getMessage(), e);
|
||||
return ResponseEntity.internalServerError()
|
||||
.body("Error retrieving file: " + e.getMessage());
|
||||
}
|
||||
}
|
||||
|
||||
return ResponseEntity.ok(result.getResult());
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user