fix: add support for bulk job enqueue; --bulk-from-file

Author: LinceMathew

client.go

@@ -3,15 +3,34 @@ package main
 import (
 	"context"
 	"encoding/json"
+	"os"
 
 	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgxpool"
 	"github.com/riverqueue/river"
 	"github.com/riverqueue/river/riverdriver/riverpgxv5"
 	"github.com/rs/zerolog/log"
 )
 
+type BulkJob struct {
+	Args     DPromptsJobArgs        `json:"args"`
+	Metadata map[string]interface{} `json:"metadata,omitempty"`
+}
+
+// RunClient enqueues a job with args and metadata as JSON strings.
+func RunClient(ctx context.Context, driver *riverpgxv5.Driver, argsJSON string, metadataJSON string, bulkFile string, dbPool *pgxpool.Pool) {
+	riverClient, err := newRiverClient(driver)
+	if err != nil {
+		log.Fatal().Err(err).Msg("Failed to create River client")
+	}
+
+	if bulkFile != "" {
+		if err := enqueueBulkJobsFromFile(ctx, riverClient, dbPool, bulkFile); err != nil {
+			log.Fatal().Err(err).Msg("Bulk insert failed")
+		}
+		return
+	}
 
-func RunClient(ctx context.Context, driver *riverpgxv5.Driver, argsJSON string, metadataJSON string) {
 	if argsJSON == "" {
 		log.Fatal().Msg("Args JSON is required in client mode")
 	}
@@ -36,17 +55,59 @@ func RunClient(ctx context.Context, driver *riverpgxv5.Driver, argsJSON string,
 		}
 	}
 
-	riverClient, err := newRiverClient(driver)
-	if err != nil {
-		log.Fatal().Err(err).Msg("Failed to create River client")
-	}
-
 	if _, err := riverClient.Insert(ctx, &args, insertOpts); err != nil {
 		log.Fatal().Err(err).Msg("Failed to enqueue job")
 	}
 	log.Info().Interface("args", args).Interface("metadata", insertOpts).Msg("Enqueued job")
 }
 
+func enqueueBulkJobsFromFile(ctx context.Context, riverClient *river.Client[pgx.Tx], dbPool *pgxpool.Pool, filename string) error {
+	file, err := os.Open(filename)
+	if err != nil {
+		return err
+	}
+	defer file.Close()
+
+	var jobs []BulkJob
+	if err := json.NewDecoder(file).Decode(&jobs); err != nil {
+		return err
+	}
+
+	tx, err := dbPool.Begin(ctx)
+	if err != nil {
+		return err
+	}
+	defer tx.Rollback(ctx)
+
+	var jobsToInsert []river.InsertManyParams
+	for _, job := range jobs {
+		var insertOpts *river.InsertOpts
+		if job.Metadata != nil {
+			metadataBytes, err := json.Marshal(job.Metadata)
+			if err != nil {
+				return err
+			}
+			insertOpts = &river.InsertOpts{Metadata: metadataBytes}
+		}
+		jobsToInsert = append(jobsToInsert, river.InsertManyParams{
+			Args:       job.Args,
+			InsertOpts: insertOpts,
+		})
+	}
+
+	results, err := riverClient.InsertManyTx(ctx, tx, jobsToInsert)
+	if err != nil {
+		return err
+	}
+
+	if err := tx.Commit(ctx); err != nil {
+		return err
+	}
+
+	log.Info().Msgf("Successfully enqueued %d jobs", len(results))
+	return nil
+}
+
 func newRiverClient(driver *riverpgxv5.Driver) (*river.Client[pgx.Tx], error) {
 	return river.NewClient[pgx.Tx](driver, &river.Config{})
 }

main.go

@@ -17,6 +17,7 @@ func main() {
 	mode := flag.String("mode", "worker", "Mode: 'client' to enqueue a job, 'worker' to run worker")
 	argsJSON := flag.String("args", "", "Job args as JSON (for client mode)")
 	metadataJSON := flag.String("metadata", "", "Job metadata as JSON (for client mode)")
+	bulkFile := flag.String("bulk-from-file", "", "Bulk insert jobs from JSON file")
 	configPath := flag.String("config", "", "Path to config file (default: $HOME/.dprompt.toml)")
 
 	flag.Parse()
@@ -35,7 +36,7 @@ func main() {
 
 	switch *mode {
 	case "client":
-		RunClient(ctx, driver, *argsJSON, *metadataJSON)
+		RunClient(ctx, driver, *argsJSON, *metadataJSON, *bulkFile, dbPool)
 	case "worker":
 		RunWorker(ctx, driver, cancel, dbPool)
 	default: