add task result files

Author: pk910

pkg/coordinator/db/schema/pgsql/20250114224233_task_summaries.sql

@@ -0,0 +1,22 @@
+-- +goose Up
+-- +goose StatementBegin
+
+CREATE TABLE IF NOT EXISTS public."task_results"
+(
+    "run_id" INTEGER NOT NULL,
+    "task_id" INTEGER NOT NULL,
+    "result_type" TEXT NOT NULL,
+    "result_index" INTEGER NOT NULL,
+    "name" TEXT NOT NULL,
+    "size" INTEGER NOT NULL,
+    "data" BYTEA NOT NULL,
+    CONSTRAINT "task_results_pkey" PRIMARY KEY ("run_id", "task_id", "result_type", "result_index")
+);
+
+CREATE INDEX IF NOT EXISTS "task_results_name_idx" ON public."task_results" ("run_id", "task_id", "name");
+
+-- +goose StatementEnd
+-- +goose Down
+-- +goose StatementBegin
+SELECT 'NOT SUPPORTED';
+-- +goose StatementEnd

pkg/coordinator/db/schema/sqlite/20250114224233_task_summaries.sql

@@ -0,0 +1,22 @@
+-- +goose Up
+-- +goose StatementBegin
+
+CREATE TABLE IF NOT EXISTS "task_results"
+(
+    "run_id" INTEGER NOT NULL,
+    "task_id" INTEGER NOT NULL,
+    "result_type" TEXT NOT NULL,
+    "result_index" INTEGER NOT NULL,
+    "name" TEXT NOT NULL,
+    "size" INTEGER NOT NULL,
+    "data" BLOB NOT NULL,
+    CONSTRAINT "task_results_pkey" PRIMARY KEY ("run_id", "task_id", "result_type", "result_index")
+);
+
+CREATE INDEX IF NOT EXISTS "task_results_name_idx" ON "task_results" ("run_id", "task_id", "name");
+
+-- +goose StatementEnd
+-- +goose Down
+-- +goose StatementBegin
+SELECT 'NOT SUPPORTED';
+-- +goose StatementEnd

pkg/coordinator/db/task_results.go

@@ -0,0 +1,88 @@
+package db
+
+import "github.com/jmoiron/sqlx"
+
+type TaskResult struct {
+	RunID  int    `db:"run_id"`
+	TaskID int    `db:"task_id"`
+	Type   string `db:"result_type"`
+	Index  int    `db:"result_index"`
+	Name   string `db:"name"`
+	Size   int    `db:"size"`
+	Data   []byte `db:"data"`
+}
+
+type TaskResultHeader struct {
+	TaskID int    `db:"task_id"`
+	Type   string `db:"result_type"`
+	Index  int    `db:"result_index"`
+	Name   string `db:"name"`
+	Size   int    `db:"size"`
+}
+
+func (db *Database) UpsertTaskResult(tx *sqlx.Tx, result *TaskResult) error {
+	_, err := tx.Exec(db.EngineQuery(map[EngineType]string{
+		EnginePgsql: `
+			INSERT INTO task_results (run_id, task_id, result_type, result_index, name, size, data)
+			VALUES ($1, $2, $3, $4, $5, $6, $7)
+			ON CONFLICT (run_id, task_id, result_type, result_index) DO UPDATE SET
+				name = excluded.name,
+				size = excluded.size,
+				data = excluded.data
+			`,
+		EngineSqlite: `
+			INSERT OR REPLACE INTO task_results (run_id, task_id, result_type, result_index, name, size, data)
+			VALUES ($1, $2, $3, $4, $5, $6, $7)`,
+	}),
+		result.RunID, result.TaskID, result.Type, result.Index, result.Name, result.Size, result.Data)
+
+	return err
+}
+
+func (db *Database) GetTaskResultByIndex(runID, taskID int, resultType string, index int) (*TaskResult, error) {
+	var result TaskResult
+	err := db.reader.Get(&result, `
+		SELECT * FROM task_results
+		WHERE run_id = $1 AND task_id = $2 AND result_type = $3 AND result_index = $4`,
+		runID, taskID, resultType, index)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &result, nil
+}
+
+func (db *Database) GetTaskResultByName(runID, taskID int, resultType, name string) (*TaskResult, error) {
+	var result TaskResult
+	err := db.reader.Get(&result, `
+		SELECT * FROM task_results
+		WHERE run_id = $1 AND task_id = $2 AND result_type = $3 AND name = $4`,
+		runID, taskID, resultType, name)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &result, nil
+}
+
+func (db *Database) GetTaskResults(runID, taskID int, summaryType string) ([]TaskResult, error) {
+	var results []TaskResult
+	err := db.reader.Select(&results, `
+		SELECT * FROM task_results
+		WHERE run_id = $1 AND task_id = $2 AND result_type = $3`,
+		runID, taskID, summaryType)
+
+	return results, err
+}
+
+func (db *Database) GetAllTaskResultHeaders(runID int) ([]TaskResultHeader, error) {
+	var headers []TaskResultHeader
+	err := db.reader.Select(&headers, `
+		SELECT task_id, result_type, result_index, name, size FROM task_results
+		WHERE run_id = $1`,
+		runID)
+
+	return headers, err
+}

