sse + ocr replanish

backend/internal/service/ocr.go

@@ -1,15 +1,21 @@
 package service
 
 import (
+	"context"
+	"database/sql"
+	"fmt"
 	"log"
 	"os"
 	"strings"
 
+	"github.com/google/uuid"
 	"github.com/vaultdrop/backend/internal/config"
+	"github.com/vaultdrop/backend/internal/db"
 	"github.com/vaultdrop/backend/internal/ocr"
 )
 
 type OCRJob struct {
+	DBID       uuid.UUID
 	ResourceID string
 	FilePath   string
 }
@@ -17,20 +23,27 @@ type OCRJob struct {
 type OCRService struct {
 	client      *ocr.Client
 	resourceSvc *ResourceService
+	broker      *EventBroker
+	queries     *db.Queries
 	jobs        chan OCRJob
 }
 
-func NewOCRService(cfg *config.Config, resourceSvc *ResourceService) *OCRService {
+func NewOCRService(database *sql.DB, queries *db.Queries, cfg *config.Config, resourceSvc *ResourceService, broker *EventBroker) *OCRService {
 	return &OCRService{
 		client:      ocr.NewClient(cfg.OCREndpoint),
 		resourceSvc: resourceSvc,
+		broker:      broker,
+		queries:     queries,
 		jobs:        make(chan OCRJob, 100),
 	}
 }
 
-func (s *OCRService) Start() {
-	go s.worker()
-	log.Println("[OCR] Worker started")
+func (s *OCRService) Start(workerCount int) {
+	s.replenish()
+	for i := range workerCount {
+		go s.worker()
+		log.Printf("[OCR] Worker %d started", i)
+	}
 }
 
 func (s *OCRService) Stop() {
@@ -38,9 +51,51 @@ func (s *OCRService) Stop() {
 	log.Println("[OCR] Worker stopped")
 }
 
-func (s *OCRService) Enqueue(resourceID, filePath string) {
-	s.jobs <- OCRJob{ResourceID: resourceID, FilePath: filePath}
-	log.Printf("[OCR] Enqueued resource %s", resourceID)
+func (s *OCRService) Enqueue(resourceID, filePath string) error {
+	ctx := context.Background()
+
+	resourceUUID, err := uuid.Parse(resourceID)
+	if err != nil {
+		return fmt.Errorf("parse resource id: %w", err)
+	}
+
+	dbJob, err := s.queries.CreateOCRJob(ctx, db.CreateOCRJobParams{
+		ResourceID: resourceUUID,
+		FilePath:   filePath,
+	})
+	if err != nil {
+		return fmt.Errorf("create ocr job: %w", err)
+	}
+
+	job := OCRJob{DBID: dbJob.ID, ResourceID: resourceID, FilePath: filePath}
+
+	select {
+	case s.jobs <- job:
+		log.Printf("[OCR] Enqueued resource %s (job %s)", resourceID, dbJob.ID)
+		return nil
+	default:
+		log.Printf("[OCR] Queue full, resource %s persisted as pending (job %s)", resourceID, dbJob.ID)
+		return nil
+	}
+}
+
+func (s *OCRService) replenish() {
+	ctx := context.Background()
+	pending, err := s.queries.ListPendingOCRJobs(ctx)
+	if err != nil {
+		log.Printf("[OCR] Failed to load pending jobs: %v", err)
+		return
+	}
+	for _, j := range pending {
+		job := OCRJob{DBID: j.ID, ResourceID: j.ResourceID.String(), FilePath: j.FilePath}
+		select {
+		case s.jobs <- job:
+			log.Printf("[OCR] Replenished job %s (resource %s)", j.ID, j.ResourceID)
+		default:
+			log.Printf("[OCR] Queue full, leaving job %s in pending", j.ID)
+			return
+		}
+	}
 }
 
 func (s *OCRService) worker() {
@@ -50,17 +105,33 @@ func (s *OCRService) worker() {
 }
 
 func (s *OCRService) process(job OCRJob) {
+	ctx := context.Background()
 	log.Printf("[OCR] Processing resource %s", job.ResourceID)
 
+	s.queries.UpdateOCRJobStatus(ctx, db.UpdateOCRJobStatusParams{
+		ID:     job.DBID,
+		Status: "processing",
+	})
+
 	data, err := os.ReadFile(job.FilePath)
 	if err != nil {
 		log.Printf("[OCR] Failed to read resource %s: %v", job.ResourceID, err)
+		s.queries.UpdateOCRJobStatus(ctx, db.UpdateOCRJobStatusParams{
+			ID:           job.DBID,
+			Status:       "failed",
+			ErrorMessage: err.Error(),
+		})
 		return
 	}
 
 	blocks, err := s.client.Recognize(data)
 	if err != nil {
 		log.Printf("[OCR] Failed to recognize resource %s: %v", job.ResourceID, err)
+		s.queries.UpdateOCRJobStatus(ctx, db.UpdateOCRJobStatusParams{
+			ID:           job.DBID,
+			Status:       "failed",
+			ErrorMessage: err.Error(),
+		})
 		return
 	}
 
@@ -68,9 +139,19 @@ func (s *OCRService) process(job OCRJob) {
 
 	if err := s.resourceSvc.UpdateOCRText(job.ResourceID, text); err != nil {
 		log.Printf("[OCR] Failed to update ocr_text for resource %s: %v", job.ResourceID, err)
+		s.queries.UpdateOCRJobStatus(ctx, db.UpdateOCRJobStatusParams{
+			ID:           job.DBID,
+			Status:       "failed",
+			ErrorMessage: err.Error(),
+		})
 		return
 	}
 
+	s.queries.UpdateOCRJobStatus(ctx, db.UpdateOCRJobStatusParams{
+		ID:     job.DBID,
+		Status: "done",
+	})
+	s.broker.Publish("ocr_done", job.ResourceID)
 	log.Printf("[OCR] Completed resource %s (%d chars)", job.ResourceID, len(text))
 }