Add input.json completion provider (#1005)

This provider suggests completions based on the nested attributes found in the `input.json` file, if such a file exists. Signed-off-by: Anders Eknert <[email protected]>
StyraInc · Aug 23, 2024 · 1fecb4c · 1fecb4c
1 parent f28ac7d
commit 1fecb4c
Show file tree

Hide file tree

Showing 8 changed files with 211 additions and 22 deletions.
diff --git a/bundle/regal/lsp/completion/providers/inputdotjson/inputdotjson.rego b/bundle/regal/lsp/completion/providers/inputdotjson/inputdotjson.rego
@@ -0,0 +1,60 @@
+package regal.lsp.completion.providers.inputdotjson
+
+import rego.v1
+
+import data.regal.lsp.completion.kind
+import data.regal.lsp.completion.location
+
+# METADATA
+# description: returns suggestions based on input.json structure (if found)
+items contains item if {
+	input.regal.context.input_dot_json_path
+
+	position := location.to_position(input.regal.context.location)
+	line := input.regal.file.lines[position.line]
+	word := location.ref_at(line, input.regal.context.location.col)
+
+	some [suggestion, type] in _matching_input_suggestions
+
+	item := {
+		"label": suggestion,
+		"kind": kind.variable,
+		"detail": type,
+		"documentation": {
+			"kind": "markdown",
+			"value": sprintf("(inferred from [`input.json`](%s))", [input.regal.context.input_dot_json_path]),
+		},
+		"textEdit": {
+			"range": location.word_range(word, position),
+			"newText": suggestion,
+		},
+	}
+}
+
+_matching_input_suggestions contains [suggestion, type] if {
+	position := location.to_position(input.regal.context.location)
+	line := input.regal.file.lines[position.line]
+
+	line != ""
+	location.in_rule_body(line)
+
+	word := location.ref_at(line, input.regal.context.location.col)
+
+	some [suggestion, type] in _input_paths
+
+	startswith(suggestion, word.text)
+}
+
+_input_paths contains [input_path, input_type] if {
+	walk(input.regal.context.input_dot_json, [path, value])
+
+	count(path) > 0
+
+	# don't traverse into arrays
+	every value in path {
+		is_string(value)
+	}
+
+	input_type := type_name(value)
+	input_path := concat(".", ["input", concat(".", path)])
+}
diff --git a/bundle/regal/lsp/completion/providers/inputdotjson/inputdotjson_test.rego b/bundle/regal/lsp/completion/providers/inputdotjson/inputdotjson_test.rego
@@ -0,0 +1,102 @@
+package regal.lsp.completion.providers.inputdotjson_test
+
+import rego.v1
+
+import data.regal.lsp.completion.providers.inputdotjson as provider
+
+# regal ignore:rule-length
+test_matching_input_suggestions if {
+	items := provider.items with input as input_obj
+	items == {
+		{
+			"detail": "object",
+			"kind": 6,
+			"label": "input.request",
+			"documentation": {
+				"kind": "markdown",
+				"value": "(inferred from [`input.json`](/foo/bar/input.json))",
+			},
+			"textEdit": {
+				"newText": "input.request",
+				"range": {
+					"end": {"character": 13, "line": 5},
+					"start": {"character": 6, "line": 5},
+				},
+			},
+		},
+		{
+			"detail": "string",
+			"kind": 6,
+			"label": "input.request.method",
+			"documentation": {
+				"kind": "markdown",
+				"value": "(inferred from [`input.json`](/foo/bar/input.json))",
+			},
+			"textEdit": {
+				"newText": "input.request.method",
+				"range": {
+					"end": {"character": 13, "line": 5},
+					"start": {"character": 6, "line": 5},
+				},
+			},
+		},
+		{
+			"detail": "string",
+			"kind": 6,
+			"label": "input.request.url",
+			"documentation": {
+				"kind": "markdown",
+				"value": "(inferred from [`input.json`](/foo/bar/input.json))",
+			},
+			"textEdit": {
+				"newText": "input.request.url",
+				"range": {
+					"end": {"character": 13, "line": 5},
+					"start": {"character": 6, "line": 5},
+				},
+			},
+		},
+	}
+}
+
+test_not_matching_input_suggestions if {
+	input_obj_new_loc := object.union(input_obj, {"regal": {"context": {"location": {
+		"row": 1,
+		"col": 1,
+	}}}})
+	items := provider.items with input as input_obj_new_loc
+	items == set()
+}
+
+input_obj := {"regal": {
+	"context": {
+		"location": {
+			"row": 6,
+			"col": 12,
+		},
+		"input_dot_json": {
+			"user": {
+				"name": {
+					"first": "John",
+					"last": "Doe",
+				},
+				"email": "[email protected]",
+				"roles": [{"name": "admin"}, {"name": "user"}],
+			},
+			"request": {
+				"method": "GET",
+				"url": "https://example.com",
+			},
+		},
+		"input_dot_json_path": "/foo/bar/input.json",
+	},
+	"file": {"lines": [
+		"package p",
+		"",
+		"import rego.v1",
+		"",
+		"allow if {",
+		"    f(input.r",
+		"}",
+	]},
+}}
diff --git a/internal/embeds/schemas/regal-ast.json b/internal/embeds/schemas/regal-ast.json
@@ -394,6 +394,12 @@
             },
             "workspace_root": {
               "type": "string"
+            },
+            "input_dot_json": {
+              "type": "object"
+            },
+            "input_dot_json_path": {
+              "type": "string"
             }
           }
         }

