From 2982a0dda0e3df56bce372efbc8879b42c52b035 Mon Sep 17 00:00:00 2001
From: Andrea Luzzardi <aluzzardi@gmail.com>
Date: Thu, 23 Dec 2021 20:23:52 +0100
Subject: [PATCH] engine: Task PreRun Hook

Tasks now have a PreRun hook that gets called before buildkit kicks in.
Allows to support local access without special casing.

Signed-off-by: Andrea Luzzardi <aluzzardi@gmail.com>
---
 plan/plan.go                | 57 ++++++++++++++++++++++---------------
 plan/task/inputdirectory.go | 17 +++++++++++
 plan/task/task.go           |  6 ++++
 tests/plan.bats             |  2 +-
 4 files changed, 58 insertions(+), 24 deletions(-)

diff --git a/plan/plan.go b/plan/plan.go
index 4e4f6e29..659b79d6 100644
--- a/plan/plan.go
+++ b/plan/plan.go
@@ -2,9 +2,7 @@ package plan
 
 import (
 	"context"
-	"errors"
 	"fmt"
-	"os"
 	"strings"
 	"time"
 
@@ -43,11 +41,11 @@ func Load(ctx context.Context, args ...string) (*Plan, error) {
 		source:  v,
 	}
 
-	if err := p.registerLocalDirs(); err != nil {
+	if err := p.configPlatform(); err != nil {
 		return nil, err
 	}
 
-	if err := p.configPlatform(); err != nil {
+	if err := p.prepare(ctx); err != nil {
 		return nil, err
 	}
 
@@ -87,26 +85,39 @@ func (p *Plan) configPlatform() error {
 	return nil
 }
 
-// registerLocalDirectories scans the context for local imports.
-// BuildKit requires to known the list of directories ahead of time.
-func (p *Plan) registerLocalDirs() error {
-	imports, err := p.source.Lookup("inputs.directories").Fields()
-	if err != nil {
-		return err
-	}
+// prepare executes the pre-run hooks of tasks
+func (p *Plan) prepare(ctx context.Context) error {
+	flow := cueflow.New(
+		&cueflow.Config{},
+		p.source.Cue(),
+		func(flowVal cue.Value) (cueflow.Runner, error) {
+			v := compiler.Wrap(flowVal)
+			t, err := task.Lookup(v)
+			if err != nil {
+				// Not a task
+				if err == task.ErrNotTask {
+					return nil, nil
+				}
+				return nil, err
+			}
+			r, ok := t.(task.PreRunner)
+			if !ok {
+				return nil, nil
+			}
 
-	for _, v := range imports {
-		dir, err := v.Value.Lookup("path").AbsPath()
-		if err != nil {
-			return err
-		}
-		if _, err := os.Stat(dir); errors.Is(err, os.ErrNotExist) {
-			return fmt.Errorf("path %q does not exist", dir)
-		}
-		p.context.LocalDirs.Add(dir)
-	}
+			return cueflow.RunnerFunc(func(t *cueflow.Task) error {
+				ctx := t.Context()
+				lg := log.Ctx(ctx).With().Str("task", t.Path().String()).Logger()
+				ctx = lg.WithContext(ctx)
 
-	return nil
+				if err := r.PreRun(ctx, p.context, compiler.Wrap(t.Value())); err != nil {
+					return fmt.Errorf("%s: %w", t.Path().String(), err)
+				}
+				return nil
+			}), nil
+		},
+	)
+	return flow.Run(ctx)
 }
 
 // Up executes the plan
@@ -175,7 +186,7 @@ func newRunner(pctx *plancontext.Context, s solver.Solver, computed *compiler.Va
 				} else {
 					lg.Error().Dur("duration", time.Since(start)).Err(err).Str("state", string(environment.StateFailed)).Msg(string(environment.StateFailed))
 				}
-				return err
+				return fmt.Errorf("%s: %w", t.Path().String(), err)
 			}
 
 			lg.Info().Dur("duration", time.Since(start)).Str("state", string(environment.StateCompleted)).Msg(string(environment.StateCompleted))
diff --git a/plan/task/inputdirectory.go b/plan/task/inputdirectory.go
index 5e1c074a..d878e213 100644
--- a/plan/task/inputdirectory.go
+++ b/plan/task/inputdirectory.go
@@ -2,6 +2,9 @@ package task
 
 import (
 	"context"
+	"errors"
+	"fmt"
+	"os"
 
 	"github.com/moby/buildkit/client/llb"
 	"github.com/rs/zerolog/log"
@@ -17,6 +20,20 @@ func init() {
 type inputDirectoryTask struct {
 }
 
+func (c *inputDirectoryTask) PreRun(ctx context.Context, pctx *plancontext.Context, v *compiler.Value) error {
+	path, err := v.Lookup("path").AbsPath()
+	if err != nil {
+		return err
+	}
+
+	if _, err := os.Stat(path); errors.Is(err, os.ErrNotExist) {
+		return fmt.Errorf("path %q does not exist", path)
+	}
+	pctx.LocalDirs.Add(path)
+
+	return nil
+}
+
 func (c *inputDirectoryTask) Run(ctx context.Context, pctx *plancontext.Context, s solver.Solver, v *compiler.Value) (*compiler.Value, error) {
 	path, err := v.Lookup("path").AbsPath()
 	if err != nil {
diff --git a/plan/task/task.go b/plan/task/task.go
index a840fdd6..9be4d546 100644
--- a/plan/task/task.go
+++ b/plan/task/task.go
@@ -29,6 +29,12 @@ type Task interface {
 	Run(ctx context.Context, pctx *plancontext.Context, s solver.Solver, v *compiler.Value) (*compiler.Value, error)
 }
 
+type PreRunner interface {
+	Task
+
+	PreRun(ctx context.Context, pctx *plancontext.Context, v *compiler.Value) error
+}
+
 // Register a task type and initializer
 func Register(typ string, f NewFunc) {
 	tasks.Store(typ, f)
diff --git a/tests/plan.bats b/tests/plan.bats
index 8540ffa6..51ee02c7 100644
--- a/tests/plan.bats
+++ b/tests/plan.bats
@@ -62,7 +62,7 @@ setup() {
   cd "$TESTDIR"
   run "$DAGGER" --europa up ./plan/inputs/directories/conflicting_values.cue
 	assert_failure
-	assert_output --partial 'failed to up environment: actions.verify.contents: conflicting values "local directory" and "local dfsadf"'
+	assert_output --partial 'conflicting values "local directory" and "local dfsadf"'
 }
 
 @test "plan/inputs/secrets exec" {