feat(handler): Phase 3E — Stage handlers (fetch/render/upload) + Job/Handler refactor

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

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
+}

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<h1>Hello</h1>"
+	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)
+	}
+}

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