byre/internal/store/workflow.go

package store

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"

	"g.fc.run/misc/byre/internal/api"
)

type FrozenWorkflow struct {
	Name       string
	Definition api.WorkflowDefinition
	Brief      string
	Title      string
}

func (s *Store) ListWorkflows(ctx context.Context) ([]api.WorkflowSummary, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT name, coalesce(json_extract(definition, '$.description'), '') FROM workflows ORDER BY name`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	workflows := []api.WorkflowSummary{}
	for rows.Next() {
		w := api.WorkflowSummary{}
		err = rows.Scan(&w.Name, &w.Description)
		if err != nil {
			return nil, err
		}
		workflows = append(workflows, w)
	}
	return workflows, rows.Err()
}

func (s *Store) GetWorkflow(ctx context.Context, name string) (api.WorkflowDefinition, error) {
	def := ""
	err := s.db.QueryRowContext(ctx, `SELECT definition FROM workflows WHERE name = ?`, name).Scan(&def)
	if errors.Is(err, sql.ErrNoRows) {
		return api.WorkflowDefinition{}, fmt.Errorf("%w: no workflow %q", ErrNotFound, name)
	}
	if err != nil {
		return api.WorkflowDefinition{}, err
	}
	w := api.WorkflowDefinition{}
	err = json.Unmarshal([]byte(def), &w)
	return w, err
}

func (s *Store) PutWorkflow(ctx context.Context, name string, def api.WorkflowDefinition, actor string) error {
	b, err := json.Marshal(def)
	if err != nil {
		return err
	}
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `INSERT INTO workflows (name, definition) VALUES (?, ?)
			ON CONFLICT (name) DO UPDATE SET definition = excluded.definition`, name, string(b))
		if err != nil {
			return err
		}
		return emit(api.Event{Kind: "workflow.updated", Actor: actor, Payload: map[string]any{"name": name}})
	})
}

func (s *Store) DeleteWorkflow(ctx context.Context, name, actor string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `DELETE FROM workflows WHERE name = ?`, name)
		if err != nil {
			return err
		}
		n, err := res.RowsAffected()
		if err != nil {
			return err
		}
		if n == 0 {
			return fmt.Errorf("%w: no workflow %q", ErrNotFound, name)
		}
		return emit(api.Event{Kind: "workflow.deleted", Actor: actor, Payload: map[string]any{"name": name}})
	})
}

func (s *Store) FreezeWorkflow(ctx context.Context, name string) (FrozenWorkflow, error) {
	def, err := s.GetWorkflow(ctx, name)
	if err != nil {
		return FrozenWorkflow{}, err
	}
	return FrozenWorkflow{Name: name, Definition: def}, nil
}

func (s *Store) TaskWorkflow(ctx context.Context, taskID string) (FrozenWorkflow, error) {
	f := FrozenWorkflow{}
	def := ""
	err := s.db.QueryRowContext(ctx, `SELECT workflow_name, workflow_definition, workflow_brief, workflow_title FROM tasks WHERE id = ?`, taskID).Scan(&f.Name, &def, &f.Brief, &f.Title)
	if errors.Is(err, sql.ErrNoRows) {
		return FrozenWorkflow{}, fmt.Errorf("%w: no task %s", ErrNotFound, taskID)
	}
	if err != nil {
		return FrozenWorkflow{}, err
	}
	if f.Name == "" {
		return FrozenWorkflow{}, fmt.Errorf("%w: task %s has no workflow", ErrNotFound, taskID)
	}
	err = json.Unmarshal([]byte(def), &f.Definition)
	return f, err
}

func (s *Store) RunTasks(ctx context.Context, run string) ([]api.Task, error) {
	return s.queryTasks(ctx, `SELECT `+taskColumns+` FROM tasks WHERE workflow_run = ? ORDER BY created_at`, run)
}

func (s *Store) RunEvents(ctx context.Context, run string) ([]api.Event, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT `+eventColumns+` FROM events
		WHERE task_id IN (SELECT id FROM tasks WHERE workflow_run = ?) AND kind IN ('task.workflow_stage', 'task.discarded', 'task.workflow_finished')
		ORDER BY seq
	`, run)
	if err != nil {
		return nil, err
	}
	return scanEvents(rows)
}

func (s *Store) RecordStage(ctx context.Context, taskID, output, commit string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE tasks SET stage_output = ?, stage_commit = ? WHERE id = ?`, output, commit, taskID)
		return err
	})
}

func (s *Store) StageResult(ctx context.Context, taskID string) (string, string, error) {
	output := ""
	commit := ""
	err := s.db.QueryRowContext(ctx, `SELECT stage_output, stage_commit FROM tasks WHERE id = ?`, taskID).Scan(&output, &commit)
	return output, commit, err
}

func (s *Store) SetStageOutput(ctx context.Context, sessionID, output string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = ? WHERE id = ?`, output, sessionID)
		return err
	})
}

func (s *Store) FinishStage(ctx context.Context, sessionID string) (bool, error) {
	finished := false
	err := s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		res, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_finished = 1 WHERE id = ? AND stage_output != '' AND NOT EXISTS (SELECT 1 FROM tasks WHERE tasks.id = agent_sessions.task_id AND tasks.discard_state != '')`, sessionID)
		if err != nil {
			return err
		}
		n, err := res.RowsAffected()
		finished = n > 0
		return err
	})
	return finished, err
}

func (s *Store) DropStageOutput(ctx context.Context, sessionID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = '' WHERE id = ? AND stage_finished = 0`, sessionID)
		return err
	})
}

func (s *Store) StageOutput(ctx context.Context, sessionID string) (string, error) {
	output := ""
	err := s.db.QueryRowContext(ctx, `SELECT stage_output FROM agent_sessions WHERE id = ?`, sessionID).Scan(&output)
	return output, err
}

func (s *Store) FinishedStages(ctx context.Context) ([]string, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT id FROM agent_sessions WHERE stage_finished = 1 ORDER BY created_at`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	ids := []string{}
	for rows.Next() {
		id := ""
		err = rows.Scan(&id)
		if err != nil {
			return nil, err
		}
		ids = append(ids, id)
	}
	return ids, rows.Err()
}

func (s *Store) ClearStage(ctx context.Context, sessionID string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		return clearStage(ctx, tx, sessionID)
	})
}

func clearStage(ctx context.Context, tx *sql.Tx, sessionID string) error {
	_, err := tx.ExecContext(ctx, `UPDATE agent_sessions SET stage_output = '', stage_finished = 0 WHERE id = ?`, sessionID)
	return err
}

func (s *Store) FinishWorkflow(ctx context.Context, task api.Task, sessionID, workflowName, output string) error {
	return s.write(ctx, func(tx *sql.Tx, emit emitFunc) error {
		err := clearStage(ctx, tx, sessionID)
		if err != nil {
			return err
		}
		return emit(api.Event{
			Kind:      "task.workflow_finished",
			ProjectId: &task.ProjectId,
			TaskId:    &task.Id,
			Actor:     "byre",
			Payload:   map[string]any{"id": task.Id, "workflow": workflowName, "output": output},
		})
	})
}

func (s *Store) WorkflowFinished(ctx context.Context, run string) (bool, error) {
	finished := false
	err := s.db.QueryRowContext(ctx, `SELECT EXISTS (SELECT 1 FROM events e JOIN tasks t ON t.id = e.task_id WHERE t.workflow_run = ? AND e.kind = 'task.workflow_finished')`, run).Scan(&finished)
	return finished, err
}