sse + ocr replanish

This commit is contained in:
m
2026-07-29 18:56:26 +02:00
parent 68a4beef37
commit 397803fb96
14 changed files with 386 additions and 26 deletions
+6 -4
View File
@@ -30,8 +30,10 @@ func main() {
queries := db.New(database) queries := db.New(database)
eventBroker := service.NewEventBroker()
resourceSvc := service.NewResourceService(database, queries, cfg) resourceSvc := service.NewResourceService(database, queries, cfg)
ocrSvc := service.NewOCRService(cfg, resourceSvc) ocrSvc := service.NewOCRService(database, queries, cfg, resourceSvc, eventBroker)
conversionSvc := service.NewConversionService(queries, cfg) conversionSvc := service.NewConversionService(queries, cfg)
urlSvc := service.NewURLService(cfg.HMACSecret, cfg.ServerHost, cfg.URLExpiryMinutes) urlSvc := service.NewURLService(cfg.HMACSecret, cfg.ServerHost, cfg.URLExpiryMinutes)
authSvc, err := auth.NewAuthService(database, queries, cfg) authSvc, err := auth.NewAuthService(database, queries, cfg)
@@ -43,13 +45,13 @@ func main() {
placementSvc := service.NewPlacementService(queries) placementSvc := service.NewPlacementService(queries)
syncSvc := service.NewSyncService(queries) syncSvc := service.NewSyncService(queries)
ocrSvc.Start() ocrSvc.Start(cfg.OCRWorkers)
defer ocrSvc.Stop() defer ocrSvc.Stop()
conversionSvc.Start() conversionSvc.Start(cfg.ConversionWorkers)
defer conversionSvc.Stop() defer conversionSvc.Stop()
h := handler.New(resourceSvc, ocrSvc, urlSvc, auth.NewAuthHandler(authSvc), conversionSvc, rebacSvc, placementSvc, syncSvc) h := handler.New(resourceSvc, ocrSvc, urlSvc, auth.NewAuthHandler(authSvc), conversionSvc, rebacSvc, placementSvc, syncSvc, eventBroker)
r := gin.Default() r := gin.Default()
handler.SetupRoutes(r, h, authSvc) handler.SetupRoutes(r, h, authSvc)
+11 -7
View File
@@ -11,12 +11,14 @@ type Config struct {
OCREndpoint string OCREndpoint string
UploadDir string UploadDir string
HMACSecret string HMACSecret string
ServerHost string ServerHost string
PASETOKey string PASETOKey string
LibreOfficePath string LibreOfficePath string
PdftoppmPath string PdftoppmPath string
ThumbnailDir string ThumbnailDir string
URLExpiryMinutes int URLExpiryMinutes int
OCRWorkers int
ConversionWorkers int
} }
func Load() *Config { func Load() *Config {
@@ -31,7 +33,9 @@ func Load() *Config {
LibreOfficePath: envOr("LIBREOFFICE_PATH", "/usr/bin/libreoffice"), LibreOfficePath: envOr("LIBREOFFICE_PATH", "/usr/bin/libreoffice"),
PdftoppmPath: envOr("PDFTOPPM_PATH", "/usr/bin/pdftoppm"), PdftoppmPath: envOr("PDFTOPPM_PATH", "/usr/bin/pdftoppm"),
ThumbnailDir: envOr("THUMBNAIL_DIR", "./uploads/thumbnails"), ThumbnailDir: envOr("THUMBNAIL_DIR", "./uploads/thumbnails"),
URLExpiryMinutes: envOrInt("URL_EXPIRY_MINUTES", 60), URLExpiryMinutes: envOrInt("URL_EXPIRY_MINUTES", 60),
OCRWorkers: envOrInt("OCR_WORKERS", 1),
ConversionWorkers: envOrInt("CONVERSION_WORKERS", 1),
} }
} }
@@ -0,0 +1 @@
DROP TABLE IF EXISTS ocr_jobs;
@@ -0,0 +1,12 @@
CREATE TABLE ocr_jobs (
id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
resource_id UUID NOT NULL REFERENCES resources(id) ON DELETE CASCADE,
file_path TEXT NOT NULL DEFAULT '',
status TEXT NOT NULL DEFAULT 'pending' CHECK (status IN ('pending', 'processing', 'done', 'failed')),
error_message TEXT NOT NULL DEFAULT '',
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
);
CREATE INDEX idx_ocr_jobs_status ON ocr_jobs(status);
CREATE INDEX idx_ocr_jobs_resource ON ocr_jobs(resource_id);
+10
View File
@@ -12,6 +12,16 @@ import (
"github.com/google/uuid" "github.com/google/uuid"
) )
type OcrJob struct {
ID uuid.UUID `json:"id"`
ResourceID uuid.UUID `json:"resource_id"`
FilePath string `json:"file_path"`
Status string `json:"status"`
ErrorMessage string `json:"error_message"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
}
type RebacRelation struct { type RebacRelation struct {
ID uuid.UUID `json:"id"` ID uuid.UUID `json:"id"`
ResourceID uuid.UUID `json:"resource_id"` ResourceID uuid.UUID `json:"resource_id"`
+122
View File
@@ -0,0 +1,122 @@
// Code generated by sqlc. DO NOT EDIT.
// versions:
// sqlc v1.31.1
// source: ocr_jobs.sql
package db
import (
"context"
"github.com/google/uuid"
)
const createOCRJob = `-- name: CreateOCRJob :one
INSERT INTO ocr_jobs (resource_id, file_path, status)
VALUES ($1, $2, 'pending')
RETURNING id, resource_id, file_path, status, error_message, created_at, updated_at
`
type CreateOCRJobParams struct {
ResourceID uuid.UUID `json:"resource_id"`
FilePath string `json:"file_path"`
}
func (q *Queries) CreateOCRJob(ctx context.Context, arg CreateOCRJobParams) (OcrJob, error) {
row := q.db.QueryRowContext(ctx, createOCRJob, arg.ResourceID, arg.FilePath)
var i OcrJob
err := row.Scan(
&i.ID,
&i.ResourceID,
&i.FilePath,
&i.Status,
&i.ErrorMessage,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const deleteOCRJob = `-- name: DeleteOCRJob :exec
DELETE FROM ocr_jobs
WHERE id = $1
`
func (q *Queries) DeleteOCRJob(ctx context.Context, id uuid.UUID) error {
_, err := q.db.ExecContext(ctx, deleteOCRJob, id)
return err
}
const getOCRJob = `-- name: GetOCRJob :one
SELECT id, resource_id, file_path, status, error_message, created_at, updated_at FROM ocr_jobs
WHERE id = $1 LIMIT 1
`
func (q *Queries) GetOCRJob(ctx context.Context, id uuid.UUID) (OcrJob, error) {
row := q.db.QueryRowContext(ctx, getOCRJob, id)
var i OcrJob
err := row.Scan(
&i.ID,
&i.ResourceID,
&i.FilePath,
&i.Status,
&i.ErrorMessage,
&i.CreatedAt,
&i.UpdatedAt,
)
return i, err
}
const listPendingOCRJobs = `-- name: ListPendingOCRJobs :many
SELECT id, resource_id, file_path, status, error_message, created_at, updated_at FROM ocr_jobs
WHERE status = 'pending'
ORDER BY created_at ASC
`
func (q *Queries) ListPendingOCRJobs(ctx context.Context) ([]OcrJob, error) {
rows, err := q.db.QueryContext(ctx, listPendingOCRJobs)
if err != nil {
return nil, err
}
defer rows.Close()
var items []OcrJob
for rows.Next() {
var i OcrJob
if err := rows.Scan(
&i.ID,
&i.ResourceID,
&i.FilePath,
&i.Status,
&i.ErrorMessage,
&i.CreatedAt,
&i.UpdatedAt,
); err != nil {
return nil, err
}
items = append(items, i)
}
if err := rows.Close(); err != nil {
return nil, err
}
if err := rows.Err(); err != nil {
return nil, err
}
return items, nil
}
const updateOCRJobStatus = `-- name: UpdateOCRJobStatus :exec
UPDATE ocr_jobs
SET status = $1, error_message = $2, updated_at = CURRENT_TIMESTAMP
WHERE id = $3
`
type UpdateOCRJobStatusParams struct {
Status string `json:"status"`
ErrorMessage string `json:"error_message"`
ID uuid.UUID `json:"id"`
}
func (q *Queries) UpdateOCRJobStatus(ctx context.Context, arg UpdateOCRJobStatusParams) error {
_, err := q.db.ExecContext(ctx, updateOCRJobStatus, arg.Status, arg.ErrorMessage, arg.ID)
return err
}
+22
View File
@@ -0,0 +1,22 @@
-- name: CreateOCRJob :one
INSERT INTO ocr_jobs (resource_id, file_path, status)
VALUES ($1, $2, 'pending')
RETURNING *;
-- name: GetOCRJob :one
SELECT * FROM ocr_jobs
WHERE id = $1 LIMIT 1;
-- name: ListPendingOCRJobs :many
SELECT * FROM ocr_jobs
WHERE status = 'pending'
ORDER BY created_at ASC;
-- name: UpdateOCRJobStatus :exec
UPDATE ocr_jobs
SET status = $1, error_message = $2, updated_at = CURRENT_TIMESTAMP
WHERE id = $3;
-- name: DeleteOCRJob :exec
DELETE FROM ocr_jobs
WHERE id = $1;
+34
View File
@@ -0,0 +1,34 @@
package handler
import (
"fmt"
"io"
"github.com/gin-gonic/gin"
"github.com/vaultdrop/backend/internal/service"
)
type EventHandler struct {
broker *service.EventBroker
}
func NewEventHandler(broker *service.EventBroker) *EventHandler {
return &EventHandler{broker: broker}
}
func (h *EventHandler) Stream(c *gin.Context) {
c.Header("Content-Type", "text/event-stream")
c.Header("Cache-Control", "no-cache")
c.Header("Connection", "keep-alive")
ch, _ := h.broker.Subscribe(c.Request.Context())
c.Stream(func(w io.Writer) bool {
event, ok := <-ch
if !ok {
return false
}
_, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event.Type, event.Data)
return err == nil
})
}
+3
View File
@@ -13,6 +13,7 @@ type Handler struct {
Share *ShareHandler Share *ShareHandler
Device *DeviceHandler Device *DeviceHandler
Sync *SyncHandler Sync *SyncHandler
Events *EventHandler
} }
func New( func New(
@@ -24,6 +25,7 @@ func New(
rebacSvc *service.RebacService, rebacSvc *service.RebacService,
placementSvc *service.PlacementService, placementSvc *service.PlacementService,
syncSvc *service.SyncService, syncSvc *service.SyncService,
eventBroker *service.EventBroker,
) *Handler { ) *Handler {
return &Handler{ return &Handler{
Resource: &ResourceHandler{ Resource: &ResourceHandler{
@@ -38,5 +40,6 @@ func New(
Share: &ShareHandler{rebac: rebacSvc}, Share: &ShareHandler{rebac: rebacSvc},
Device: &DeviceHandler{placement: placementSvc}, Device: &DeviceHandler{placement: placementSvc},
Sync: &SyncHandler{sync: syncSvc}, Sync: &SyncHandler{sync: syncSvc},
Events: NewEventHandler(eventBroker),
} }
} }
+6 -2
View File
@@ -43,10 +43,14 @@ func (h *ResourceHandler) Upload(c *gin.Context) {
return return
} }
h.ocr.Enqueue(result.ID, result.Path) if err := h.ocr.Enqueue(result.ID, result.Path); err != nil {
log.Printf("WARN %v", err)
}
if service.IsConvertible(result.MimeType) { if service.IsConvertible(result.MimeType) {
h.conversion.Enqueue(result.ID, result.Path, result.MimeType) if err := h.conversion.Enqueue(result.ID, result.Path, result.MimeType); err != nil {
log.Printf("WARN %v", err)
}
} }
results = append(results, gin.H{ results = append(results, gin.H{
+3
View File
@@ -38,6 +38,9 @@ func SetupRoutes(r *gin.Engine, h *Handler, authMiddleware *auth.AuthService) {
// Variants // Variants
protected.GET("/resources/:id/variants", h.Resource.GetVariants) protected.GET("/resources/:id/variants", h.Resource.GetVariants)
// Events (SSE)
protected.GET("/events", h.Events.Stream)
// Dedup // Dedup
protected.POST("/resources/dedup-check", h.Resource.CheckDuplicates) protected.POST("/resources/dedup-check", h.Resource.CheckDuplicates)
+14 -6
View File
@@ -36,9 +36,11 @@ func NewConversionService(queries *db.Queries, cfg *config.Config) *ConversionSe
} }
} }
func (s *ConversionService) Start() { func (s *ConversionService) Start(workerCount int) {
go s.worker() for i := range workerCount {
log.Println("[Conversion] Worker started") go s.worker()
log.Printf("[Conversion] Worker %d started", i)
}
} }
func (s *ConversionService) Stop() { func (s *ConversionService) Stop() {
@@ -46,9 +48,15 @@ func (s *ConversionService) Stop() {
log.Println("[Conversion] Worker stopped") log.Println("[Conversion] Worker stopped")
} }
func (s *ConversionService) Enqueue(resourceID, filePath, mimeType string) { func (s *ConversionService) Enqueue(resourceID, filePath, mimeType string) error {
s.jobs <- ConversionJob{ResourceID: resourceID, FilePath: filePath, MimeType: mimeType} select {
log.Printf("[Conversion] Enqueued resource %s", resourceID) case s.jobs <- ConversionJob{ResourceID: resourceID, FilePath: filePath, MimeType: mimeType}:
log.Printf("[Conversion] Enqueued resource %s", resourceID)
return nil
default:
log.Printf("[Conversion] Queue full, dropping resource %s", resourceID)
return fmt.Errorf("conversion queue full (%d pending)", len(s.jobs))
}
} }
func (s *ConversionService) worker() { func (s *ConversionService) worker() {
+54
View File
@@ -0,0 +1,54 @@
package service
import (
"context"
"sync"
"github.com/google/uuid"
)
type Event struct {
Type string
Data string
}
type EventBroker struct {
mu sync.RWMutex
subscribers map[string]chan Event
}
func NewEventBroker() *EventBroker {
return &EventBroker{
subscribers: make(map[string]chan Event),
}
}
func (b *EventBroker) Subscribe(ctx context.Context) (<-chan Event, string) {
id := uuid.New().String()
ch := make(chan Event, 16)
b.mu.Lock()
b.subscribers[id] = ch
b.mu.Unlock()
go func() {
<-ctx.Done()
b.mu.Lock()
delete(b.subscribers, id)
close(ch)
b.mu.Unlock()
}()
return ch, id
}
func (b *EventBroker) Publish(eventType, data string) {
b.mu.RLock()
defer b.mu.RUnlock()
for _, ch := range b.subscribers {
select {
case ch <- Event{Type: eventType, Data: data}:
default:
}
}
}
+88 -7
View File
@@ -1,15 +1,21 @@
package service package service
import ( import (
"context"
"database/sql"
"fmt"
"log" "log"
"os" "os"
"strings" "strings"
"github.com/google/uuid"
"github.com/vaultdrop/backend/internal/config" "github.com/vaultdrop/backend/internal/config"
"github.com/vaultdrop/backend/internal/db"
"github.com/vaultdrop/backend/internal/ocr" "github.com/vaultdrop/backend/internal/ocr"
) )
type OCRJob struct { type OCRJob struct {
DBID uuid.UUID
ResourceID string ResourceID string
FilePath string FilePath string
} }
@@ -17,20 +23,27 @@ type OCRJob struct {
type OCRService struct { type OCRService struct {
client *ocr.Client client *ocr.Client
resourceSvc *ResourceService resourceSvc *ResourceService
broker *EventBroker
queries *db.Queries
jobs chan OCRJob 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{ return &OCRService{
client: ocr.NewClient(cfg.OCREndpoint), client: ocr.NewClient(cfg.OCREndpoint),
resourceSvc: resourceSvc, resourceSvc: resourceSvc,
broker: broker,
queries: queries,
jobs: make(chan OCRJob, 100), jobs: make(chan OCRJob, 100),
} }
} }
func (s *OCRService) Start() { func (s *OCRService) Start(workerCount int) {
go s.worker() s.replenish()
log.Println("[OCR] Worker started") for i := range workerCount {
go s.worker()
log.Printf("[OCR] Worker %d started", i)
}
} }
func (s *OCRService) Stop() { func (s *OCRService) Stop() {
@@ -38,9 +51,51 @@ func (s *OCRService) Stop() {
log.Println("[OCR] Worker stopped") log.Println("[OCR] Worker stopped")
} }
func (s *OCRService) Enqueue(resourceID, filePath string) { func (s *OCRService) Enqueue(resourceID, filePath string) error {
s.jobs <- OCRJob{ResourceID: resourceID, FilePath: filePath} ctx := context.Background()
log.Printf("[OCR] Enqueued resource %s", resourceID)
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() { func (s *OCRService) worker() {
@@ -50,17 +105,33 @@ func (s *OCRService) worker() {
} }
func (s *OCRService) process(job OCRJob) { func (s *OCRService) process(job OCRJob) {
ctx := context.Background()
log.Printf("[OCR] Processing resource %s", job.ResourceID) 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) data, err := os.ReadFile(job.FilePath)
if err != nil { if err != nil {
log.Printf("[OCR] Failed to read resource %s: %v", job.ResourceID, err) 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 return
} }
blocks, err := s.client.Recognize(data) blocks, err := s.client.Recognize(data)
if err != nil { if err != nil {
log.Printf("[OCR] Failed to recognize resource %s: %v", job.ResourceID, err) 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 return
} }
@@ -68,9 +139,19 @@ func (s *OCRService) process(job OCRJob) {
if err := s.resourceSvc.UpdateOCRText(job.ResourceID, text); err != nil { 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) 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 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)) log.Printf("[OCR] Completed resource %s (%d chars)", job.ResourceID, len(text))
} }