feat: add file export
This commit is contained in:
@@ -34,3 +34,8 @@ type TaskListResponse struct {
|
||||
TaskList []*TaskListDTO `json:"task_list"`
|
||||
Total int64 `json:"total"`
|
||||
}
|
||||
|
||||
type ExportTaskRequest struct {
|
||||
TaskID int64 `json:"task_id" binding:"required"`
|
||||
Type string `json:"type" binding:"required,oneof=pdf docx"`
|
||||
}
|
||||
|
||||
@@ -1,8 +1,13 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"gitea.com/bitwsd/document_ai/internal/model/task"
|
||||
@@ -111,3 +116,84 @@ func (svc *TaskService) GetTaskList(ctx context.Context, req *task.TaskListReque
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (svc *TaskService) ExportTask(ctx context.Context, req *task.ExportTaskRequest) ([]byte, string, error) {
|
||||
recognitionTask, err := svc.recognitionTaskDao.GetTaskByID(dao.DB.WithContext(ctx), req.TaskID)
|
||||
if err != nil {
|
||||
log.Error(ctx, "func", "ExportTask", "msg", "get task by task id failed", "error", err)
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
if recognitionTask == nil {
|
||||
log.Error(ctx, "func", "ExportTask", "msg", "task not found")
|
||||
return nil, "", errors.New("task not found")
|
||||
}
|
||||
|
||||
if recognitionTask.Status != dao.TaskStatusCompleted {
|
||||
log.Error(ctx, "func", "ExportTask", "msg", "task not finished")
|
||||
return nil, "", errors.New("task not finished")
|
||||
}
|
||||
|
||||
recognitionResult, err := svc.recognitionResultDao.GetByTaskID(dao.DB.WithContext(ctx), req.TaskID)
|
||||
if err != nil {
|
||||
log.Error(ctx, "func", "ExportTask", "msg", "get recognition result by task id failed", "error", err)
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
if recognitionResult == nil {
|
||||
log.Error(ctx, "func", "ExportTask", "msg", "recognition result not found")
|
||||
return nil, "", errors.New("recognition result not found")
|
||||
}
|
||||
|
||||
markdown := recognitionResult.Markdown
|
||||
if markdown == "" {
|
||||
log.Error(ctx, "func", "ExportTask", "msg", "markdown not found")
|
||||
return nil, "", errors.New("markdown not found")
|
||||
}
|
||||
|
||||
// call http://localhost:8055/export
|
||||
body := &bytes.Buffer{}
|
||||
writer := multipart.NewWriter(body)
|
||||
_ = writer.WriteField("markdown", markdown)
|
||||
_ = writer.WriteField("type", req.Type)
|
||||
writer.Close()
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, "http://localhost:8055/export", body)
|
||||
if err != nil {
|
||||
log.Error(ctx, "func", "ExportTask", "msg", "create http request failed", "error", err)
|
||||
return nil, "", err
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
|
||||
client := &http.Client{}
|
||||
resp, err := client.Do(httpReq)
|
||||
if err != nil {
|
||||
log.Error(ctx, "func", "ExportTask", "msg", "http request failed", "error", err)
|
||||
return nil, "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Error(ctx, "func", "ExportTask", "msg", "http request failed", "status", resp.StatusCode)
|
||||
return nil, "", fmt.Errorf("export service returned status: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
fileData, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
log.Error(ctx, "func", "ExportTask", "msg", "read response body failed", "error", err)
|
||||
return nil, "", err
|
||||
}
|
||||
|
||||
// determine content type based on export type
|
||||
var contentType string
|
||||
switch req.Type {
|
||||
case "pdf":
|
||||
contentType = "application/pdf"
|
||||
case "docx":
|
||||
contentType = "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
||||
default:
|
||||
contentType = "application/octet-stream"
|
||||
}
|
||||
|
||||
return fileData, contentType, nil
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user