pkg/coordinator/scheduler/scheduler.go

@@ -43,6 +43,10 @@ func (ts *TaskScheduler) GetServices() types.TaskServices {
 	return ts.services
 }
 
+func (ts *TaskScheduler) GetTestRunID() uint64 {
+	return ts.testRunID
+}
+
 func (ts *TaskScheduler) AddRootTask(options *types.TaskOptions) (types.TaskIndex, error) {
 	task, err := ts.newTaskState(options, nil, nil, false)
 	if err != nil {

pkg/coordinator/tasks/run_shell/result_file.go

@@ -0,0 +1,40 @@
+package runshell
+
+import (
+	"fmt"
+	"os"
+)
+
+type resultFile struct {
+	filePath string
+}
+
+func newResultFile(filePath string) (*resultFile, error) {
+	// Touch the file
+	if err := os.WriteFile(filePath, []byte{}, 0o600); err != nil {
+		return nil, fmt.Errorf("failed to create result file: %w", err)
+	}
+
+	return &resultFile{
+		filePath: filePath,
+	}, nil
+}
+
+func (f *resultFile) FilePath() string {
+	return f.filePath
+}
+
+func (f *resultFile) Cleanup() ([]byte, error) {
+	// Read file contents
+	data, err := os.ReadFile(f.filePath)
+	if err != nil && !os.IsNotExist(err) {
+		return nil, fmt.Errorf("failed to read result file: %w", err)
+	}
+
+	// Remove temp file
+	if err := os.Remove(f.filePath); err != nil {
+		return nil, fmt.Errorf("failed to remove temp file: %w", err)
+	}
+
+	return data, nil
+}

pkg/coordinator/tasks/run_shell/task.go

@@ -6,12 +6,16 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"os"
 	"os/exec"
+	"path/filepath"
 	"regexp"
 	"syscall"
 	"time"
 
+	"github.com/ethpandaops/assertoor/pkg/coordinator/db"
 	"github.com/ethpandaops/assertoor/pkg/coordinator/types"
+	"github.com/jmoiron/sqlx"
 	"github.com/sirupsen/logrus"
 )
 
@@ -78,6 +82,19 @@ func (t *Task) Execute(ctx context.Context) error {
 	cmdLogger := t.logger.WithField("shell", t.config.Shell)
 	cmdLogger.Info("running command")
 
+	// create temp dir for task
+	taskDir, err := os.MkdirTemp(os.TempDir(), fmt.Sprintf("assertoor_%v_%v_", t.ctx.Scheduler.GetTestRunID(), t.ctx.Index))
+	if err != nil {
+		cmdLogger.Errorf("failed creating task dir: %v", err)
+		return err
+	}
+
+	defer func() {
+		if err2 := os.RemoveAll(taskDir); err2 != nil {
+			cmdLogger.Errorf("failed cleaning up task dir: %v", err2)
+		}
+	}()
+
 	//nolint:gosec // ignore
 	command := exec.CommandContext(ctx, t.config.Shell, t.config.ShellArgs...)
 
@@ -124,6 +141,28 @@ func (t *Task) Execute(ctx context.Context) error {
 		command.Env = append(command.Env, fmt.Sprintf("%v=%v", envName, string(varJSON)))
 	}
 
+	// create summaries file
+	summaryFile, err := newResultFile(filepath.Join(taskDir, "summary"))
+	if err != nil {
+		cmdLogger.Errorf("failed creating summary file: %v", err)
+		return err
+	}
+
+	command.Env = append(command.Env, fmt.Sprintf("ASSERTOOR_SUMMARY=%v", summaryFile.FilePath()))
+
+	// create folder for result files
+	resultDir := filepath.Join(taskDir, "results")
+	if err = os.MkdirAll(resultDir, 0o700); err != nil {
+		cmdLogger.Errorf("failed creating result dir: %v", err)
+		return err
+	}
+
+	command.Env = append(command.Env, fmt.Sprintf("ASSERTOOR_RESULT_DIR=%v", resultDir))
+
+	defer func() {
+		t.storeTaskResults(summaryFile, resultDir)
+	}()
+
 	// start shell
 	err = command.Start()
 	if err != nil {
@@ -305,3 +344,70 @@ func (t *Task) parseOutputVars(line string) bool {
 
 	return false
 }
+
+func (t *Task) storeTaskResults(summaryFile *resultFile, resultDir string) {
+	// store files to db
+	database := t.ctx.Scheduler.GetServices().Database()
+	if err2 := database.RunTransaction(func(tx *sqlx.Tx) error {
+		// store summary file
+		data, err3 := summaryFile.Cleanup()
+		if err3 != nil {
+			t.logger.Errorf("failed cleaning up summary file: %v", err3)
+		} else if err3 = database.UpsertTaskResult(tx, &db.TaskResult{
+			RunID:  int(t.ctx.Scheduler.GetTestRunID()),
+			TaskID: int(t.ctx.Index),
+			Type:   "summary",
+			Index:  0,
+			Name:   "",
+			Size:   len(data),
+			Data:   data,
+		}); err3 != nil {
+			t.logger.Errorf("failed storing summary file to db: %v", err3)
+		}
+
+		// store result files
+		fileIdx := 0
+
+		var storeResultFilesFn func(path string, prefix string)
+		storeResultFilesFn = func(path string, prefix string) {
+			if prefix != "" {
+				prefix += "/"
+			}
+
+			files, err3 := os.ReadDir(path)
+			if err3 != nil {
+				t.logger.Errorf("failed reading result dir: %v", err3)
+			} else {
+				for _, file := range files {
+					if file.IsDir() {
+						storeResultFilesFn(filepath.Join(path, file.Name()), fmt.Sprintf("%v%v", prefix, file.Name()))
+						continue
+					}
+
+					data, err3 := os.ReadFile(filepath.Join(path, file.Name()))
+					if err3 != nil {
+						t.logger.Errorf("failed reading result file: %v", err3)
+					} else if err3 = database.UpsertTaskResult(tx, &db.TaskResult{
+						RunID:  int(t.ctx.Scheduler.GetTestRunID()),
+						TaskID: int(t.ctx.Index),
+						Type:   "result",
+						Index:  fileIdx,
+						Name:   file.Name(),
+						Size:   len(data),
+						Data:   data,
+					}); err3 != nil {
+						t.logger.Errorf("failed storing result file to db: %v", err3)
+					}
+
+					fileIdx++
+				}
+			}
+		}
+
+		storeResultFilesFn(resultDir, "")
+
+		return nil
+	}); err2 != nil {
+		t.logger.Errorf("failed storing task results to db: %v", err2)
+	}
+}

pkg/coordinator/types/scheduler.go

@@ -13,6 +13,7 @@ import (
 type TaskSchedulerRunner interface {
 	TaskScheduler
 	GetServices() TaskServices
+	GetTestRunID() uint64
 	ParseTaskOptions(rawtask helper.IRawMessage) (*TaskOptions, error)
 	ExecuteTask(ctx context.Context, taskIndex TaskIndex, taskWatchFn func(ctx context.Context, cancelFn context.CancelFunc, taskIndex TaskIndex)) error
 	WatchTaskPass(ctx context.Context, cancelFn context.CancelFunc, taskIndex TaskIndex)