Batch create AI jobs efficiently

internal/store/store.go

@@ -69,16 +69,7 @@ func (s *Store) CreateJob(ctx context.Context, in model.CreateJob) (*model.Job,
 	if err := validateCreateJob(in); err != nil {
 		return nil, err
 	}
-	if in.MaxAttempts <= 0 {
-		in.MaxAttempts = 3
-	}
-	if len(in.Input) == 0 {
-		in.Input = json.RawMessage(`{}`)
-	}
-	scheduledAt := time.Now().UTC()
-	if in.ScheduledAt != nil {
-		scheduledAt = in.ScheduledAt.UTC()
-	}
+	normalizeCreateJob(&in)
 
 	const q = `
 INSERT INTO ai_jobs (
@@ -98,12 +89,81 @@ RETURNING ` + jobSelectColumns + `
 		in.Priority,
 		in.MaxAttempts,
 		in.Input,
-		scheduledAt,
+		*in.ScheduledAt,
 		in.IdempotencyKey,
 	)
 	return scanJob(row)
 }
 
+func (s *Store) CreateJobs(ctx context.Context, items []model.CreateJob) ([]*model.Job, error) {
+	if len(items) == 0 {
+		return []*model.Job{}, nil
+	}
+	const q = `
+INSERT INTO ai_jobs (
+  owner_service, owner_ref, task_type, model_profile, priority, max_attempts,
+  input, scheduled_at, idempotency_key
+)
+VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9)
+ON CONFLICT (idempotency_key) WHERE idempotency_key IS NOT NULL
+DO UPDATE SET updated_at = ai_jobs.updated_at
+RETURNING ` + jobSelectColumns + `
+`
+	var batch pgx.Batch
+	for i := range items {
+		if err := validateCreateJob(items[i]); err != nil {
+			return nil, err
+		}
+		normalizeCreateJob(&items[i])
+		batch.Queue(q,
+			items[i].OwnerService,
+			items[i].OwnerRef,
+			items[i].TaskType,
+			items[i].ModelProfile,
+			items[i].Priority,
+			items[i].MaxAttempts,
+			items[i].Input,
+			*items[i].ScheduledAt,
+			items[i].IdempotencyKey,
+		)
+	}
+	br := s.pool.SendBatch(ctx, &batch)
+	batchClosed := false
+	defer func() {
+		if !batchClosed {
+			_ = br.Close()
+		}
+	}()
+	out := make([]*model.Job, 0, len(items))
+	for range items {
+		job, err := scanJob(br.QueryRow())
+		if err != nil {
+			return nil, err
+		}
+		out = append(out, job)
+	}
+	err := br.Close()
+	batchClosed = true
+	if err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func normalizeCreateJob(in *model.CreateJob) {
+	if in.MaxAttempts <= 0 {
+		in.MaxAttempts = 3
+	}
+	if len(in.Input) == 0 {
+		in.Input = json.RawMessage(`{}`)
+	}
+	scheduledAt := time.Now().UTC()
+	if in.ScheduledAt != nil {
+		scheduledAt = in.ScheduledAt.UTC()
+	}
+	in.ScheduledAt = &scheduledAt
+}
+
 func validateCreateJob(in model.CreateJob) error {
 	switch {
 	case strings.TrimSpace(in.OwnerService) == "":