byre/internal/store/schedule.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"slices"
	"time"

	"g.fc.run/ai/byre/internal/api"
	"g.fc.run/ai/byre/internal/schedule"
)

const scheduleColumns = `id, project_id, owner, title, every, hour, minute, day_of_month, prompt, model, workflow, next_run_at`

func querySchedules(ctx context.Context, q querier, where string, args ...any) ([]api.Schedule, error) {
	rows, err := q.QueryContext(ctx, `SELECT `+scheduleColumns+` FROM schedules `+where+` ORDER BY next_run_at, id`, args...)
	if err != nil {
		return nil, err
	}
	out := []api.Schedule{}
	for rows.Next() {
		sc := api.Schedule{Weekdays: []api.Weekday{}}
		next := ""
		err = rows.Scan(&sc.Id, &sc.ProjectId, &sc.Owner, &sc.Title, &sc.Every, &sc.Hour, &sc.Minute, &sc.DayOfMonth, &sc.Prompt, &sc.Model, &sc.Workflow, &next)
		if err != nil {
			rows.Close()
			return nil, err
		}
		sc.NextRun, err = parseTime(next)
		if err != nil {
			rows.Close()
			return nil, err
		}
		out = append(out, sc)
	}
	err = errors.Join(rows.Err(), rows.Close())
	if err != nil {
		return nil, err
	}
	for i := range out {
		days, err := queryIDs(ctx, q, `SELECT weekday FROM schedule_weekdays WHERE schedule_id = ?`, out[i].Id)
		if err != nil {
			return nil, err
		}
		for _, d := range schedule.Weekdays {
			if slices.Contains(days, string(d)) {
				out[i].Weekdays = append(out[i].Weekdays, d)
			}
		}
	}
	return out, nil
}

func (s *Store) ListSchedules(ctx context.Context, projectID string) ([]api.Schedule, error) {
	if projectID == "" {
		return querySchedules(ctx, s.db, ``)
	}
	_, err := s.GetProject(ctx, projectID)
	if err != nil {
		return nil, err
	}
	return querySchedules(ctx, s.db, `WHERE project_id = ?`, projectID)
}

func (s *Store) GetSchedule(ctx context.Context, id string) (api.Schedule, error) {
	list, err := querySchedules(ctx, s.db, `WHERE id = ?`, id)
	if err != nil {
		return api.Schedule{}, err
	}
	if len(list) == 0 {
		return api.Schedule{}, fmt.Errorf("schedule %s: %w", id, ErrNotFound)
	}
	return list[0], nil
}

func writeWeekdays(ctx context.Context, tx *sql.Tx, sc api.Schedule) error {
	_, err := tx.ExecContext(ctx, `DELETE FROM schedule_weekdays WHERE schedule_id = ?`, sc.Id)
	if err != nil {
		return err
	}
	for _, d := range sc.Weekdays {
		_, err = tx.ExecContext(ctx, `INSERT INTO schedule_weekdays (schedule_id, weekday) VALUES (?, ?)`, sc.Id, string(d))
		if err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) CreateSchedule(ctx context.Context, sc api.Schedule, actor string) (api.Schedule, error) {
	sc.NextRun = schedule.Next(sc, time.Now())
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `INSERT INTO schedules (`+scheduleColumns+`, created_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
			sc.Id, sc.ProjectId, sc.Owner, sc.Title, string(sc.Every), sc.Hour, sc.Minute, sc.DayOfMonth, sc.Prompt, sc.Model, sc.Workflow, formatTime(sc.NextRun), now())
		if err != nil {
			return err
		}
		err = writeWeekdays(ctx, tx, sc)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "schedule.saved", ProjectId: &sc.ProjectId, Actor: actor, Payload: ToPayload(sc)})
	})
	return sc, err
}

func (s *Store) UpdateSchedule(ctx context.Context, sc api.Schedule, actor string) (api.Schedule, error) {
	sc.NextRun = schedule.Next(sc, time.Now())
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE schedules SET title = ?, every = ?, hour = ?, minute = ?, day_of_month = ?, prompt = ?, model = ?, workflow = ?, next_run_at = ? WHERE id = ?`,
			sc.Title, string(sc.Every), sc.Hour, sc.Minute, sc.DayOfMonth, sc.Prompt, sc.Model, sc.Workflow, formatTime(sc.NextRun), sc.Id)
		if err != nil {
			return err
		}
		err = writeWeekdays(ctx, tx, sc)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "schedule.saved", ProjectId: &sc.ProjectId, Actor: actor, Payload: ToPayload(sc)})
	})
	return sc, err
}

func (s *Store) DeleteSchedule(ctx context.Context, id, actor string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		projectID := ""
		err := tx.QueryRowContext(ctx, `SELECT project_id FROM schedules WHERE id = ?`, id).Scan(&projectID)
		if errors.Is(err, sql.ErrNoRows) {
			return fmt.Errorf("schedule %s: %w", id, ErrNotFound)
		}
		if err != nil {
			return err
		}
		_, err = tx.ExecContext(ctx, `DELETE FROM schedule_weekdays WHERE schedule_id = ?`, id)
		if err != nil {
			return err
		}
		_, err = tx.ExecContext(ctx, `DELETE FROM schedules WHERE id = ?`, id)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "schedule.deleted", ProjectId: &projectID, Actor: actor, Payload: map[string]any{"id": id}})
	})
}

func (s *Store) AdvanceSchedule(ctx context.Context, sc api.Schedule) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		list, err := querySchedules(ctx, tx, `WHERE id = ? AND next_run_at = ?`, sc.Id, formatTime(sc.NextRun))
		if err != nil || len(list) == 0 {
			return err
		}
		next := list[0]
		next.NextRun = schedule.Next(next, time.Now())
		_, err = tx.ExecContext(ctx, `UPDATE schedules SET next_run_at = ? WHERE id = ?`, formatTime(next.NextRun), sc.Id)
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "schedule.saved", ProjectId: &next.ProjectId, Actor: "byre", Payload: ToPayload(next)})
	})
}

func dueSchedules(ctx context.Context, tx *sql.Tx) ([]api.Schedule, *time.Time, error) {
	at := formatTime(time.Now())
	due, err := querySchedules(ctx, tx, `WHERE next_run_at <= ?`, at)
	if err != nil {
		return nil, nil, err
	}
	soonest := ""
	err = tx.QueryRowContext(ctx, `SELECT next_run_at FROM schedules WHERE next_run_at > ? ORDER BY next_run_at LIMIT 1`, at).Scan(&soonest)
	if errors.Is(err, sql.ErrNoRows) {
		return due, nil, nil
	}
	if err != nil {
		return nil, nil, err
	}
	t, err := parseTime(soonest)
	if err != nil {
		return nil, nil, err
	}
	return due, &t, nil
}