feat: Add database migration service with Alembic for PostgreSQL, integrating job management and counting functionalities. Update Docker Compose to include db-migrations service, enhancing the microservices architecture. Revise documentation for clarity on database schema and migration processes.

This commit is contained in:
erik
2026-07-06 00:50:50 -05:00
parent d95f8f6ad8
commit 081da999d4
43 changed files with 2911 additions and 2271 deletions
@@ -1,53 +0,0 @@
package api
import (
"context"
"github.com/pys/pdf-count/internal/pipeline"
"github.com/pys/pdf-count/internal/store"
)
// FilePayload holds a PDF ready for inline job processing.
type FilePayload struct {
Filename string
Data []byte
}
// JobRunner processes uploaded job files in-process using the API OCR pool.
type JobRunner struct {
store *store.Store
newProc func(includeRejected bool) *pipeline.Processor
sem chan struct{}
}
func NewJobRunner(st *store.Store, newProc func(includeRejected bool) *pipeline.Processor, concurrency int) *JobRunner {
if concurrency < 1 {
concurrency = 1
}
return &JobRunner{
store: st,
newProc: newProc,
sem: make(chan struct{}, concurrency),
}
}
// ProcessOne schedules a single PDF for inline OCR (non-blocking).
func (jr *JobRunner) ProcessOne(jobID string, file FilePayload, includeRejected bool) {
go func() {
jr.sem <- struct{}{}
defer func() { <-jr.sem }()
ctx := context.Background()
proc := jr.newProc(includeRejected)
result := proc.ProcessPDF(file.Filename, file.Data)
_ = jr.store.SaveResult(ctx, jobID, result)
_ = jr.store.TryCompleteJob(ctx, jobID)
}()
}
// ProcessFiles runs OCR in background goroutines and updates job progress in Redis.
func (jr *JobRunner) ProcessFiles(jobID string, files []FilePayload, includeRejected bool) {
for _, file := range files {
jr.ProcessOne(jobID, file, includeRejected)
}
}
@@ -4,87 +4,47 @@ import (
"context"
"encoding/json"
"io"
"mime/multipart"
"net/http"
"os"
"path/filepath"
"strconv"
"strings"
"sync"
"time"
"github.com/go-chi/chi/v5"
"github.com/go-chi/chi/v5/middleware"
"github.com/google/uuid"
"github.com/pys/pdf-count/internal/config"
"github.com/pys/pdf-count/internal/domain"
"github.com/pys/pdf-count/internal/model"
"github.com/pys/pdf-count/internal/ocr"
"github.com/pys/pdf-count/internal/pipeline"
"github.com/pys/pdf-count/internal/queue"
"github.com/pys/pdf-count/internal/storage"
"github.com/pys/pdf-count/internal/store"
"github.com/pys/pdf-count/web"
"gopkg.in/yaml.v3"
)
// Server holds the shared, stateless dependencies of the count-api.
type Server struct {
cfg config.Config
predictor *model.Predictor
predictor *model.Predictor // nil → heuristic fallback
ocrEngine *ocr.Engine
store *store.Store
storage *storage.Client
publisher *queue.Publisher
jobRunner *JobRunner
openAPI []byte
ready bool
readyMu sync.RWMutex
}
func NewServer(cfg config.Config, predictor *model.Predictor, ocrEngine *ocr.Engine, st *store.Store, stg *storage.Client, pub *queue.Publisher, openAPI []byte) *Server {
srv := &Server{
// NewServer constructs a Server. predictor may be nil when the ONNX model is
// unavailable; the pipeline will fall back to a counting heuristic.
func NewServer(
cfg config.Config,
predictor *model.Predictor,
ocrEngine *ocr.Engine,
openAPI []byte,
) *Server {
return &Server{
cfg: cfg,
predictor: predictor,
ocrEngine: ocrEngine,
store: st,
storage: stg,
publisher: pub,
openAPI: openAPI,
ready: true,
}
srv.jobRunner = NewJobRunner(st, srv.newProcessor, srv.ocrConcurrency())
return srv
}
func (s *Server) processConcurrency() int {
c := s.cfg.SyncProcessConcurrency
if c <= 0 {
c = s.cfg.OCRPoolSize
}
if c <= 0 {
c = 4
}
return c
}
func (s *Server) ocrConcurrency() int {
c := s.processConcurrency()
pool := s.cfg.OCRPoolSize
if pool <= 0 {
pool = 4
}
if c > pool {
c = pool
}
return c
}
func (s *Server) SetReady(v bool) {
s.readyMu.Lock()
s.ready = v
s.readyMu.Unlock()
}
// Router registers all HTTP routes.
func (s *Server) Router() http.Handler {
r := chi.NewRouter()
r.Use(middleware.RequestID)
@@ -92,46 +52,41 @@ func (s *Server) Router() http.Handler {
r.Use(middleware.Recoverer)
r.Get("/health", s.health)
r.Get("/ready", s.readyHandler)
r.Get("/ready", s.ready)
r.Get("/openapi.json", s.openAPIHandler)
r.Get("/demo", func(w http.ResponseWriter, r *http.Request) {
http.Redirect(w, r, "/demo/", http.StatusMovedPermanently)
})
r.Handle("/demo/*", http.StripPrefix("/demo/", http.FileServer(http.FS(web.FS))))
r.Route("/v1", func(r chi.Router) {
r.Post("/count", s.countSync)
r.Post("/jobs", s.createJob)
r.Post("/jobs/{jobID}/upload", s.uploadChunk)
r.Get("/jobs/{jobID}", s.getJob)
r.Get("/jobs/{jobID}/results", s.getJobResults)
r.Delete("/jobs/{jobID}", s.cancelJob)
})
return r
}
func (s *Server) health(w http.ResponseWriter, r *http.Request) {
// health returns a simple liveness response.
func (s *Server) health(w http.ResponseWriter, _ *http.Request) {
writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
}
func (s *Server) readyHandler(w http.ResponseWriter, r *http.Request) {
s.readyMu.RLock()
ok := s.ready
s.readyMu.RUnlock()
if !ok {
writeJSON(w, http.StatusServiceUnavailable, map[string]string{"status": "not_ready"})
return
}
// ready returns 200 once the server is accepting traffic (always ready after init).
func (s *Server) ready(w http.ResponseWriter, _ *http.Request) {
writeJSON(w, http.StatusOK, map[string]string{"status": "ready"})
}
func (s *Server) openAPIHandler(w http.ResponseWriter, r *http.Request) {
// openAPIHandler serves the OpenAPI specification as JSON.
func (s *Server) openAPIHandler(w http.ResponseWriter, _ *http.Request) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(http.StatusOK)
_, _ = w.Write(s.openAPI)
}
// countSync processes a multipart batch of PDFs synchronously and returns the
// counting results in a single JSON response. No state is persisted.
//
// POST /v1/count
//
// Query: include_rejected=true|false
// Body: multipart/form-data, field "files" (one or more PDFs)
// Limits: SYNC_MAX_FILES (default 50), MAX_PDF_SIZE_MB (default 20)
func (s *Server) countSync(w http.ResponseWriter, r *http.Request) {
includeRejected := r.URL.Query().Get("include_rejected") == "true"
maxBytes := int64(s.cfg.MaxPDFSizeMB) * 1024 * 1024
@@ -151,61 +106,51 @@ func (s *Server) countSync(w http.ResponseWriter, r *http.Request) {
return
}
proc := s.newProcessor(includeRejected)
type fileInput struct {
filename string
data []byte
errRes *domain.PDFResult
}
inputs := make([]fileInput, 0, len(files))
inputs := make([]fileInput, 0, len(files))
for _, fh := range files {
name := filepath.Base(fh.Filename)
f, err := fh.Open()
if err != nil {
errRes := domain.PDFResult{
Filename: filepath.Base(fh.Filename),
Status: domain.StatusError,
Error: err.Error(),
}
errRes := domain.PDFResult{Filename: name, Status: domain.StatusError, Error: err.Error()}
inputs = append(inputs, fileInput{errRes: &errRes})
continue
}
data, err := io.ReadAll(io.LimitReader(f, maxBytes+1))
_ = f.Close()
if err != nil {
errRes := domain.PDFResult{
Filename: filepath.Base(fh.Filename),
Status: domain.StatusError,
Error: err.Error(),
}
errRes := domain.PDFResult{Filename: name, Status: domain.StatusError, Error: err.Error()}
inputs = append(inputs, fileInput{errRes: &errRes})
continue
}
if int64(len(data)) > maxBytes {
errRes := domain.PDFResult{
Filename: filepath.Base(fh.Filename),
Status: domain.StatusError,
Error: "file too large",
}
errRes := domain.PDFResult{Filename: name, Status: domain.StatusError, Error: "file too large"}
inputs = append(inputs, fileInput{errRes: &errRes})
continue
}
inputs = append(inputs, fileInput{
filename: filepath.Base(fh.Filename),
data: data,
})
inputs = append(inputs, fileInput{filename: name, data: data})
}
proc := s.newProcessor(includeRejected)
results := make([]domain.PDFResult, len(inputs))
failed := 0
sem := make(chan struct{}, s.ocrConcurrency())
concurrency := s.syncConcurrency()
sem := make(chan struct{}, concurrency)
var wg sync.WaitGroup
var mu sync.Mutex
for i, in := range inputs {
if in.errRes != nil {
results[i] = *in.errRes
mu.Lock()
failed++
mu.Unlock()
continue
}
wg.Add(1)
@@ -231,314 +176,23 @@ func (s *Server) countSync(w http.ResponseWriter, r *http.Request) {
})
}
type createJobRequest struct {
Files []string `json:"files"`
Prefix string `json:"prefix"`
IncludeRejected bool `json:"include_rejected"`
CallbackURL string `json:"callback_url"`
}
func (s *Server) createJob(w http.ResponseWriter, r *http.Request) {
var req createJobRequest
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
httpError(w, http.StatusBadRequest, "invalid json body")
return
// syncConcurrency returns the effective parallelism for /v1/count requests.
func (s *Server) syncConcurrency() int {
c := s.cfg.SyncProcessConcurrency
if c <= 0 {
c = s.cfg.OCRPoolSize
}
ctx := r.Context()
pending, err := s.store.PendingCount(ctx)
if err == nil && int(pending) > s.cfg.QueueMaxPending {
httpError(w, http.StatusTooManyRequests, "queue full")
return
if c <= 0 {
c = 4
}
jobID := uuid.NewString()
var objectKeys []string
var filenames []string
if req.Prefix != "" {
prefix := strings.TrimSuffix(req.Prefix, "/")
keys, err := s.storage.ListByPrefix(ctx, prefix)
if err != nil {
httpError(w, http.StatusBadRequest, err.Error())
return
}
objectKeys = keys
for _, k := range keys {
filenames = append(filenames, filepath.Base(k))
}
} else if len(req.Files) > 0 {
for _, f := range req.Files {
name := filepath.Base(f)
objectKeys = append(objectKeys, storage.JobObjectKey(jobID, name))
filenames = append(filenames, name)
}
} else {
// Empty job shell for chunked upload flow.
if err := s.store.CreateJob(ctx, jobID, 0, req.IncludeRejected); err != nil {
httpError(w, http.StatusInternalServerError, err.Error())
return
}
writeJSON(w, http.StatusAccepted, domain.JobStatus{
JobID: jobID,
Status: domain.JobQueued,
Total: 0,
})
return
pool := s.cfg.OCRPoolSize
if pool <= 0 {
pool = 4
}
if len(objectKeys) > s.cfg.JobMaxFiles {
httpError(w, http.StatusBadRequest, "too many files in job")
return
if c > pool {
c = pool
}
if err := s.store.CreateJob(ctx, jobID, len(objectKeys), req.IncludeRejected); err != nil {
httpError(w, http.StatusInternalServerError, err.Error())
return
}
if err := s.store.AddJobFiles(ctx, jobID, filenames); err != nil {
httpError(w, http.StatusInternalServerError, err.Error())
return
}
_ = s.store.MarkRunning(ctx, jobID)
for i, key := range objectKeys {
msg := domain.QueueMessage{
JobID: jobID,
ObjectKey: key,
Filename: filenames[i],
IncludeRejected: req.IncludeRejected,
}
if err := s.publisher.Publish(ctx, msg); err != nil {
httpError(w, http.StatusInternalServerError, err.Error())
return
}
}
writeJSON(w, http.StatusAccepted, domain.JobStatus{
JobID: jobID,
Status: domain.JobRunning,
Total: len(objectKeys),
})
}
func (s *Server) uploadChunk(w http.ResponseWriter, r *http.Request) {
jobID := chi.URLParam(r, "jobID")
includeRejected := r.URL.Query().Get("include_rejected") == "true"
maxBytes := int64(s.cfg.MaxPDFSizeMB) * 1024 * 1024
if err := r.ParseMultipartForm(maxBytes * 500); err != nil {
httpError(w, http.StatusBadRequest, "invalid multipart form")
return
}
files := r.MultipartForm.File["files"]
if len(files) == 0 {
httpError(w, http.StatusBadRequest, "field 'files' required")
return
}
if len(files) > 500 {
httpError(w, http.StatusBadRequest, "max 500 files per chunk")
return
}
ctx := r.Context()
status, err := s.store.GetJob(ctx, jobID)
if err != nil {
_ = s.store.CreateJob(ctx, jobID, 0, includeRejected)
status = domain.JobStatus{JobID: jobID, Status: domain.JobQueued}
}
if status.Status == domain.JobCancelled {
httpError(w, http.StatusConflict, "job cancelled")
return
}
var filenames []string
if s.cfg.JobInlineProcessing {
_ = s.store.MarkRunning(ctx, jobID)
for _, fh := range files {
f, err := fh.Open()
if err != nil {
continue
}
data, err := io.ReadAll(io.LimitReader(f, maxBytes+1))
_ = f.Close()
if err != nil || int64(len(data)) > maxBytes {
continue
}
name := filepath.Base(fh.Filename)
if err := s.store.AddJobFiles(ctx, jobID, []string{name}); err != nil {
httpError(w, http.StatusInternalServerError, err.Error())
return
}
if err := s.store.IncrementTotal(ctx, jobID, 1); err != nil {
httpError(w, http.StatusInternalServerError, err.Error())
return
}
filenames = append(filenames, name)
s.jobRunner.ProcessOne(jobID, FilePayload{Filename: name, Data: data}, includeRejected)
}
} else {
var err error
filenames, err = s.uploadChunkMinIO(ctx, jobID, files, includeRejected, maxBytes)
if err != nil {
httpError(w, http.StatusInternalServerError, err.Error())
return
}
}
if len(filenames) == 0 {
httpError(w, http.StatusBadRequest, "no valid files uploaded")
return
}
writeJSON(w, http.StatusAccepted, map[string]interface{}{
"job_id": jobID,
"uploaded": len(filenames),
})
}
func (s *Server) uploadChunkMinIO(ctx context.Context, jobID string, files []*multipart.FileHeader, includeRejected bool, maxBytes int64) ([]string, error) {
filenames := make([]string, 0, len(files))
for _, fh := range files {
filenames = append(filenames, filepath.Base(fh.Filename))
}
if len(filenames) == 0 {
return nil, nil
}
if err := s.store.AddJobFiles(ctx, jobID, filenames); err != nil {
return nil, err
}
if err := s.store.IncrementTotal(ctx, jobID, len(filenames)); err != nil {
return nil, err
}
if err := s.store.MarkRunning(ctx, jobID); err != nil {
return nil, err
}
ioLimit := 8
if len(files) < ioLimit {
ioLimit = len(files)
}
ioSem := make(chan struct{}, ioLimit)
var wg sync.WaitGroup
var mu sync.Mutex
uploadedNames := make([]string, 0, len(files))
for _, fh := range files {
wg.Add(1)
go func(fh *multipart.FileHeader) {
defer wg.Done()
ioSem <- struct{}{}
defer func() { <-ioSem }()
name := filepath.Base(fh.Filename)
f, err := fh.Open()
if err != nil {
s.recordUploadError(ctx, jobID, name, err.Error())
return
}
data, err := io.ReadAll(io.LimitReader(f, maxBytes+1))
_ = f.Close()
if err != nil {
s.recordUploadError(ctx, jobID, name, err.Error())
return
}
if int64(len(data)) > maxBytes {
s.recordUploadError(ctx, jobID, name, "file too large")
return
}
key := storage.JobObjectKey(jobID, name)
if err := s.storage.Put(ctx, key, data, "application/pdf"); err != nil {
s.recordUploadError(ctx, jobID, name, err.Error())
return
}
msg := domain.QueueMessage{
JobID: jobID,
ObjectKey: key,
Filename: name,
IncludeRejected: includeRejected,
}
if err := s.publisher.Publish(ctx, msg); err != nil {
s.recordUploadError(ctx, jobID, name, err.Error())
return
}
mu.Lock()
uploadedNames = append(uploadedNames, name)
mu.Unlock()
}(fh)
}
wg.Wait()
if len(uploadedNames) == 0 {
return nil, nil
}
return uploadedNames, nil
}
func (s *Server) recordUploadError(ctx context.Context, jobID, filename, errMsg string) {
_ = s.store.SaveResult(ctx, jobID, domain.PDFResult{
Filename: filename,
Status: domain.StatusError,
Error: errMsg,
})
_ = s.store.TryCompleteJob(ctx, jobID)
}
func (s *Server) getJob(w http.ResponseWriter, r *http.Request) {
jobID := chi.URLParam(r, "jobID")
waitMs, _ := strconv.Atoi(r.URL.Query().Get("wait_ms"))
if waitMs > 30000 {
waitMs = 30000
}
deadline := time.Now()
if waitMs > 0 {
deadline = deadline.Add(time.Duration(waitMs) * time.Millisecond)
}
for {
status, err := s.store.GetJob(r.Context(), jobID)
if err != nil {
httpError(w, http.StatusNotFound, "job not found")
return
}
if isTerminalJobStatus(status.Status) || waitMs <= 0 || !time.Now().Before(deadline) {
writeJSON(w, http.StatusOK, status)
return
}
time.Sleep(100 * time.Millisecond)
}
}
func isTerminalJobStatus(status string) bool {
return status == domain.JobCompleted || status == domain.JobCancelled
}
func (s *Server) getJobResults(w http.ResponseWriter, r *http.Request) {
jobID := chi.URLParam(r, "jobID")
cursor := r.URL.Query().Get("cursor")
limit, _ := strconv.Atoi(r.URL.Query().Get("limit"))
page, err := s.store.GetResultsPage(r.Context(), jobID, cursor, limit)
if err != nil {
httpError(w, http.StatusNotFound, err.Error())
return
}
writeJSON(w, http.StatusOK, page)
}
func (s *Server) cancelJob(w http.ResponseWriter, r *http.Request) {
jobID := chi.URLParam(r, "jobID")
if err := s.store.CancelJob(r.Context(), jobID); err != nil {
httpError(w, http.StatusInternalServerError, err.Error())
return
}
writeJSON(w, http.StatusOK, map[string]string{"job_id": jobID, "status": domain.JobCancelled})
return c
}
func (s *Server) newProcessor(includeRejected bool) *pipeline.Processor {
@@ -552,20 +206,24 @@ func (s *Server) newProcessor(includeRejected bool) *pipeline.Processor {
})
}
// writeJSON encodes v as JSON and writes it with the given HTTP status.
func writeJSON(w http.ResponseWriter, status int, v interface{}) {
w.Header().Set("Content-Type", "application/json")
w.WriteHeader(status)
_ = json.NewEncoder(w).Encode(v)
}
// httpError writes a JSON error response.
func httpError(w http.ResponseWriter, status int, msg string) {
writeJSON(w, status, map[string]string{"error": msg})
}
// LoadOpenAPI reads an OpenAPI spec from path.
// YAML files are converted to JSON on the fly. Returns a minimal stub on error.
func LoadOpenAPI(path string) []byte {
data, err := os.ReadFile(path)
if err != nil {
return []byte(`{"openapi":"3.1.0","info":{"title":"pdf-count","version":"1.0.0"}}`)
return []byte(`{"openapi":"3.1.0","info":{"title":"count-api","version":"1.0.0"}}`)
}
if strings.HasSuffix(strings.ToLower(path), ".yaml") || strings.HasSuffix(strings.ToLower(path), ".yml") {
var doc interface{}
@@ -578,6 +236,7 @@ func LoadOpenAPI(path string) []byte {
return data
}
// ListenAndServe starts the HTTP server and blocks until ctx is cancelled.
func ListenAndServe(ctx context.Context, addr string, handler http.Handler) error {
srv := &http.Server{Addr: addr, Handler: handler}
go func() {
@@ -6,67 +6,39 @@ import (
"time"
)
// Config holds all runtime parameters for the count-api service.
// All values are loaded from environment variables with safe defaults.
type Config struct {
Port string
RedisURL string
NATSURL string
MinIOEndpoint string
MinIOAccessKey string
MinIOSecretKey string
MinIOBucket string
MinIOUseSSL bool
ModelPath string
InferencePath string
OCRDPI int
OCRLang string
DBSCANEps float64
DBSCANMinSamples int
SyncMaxFiles int
JobMaxFiles int
QueueMaxPending int
WorkerConcurrency int
MaxPDFSizeMB int
PDFProcessTimeout time.Duration
JobResultTTL time.Duration
StreamName string
SubjectPrefix string
Port string
ModelPath string
InferencePath string
OCRDPI int
OCRLang string
OCRPoolSize int
LogPipelineTiming bool
NATSFetchMaxWait time.Duration
JobInlineProcessing bool
DBSCANEps float64
DBSCANMinSamples int
SyncMaxFiles int
MaxPDFSizeMB int
PDFProcessTimeout time.Duration
SyncProcessConcurrency int
LogPipelineTiming bool
}
func Load() Config {
return Config{
Port: getenv("PORT", "8080"),
RedisURL: getenv("REDIS_URL", "redis://redis:6379/0"),
NATSURL: getenv("NATS_URL", "nats://nats:4222"),
MinIOEndpoint: getenv("MINIO_ENDPOINT", "minio:9000"),
MinIOAccessKey: getenv("MINIO_ACCESS_KEY", "minioadmin"),
MinIOSecretKey: getenv("MINIO_SECRET_KEY", "minioadmin"),
MinIOBucket: getenv("MINIO_BUCKET", "pdf-count"),
MinIOUseSSL: getenvBool("MINIO_USE_SSL", false),
ModelPath: getenv("MODEL_PATH", "/app/models/model.onnx"),
InferencePath: getenv("INFERENCE_CONFIG", "/app/models/inference.json"),
OCRDPI: getenvInt("OCR_DPI", 300),
OCRLang: getenv("OCR_LANG", "spa"),
DBSCANEps: getenvFloat("DBSCAN_EPS", 10),
DBSCANMinSamples: getenvInt("DBSCAN_MIN_SAMPLES", 1),
SyncMaxFiles: getenvInt("SYNC_MAX_FILES", 50),
JobMaxFiles: getenvInt("JOB_MAX_FILES", 50000),
QueueMaxPending: getenvInt("QUEUE_MAX_PENDING", 100000),
WorkerConcurrency: getenvInt("WORKER_CONCURRENCY", 4),
MaxPDFSizeMB: getenvInt("MAX_PDF_SIZE_MB", 20),
PDFProcessTimeout: time.Duration(getenvInt("PDF_PROCESS_TIMEOUT_SEC", 120)) * time.Second,
JobResultTTL: time.Duration(getenvInt("JOB_RESULT_TTL_HOURS", 72)) * time.Hour,
StreamName: getenv("NATS_STREAM", "PDF_JOBS"),
SubjectPrefix: getenv("NATS_SUBJECT_PREFIX", "pdf.process"),
Port: getenv("PORT", "8080"),
ModelPath: getenv("MODEL_PATH", "/app/models/model.onnx"),
InferencePath: getenv("INFERENCE_CONFIG", "/app/models/inference.json"),
OCRDPI: getenvInt("OCR_DPI", 300),
OCRLang: getenv("OCR_LANG", "spa"),
OCRPoolSize: getenvInt("OCR_POOL_SIZE", 0),
LogPipelineTiming: getenvBool("LOG_PIPELINE_TIMING", false),
NATSFetchMaxWait: time.Duration(getenvInt("NATS_FETCH_MAX_WAIT_MS", 300)) * time.Millisecond,
JobInlineProcessing: getenvBool("JOB_INLINE_PROCESSING", false),
DBSCANEps: getenvFloat("DBSCAN_EPS", 10),
DBSCANMinSamples: getenvInt("DBSCAN_MIN_SAMPLES", 1),
SyncMaxFiles: getenvInt("SYNC_MAX_FILES", 50),
MaxPDFSizeMB: getenvInt("MAX_PDF_SIZE_MB", 20),
PDFProcessTimeout: time.Duration(getenvInt("PDF_PROCESS_TIMEOUT_SEC", 120)) * time.Second,
SyncProcessConcurrency: getenvInt("SYNC_PROCESS_CONCURRENCY", 0),
LogPipelineTiming: getenvBool("LOG_PIPELINE_TIMING", false),
}
}
@@ -1,13 +1,13 @@
package domain
// ServiceRow is one detected invoice line.
// ServiceRow represents a single detected invoice line in a PDF page.
type ServiceRow struct {
Text string `json:"text"`
Probability float64 `json:"probability"`
IsService *bool `json:"is_service,omitempty"`
}
// PDFResult is the processing outcome for a single PDF.
// PDFResult is the counting outcome for a single PDF document.
type PDFResult struct {
Filename string `json:"filename"`
ServiceCount int `json:"service_count"`
@@ -16,44 +16,14 @@ type PDFResult struct {
Error string `json:"error,omitempty"`
}
// CountResponse wraps sync batch results.
// CountResponse wraps the results of a synchronous batch request to POST /v1/count.
type CountResponse struct {
Results []PDFResult `json:"results"`
Processed int `json:"processed"`
Failed int `json:"failed"`
}
// JobStatus tracks async job progress.
type JobStatus struct {
JobID string `json:"job_id"`
Status string `json:"status"`
Total int `json:"total"`
Processed int `json:"processed"`
Failed int `json:"failed"`
}
// JobResultsPage is a paginated slice of PDF results.
type JobResultsPage struct {
JobID string `json:"job_id"`
Results []PDFResult `json:"results"`
NextCursor string `json:"next_cursor,omitempty"`
HasMore bool `json:"has_more"`
}
// QueueMessage is published per PDF for workers.
type QueueMessage struct {
JobID string `json:"job_id"`
ObjectKey string `json:"object_key"`
Filename string `json:"filename"`
IncludeRejected bool `json:"include_rejected"`
}
const (
StatusOK = "ok"
StatusError = "error"
JobQueued = "queued"
JobRunning = "running"
JobCompleted = "completed"
JobCancelled = "cancelled"
)
@@ -73,6 +73,21 @@ func (p *Processor) ProcessPDF(filename string, pdfData []byte) domain.PDFResult
return result
}
// Sin modelo ML: contar todas las filas agrupadas como servicios (heurística).
if p.predictor == nil {
services := make([]domain.ServiceRow, 0, len(rows))
for _, row := range rows {
services = append(services, domain.ServiceRow{
Text: row.TextoFila,
Probability: 1.0,
})
}
result.Services = services
result.ServiceCount = len(rows)
p.logTiming(filename, renderDur, ocrDur, clusterDur, 0, time.Since(start))
return result
}
features := make([][8]float32, len(rows))
for i, r := range rows {
features[i] = r.Features
@@ -1,109 +0,0 @@
package queue
import (
"context"
"encoding/json"
"fmt"
"time"
"github.com/nats-io/nats.go"
"github.com/nats-io/nats.go/jetstream"
"github.com/pys/pdf-count/internal/domain"
)
type Publisher struct {
js jetstream.JetStream
stream string
subject string
}
func NewPublisher(natsURL, streamName, subjectPrefix string) (*Publisher, error) {
nc, err := nats.Connect(natsURL)
if err != nil {
return nil, fmt.Errorf("nats connect: %w", err)
}
js, err := jetstream.New(nc)
if err != nil {
nc.Close()
return nil, err
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
_, err = js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{
Name: streamName,
Subjects: []string{subjectPrefix + ".>"},
Retention: jetstream.WorkQueuePolicy,
MaxMsgs: 200000,
})
if err != nil {
nc.Close()
return nil, fmt.Errorf("create stream: %w", err)
}
return &Publisher{js: js, stream: streamName, subject: subjectPrefix}, nil
}
func (p *Publisher) Publish(ctx context.Context, msg domain.QueueMessage) error {
data, err := json.Marshal(msg)
if err != nil {
return err
}
subject := fmt.Sprintf("%s.%s", p.subject, msg.JobID)
_, err = p.js.Publish(ctx, subject, data)
return err
}
type Consumer struct {
js jetstream.JetStream
subject string
consumerName string
}
func NewConsumer(natsURL, streamName, subjectPrefix, consumerName string) (*Consumer, jetstream.Consumer, *nats.Conn, error) {
nc, err := nats.Connect(natsURL)
if err != nil {
return nil, nil, nil, err
}
js, err := jetstream.New(nc)
if err != nil {
nc.Close()
return nil, nil, nil, err
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
_, err = js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{
Name: streamName,
Subjects: []string{subjectPrefix + ".>"},
Retention: jetstream.WorkQueuePolicy,
MaxMsgs: 200000,
})
if err != nil {
nc.Close()
return nil, nil, nil, err
}
cons, err := js.CreateOrUpdateConsumer(ctx, streamName, jetstream.ConsumerConfig{
Name: consumerName,
FilterSubject: subjectPrefix + ".>",
AckPolicy: jetstream.AckExplicitPolicy,
MaxDeliver: 3,
AckWait: 2 * time.Minute,
MaxAckPending: 256,
})
if err != nil {
nc.Close()
return nil, nil, nil, err
}
return &Consumer{js: js, subject: subjectPrefix, consumerName: consumerName}, cons, nc, nil
}
func DecodeMessage(data []byte) (domain.QueueMessage, error) {
var msg domain.QueueMessage
err := json.Unmarshal(data, &msg)
return msg, err
}
@@ -1,86 +0,0 @@
package storage
import (
"bytes"
"context"
"fmt"
"io"
"strings"
"github.com/minio/minio-go/v7"
"github.com/minio/minio-go/v7/pkg/credentials"
)
type Client struct {
client *minio.Client
bucket string
}
func New(endpoint, accessKey, secretKey, bucket string, useSSL bool) (*Client, error) {
client, err := minio.New(endpoint, &minio.Options{
Creds: credentials.NewStaticV4(accessKey, secretKey, ""),
Secure: useSSL,
})
if err != nil {
return nil, err
}
return &Client{client: client, bucket: bucket}, nil
}
func (c *Client) EnsureBucket(ctx context.Context) error {
exists, err := c.client.BucketExists(ctx, c.bucket)
if err != nil {
return err
}
if exists {
return nil
}
err = c.client.MakeBucket(ctx, c.bucket, minio.MakeBucketOptions{})
if err == nil {
return nil
}
resp := minio.ToErrorResponse(err)
if resp.Code == "BucketAlreadyOwnedByYou" || resp.Code == "BucketAlreadyExists" {
return nil
}
return err
}
func (c *Client) Put(ctx context.Context, objectKey string, data []byte, contentType string) error {
if contentType == "" {
contentType = "application/pdf"
}
_, err := c.client.PutObject(ctx, c.bucket, objectKey, bytes.NewReader(data), int64(len(data)), minio.PutObjectOptions{
ContentType: contentType,
})
return err
}
func (c *Client) Get(ctx context.Context, objectKey string) ([]byte, error) {
obj, err := c.client.GetObject(ctx, c.bucket, objectKey, minio.GetObjectOptions{})
if err != nil {
return nil, err
}
defer obj.Close()
return io.ReadAll(obj)
}
func (c *Client) ListByPrefix(ctx context.Context, prefix string) ([]string, error) {
var keys []string
for obj := range c.client.ListObjects(ctx, c.bucket, minio.ListObjectsOptions{
Prefix: prefix,
Recursive: true,
}) {
if obj.Err != nil {
return nil, obj.Err
}
if strings.HasSuffix(strings.ToLower(obj.Key), ".pdf") {
keys = append(keys, obj.Key)
}
}
return keys, nil
}
func JobObjectKey(jobID, filename string) string {
return fmt.Sprintf("jobs/%s/%s", jobID, filename)
}
@@ -1,276 +0,0 @@
package store
import (
"context"
"encoding/json"
"fmt"
"strconv"
"time"
"github.com/pys/pdf-count/internal/domain"
"github.com/redis/go-redis/v9"
)
type Store struct {
client *redis.Client
ttl time.Duration
}
func New(redisURL string, ttl time.Duration) (*Store, error) {
opts, err := redis.ParseURL(redisURL)
if err != nil {
return nil, err
}
client := redis.NewClient(opts)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := client.Ping(ctx).Err(); err != nil {
return nil, fmt.Errorf("redis ping: %w", err)
}
return &Store{client: client, ttl: ttl}, nil
}
func (s *Store) Close() error {
return s.client.Close()
}
func metaKey(jobID string) string { return fmt.Sprintf("job:%s:meta", jobID) }
func filesKey(jobID string) string { return fmt.Sprintf("job:%s:files", jobID) }
func resultsKey(jobID string) string {
return fmt.Sprintf("job:%s:results", jobID)
}
func (s *Store) CreateJob(ctx context.Context, jobID string, total int, includeRejected bool) error {
pipe := s.client.Pipeline()
pipe.HSet(ctx, metaKey(jobID), map[string]interface{}{
"status": domain.JobQueued,
"total": total,
"processed": 0,
"failed": 0,
"include_rejected": boolStr(includeRejected),
})
pipe.Expire(ctx, metaKey(jobID), s.ttl)
_, err := pipe.Exec(ctx)
return err
}
func (s *Store) AddJobFiles(ctx context.Context, jobID string, filenames []string) error {
if len(filenames) == 0 {
return nil
}
members := make([]interface{}, len(filenames))
for i, f := range filenames {
members[i] = f
}
pipe := s.client.Pipeline()
pipe.RPush(ctx, filesKey(jobID), members...)
pipe.Expire(ctx, filesKey(jobID), s.ttl)
_, err := pipe.Exec(ctx)
return err
}
func (s *Store) MarkRunning(ctx context.Context, jobID string) error {
return s.client.HSet(ctx, metaKey(jobID), "status", domain.JobRunning).Err()
}
var saveResultScript = redis.NewScript(`
local meta = KEYS[1]
local results = KEYS[2]
redis.call('HSET', results, ARGV[1], ARGV[2])
local processed = redis.call('HINCRBY', meta, 'processed', 1)
if ARGV[3] == '1' then
redis.call('HINCRBY', meta, 'failed', 1)
end
redis.call('EXPIRE', results, ARGV[4])
local total = tonumber(redis.call('HGET', meta, 'total') or '0')
local status = redis.call('HGET', meta, 'status')
if status ~= 'cancelled' and status ~= 'completed' and total > 0 and processed >= total then
redis.call('HSET', meta, 'status', 'completed')
end
return processed
`)
func (s *Store) SaveResult(ctx context.Context, jobID string, result domain.PDFResult) error {
data, err := json.Marshal(result)
if err != nil {
return err
}
failFlag := "0"
if result.Status == domain.StatusError {
failFlag = "1"
}
ttlSec := int64(s.ttl.Seconds())
if ttlSec <= 0 {
ttlSec = 72 * 3600
}
_, err = saveResultScript.Run(ctx, s.client,
[]string{metaKey(jobID), resultsKey(jobID)},
result.Filename, string(data), failFlag, ttlSec,
).Result()
return err
}
func (s *Store) TryCompleteJob(ctx context.Context, jobID string) error {
meta, err := s.client.HGetAll(ctx, metaKey(jobID)).Result()
if err != nil {
return err
}
_, err = s.maybeCompleteJob(ctx, jobID, meta)
return err
}
func (s *Store) maybeCompleteJob(ctx context.Context, jobID string, meta map[string]string) (string, error) {
status := meta["status"]
if status == domain.JobCancelled || status == domain.JobCompleted {
return status, nil
}
total, _ := strconv.Atoi(meta["total"])
processed, _ := strconv.Atoi(meta["processed"])
if processed >= total && total > 0 {
if err := s.client.HSet(ctx, metaKey(jobID), "status", domain.JobCompleted).Err(); err != nil {
return status, err
}
return domain.JobCompleted, nil
}
return status, nil
}
func (s *Store) GetJob(ctx context.Context, jobID string) (domain.JobStatus, error) {
meta, err := s.client.HGetAll(ctx, metaKey(jobID)).Result()
if err != nil {
return domain.JobStatus{}, err
}
if len(meta) == 0 {
return domain.JobStatus{}, fmt.Errorf("job not found")
}
status, err := s.maybeCompleteJob(ctx, jobID, meta)
if err != nil {
return domain.JobStatus{}, err
}
if status == domain.JobCompleted {
meta["status"] = domain.JobCompleted
}
total, _ := strconv.Atoi(meta["total"])
processed, _ := strconv.Atoi(meta["processed"])
failed, _ := strconv.Atoi(meta["failed"])
return domain.JobStatus{
JobID: jobID,
Status: status,
Total: total,
Processed: processed,
Failed: failed,
}, nil
}
func (s *Store) GetResultsPage(ctx context.Context, jobID, cursor string, limit int) (domain.JobResultsPage, error) {
if limit <= 0 || limit > 1000 {
limit = 500
}
allFiles, err := s.client.LRange(ctx, filesKey(jobID), 0, -1).Result()
if err != nil {
return domain.JobResultsPage{}, err
}
start := 0
if cursor != "" {
for i, f := range allFiles {
if f == cursor {
start = i + 1
break
}
}
}
end := start + limit
if end > len(allFiles) {
end = len(allFiles)
}
pageFiles := allFiles[start:end]
results := make([]domain.PDFResult, 0, len(pageFiles))
if len(pageFiles) > 0 {
raw, err := s.client.HMGet(ctx, resultsKey(jobID), pageFiles...).Result()
if err != nil {
return domain.JobResultsPage{}, err
}
for i, v := range raw {
if v == nil {
results = append(results, domain.PDFResult{
Filename: pageFiles[i],
Status: "pending",
})
continue
}
var r domain.PDFResult
if err := json.Unmarshal([]byte(v.(string)), &r); err != nil {
continue
}
results = append(results, r)
}
}
page := domain.JobResultsPage{
JobID: jobID,
Results: results,
HasMore: end < len(allFiles),
}
if page.HasMore && len(pageFiles) > 0 {
page.NextCursor = pageFiles[len(pageFiles)-1]
}
return page, nil
}
func (s *Store) CancelJob(ctx context.Context, jobID string) error {
return s.client.HSet(ctx, metaKey(jobID), "status", domain.JobCancelled).Err()
}
func (s *Store) IsCancelled(ctx context.Context, jobID string) (bool, error) {
status, err := s.client.HGet(ctx, metaKey(jobID), "status").Result()
if err != nil {
return false, err
}
return status == domain.JobCancelled, nil
}
func (s *Store) PendingCount(ctx context.Context) (int64, error) {
// Approximate: sum of (total - processed) for running jobs
var cursor uint64
var pending int64
for {
keys, next, err := s.client.Scan(ctx, cursor, "job:*:meta", 100).Result()
if err != nil {
return pending, err
}
for _, key := range keys {
meta, err := s.client.HGetAll(ctx, key).Result()
if err != nil {
continue
}
if meta["status"] != domain.JobQueued && meta["status"] != domain.JobRunning {
continue
}
total, _ := strconv.Atoi(meta["total"])
processed, _ := strconv.Atoi(meta["processed"])
pending += int64(total - processed)
}
cursor = next
if cursor == 0 {
break
}
}
return pending, nil
}
func (s *Store) IncrementTotal(ctx context.Context, jobID string, delta int) error {
pipe := s.client.Pipeline()
pipe.HIncrBy(ctx, metaKey(jobID), "total", int64(delta))
pipe.Expire(ctx, metaKey(jobID), s.ttl)
_, err := pipe.Exec(ctx)
return err
}
func boolStr(v bool) string {
if v {
return "true"
}
return "false"
}
@@ -1,212 +0,0 @@
package worker
import (
"context"
"log"
"os/signal"
"sync"
"syscall"
"time"
"github.com/nats-io/nats.go/jetstream"
"github.com/pys/pdf-count/internal/config"
"github.com/pys/pdf-count/internal/domain"
"github.com/pys/pdf-count/internal/model"
"github.com/pys/pdf-count/internal/ocr"
"github.com/pys/pdf-count/internal/pipeline"
"github.com/pys/pdf-count/internal/queue"
"github.com/pys/pdf-count/internal/storage"
"github.com/pys/pdf-count/internal/store"
)
type Runner struct {
cfg config.Config
predictor *model.Predictor
ocrEngine *ocr.Engine
store *store.Store
storage *storage.Client
consumer jetstream.Consumer
ncClose func()
}
func NewRunner(cfg config.Config) (*Runner, error) {
predictor, err := model.NewPredictor(cfg.ModelPath, cfg.InferencePath)
if err != nil {
return nil, err
}
poolSize := cfg.OCRPoolSize
if poolSize <= 0 {
poolSize = cfg.WorkerConcurrency
}
ocrEngine, err := ocr.NewEngine(cfg.OCRLang, poolSize)
if err != nil {
predictor.Close()
return nil, err
}
st, err := store.New(cfg.RedisURL, cfg.JobResultTTL)
if err != nil {
ocrEngine.Close()
predictor.Close()
return nil, err
}
stg, err := storage.New(cfg.MinIOEndpoint, cfg.MinIOAccessKey, cfg.MinIOSecretKey, cfg.MinIOBucket, cfg.MinIOUseSSL)
if err != nil {
ocrEngine.Close()
st.Close()
predictor.Close()
return nil, err
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
if err := stg.EnsureBucket(ctx); err != nil {
ocrEngine.Close()
st.Close()
predictor.Close()
return nil, err
}
_, cons, nc, err := queue.NewConsumer(cfg.NATSURL, cfg.StreamName, cfg.SubjectPrefix, "pdf-workers")
if err != nil {
ocrEngine.Close()
st.Close()
predictor.Close()
return nil, err
}
return &Runner{
cfg: cfg,
predictor: predictor,
ocrEngine: ocrEngine,
store: st,
storage: stg,
consumer: cons,
ncClose: func() { nc.Close() },
}, nil
}
func (r *Runner) Close() {
r.ncClose()
_ = r.ocrEngine.Close()
_ = r.store.Close()
_ = r.predictor.Close()
}
func (r *Runner) processSlots() int {
c := r.cfg.WorkerConcurrency
pool := r.cfg.OCRPoolSize
if pool <= 0 {
pool = r.cfg.WorkerConcurrency
}
if pool <= 0 {
pool = 4
}
if c > pool {
c = pool
}
if c <= 0 {
c = 4
}
return c
}
func (r *Runner) Run(ctx context.Context) error {
slots := r.processSlots()
sem := make(chan struct{}, slots)
var wg sync.WaitGroup
go func() {
<-ctx.Done()
wg.Wait()
}()
for {
select {
case <-ctx.Done():
return nil
default:
}
msgs, err := r.consumer.Fetch(slots, jetstream.FetchMaxWait(r.cfg.NATSFetchMaxWait))
if err != nil {
continue
}
for msg := range msgs.Messages() {
sem <- struct{}{}
wg.Add(1)
m := msg
go func() {
defer func() {
<-sem
wg.Done()
}()
r.handleMessage(ctx, m)
}()
}
_ = msgs.Error()
}
}
func (r *Runner) handleMessage(ctx context.Context, msg jetstream.Msg) {
processCtx, cancel := context.WithTimeout(ctx, r.cfg.PDFProcessTimeout)
defer cancel()
qmsg, err := queue.DecodeMessage(msg.Data())
if err != nil {
_ = msg.Term()
return
}
cancelled, err := r.store.IsCancelled(processCtx, qmsg.JobID)
if err != nil || cancelled {
_ = msg.Ack()
return
}
data, err := r.storage.Get(processCtx, qmsg.ObjectKey)
if err != nil {
_ = r.store.SaveResult(processCtx, qmsg.JobID, domain.PDFResult{
Filename: qmsg.Filename,
Status: domain.StatusError,
Error: err.Error(),
})
_ = r.store.TryCompleteJob(processCtx, qmsg.JobID)
_ = msg.Ack()
return
}
proc := pipeline.NewProcessor(r.predictor, r.ocrEngine, pipeline.Options{
DPI: r.cfg.OCRDPI,
Lang: r.cfg.OCRLang,
DBSCANEps: r.cfg.DBSCANEps,
DBSCANMinSamples: r.cfg.DBSCANMinSamples,
IncludeRejected: qmsg.IncludeRejected,
LogPipelineTiming: r.cfg.LogPipelineTiming,
})
result := proc.ProcessPDF(qmsg.Filename, data)
data = nil
if err := r.store.SaveResult(processCtx, qmsg.JobID, result); err != nil {
_ = msg.NakWithDelay(5 * time.Second)
return
}
_ = r.store.TryCompleteJob(processCtx, qmsg.JobID)
_ = msg.Ack()
}
func RunMain() {
cfg := config.Load()
runner, err := NewRunner(cfg)
if err != nil {
log.Fatalf("worker init: %v", err)
}
defer runner.Close()
ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
defer cancel()
log.Printf("worker started slots=%d", runner.processSlots())
if err := runner.Run(ctx); err != nil {
log.Fatalf("worker: %v", err)
}
log.Println("worker stopped")
}