294 lines
8.9 KiB
Go
294 lines
8.9 KiB
Go
package extraction
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"path/filepath"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// ── Status Constants ───────────────────────
|
|
|
|
const (
|
|
StatusPending = "pending"
|
|
StatusProcessing = "processing"
|
|
StatusComplete = "complete"
|
|
StatusFailed = "failed"
|
|
StatusSkipped = "skipped" // not extractable (images, etc.)
|
|
)
|
|
|
|
// ── Queue Item ─────────────────────────────
|
|
|
|
// QueueItem is the status.json written to the processing directory.
|
|
// The sidecar extractor watches this directory for pending items.
|
|
type QueueItem struct {
|
|
AttachmentID string `json:"attachment_id"`
|
|
StorageKey string `json:"storage_key"`
|
|
ContentType string `json:"content_type"`
|
|
Filename string `json:"filename"`
|
|
Status string `json:"status"`
|
|
Error string `json:"error,omitempty"`
|
|
OutputKey string `json:"output_key,omitempty"` // path to extracted text
|
|
QueuedAt string `json:"queued_at"`
|
|
StartedAt string `json:"started_at,omitempty"`
|
|
CompletedAt string `json:"completed_at,omitempty"`
|
|
}
|
|
|
|
// ── Queue Manager ──────────────────────────
|
|
|
|
// Queue manages the filesystem-based extraction queue.
|
|
// Processing state lives in {storagePath}/processing/{attachment_id}/status.json.
|
|
// The sidecar extractor watches for pending items and updates status.
|
|
type Queue struct {
|
|
storagePath string
|
|
concurrency int
|
|
sem chan struct{} // semaphore for concurrent extraction limit
|
|
mu sync.Mutex
|
|
}
|
|
|
|
// NewQueue creates a new extraction queue manager.
|
|
// storagePath is the PVC mount (e.g., /data/storage).
|
|
// concurrency limits parallel extractions (default 3).
|
|
func NewQueue(storagePath string, concurrency int) (*Queue, error) {
|
|
if concurrency <= 0 {
|
|
concurrency = 3
|
|
}
|
|
|
|
procDir := filepath.Join(storagePath, "processing")
|
|
if err := os.MkdirAll(procDir, 0750); err != nil {
|
|
return nil, fmt.Errorf("extraction: create processing dir: %w", err)
|
|
}
|
|
|
|
return &Queue{
|
|
storagePath: storagePath,
|
|
concurrency: concurrency,
|
|
sem: make(chan struct{}, concurrency),
|
|
}, nil
|
|
}
|
|
|
|
// Enqueue adds an attachment to the extraction queue.
|
|
// Creates {storagePath}/processing/{id}/status.json with status "pending".
|
|
func (q *Queue) Enqueue(attachmentID, storageKey, contentType, filename string) error {
|
|
q.mu.Lock()
|
|
defer q.mu.Unlock()
|
|
|
|
itemDir := filepath.Join(q.storagePath, "processing", attachmentID)
|
|
if err := os.MkdirAll(itemDir, 0750); err != nil {
|
|
return fmt.Errorf("extraction: create item dir: %w", err)
|
|
}
|
|
|
|
item := QueueItem{
|
|
AttachmentID: attachmentID,
|
|
StorageKey: storageKey,
|
|
ContentType: contentType,
|
|
Filename: filename,
|
|
Status: StatusPending,
|
|
QueuedAt: time.Now().UTC().Format(time.RFC3339),
|
|
}
|
|
|
|
return q.writeStatus(attachmentID, &item)
|
|
}
|
|
|
|
// GetStatus reads the current extraction status for an attachment.
|
|
func (q *Queue) GetStatus(attachmentID string) (*QueueItem, error) {
|
|
statusPath := filepath.Join(q.storagePath, "processing", attachmentID, "status.json")
|
|
data, err := os.ReadFile(statusPath)
|
|
if err != nil {
|
|
if os.IsNotExist(err) {
|
|
return nil, nil // not queued
|
|
}
|
|
return nil, fmt.Errorf("extraction: read status: %w", err)
|
|
}
|
|
|
|
var item QueueItem
|
|
if err := json.Unmarshal(data, &item); err != nil {
|
|
return nil, fmt.Errorf("extraction: parse status: %w", err)
|
|
}
|
|
return &item, nil
|
|
}
|
|
|
|
// MarkComplete updates status to complete and records the output key.
|
|
func (q *Queue) MarkComplete(attachmentID, outputKey string) error {
|
|
q.mu.Lock()
|
|
defer q.mu.Unlock()
|
|
|
|
item, err := q.GetStatus(attachmentID)
|
|
if err != nil || item == nil {
|
|
return fmt.Errorf("extraction: item not found: %s", attachmentID)
|
|
}
|
|
|
|
item.Status = StatusComplete
|
|
item.OutputKey = outputKey
|
|
item.CompletedAt = time.Now().UTC().Format(time.RFC3339)
|
|
return q.writeStatus(attachmentID, item)
|
|
}
|
|
|
|
// MarkFailed updates status to failed with an error message.
|
|
func (q *Queue) MarkFailed(attachmentID, errMsg string) error {
|
|
q.mu.Lock()
|
|
defer q.mu.Unlock()
|
|
|
|
item, err := q.GetStatus(attachmentID)
|
|
if err != nil || item == nil {
|
|
return fmt.Errorf("extraction: item not found: %s", attachmentID)
|
|
}
|
|
|
|
item.Status = StatusFailed
|
|
item.Error = errMsg
|
|
item.CompletedAt = time.Now().UTC().Format(time.RFC3339)
|
|
return q.writeStatus(attachmentID, item)
|
|
}
|
|
|
|
// Cleanup removes the processing directory for a completed/failed item.
|
|
func (q *Queue) Cleanup(attachmentID string) error {
|
|
itemDir := filepath.Join(q.storagePath, "processing", attachmentID)
|
|
return os.RemoveAll(itemDir)
|
|
}
|
|
|
|
// RecoverStale scans for items stuck in "processing" state (crash recovery).
|
|
// Items older than maxAge are reset to "pending" for re-processing.
|
|
func (q *Queue) RecoverStale(maxAge time.Duration) (int, error) {
|
|
procDir := filepath.Join(q.storagePath, "processing")
|
|
entries, err := os.ReadDir(procDir)
|
|
if err != nil {
|
|
if os.IsNotExist(err) {
|
|
return 0, nil
|
|
}
|
|
return 0, err
|
|
}
|
|
|
|
recovered := 0
|
|
cutoff := time.Now().Add(-maxAge)
|
|
|
|
for _, entry := range entries {
|
|
if !entry.IsDir() {
|
|
continue
|
|
}
|
|
item, err := q.GetStatus(entry.Name())
|
|
if err != nil || item == nil {
|
|
continue
|
|
}
|
|
|
|
if item.Status != StatusProcessing {
|
|
continue
|
|
}
|
|
|
|
// Check if started_at is older than maxAge
|
|
if item.StartedAt != "" {
|
|
startedAt, err := time.Parse(time.RFC3339, item.StartedAt)
|
|
if err == nil && startedAt.Before(cutoff) {
|
|
q.mu.Lock()
|
|
item.Status = StatusPending
|
|
item.StartedAt = ""
|
|
item.Error = "recovered from stale processing state"
|
|
q.writeStatus(entry.Name(), item)
|
|
q.mu.Unlock()
|
|
recovered++
|
|
log.Printf("extraction: recovered stale item %s", entry.Name())
|
|
}
|
|
}
|
|
}
|
|
|
|
return recovered, nil
|
|
}
|
|
|
|
// ListPending returns all items with status "pending".
|
|
func (q *Queue) ListPending() ([]QueueItem, error) {
|
|
return q.listByStatus(StatusPending)
|
|
}
|
|
|
|
// ListAll returns all items in the processing directory.
|
|
func (q *Queue) ListAll() ([]QueueItem, error) {
|
|
procDir := filepath.Join(q.storagePath, "processing")
|
|
entries, err := os.ReadDir(procDir)
|
|
if err != nil {
|
|
if os.IsNotExist(err) {
|
|
return nil, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
|
|
var items []QueueItem
|
|
for _, entry := range entries {
|
|
if !entry.IsDir() {
|
|
continue
|
|
}
|
|
item, err := q.GetStatus(entry.Name())
|
|
if err != nil || item == nil {
|
|
continue
|
|
}
|
|
items = append(items, *item)
|
|
}
|
|
return items, nil
|
|
}
|
|
|
|
// ── Internal Helpers ───────────────────────
|
|
|
|
func (q *Queue) writeStatus(attachmentID string, item *QueueItem) error {
|
|
statusPath := filepath.Join(q.storagePath, "processing", attachmentID, "status.json")
|
|
data, err := json.MarshalIndent(item, "", " ")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Atomic write: temp + rename
|
|
tmpPath := statusPath + ".tmp"
|
|
if err := os.WriteFile(tmpPath, data, 0640); err != nil {
|
|
return err
|
|
}
|
|
return os.Rename(tmpPath, statusPath)
|
|
}
|
|
|
|
func (q *Queue) listByStatus(status string) ([]QueueItem, error) {
|
|
all, err := q.ListAll()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var filtered []QueueItem
|
|
for _, item := range all {
|
|
if item.Status == status {
|
|
filtered = append(filtered, item)
|
|
}
|
|
}
|
|
return filtered, nil
|
|
}
|
|
|
|
// ── Extractable Check ──────────────────────
|
|
|
|
// IsExtractable returns true if the content type supports text extraction.
|
|
// Images and plain text do not need extraction (handled inline by upload handler).
|
|
var extractableTypes = map[string]bool{
|
|
"application/pdf": true,
|
|
"application/vnd.openxmlformats-officedocument.wordprocessingml.document": true,
|
|
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet": true,
|
|
"application/vnd.openxmlformats-officedocument.presentationml.presentation": true,
|
|
"application/msword": true,
|
|
"application/vnd.ms-excel": true,
|
|
"application/vnd.oasis.opendocument.text": true,
|
|
"application/vnd.oasis.opendocument.spreadsheet": true,
|
|
"application/vnd.oasis.opendocument.presentation": true,
|
|
"application/rtf": true,
|
|
}
|
|
|
|
// IsExtractable returns true if the given content type requires sidecar extraction.
|
|
func IsExtractable(contentType string) bool {
|
|
return extractableTypes[contentType]
|
|
}
|
|
|
|
// ── Null Queue ─────────────────────────────
|
|
|
|
// NullQueue is a no-op implementation for when extraction is disabled.
|
|
type NullQueue struct{}
|
|
|
|
func (NullQueue) Enqueue(_, _, _, _ string) error { return nil }
|
|
func (NullQueue) GetStatus(_ string) (*QueueItem, error) { return nil, nil }
|
|
func (NullQueue) MarkComplete(_, _ string) error { return nil }
|
|
func (NullQueue) MarkFailed(_, _ string) error { return nil }
|
|
func (NullQueue) Cleanup(_ string) error { return nil }
|
|
func (NullQueue) RecoverStale(_ time.Duration) (int, error) { return 0, nil }
|
|
func (NullQueue) ListPending() ([]QueueItem, error) { return nil, nil }
|
|
func (NullQueue) ListAll() ([]QueueItem, error) { return nil, nil } |