diff --git a/cmd/worker/main.go b/cmd/worker/main.go index cdce4fe..f1c6532 100644 --- a/cmd/worker/main.go +++ b/cmd/worker/main.go @@ -2,6 +2,7 @@ package main import ( "context" + "encoding/base64" "flag" "log/slog" "os" @@ -11,7 +12,10 @@ import ( "github.com/jackc/pgx/v5/pgxpool" goredis "github.com/redis/go-redis/v9" + "golang.org/x/oauth2" + "golang.org/x/oauth2/google" + "github.com/smallchungus/disttaskqueue/internal/handler" "github.com/smallchungus/disttaskqueue/internal/queue" "github.com/smallchungus/disttaskqueue/internal/store" "github.com/smallchungus/disttaskqueue/internal/worker" @@ -55,12 +59,44 @@ func main() { redis := goredis.NewClient(opts) defer func() { _ = redis.Close() }() + var h worker.Handler + switch *stage { + case "test": + h = worker.NoopHandler{} + case "fetch": + cfg, key := loadGoogleCfgAndKey() + h = handler.NewFetchHandler(handler.FetchConfig{ + Store: store.New(pool), + EncryptionKey: key, + OAuth2: cfg, + DataDir: envOr("DATA_DIR", "/data"), + }) + case "render": + h = handler.NewRenderHandler(handler.RenderConfig{ + DataDir: envOr("DATA_DIR", "/data"), + PDFEndpoint: envOr("GOTENBERG_URL", "http://gotenberg:3000"), + }) + case "upload": + cfg, key := loadGoogleCfgAndKey() + h = handler.NewUploadHandler(handler.UploadConfig{ + Store: store.New(pool), + EncryptionKey: key, + OAuth2: cfg, + Redis: redis, + DataDir: envOr("DATA_DIR", "/data"), + RootFolderID: envOr("DRIVE_ROOT_FOLDER_ID", ""), + }) + default: + slog.Error("unknown stage", "stage", *stage) + os.Exit(2) + } + w := worker.New(worker.Config{ Stage: *stage, WorkerID: worker.NewWorkerID(), Store: store.New(pool), Queue: queue.New(redis), - Handler: worker.NoopHandler{}, + Handler: h, PopTimeout: 30 * time.Second, HeartbeatTTL: 15 * time.Second, HeartbeatInterval: 5 * time.Second, @@ -74,6 +110,36 @@ func main() { slog.Info("worker stopped") } +func loadGoogleCfgAndKey() (*oauth2.Config, []byte) { + clientID := mustEnv("GOOGLE_OAUTH_CLIENT_ID") + clientSecret := mustEnv("GOOGLE_OAUTH_CLIENT_SECRET") + keyB64 := mustEnv("TOKEN_ENCRYPTION_KEY") + key, err := base64.StdEncoding.DecodeString(keyB64) + if err != nil { + slog.Error("decode TOKEN_ENCRYPTION_KEY", "err", err) + os.Exit(1) + } + if len(key) != 32 { + slog.Error("TOKEN_ENCRYPTION_KEY must decode to 32 bytes", "got", len(key)) + os.Exit(1) + } + return &oauth2.Config{ + ClientID: clientID, + ClientSecret: clientSecret, + Endpoint: google.Endpoint, + Scopes: []string{"https://www.googleapis.com/auth/gmail.readonly", "https://www.googleapis.com/auth/drive.file"}, + }, key +} + +func mustEnv(k string) string { + v := os.Getenv(k) + if v == "" { + slog.Error("missing env var", "key", k) + os.Exit(1) + } + return v +} + func envOr(k, def string) string { if v := os.Getenv(k); v != "" { return v diff --git a/internal/handler/fetch.go b/internal/handler/fetch.go new file mode 100644 index 0000000..af04a5f --- /dev/null +++ b/internal/handler/fetch.go @@ -0,0 +1,66 @@ +package handler + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + + "golang.org/x/oauth2" + + "github.com/smallchungus/disttaskqueue/internal/gmail" + "github.com/smallchungus/disttaskqueue/internal/store" +) + +type FetchConfig struct { + Store *store.Store + EncryptionKey []byte + OAuth2 *oauth2.Config + GmailEndpoint string + DataDir string +} + +type FetchHandler struct { + cfg FetchConfig +} + +func NewFetchHandler(cfg FetchConfig) *FetchHandler { + return &FetchHandler{cfg: cfg} +} + +func (h *FetchHandler) Process(ctx context.Context, job store.Job) (string, error) { + if job.UserID == nil { + return "", errors.New("fetch: job missing user_id") + } + if job.GmailMessageID == nil { + return "", errors.New("fetch: job missing gmail_message_id") + } + + client, err := gmail.New(ctx, gmail.Config{ + Store: h.cfg.Store, + UserID: *job.UserID, + EncryptionKey: h.cfg.EncryptionKey, + OAuth2: h.cfg.OAuth2, + Endpoint: h.cfg.GmailEndpoint, + }) + if err != nil { + return "", fmt.Errorf("gmail client: %w", err) + } + + raw, err := client.FetchMessage(ctx, *job.GmailMessageID) + if err != nil { + return "", fmt.Errorf("fetch: %w", err) + } + + dir := filepath.Join(h.cfg.DataDir, "mime") + if err := os.MkdirAll(dir, 0o750); err != nil { + return "", fmt.Errorf("mkdir: %w", err) + } + path := filepath.Join(dir, job.ID.String()+".eml") + if err := os.WriteFile(path, raw, 0o600); err != nil { + return "", fmt.Errorf("write mime: %w", err) + } + + return "render", nil +} diff --git a/internal/handler/fetch_integration_test.go b/internal/handler/fetch_integration_test.go new file mode 100644 index 0000000..bfa1a16 --- /dev/null +++ b/internal/handler/fetch_integration_test.go @@ -0,0 +1,206 @@ +//go:build integration + +package handler_test + +import ( + "context" + "encoding/base64" + "encoding/json" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "testing" + "time" + + "github.com/google/uuid" + "golang.org/x/oauth2" + + "github.com/smallchungus/disttaskqueue/internal/handler" + "github.com/smallchungus/disttaskqueue/internal/oauth" + "github.com/smallchungus/disttaskqueue/internal/store" + "github.com/smallchungus/disttaskqueue/internal/testutil" +) + +func TestRenderHandler_WritesPDFAndReturnsUpload(t *testing.T) { + const fakeMime = "From: a@b.com\r\nSubject: hi\r\nContent-Type: text/html\r\n\r\n

Hello

" + const fakePDF = "%PDF-1.7\nfake\n%%EOF" + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/pdf") + _, _ = w.Write([]byte(fakePDF)) + })) + defer srv.Close() + + dataDir := t.TempDir() + jobID := uuid.New() + + mimeDir := filepath.Join(dataDir, "mime") + if err := os.MkdirAll(mimeDir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(mimeDir, jobID.String()+".eml"), []byte(fakeMime), 0o600); err != nil { + t.Fatal(err) + } + + h := handler.NewRenderHandler(handler.RenderConfig{ + DataDir: dataDir, + PDFEndpoint: srv.URL, + }) + + job := store.Job{ID: jobID, Stage: "render"} + next, err := h.Process(context.Background(), job) + if err != nil { + t.Fatalf("process: %v", err) + } + if next != "upload" { + t.Fatalf("next: %q, want upload", next) + } + + written, err := os.ReadFile(filepath.Join(dataDir, "pdf", jobID.String()+".pdf")) + if err != nil { + t.Fatalf("read pdf: %v", err) + } + if string(written) != fakePDF { + t.Fatalf("pdf: got %q, want %q", written, fakePDF) + } +} + +func newKey32() []byte { + k := make([]byte, 32) + for i := range k { + k[i] = byte(i) + } + return k +} + +func saveFakeToken(t *testing.T, s *store.Store, uid uuid.UUID, key []byte) { + t.Helper() + tok := &oauth2.Token{AccessToken: "x", RefreshToken: "y", Expiry: time.Now().Add(time.Hour)} + if err := oauth.SaveToken(context.Background(), s, uid, key, "google", tok); err != nil { + t.Fatal(err) + } +} + +func encodeURL(s string) string { + return base64.URLEncoding.WithPadding(base64.NoPadding).EncodeToString([]byte(s)) +} + +func TestFetchHandler_WritesMimeAndReturnsRender(t *testing.T) { + const rawMime = "From: alice@example.com\r\nSubject: hi\r\n\r\nbody" + + encB64 := encodeURL(rawMime) + msgJSON, _ := json.Marshal(map[string]any{"raw": encB64}) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write(msgJSON) + })) + defer srv.Close() + + pool := testutil.StartPostgres(t) + if err := store.Migrate(context.Background(), pool.Config().ConnString()); err != nil { + t.Fatal(err) + } + s := store.New(pool) + u, _ := s.CreateUser(context.Background(), fmt.Sprintf("u+%d@example.com", time.Now().UnixNano())) + key := newKey32() + saveFakeToken(t, s, u.ID, key) + + dataDir := t.TempDir() + + h := handler.NewFetchHandler(handler.FetchConfig{ + Store: s, + EncryptionKey: key, + OAuth2: &oauth2.Config{ClientID: "x", ClientSecret: "y"}, + GmailEndpoint: srv.URL, + DataDir: dataDir, + }) + + gmailMsgID := "m1" + job, _ := s.EnqueueJob(context.Background(), store.NewJob{ + Stage: "fetch", + UserID: &u.ID, + GmailMessageID: &gmailMsgID, + }) + _ = s.ClaimJob(context.Background(), job.ID, "w1") + + job, _ = s.GetJob(context.Background(), job.ID) + next, err := h.Process(context.Background(), job) + if err != nil { + t.Fatalf("process: %v", err) + } + if next != "render" { + t.Fatalf("next: %q, want render", next) + } + + written, err := os.ReadFile(filepath.Join(dataDir, "mime", job.ID.String()+".eml")) + if err != nil { + t.Fatalf("read: %v", err) + } + if string(written) != rawMime { + t.Fatalf("mime: %q, want %q", written, rawMime) + } +} + +func TestUploadHandler_UploadsToDateTreeAndReturnsTerminal(t *testing.T) { + uploadJSON, _ := json.Marshal(map[string]any{"id": "uploaded-pdf-id"}) + listJSON, _ := json.Marshal(map[string]any{"files": []map[string]any{}}) + createJSON, _ := json.Marshal(map[string]any{"id": "folder-id"}) + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + switch { + case r.URL.Path == "/files" && r.Method == http.MethodGet: + _, _ = w.Write(listJSON) + case r.URL.Path == "/files" && r.Method == http.MethodPost: + _, _ = w.Write(createJSON) + case r.URL.Path == "/upload/drive/v3/files": + _, _ = w.Write(uploadJSON) + default: + http.NotFound(w, r) + } + })) + defer srv.Close() + + pool := testutil.StartPostgres(t) + if err := store.Migrate(context.Background(), pool.Config().ConnString()); err != nil { + t.Fatal(err) + } + s := store.New(pool) + u, _ := s.CreateUser(context.Background(), fmt.Sprintf("up+%d@example.com", time.Now().UnixNano())) + key := newKey32() + saveFakeToken(t, s, u.ID, key) + + redis := testutil.StartRedis(t) + + dataDir := t.TempDir() + jobID := uuid.New() + pdfDir := filepath.Join(dataDir, "pdf") + if err := os.MkdirAll(pdfDir, 0o755); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(pdfDir, jobID.String()+".pdf"), []byte("PDF-DATA"), 0o600); err != nil { + t.Fatal(err) + } + + h := handler.NewUploadHandler(handler.UploadConfig{ + Store: s, + EncryptionKey: key, + OAuth2: &oauth2.Config{ClientID: "x", ClientSecret: "y"}, + DriveEndpoint: srv.URL, + Redis: redis, + DataDir: dataDir, + RootFolderID: "root-folder-id", + }) + + job := store.Job{ID: jobID, UserID: &u.ID, Stage: "upload", CreatedAt: time.Date(2026, 4, 17, 0, 0, 0, 0, time.UTC)} + next, err := h.Process(context.Background(), job) + if err != nil { + t.Fatalf("process: %v", err) + } + if next != "" { + t.Fatalf("next: %q, want '' (terminal)", next) + } +} diff --git a/internal/handler/paths.go b/internal/handler/paths.go new file mode 100644 index 0000000..affe7ea --- /dev/null +++ b/internal/handler/paths.go @@ -0,0 +1,16 @@ +package handler + +import ( + "fmt" + "time" +) + +// DateTreeFolders returns the year/month/day folder names for the given time. +func DateTreeFolders(t time.Time) []string { + t = t.UTC() + return []string{ + fmt.Sprintf("%04d", t.Year()), + fmt.Sprintf("%02d", int(t.Month())), + fmt.Sprintf("%02d", t.Day()), + } +} diff --git a/internal/handler/paths_test.go b/internal/handler/paths_test.go new file mode 100644 index 0000000..b3c2234 --- /dev/null +++ b/internal/handler/paths_test.go @@ -0,0 +1,15 @@ +package handler + +import ( + "reflect" + "testing" + "time" +) + +func TestDateTreeFolders(t *testing.T) { + got := DateTreeFolders(time.Date(2026, 4, 17, 10, 30, 0, 0, time.UTC)) + want := []string{"2026", "04", "17"} + if !reflect.DeepEqual(got, want) { + t.Fatalf("got %v, want %v", got, want) + } +} diff --git a/internal/handler/render.go b/internal/handler/render.go new file mode 100644 index 0000000..c7fb49b --- /dev/null +++ b/internal/handler/render.go @@ -0,0 +1,112 @@ +package handler + +import ( + "bytes" + "context" + "errors" + "fmt" + "net/mail" + "os" + "path/filepath" + "strings" + + "github.com/smallchungus/disttaskqueue/internal/pdf" + "github.com/smallchungus/disttaskqueue/internal/store" +) + +type RenderConfig struct { + DataDir string + PDFEndpoint string +} + +type RenderHandler struct { + cfg RenderConfig + client *pdf.Client +} + +func NewRenderHandler(cfg RenderConfig) *RenderHandler { + return &RenderHandler{cfg: cfg, client: pdf.New(cfg.PDFEndpoint)} +} + +func (h *RenderHandler) Process(ctx context.Context, job store.Job) (string, error) { + mimePath := filepath.Join(h.cfg.DataDir, "mime", job.ID.String()+".eml") + rawMime, err := os.ReadFile(mimePath) //nolint:gosec // path derived from trusted job ID, not user input + if err != nil { + return "", fmt.Errorf("read mime: %w", err) + } + + html, err := htmlFromMime(rawMime) + if err != nil { + return "", fmt.Errorf("parse mime: %w", err) + } + + pdfBytes, err := h.client.RenderHTML(ctx, html) + if err != nil { + return "", fmt.Errorf("render: %w", err) + } + + pdfDir := filepath.Join(h.cfg.DataDir, "pdf") + if err := os.MkdirAll(pdfDir, 0o750); err != nil { + return "", fmt.Errorf("mkdir: %w", err) + } + pdfPath := filepath.Join(pdfDir, job.ID.String()+".pdf") + if err := os.WriteFile(pdfPath, pdfBytes, 0o600); err != nil { + return "", fmt.Errorf("write pdf: %w", err) + } + + return "upload", nil +} + +func htmlFromMime(raw []byte) ([]byte, error) { + msg, err := mail.ReadMessage(bytes.NewReader(raw)) + if err != nil { + return nil, fmt.Errorf("read message: %w", err) + } + subject := msg.Header.Get("Subject") + from := msg.Header.Get("From") + date := msg.Header.Get("Date") + + body := &bytes.Buffer{} + if _, err := body.ReadFrom(msg.Body); err != nil { + return nil, fmt.Errorf("read body: %w", err) + } + + if subject == "" && from == "" && body.Len() == 0 { + return nil, errors.New("empty message") + } + + out := &bytes.Buffer{} + fmt.Fprintf(out, "%s", htmlEscape(subject)) + fmt.Fprintf(out, "
") + fmt.Fprintf(out, "
From: %s
", htmlEscape(from)) + fmt.Fprintf(out, "
Date: %s
", htmlEscape(date)) + fmt.Fprintf(out, "
Subject: %s
", htmlEscape(subject)) + fmt.Fprintf(out, "
") + contentType := msg.Header.Get("Content-Type") + if contentType == "" || strings.HasPrefix(contentType, "text/html") { + out.Write(body.Bytes()) + } else { + fmt.Fprintf(out, "
%s
", htmlEscape(body.String())) + } + out.WriteString("") + return out.Bytes(), nil +} + +func htmlEscape(s string) string { + r := bytes.NewBuffer(nil) + for _, c := range []byte(s) { + switch c { + case '<': + r.WriteString("<") + case '>': + r.WriteString(">") + case '&': + r.WriteString("&") + case '"': + r.WriteString(""") + default: + r.WriteByte(c) + } + } + return r.String() +} diff --git a/internal/handler/upload.go b/internal/handler/upload.go new file mode 100644 index 0000000..e82c456 --- /dev/null +++ b/internal/handler/upload.go @@ -0,0 +1,88 @@ +package handler + +import ( + "context" + "errors" + "fmt" + "os" + "path" + "path/filepath" + "time" + + goredis "github.com/redis/go-redis/v9" + "golang.org/x/oauth2" + + "github.com/smallchungus/disttaskqueue/internal/drive" + "github.com/smallchungus/disttaskqueue/internal/store" +) + +type UploadConfig struct { + Store *store.Store + EncryptionKey []byte + OAuth2 *oauth2.Config + DriveEndpoint string + Redis *goredis.Client + DataDir string + RootFolderID string +} + +type UploadHandler struct { + cfg UploadConfig + cache *drive.FolderCache +} + +func NewUploadHandler(cfg UploadConfig) *UploadHandler { + return &UploadHandler{cfg: cfg, cache: drive.NewFolderCache(cfg.Redis)} +} + +func (h *UploadHandler) Process(ctx context.Context, job store.Job) (string, error) { + if job.UserID == nil { + return "", errors.New("upload: job missing user_id") + } + + client, err := drive.New(ctx, drive.Config{ + Store: h.cfg.Store, + UserID: *job.UserID, + EncryptionKey: h.cfg.EncryptionKey, + OAuth2: h.cfg.OAuth2, + Endpoint: h.cfg.DriveEndpoint, + }) + if err != nil { + return "", fmt.Errorf("drive client: %w", err) + } + + folders := DateTreeFolders(job.CreatedAt) + parent := h.cfg.RootFolderID + pathSoFar := "" + for _, f := range folders { + if pathSoFar == "" { + pathSoFar = f + } else { + pathSoFar = path.Join(pathSoFar, f) + } + + if cached, ok, err := h.cache.Get(ctx, *job.UserID, pathSoFar); err == nil && ok { + parent = cached + continue + } + + folderID, err := client.EnsureFolder(ctx, parent, f) + if err != nil { + return "", fmt.Errorf("ensure folder %s: %w", f, err) + } + _ = h.cache.Set(ctx, *job.UserID, pathSoFar, folderID, 24*time.Hour) + parent = folderID + } + + pdfBytes, err := os.ReadFile(filepath.Join(h.cfg.DataDir, "pdf", job.ID.String()+".pdf")) + if err != nil { + return "", fmt.Errorf("read pdf: %w", err) + } + + name := job.ID.String() + ".pdf" + if _, err := client.Upload(ctx, parent, name, "application/pdf", pdfBytes); err != nil { + return "", fmt.Errorf("upload: %w", err) + } + + return "", nil +} diff --git a/internal/store/store.go b/internal/store/store.go index 9f1b9fa..80c2ba2 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -30,15 +30,17 @@ func (s *Store) EnqueueJob(ctx context.Context, nj NewJob) (Job, error) { } const q = ` - INSERT INTO pipeline_jobs (id, stage, status, payload) - VALUES ($1, $2, $3, $4) - RETURNING id, stage, status, payload, worker_id, attempts, max_attempts, - last_error, next_run_at, claimed_at, completed_at, created_at, updated_at` + INSERT INTO pipeline_jobs (id, user_id, gmail_message_id, stage, status, payload, is_synthetic) + VALUES ($1, $2, $3, $4, $5, $6, $7) + RETURNING id, user_id, gmail_message_id, stage, status, payload, is_synthetic, + worker_id, attempts, max_attempts, last_error, next_run_at, + claimed_at, completed_at, created_at, updated_at` - row := s.pool.QueryRow(ctx, q, id, nj.Stage, StatusQueued, payload) + row := s.pool.QueryRow(ctx, q, id, nj.UserID, nj.GmailMessageID, nj.Stage, StatusQueued, payload, nj.IsSynthetic) var j Job - if err := row.Scan(&j.ID, &j.Stage, &j.Status, &j.Payload, &j.WorkerID, &j.Attempts, - &j.MaxAttempts, &j.LastError, &j.NextRunAt, &j.ClaimedAt, &j.CompletedAt, &j.CreatedAt, &j.UpdatedAt); err != nil { + if err := row.Scan(&j.ID, &j.UserID, &j.GmailMessageID, &j.Stage, &j.Status, &j.Payload, &j.IsSynthetic, + &j.WorkerID, &j.Attempts, &j.MaxAttempts, &j.LastError, &j.NextRunAt, + &j.ClaimedAt, &j.CompletedAt, &j.CreatedAt, &j.UpdatedAt); err != nil { return Job{}, fmt.Errorf("enqueue: %w", err) } return j, nil @@ -103,8 +105,9 @@ func (s *Store) MarkFailed(ctx context.Context, id uuid.UUID, errMsg string, nex func (s *Store) ListRunningJobs(ctx context.Context) ([]Job, error) { const q = ` - SELECT id, stage, status, payload, worker_id, attempts, max_attempts, - last_error, next_run_at, claimed_at, completed_at, created_at, updated_at + SELECT id, user_id, gmail_message_id, stage, status, payload, is_synthetic, + worker_id, attempts, max_attempts, last_error, next_run_at, + claimed_at, completed_at, created_at, updated_at FROM pipeline_jobs WHERE status = $1` rows, err := s.pool.Query(ctx, q, StatusRunning) if err != nil { @@ -116,8 +119,9 @@ func (s *Store) ListRunningJobs(ctx context.Context) ([]Job, error) { func (s *Store) ListReadyRetryJobs(ctx context.Context) ([]Job, error) { const q = ` - SELECT id, stage, status, payload, worker_id, attempts, max_attempts, - last_error, next_run_at, claimed_at, completed_at, created_at, updated_at + SELECT id, user_id, gmail_message_id, stage, status, payload, is_synthetic, + worker_id, attempts, max_attempts, last_error, next_run_at, + claimed_at, completed_at, created_at, updated_at FROM pipeline_jobs WHERE status = $1 AND last_error IS NOT NULL AND next_run_at <= now()` rows, err := s.pool.Query(ctx, q, StatusQueued) @@ -132,8 +136,9 @@ func scanJobs(rows pgx.Rows) ([]Job, error) { var out []Job for rows.Next() { var j Job - if err := rows.Scan(&j.ID, &j.Stage, &j.Status, &j.Payload, &j.WorkerID, &j.Attempts, - &j.MaxAttempts, &j.LastError, &j.NextRunAt, &j.ClaimedAt, &j.CompletedAt, &j.CreatedAt, &j.UpdatedAt); err != nil { + if err := rows.Scan(&j.ID, &j.UserID, &j.GmailMessageID, &j.Stage, &j.Status, &j.Payload, &j.IsSynthetic, + &j.WorkerID, &j.Attempts, &j.MaxAttempts, &j.LastError, &j.NextRunAt, + &j.ClaimedAt, &j.CompletedAt, &j.CreatedAt, &j.UpdatedAt); err != nil { return nil, fmt.Errorf("scan: %w", err) } out = append(out, j) @@ -144,16 +149,34 @@ func scanJobs(rows pgx.Rows) ([]Job, error) { return out, nil } +func (s *Store) AdvanceJob(ctx context.Context, id uuid.UUID, nextStage string) error { + const q = ` + UPDATE pipeline_jobs + SET stage = $1, status = $2, worker_id = NULL, claimed_at = NULL, updated_at = now() + WHERE id = $3 AND status = $4` + + tag, err := s.pool.Exec(ctx, q, nextStage, StatusQueued, id, StatusRunning) + if err != nil { + return fmt.Errorf("advance: %w", err) + } + if tag.RowsAffected() == 0 { + return ErrJobNotFound + } + return nil +} + func (s *Store) GetJob(ctx context.Context, id uuid.UUID) (Job, error) { const q = ` - SELECT id, stage, status, payload, worker_id, attempts, max_attempts, - last_error, next_run_at, claimed_at, completed_at, created_at, updated_at + SELECT id, user_id, gmail_message_id, stage, status, payload, is_synthetic, + worker_id, attempts, max_attempts, last_error, next_run_at, + claimed_at, completed_at, created_at, updated_at FROM pipeline_jobs WHERE id = $1` row := s.pool.QueryRow(ctx, q, id) var j Job - err := row.Scan(&j.ID, &j.Stage, &j.Status, &j.Payload, &j.WorkerID, &j.Attempts, - &j.MaxAttempts, &j.LastError, &j.NextRunAt, &j.ClaimedAt, &j.CompletedAt, &j.CreatedAt, &j.UpdatedAt) + err := row.Scan(&j.ID, &j.UserID, &j.GmailMessageID, &j.Stage, &j.Status, &j.Payload, &j.IsSynthetic, + &j.WorkerID, &j.Attempts, &j.MaxAttempts, &j.LastError, &j.NextRunAt, + &j.ClaimedAt, &j.CompletedAt, &j.CreatedAt, &j.UpdatedAt) if errors.Is(err, pgx.ErrNoRows) { return Job{}, ErrJobNotFound } diff --git a/internal/store/store_integration_test.go b/internal/store/store_integration_test.go index b52f985..c6fefb4 100644 --- a/internal/store/store_integration_test.go +++ b/internal/store/store_integration_test.go @@ -293,6 +293,28 @@ func TestGetOAuthToken_ReturnsErrOAuthTokenNotFound(t *testing.T) { } } +func TestAdvanceJob_TransitionsToNextStage(t *testing.T) { + s := newStore(t) + ctx := context.Background() + j, _ := s.EnqueueJob(ctx, store.NewJob{Stage: "fetch"}) + _ = s.ClaimJob(ctx, j.ID, "w1") + + if err := s.AdvanceJob(ctx, j.ID, "render"); err != nil { + t.Fatalf("advance: %v", err) + } + + got, _ := s.GetJob(ctx, j.ID) + if got.Stage != "render" { + t.Fatalf("stage: %s, want render", got.Stage) + } + if got.Status != store.StatusQueued { + t.Fatalf("status: %s, want queued", got.Status) + } + if got.WorkerID != nil { + t.Fatalf("worker_id: %v, want nil", got.WorkerID) + } +} + func jobIDs(js []store.Job) []string { out := make([]string, len(js)) for i, j := range js { diff --git a/internal/store/types.go b/internal/store/types.go index 50d50ea..ecb370b 100644 --- a/internal/store/types.go +++ b/internal/store/types.go @@ -17,24 +17,30 @@ const ( ) type Job struct { - ID uuid.UUID - Stage string - Status JobStatus - Payload json.RawMessage - WorkerID *string - Attempts int - MaxAttempts int - LastError *string - NextRunAt time.Time - ClaimedAt *time.Time - CompletedAt *time.Time - CreatedAt time.Time - UpdatedAt time.Time + ID uuid.UUID + UserID *uuid.UUID + GmailMessageID *string + Stage string + Status JobStatus + Payload json.RawMessage + IsSynthetic bool + WorkerID *string + Attempts int + MaxAttempts int + LastError *string + NextRunAt time.Time + ClaimedAt *time.Time + CompletedAt *time.Time + CreatedAt time.Time + UpdatedAt time.Time } type NewJob struct { - Stage string - Payload json.RawMessage + Stage string + Payload json.RawMessage + UserID *uuid.UUID + GmailMessageID *string + IsSynthetic bool } type User struct { diff --git a/internal/worker/handler.go b/internal/worker/handler.go index d23cb1b..c8f764a 100644 --- a/internal/worker/handler.go +++ b/internal/worker/handler.go @@ -9,12 +9,12 @@ import ( ) type Handler interface { - Process(ctx context.Context, job store.Job) error + Process(ctx context.Context, job store.Job) (nextStage string, err error) } type NoopHandler struct{} -func (NoopHandler) Process(ctx context.Context, job store.Job) error { +func (NoopHandler) Process(ctx context.Context, job store.Job) (string, error) { var p struct { SleepMs int `json:"sleepMs"` } @@ -22,12 +22,12 @@ func (NoopHandler) Process(ctx context.Context, job store.Job) error { _ = json.Unmarshal(job.Payload, &p) } if p.SleepMs <= 0 { - return nil + return "", nil } select { case <-ctx.Done(): - return ctx.Err() + return "", ctx.Err() case <-time.After(time.Duration(p.SleepMs) * time.Millisecond): - return nil + return "", nil } } diff --git a/internal/worker/handler_test.go b/internal/worker/handler_test.go index 2d94fb4..0ca8724 100644 --- a/internal/worker/handler_test.go +++ b/internal/worker/handler_test.go @@ -15,7 +15,7 @@ func TestNoopHandler_SleepsThenReturnsNil(t *testing.T) { job := store.Job{ID: uuid.New(), Payload: json.RawMessage(`{"sleepMs":50}`)} start := time.Now() - err := h.Process(context.Background(), job) + _, err := h.Process(context.Background(), job) elapsed := time.Since(start) if err != nil { @@ -31,7 +31,7 @@ func TestNoopHandler_DefaultsToZeroSleep(t *testing.T) { job := store.Job{ID: uuid.New(), Payload: json.RawMessage(`{}`)} start := time.Now() - if err := h.Process(context.Background(), job); err != nil { + if _, err := h.Process(context.Background(), job); err != nil { t.Fatal(err) } if time.Since(start) > 50*time.Millisecond { diff --git a/internal/worker/worker.go b/internal/worker/worker.go index a50acfb..b6fab08 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -74,7 +74,7 @@ func (w *Worker) ProcessOne(ctx context.Context) (didWork bool, err error) { hbCtx, hbCancel := context.WithCancel(ctx) go w.heartbeatLoop(hbCtx) - handlerErr := w.cfg.Handler.Process(ctx, job) + nextStage, handlerErr := w.cfg.Handler.Process(ctx, job) hbCancel() @@ -87,6 +87,17 @@ func (w *Worker) ProcessOne(ctx context.Context) (didWork bool, err error) { return true, nil } + if nextStage != "" { + if err := w.cfg.Store.AdvanceJob(ctx, jobID, nextStage); err != nil { + return true, fmt.Errorf("advance: %w", err) + } + if err := w.cfg.Queue.Push(ctx, nextStage, jobID.String()); err != nil { + return true, fmt.Errorf("push next stage: %w", err) + } + slog.Info("job advanced", "job_id", jobID, "from", job.Stage, "to", nextStage) + return true, nil + } + if err := w.cfg.Store.MarkDone(ctx, jobID); err != nil { return true, fmt.Errorf("mark done: %w", err) } diff --git a/internal/worker/worker_integration_test.go b/internal/worker/worker_integration_test.go index 999ad8d..77d4e07 100644 --- a/internal/worker/worker_integration_test.go +++ b/internal/worker/worker_integration_test.go @@ -80,14 +80,41 @@ func TestProcessOne_ReturnsNoWorkOnEmptyQueue(t *testing.T) { type errHandler struct{} -func (errHandler) Process(ctx context.Context, job store.Job) error { - return &handlerErr{"boom"} +func (errHandler) Process(ctx context.Context, job store.Job) (string, error) { + return "", &handlerErr{"boom"} } type handlerErr struct{ msg string } func (e *handlerErr) Error() string { return e.msg } +type advancingHandler struct{ next string } + +func (h advancingHandler) Process(ctx context.Context, job store.Job) (string, error) { + return h.next, nil +} + +func setupHarnessWithStage(t *testing.T, h worker.Handler, stage string) *harness { + t.Helper() + pool := testutil.StartPostgres(t) + if err := store.Migrate(context.Background(), pool.Config().ConnString()); err != nil { + t.Fatal(err) + } + s := store.New(pool) + q := queue.New(testutil.StartRedis(t)) + w := worker.New(worker.Config{ + Stage: stage, + WorkerID: "worker-test-1", + Store: s, + Queue: q, + Handler: h, + PopTimeout: 500 * time.Millisecond, + HeartbeatTTL: 5 * time.Second, + HeartbeatInterval: 1 * time.Second, + }) + return &harness{store: s, queue: q, w: w} +} + func TestProcessOne_MarksFailedOnHandlerError(t *testing.T) { h := setupHarness(t, errHandler{}) ctx := context.Background() @@ -114,3 +141,31 @@ func TestProcessOne_MarksFailedOnHandlerError(t *testing.T) { t.Fatalf("last_error: %v", got.LastError) } } + +func TestProcessOne_AdvancesToNextStage(t *testing.T) { + h := setupHarnessWithStage(t, advancingHandler{next: "render"}, "fetch") + ctx := context.Background() + job, _ := h.store.EnqueueJob(ctx, store.NewJob{Stage: "fetch"}) + _ = h.queue.Push(ctx, "fetch", job.ID.String()) + + didWork, err := h.w.ProcessOne(ctx) + if err != nil { + t.Fatalf("ProcessOne: %v", err) + } + if !didWork { + t.Fatal("expected didWork=true") + } + + got, _ := h.store.GetJob(ctx, job.ID) + if got.Stage != "render" || got.Status != store.StatusQueued { + t.Fatalf("stage/status: %s/%s, want render/queued", got.Stage, got.Status) + } + + popped, err := h.queue.BlockingPop(ctx, "render", 500*time.Millisecond) + if err != nil { + t.Fatalf("expected pushed to render: %v", err) + } + if popped != job.ID.String() { + t.Fatalf("popped %q, want %q", popped, job.ID.String()) + } +}