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
%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()) + } +}