diff --git a/internal/io/io.go b/internal/io/io.go
@@ -2,9 +2,12 @@ package io
 
 import (
 	"fmt"
+	"io"
 	files "io/fs"
 	"log"
 	"os"
+	"path"
+	"path/filepath"
 	"strings"
 
 	"github.com/anderseknert/roast/pkg/encoding"
@@ -101,3 +104,21 @@ func ExcludeTestFilter() filter.LoaderFilter {
 			info.Name() != "todo_test.rego"
 	}
 }
+
+// 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) {
+	relative := strings.TrimPrefix(file, workspacePath)
+	components := strings.Split(path.Dir(relative), string(filepath.Separator))
+
+	for i := range len(components) {
+		inputPath := path.Join(workspacePath, path.Join(components[:len(components)-i]...), "input.json")
+
+		f, err := os.Open(inputPath)
+		if err == nil {
+			return inputPath, f
+		}
+	}
+
+	return "", nil
+}
diff --git a/internal/lsp/completions/providers/policy.go b/internal/lsp/completions/providers/policy.go
@@ -2,8 +2,10 @@ package providers
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"os"
 
 	"github.com/open-policy-agent/opa/ast"
@@ -71,6 +73,23 @@ func (p *Policy) Run(c *cache.Cache, params types.CompletionParams, opts *Option
 	inputContext["workspace_root"] = uri.ToPath(opts.ClientIdentifier, opts.RootURI)
 	inputContext["path_separator"] = string(os.PathSeparator)
 
+	workspacePath := uri.ToPath(opts.ClientIdentifier, opts.RootURI)
+	inputDotJSONPath, inputDotJSONReader := 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
+			}
+		}
+	}
+
 	input, err := rego2.ToInput(
 		params.TextDocument.URI,
 		opts.ClientIdentifier,

diff --git a/internal/lsp/eval.go b/internal/lsp/eval.go
@@ -5,10 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"os"
-	"path"
-	"path/filepath"
-	"strings"
 
 	"github.com/anderseknert/roast/pkg/encoding"
 
@@ -97,22 +93,6 @@ type EvalPathResult struct {
 	PrintOutput map[int][]string `json:"printOutput"`
 }
 
-func FindInput(file string, workspacePath string) io.Reader {
-	relative := strings.TrimPrefix(file, workspacePath)
-	components := strings.Split(path.Dir(relative), string(filepath.Separator))
-
-	for i := range len(components) {
-		inputPath := path.Join(workspacePath, path.Join(components[:len(components)-i]...), "input.json")
-
-		f, err := os.Open(inputPath)
-		if err == nil {
-			return f
-		}
-	}
-
-	return nil
-}
-
 func (l *LanguageServer) EvalWorkspacePath(
 	ctx context.Context,
 	query string,

diff --git a/internal/lsp/eval_test.go b/internal/lsp/eval_test.go
@@ -7,6 +7,7 @@ import (
 	"strings"
 	"testing"
 
+	rio "github.com/styrainc/regal/internal/io"
 	"github.com/styrainc/regal/internal/parse"
 )
 
@@ -108,7 +109,7 @@ func createWithContent(t *testing.T, path string, content string) {
 func readInputString(t *testing.T, file, workspacePath string) string {
 	t.Helper()
 
-	input := FindInput(file, workspacePath)
+	_, input := rio.FindInput(file, workspacePath)
 
 	if input == nil {
 		return ""

diff --git a/internal/lsp/server.go b/internal/lsp/server.go
@@ -506,7 +506,7 @@ func (l *LanguageServer) StartCommandWorker(ctx context.Context) {
 				ruleHeadLocations := allRuleHeadLocations[path]
 
 				workspacePath := uri.ToPath(l.clientIdentifier, l.workspaceRootURI)
-				input := FindInput(uri.ToPath(l.clientIdentifier, file), workspacePath)
+				_, input := rio.FindInput(uri.ToPath(l.clientIdentifier, file), workspacePath)
 
 				result, err := l.EvalWorkspacePath(ctx, path, input)
 				if err != nil {