Add support for using input.yaml in Evaluate code lens (#1269)

This addresses a request filed in the VS Code extension:
open-policy-agent/vscode-opa#308

Sadly this doesn't yet work for the debug feature as OPA currently
only will do JSON decoding in that path, so next step is to submit
a fix for that there.

Signed-off-by: Anders Eknert <[email protected]>

Author: anderseknert

.gitignore

@@ -6,9 +6,12 @@ dist/
 /regal.exe
 
 # These two files are used by the Regal evaluation Code Lens, where input.json
-# defines the input to use for evaluation, and output.json is where the output
-# ends up unless the client supports presenting it in a different way.
+# (or input.yaml) defines the input to use for evaluation, and output.json is
+# where the output ends up unless the client supports presenting it in a
+# different way.
 input.json
+input.yaml
+
 output.json
 
 build/node_modules/

docs/language-server.md

@@ -173,10 +173,11 @@ the way it did, or where rule evaluation failed.
   src={require('./assets/lsp/evalcodelensprint.png').default}
   alt="Screenshot of evaluation with print call performed via code lens"/>
 
-Policy evaluation often depends on **input**. This can be provided via an `input.json` file which Regal will search
-for first in the same directory as the policy file evaluated. If not found there, Regal will proceed to search each
-parent directory up until the workspace root directory. It is recommended to add `input.json` to your `.gitignore`
-file so that you can work freely with evaluation in any directory without having your input accidentally committed.
+Policy evaluation often depends on **input**. This can be provided via an `input.json` or `input.yaml` file which
+Regal will search for first in the same directory as the policy file evaluated. If not found there, Regal will proceed
+to search each parent directory up until the workspace root directory. It is recommended to add `input.json/yaml` to
+your `.gitignore` file so that you can work freely with evaluation in any directory without having your input
+accidentally committed.
 
 #### Editor support
 

internal/io/io.go

@@ -10,6 +10,7 @@ import (
 	"strings"
 
 	"github.com/anderseknert/roast/pkg/encoding"
+	"gopkg.in/yaml.v3"
 
 	"github.com/open-policy-agent/opa/bundle"
 	"github.com/open-policy-agent/opa/loader/filter"
@@ -94,22 +95,61 @@ func ExcludeTestFilter() filter.LoaderFilter {
 	}
 }
 
-// FindInput finds input.json file in workspace closest to the file, and returns
-// both the location and the reader.
-func FindInput(file string, workspacePath string) (string, io.Reader) {
+// FindInput finds input.json or input.yaml file in workspace closest to the file, and returns
+// both the location and the contents of the file (as map), or an empty string and nil if not found.
+// Note that:
+// - This function doesn't do error handling. If the file can't be read, nothing is returned.
+// - While the input data theoritcally could be anything JSON/YAML value, we only support an object.
+func FindInput(file string, workspacePath string) (string, map[string]any) {
 	relative := strings.TrimPrefix(file, workspacePath)
 	components := strings.Split(filepath.Dir(relative), string(filepath.Separator))
 
+	var (
+		inputPath string
+		content   []byte
+	)
+
 	for i := range components {
-		inputPath := filepath.Join(workspacePath, filepath.Join(components[:len(components)-i]...), "input.json")
+		current := components[:len(components)-i]
+
+		inputPathJSON := filepath.Join(workspacePath, filepath.Join(current...), "input.json")
+
+		f, err := os.Open(inputPathJSON)
+		if err == nil {
+			inputPath = inputPathJSON
+			content, _ = io.ReadAll(f)
 
-		f, err := os.Open(inputPath)
+			break
+		}
+
+		inputPathYAML := filepath.Join(workspacePath, filepath.Join(current...), "input.yaml")
+
+		f, err = os.Open(inputPathYAML)
 		if err == nil {
-			return inputPath, f
+			inputPath = inputPathYAML
+			content, _ = io.ReadAll(f)
+
+			break
+		}
+	}
+
+	if inputPath == "" || content == nil {
+		return "", nil
+	}
+
+	var input map[string]any
+
+	if strings.HasSuffix(inputPath, ".json") {
+		if err := encoding.JSON().Unmarshal(content, &input); err != nil {
+			return "", nil
+		}
+	} else if strings.HasSuffix(inputPath, ".yaml") {
+		if err := yaml.Unmarshal(content, &input); err != nil {
+			return "", nil
 		}
 	}
 
-	return "", nil
+	return inputPath, input
 }
 
 func IsSkipWalkDirectory(info files.DirEntry) bool {

internal/lsp/completions/providers/policy.go

@@ -2,10 +2,8 @@ package providers
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"os"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -70,20 +68,15 @@ func (p *Policy) Run(
 	inputContext["path_separator"] = string(os.PathSeparator)
 
 	workspacePath := uri.ToPath(opts.ClientIdentifier, opts.RootURI)
-	inputDotJSONPath, inputDotJSONReader := rio.FindInput(
+
+	inputDotJSONPath, inputDotJSONContent := rio.FindInput(
 		uri.ToPath(opts.ClientIdentifier, params.TextDocument.URI),
 		workspacePath,
 	)
 
-	if inputDotJSONReader != nil {
-		inputDotJSON := make(map[string]any)
-
-		if bs, err := io.ReadAll(inputDotJSONReader); err == nil {
-			if err = json.Unmarshal(bs, &inputDotJSON); err == nil {
-				inputContext["input_dot_json_path"] = inputDotJSONPath
-				inputContext["input_dot_json"] = inputDotJSON
-			}
-		}
+	if inputDotJSONPath != "" && inputDotJSONContent != nil {
+		inputContext["input_dot_json_path"] = inputDotJSONPath
+		inputContext["input_dot_json"] = inputDotJSONContent
 	}
 
 	input, err := rego2.ToInput(

internal/lsp/eval.go

@@ -4,11 +4,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"strings"
 
-	"github.com/anderseknert/roast/pkg/encoding"
-
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/bundle"
 	"github.com/open-policy-agent/opa/rego"
@@ -24,7 +21,7 @@ import (
 func (l *LanguageServer) Eval(
 	ctx context.Context,
 	query string,
-	input io.Reader,
+	input map[string]any,
 	printHook print.Hook,
 	dataBundles map[string]bundle.Bundle,
 ) (rego.ResultSet, error) {
@@ -88,20 +85,7 @@ func (l *LanguageServer) Eval(
 	}
 
 	if input != nil {
-		inputMap := make(map[string]any)
-
-		in, err := io.ReadAll(input)
-		if err != nil {
-			return nil, fmt.Errorf("failed reading input: %w", err)
-		}
-
-		json := encoding.JSON()
-
-		if err = json.Unmarshal(in, &inputMap); err != nil {
-			return nil, fmt.Errorf("failed unmarshalling input: %w", err)
-		}
-
-		return pq.Eval(ctx, rego.EvalInput(inputMap)) //nolint:wrapcheck
+		return pq.Eval(ctx, rego.EvalInput(input)) //nolint:wrapcheck
 	}
 
 	return pq.Eval(ctx) //nolint:wrapcheck
@@ -116,7 +100,7 @@ type EvalPathResult struct {
 func (l *LanguageServer) EvalWorkspacePath(
 	ctx context.Context,
 	query string,
-	input io.Reader,
+	input map[string]any,
 ) (EvalPathResult, error) {
 	resultQuery := "result := " + query
 

internal/lsp/eval_test.go

@@ -2,10 +2,9 @@ package lsp
 
 import (
 	"context"
-	"io"
+	"maps"
 	"os"
 	"slices"
-	"strings"
 	"testing"
 
 	rio "github.com/styrainc/regal/internal/io"
@@ -49,7 +48,9 @@ func TestEvalWorkspacePath(t *testing.T) {
 	ls.cache.SetModule("file://policy1.rego", module1)
 	ls.cache.SetModule("file://policy2.rego", module2)
 
-	input := strings.NewReader(`{"exists": true}`)
+	input := map[string]any{
+		"exists": true,
+	}
 
 	res, err := ls.EvalWorkspacePath(context.TODO(), "data.policy1.allow", input)
 	if err != nil {
@@ -71,7 +72,7 @@ func TestEvalWorkspacePathInternalData(t *testing.T) {
 		&LanguageServerOptions{LogWriter: logger, LogLevel: log.LevelDebug},
 	)
 
-	res, err := ls.EvalWorkspacePath(context.TODO(), "object.keys(data.internal)", strings.NewReader("{}"))
+	res, err := ls.EvalWorkspacePath(context.TODO(), "object.keys(data.internal)", map[string]any{})
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -104,35 +105,50 @@ func TestEvalWorkspacePathInternalData(t *testing.T) {
 func TestFindInput(t *testing.T) {
 	t.Parallel()
 
-	tmpDir := t.TempDir()
+	cases := []struct {
+		fileType    string
+		fileContent string
+	}{
+		{"json", `{"x": true}`},
+		{"yaml", "x: true"},
+	}
 
-	workspacePath := tmpDir + "/workspace"
-	file := tmpDir + "/workspace/foo/bar/baz.rego"
+	for _, tc := range cases {
+		t.Run(tc.fileType, func(t *testing.T) {
+			t.Parallel()
 
-	if err := os.MkdirAll(workspacePath+"/foo/bar", 0o755); err != nil {
-		t.Fatal(err)
-	}
+			tmpDir := t.TempDir()
 
-	if readInputString(t, file, workspacePath) != "" {
-		t.Fatalf("did not expect to find input.json")
-	}
+			workspacePath := tmpDir + "/workspace"
+			file := tmpDir + "/workspace/foo/bar/baz.rego"
 
-	content := `{"x": 1}`
+			if err := os.MkdirAll(workspacePath+"/foo/bar", 0o755); err != nil {
+				t.Fatal(err)
+			}
 
-	createWithContent(t, tmpDir+"/workspace/foo/bar/input.json", content)
+			path, content := rio.FindInput(file, workspacePath)
+			if path != "" || content != nil {
+				t.Fatalf("did not expect to find input.%s", tc.fileType)
+			}
 
-	if res := readInputString(t, file, workspacePath); res != content {
-		t.Errorf("expected input at %s, got %s", content, res)
-	}
+			createWithContent(t, tmpDir+"/workspace/foo/bar/input."+tc.fileType, tc.fileContent)
 
-	if err := os.Remove(tmpDir + "/workspace/foo/bar/input.json"); err != nil {
-		t.Fatal(err)
-	}
+			path, content = rio.FindInput(file, workspacePath)
+			if path != workspacePath+"/foo/bar/input."+tc.fileType || !maps.Equal(content, map[string]any{"x": true}) {
+				t.Errorf(`expected input {"x": true} at, got %s`, content)
+			}
+
+			if err := os.Remove(tmpDir + "/workspace/foo/bar/input." + tc.fileType); err != nil {
+				t.Fatal(err)
+			}
 
-	createWithContent(t, tmpDir+"/workspace/input.json", content)
+			createWithContent(t, tmpDir+"/workspace/input."+tc.fileType, tc.fileContent)
 
-	if res := readInputString(t, file, workspacePath); res != content {
-		t.Errorf("expected input at %s, got %s", content, res)
+			path, content = rio.FindInput(file, workspacePath)
+			if path != workspacePath+"/input."+tc.fileType || !maps.Equal(content, map[string]any{"x": true}) {
+				t.Errorf(`expected input {"x": true} at, got %s`, content)
+			}
+		})
 	}
 }
 
@@ -150,24 +166,3 @@ func createWithContent(t *testing.T, path string, content string) {
 		t.Fatal(err)
 	}
 }
-
-func readInputString(t *testing.T, file, workspacePath string) string {
-	t.Helper()
-
-	_, input := rio.FindInput(file, workspacePath)
-
-	if input == nil {
-		return ""
-	}
-
-	bs, err := io.ReadAll(input)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	if bs == nil {
-		return ""
-	}
-
-	return string(bs)
-}

internal/lsp/server.go

@@ -2,7 +2,6 @@
 package lsp
 
 import (
-	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -820,33 +819,28 @@ func (l *LanguageServer) StartCommandWorker(ctx context.Context) { // nolint:mai
 				// if there are none, then it's a package evaluation
 				ruleHeadLocations := allRuleHeadLocations[path]
 
-				var input io.Reader
+				var inputMap map[string]any
 
 				// When the first comment in the file is `regal eval: use-as-input`, the AST of that module is
-				// used as the input rather than the contents of input.json. This is a development feature for
+				// used as the input rather than the contents of input.json/yaml. This is a development feature for
 				// working on rules (built-in or custom), allowing querying the AST of the module directly.
 				if len(currentModule.Comments) > 0 && regalEvalUseAsInputComment.Match(currentModule.Comments[0].Text) {
-					inputMap, err := rparse.PrepareAST(file, currentContents, currentModule)
+					inputMap, err = rparse.PrepareAST(file, currentContents, currentModule)
 					if err != nil {
 						l.logf(log.LevelMessage, "failed to prepare module: %s", err)
 
 						break
 					}
+				} else {
+					// Normal mode — try to find the input.json/yaml file in the workspace and use as input
+					_, inputMap = rio.FindInput(uri.ToPath(l.clientIdentifier, file), l.workspacePath())
 
-					bs, err := encoding.JSON().Marshal(inputMap)
-					if err != nil {
-						l.logf(log.LevelMessage, "failed to marshal module: %s", err)
-
+					if inputMap == nil {
 						break
 					}
-
-					input = bytes.NewReader(bs)
-				} else {
-					// Normal mode — try to find the input.json file in the workspace and use as input
-					_, input = rio.FindInput(uri.ToPath(l.clientIdentifier, file), l.workspacePath())
 				}
 
-				result, err := l.EvalWorkspacePath(ctx, path, input)
+				result, err := l.EvalWorkspacePath(ctx, path, inputMap)
 				if err != nil {
 					fmt.Fprintf(os.Stderr, "failed to evaluate workspace path: %v\n", err)