Skip to content

Commit a318e6c

Browse files
anderseknertanderseknert
authoredAug 7, 2024··
lsp: Add code lens support for evaluating rules (#968)
- This is currently limited to VS Code clients, but may be extended to other clients too in the furure, if they can support it on their side. - Implementation wil recursively look for an `input.json` file and pick the one closest to the file evaluated, down to the level of the workspace root (if any file is found). Signed-off-by: Anders Eknert <anders@styra.com>
·
v0.32.0v0.25.0
1 parent 95d1eb1 commit a318e6c

File tree

7 files changed

+378
-11
lines changed

7 files changed

+378
-11
lines changed
 

‎README.md

+1-1
Original file line numberDiff line numberDiff line change
@@ -702,7 +702,7 @@ in the near future:
702702

703703
- [ ] Make "Check on save" unnecessary by allowing diagnostics to include
704704
[compilation errors](https://github.com/StyraInc/regal/issues/745)
705-
- [ ] Add Code Lens to "Evaluate" any rule or package (VS Code only, initially)
705+
- [x] Add Code Lens to "Evaluate" any rule or package (VS Code only, initially)
706706
- [ ] Implement [Signature Help](https://github.com/StyraInc/regal/issues/695) feature
707707

708708
The roadmap is updated when all the current items have been completed.

‎bundle/regal/ast/search.rego

-5
Original file line numberDiff line numberDiff line change
@@ -287,11 +287,6 @@ find_some_decl_names_in_scope(rule, location) := {some_var.value |
287287
_before_location(rule, some_var, location)
288288
}
289289

290-
# _rules_with_bodies[rule_index] := rule if {
291-
# some rule_index, rule in input.rules
292-
# not generated_body(rule)
293-
# }
294-
295290
exprs[rule_index][expr_index] := expr if {
296291
some rule_index, rule in input.rules
297292
some expr_index, expr in rule.body

‎bundle/regal/main.rego

+2-2
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,10 @@ lint.aggregates := aggregate
1212

1313
lint.ignore_directives[input.regal.file.name] := ast.ignore_directives
1414

15-
lint.violations := report
16-
1715
lint_aggregate.violations := aggregate_report
1816

17+
lint.violations := report
18+
1919
rules_to_run[category][title] if {
2020
some category, title
2121
config.merged_config.rules[category][title]

‎internal/lsp/eval.go

+113
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
package lsp
2+
3+
import (
4+
"context"
5+
"encoding/json"
6+
"errors"
7+
"fmt"
8+
"os"
9+
"path"
10+
"path/filepath"
11+
"strings"
12+
13+
"github.com/open-policy-agent/opa/ast"
14+
"github.com/open-policy-agent/opa/bundle"
15+
"github.com/open-policy-agent/opa/rego"
16+
"github.com/open-policy-agent/opa/topdown"
17+
18+
"github.com/styrainc/regal/internal/lsp/uri"
19+
"github.com/styrainc/regal/pkg/builtins"
20+
)
21+
22+
func (l *LanguageServer) Eval(ctx context.Context, query string, input string) (rego.ResultSet, error) {
23+
modules := l.cache.GetAllModules()
24+
moduleFiles := make([]bundle.ModuleFile, 0, len(modules))
25+
26+
for fileURI, module := range modules {
27+
moduleFiles = append(moduleFiles, bundle.ModuleFile{
28+
URL: fileURI,
29+
Parsed: module,
30+
Path: uri.ToPath(l.clientIdentifier, fileURI),
31+
})
32+
}
33+
34+
bd := bundle.Bundle{
35+
Data: make(map[string]any),
36+
Manifest: bundle.Manifest{
37+
Roots: &[]string{""},
38+
Metadata: map[string]any{"name": "workspace"},
39+
},
40+
Modules: moduleFiles,
41+
}
42+
43+
regoArgs := prepareRegoArgs(ast.MustParseBody(query), bd)
44+
45+
pq, err := rego.New(regoArgs...).PrepareForEval(ctx)
46+
if err != nil {
47+
return nil, fmt.Errorf("failed preparing query: %w", err)
48+
}
49+
50+
if input != "" {
51+
inputMap := make(map[string]any)
52+
53+
err = json.Unmarshal([]byte(input), &inputMap)
54+
if err != nil {
55+
return nil, fmt.Errorf("failed unmarshalling input: %w", err)
56+
}
57+
58+
return pq.Eval(ctx, rego.EvalInput(inputMap)) //nolint:wrapcheck
59+
}
60+
61+
return pq.Eval(ctx) //nolint:wrapcheck
62+
}
63+
64+
type EvalPathResult struct {
65+
Value any `json:"value"`
66+
IsUndefined bool `json:"isUndefined"`
67+
}
68+
69+
func FindInput(file string, workspacePath string) string {
70+
relative := strings.TrimPrefix(file, workspacePath)
71+
components := strings.Split(path.Dir(relative), string(filepath.Separator))
72+
73+
for i := range len(components) {
74+
inputPath := path.Join(workspacePath, path.Join(components[:len(components)-i]...), "input.json")
75+
76+
if input, err := os.ReadFile(inputPath); err == nil {
77+
return string(input)
78+
}
79+
}
80+
81+
return ""
82+
}
83+
84+
func (l *LanguageServer) EvalWorkspacePath(ctx context.Context, query string, input string) (EvalPathResult, error) {
85+
resultQuery := "result := " + query
86+
87+
result, err := l.Eval(ctx, resultQuery, input)
88+
if err != nil {
89+
return EvalPathResult{}, fmt.Errorf("failed evaluating query: %w", err)
90+
}
91+
92+
if len(result) == 0 {
93+
return EvalPathResult{IsUndefined: true}, nil
94+
}
95+
96+
res, ok := result[0].Bindings["result"]
97+
if !ok {
98+
return EvalPathResult{}, errors.New("expected result in bindings, didn't get it")
99+
}
100+
101+
return EvalPathResult{Value: res}, nil
102+
}
103+
104+
func prepareRegoArgs(query ast.Body, bd bundle.Bundle) []func(*rego.Rego) {
105+
return []func(*rego.Rego){
106+
rego.ParsedQuery(query),
107+
rego.ParsedBundle("workspace", &bd),
108+
rego.Function2(builtins.RegalParseModuleMeta, builtins.RegalParseModule),
109+
rego.Function1(builtins.RegalLastMeta, builtins.RegalLast),
110+
rego.EnablePrintStatements(true),
111+
rego.PrintHook(topdown.NewPrintHook(os.Stderr)),
112+
}
113+
}

‎internal/lsp/eval_test.go

+108
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
package lsp
2+
3+
import (
4+
"context"
5+
"os"
6+
"testing"
7+
8+
"github.com/styrainc/regal/internal/parse"
9+
)
10+
11+
func TestEvalWorkspacePath(t *testing.T) {
12+
t.Parallel()
13+
14+
ls := NewLanguageServer(&LanguageServerOptions{ErrorLog: os.Stderr})
15+
16+
policy1 := `package policy1
17+
18+
import rego.v1
19+
20+
import data.policy2
21+
22+
default allow := false
23+
24+
allow if policy2.allow
25+
`
26+
27+
policy2 := `package policy2
28+
29+
import rego.v1
30+
31+
allow if input.exists
32+
`
33+
34+
module1 := parse.MustParseModule(policy1)
35+
module2 := parse.MustParseModule(policy2)
36+
37+
ls.cache.SetFileContents("file://policy1.rego", policy1)
38+
ls.cache.SetFileContents("file://policy2.rego", policy2)
39+
ls.cache.SetModule("file://policy1.rego", module1)
40+
ls.cache.SetModule("file://policy2.rego", module2)
41+
42+
res, err := ls.EvalWorkspacePath(context.TODO(), "data.policy1.allow", `{"exists": true}`)
43+
if err != nil {
44+
t.Fatal(err)
45+
}
46+
47+
empty := EvalPathResult{}
48+
49+
if res == empty {
50+
t.Fatal("expected result, got nil")
51+
}
52+
53+
if val, ok := res.Value.(bool); !ok || val != true {
54+
t.Fatalf("expected true, got false")
55+
}
56+
}
57+
58+
func TestFindInput(t *testing.T) {
59+
t.Parallel()
60+
61+
tmpDir := t.TempDir()
62+
63+
workspacePath := tmpDir + "/workspace"
64+
file := tmpDir + "/workspace/foo/bar/baz.rego"
65+
66+
if err := os.MkdirAll(workspacePath+"/foo/bar", 0o755); err != nil {
67+
t.Fatal(err)
68+
}
69+
70+
if FindInput(file, workspacePath) != "" {
71+
t.Fatalf("did not expect to find input.json")
72+
}
73+
74+
content := `{"x": 1}`
75+
76+
createWithContent(t, tmpDir+"/workspace/foo/bar/input.json", content)
77+
78+
if res := FindInput(file, workspacePath); res != content {
79+
t.Errorf("expected input at %s, got %s", content, res)
80+
}
81+
82+
err := os.Remove(tmpDir + "/workspace/foo/bar/input.json")
83+
if err != nil {
84+
t.Fatal(err)
85+
}
86+
87+
createWithContent(t, tmpDir+"/workspace/input.json", content)
88+
89+
if res := FindInput(file, workspacePath); res != content {
90+
t.Errorf("expected input at %s, got %s", content, res)
91+
}
92+
}
93+
94+
func createWithContent(t *testing.T, path string, content string) {
95+
t.Helper()
96+
97+
f, err := os.Create(path)
98+
if err != nil {
99+
t.Fatal(err)
100+
}
101+
102+
defer f.Close()
103+
104+
_, err = f.WriteString(content)
105+
if err != nil {
106+
t.Fatal(err)
107+
}
108+
}

‎internal/lsp/server.go

+139-3
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ type fileUpdateEvent struct {
109109
Content string
110110
}
111111

112+
//nolint:gocyclo
112113
func (l *LanguageServer) Handle(
113114
ctx context.Context,
114115
conn *jsonrpc2.Conn,
@@ -148,6 +149,8 @@ func (l *LanguageServer) Handle(
148149
return l.handleTextDocumentHover(ctx, conn, req)
149150
case "textDocument/inlayHint":
150151
return l.handleTextDocumentInlayHint(ctx, conn, req)
152+
case "textDocument/codeLens":
153+
return l.handleTextDocumentCodeLens(ctx, conn, req)
151154
case "textDocument/completion":
152155
return l.handleTextDocumentCompletion(ctx, conn, req)
153156
case "workspace/didChangeWatchedFiles":
@@ -447,6 +450,59 @@ func (l *LanguageServer) StartCommandWorker(ctx context.Context) {
447450
commands.ParseOptions{TargetArgIndex: 0, RowArgIndex: 1, ColArgIndex: 2},
448451
params,
449452
)
453+
case "regal.eval":
454+
fmt.Fprintf(os.Stderr, "regal.eval called with params: %v\n", params)
455+
456+
if len(params.Arguments) != 3 {
457+
l.logError(fmt.Errorf("expected three arguments, got %d", len(params.Arguments)))
458+
459+
break
460+
}
461+
462+
file, ok := params.Arguments[0].(string)
463+
if !ok {
464+
l.logError(fmt.Errorf("expected first argument to be a string, got %T", params.Arguments[0]))
465+
466+
break
467+
}
468+
469+
path, ok := params.Arguments[1].(string)
470+
if !ok {
471+
l.logError(fmt.Errorf("expected second argument to be a string, got %T", params.Arguments[1]))
472+
473+
break
474+
}
475+
476+
line, ok := params.Arguments[2].(float64)
477+
if !ok {
478+
l.logError(fmt.Errorf("expected third argument to be a number, got %T", params.Arguments[2]))
479+
480+
break
481+
}
482+
483+
input := FindInput(
484+
uri.ToPath(l.clientIdentifier, file),
485+
uri.ToPath(l.clientIdentifier, l.workspaceRootURI),
486+
)
487+
488+
result, err := l.EvalWorkspacePath(ctx, path, input)
489+
if err != nil {
490+
fmt.Fprintf(os.Stderr, "failed to evaluate workspace path: %v\n", err)
491+
492+
break
493+
}
494+
495+
responseParams := map[string]any{
496+
"result": result,
497+
"line": line,
498+
}
499+
500+
responseResult := map[string]any{}
501+
502+
err = l.conn.Call(ctx, "regal/showEvalResult", responseParams, &responseResult)
503+
if err != nil {
504+
l.logError(fmt.Errorf("failed %s notify: %v", "regal/hello", err.Error()))
505+
}
450506
}
451507

452508
if err != nil {
@@ -934,6 +990,78 @@ func (l *LanguageServer) handleTextDocumentInlayHint(
934990
return inlayHints, nil
935991
}
936992

993+
func (l *LanguageServer) handleTextDocumentCodeLens(
994+
_ context.Context,
995+
_ *jsonrpc2.Conn,
996+
req *jsonrpc2.Request,
997+
) (result any, err error) {
998+
var params types.CodeLensParams
999+
if err := json.Unmarshal(*req.Params, &params); err != nil {
1000+
return nil, fmt.Errorf("failed to unmarshal params: %w", err)
1001+
}
1002+
1003+
if l.clientIdentifier != clients.IdentifierVSCode {
1004+
// only VSCode has the client side capability to handle the callback request
1005+
// to handle the result of evaluation, so displaying code lenses for any other
1006+
// editor is likely just going to result in a bad experience
1007+
return nil, nil // return a null response, as per the spec
1008+
}
1009+
1010+
module, ok := l.cache.GetModule(params.TextDocument.URI)
1011+
if !ok {
1012+
l.logError(fmt.Errorf("failed to get module for uri %q", params.TextDocument.URI))
1013+
1014+
// return a null response, as per the spec
1015+
return nil, nil
1016+
}
1017+
1018+
codeLenses := make([]types.CodeLens, 0)
1019+
1020+
// Package
1021+
1022+
pkgLens := types.CodeLens{
1023+
Range: locationToRange(module.Package.Location),
1024+
Command: &types.Command{
1025+
Title: "Evaluate",
1026+
Command: "regal.eval",
1027+
Arguments: &[]any{
1028+
module.Package.Location.File,
1029+
module.Package.Path.String(),
1030+
module.Package.Location.Row,
1031+
},
1032+
},
1033+
}
1034+
1035+
codeLenses = append(codeLenses, pkgLens)
1036+
1037+
// Rules
1038+
1039+
for _, rule := range module.Rules {
1040+
if rule.Head.Args != nil {
1041+
// Skip functions for now, as it's not clear how to best
1042+
// provide inputs for them.
1043+
continue
1044+
}
1045+
1046+
ruleLens := types.CodeLens{
1047+
Range: locationToRange(rule.Location),
1048+
Command: &types.Command{
1049+
Title: "Evaluate",
1050+
Command: "regal.eval",
1051+
Arguments: &[]any{
1052+
module.Package.Location.File,
1053+
module.Package.Path.String() + "." + string(rule.Head.Name),
1054+
rule.Head.Location.Row,
1055+
},
1056+
},
1057+
}
1058+
1059+
codeLenses = append(codeLenses, ruleLens)
1060+
}
1061+
1062+
return codeLenses, nil
1063+
}
1064+
9371065
func (l *LanguageServer) handleTextDocumentCompletion(
9381066
_ context.Context,
9391067
_ *jsonrpc2.Conn,
@@ -1527,7 +1655,7 @@ func (l *LanguageServer) handleInitialize(
15271655
ctx context.Context,
15281656
_ *jsonrpc2.Conn,
15291657
req *jsonrpc2.Request,
1530-
) (result any, err error) {
1658+
) (any, error) {
15311659
var params types.InitializeParams
15321660
if err := json.Unmarshal(*req.Params, &params); err != nil {
15331661
return nil, fmt.Errorf("failed to unmarshal params: %w", err)
@@ -1554,7 +1682,7 @@ func (l *LanguageServer) handleInitialize(
15541682
},
15551683
}
15561684

1557-
result = types.InitializeResult{
1685+
initializeResult := types.InitializeResult{
15581686
Capabilities: types.ServerCapabilities{
15591687
TextDocumentSyncOptions: types.TextDocumentSyncOptions{
15601688
OpenClose: true,
@@ -1590,6 +1718,7 @@ func (l *LanguageServer) handleInitialize(
15901718
},
15911719
ExecuteCommandProvider: types.ExecuteCommandOptions{
15921720
Commands: []string{
1721+
"regal.eval",
15931722
"regal.fix.opa-fmt",
15941723
"regal.fix.use-rego-v1",
15951724
"regal.fix.use-assignment-operator",
@@ -1610,6 +1739,13 @@ func (l *LanguageServer) handleInitialize(
16101739
},
16111740
}
16121741

1742+
// Since evaluation requires some client side handling, this can't be supported
1743+
// purely by the LSP. Clients that are capable of handling the code lens callback
1744+
// should be added here though.
1745+
if l.clientIdentifier == clients.IdentifierVSCode {
1746+
initializeResult.Capabilities.CodeLensProvider = &types.CodeLensOptions{}
1747+
}
1748+
16131749
if l.workspaceRootURI != "" {
16141750
configFile, err := config.FindConfig(uri.ToPath(l.clientIdentifier, l.workspaceRootURI))
16151751
if err == nil {
@@ -1624,7 +1760,7 @@ func (l *LanguageServer) handleInitialize(
16241760
l.diagnosticRequestWorkspace <- "server initialize"
16251761
}
16261762

1627-
return result, nil
1763+
return initializeResult, nil
16281764
}
16291765

16301766
func (l *LanguageServer) loadWorkspaceContents(ctx context.Context, newOnly bool) ([]string, error) {

‎internal/lsp/types/types.go

+15
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,7 @@ type ServerCapabilities struct {
105105
WorkspaceSymbolProvider bool `json:"workspaceSymbolProvider"`
106106
DefinitionProvider bool `json:"definitionProvider"`
107107
CompletionProvider CompletionOptions `json:"completionProvider"`
108+
CodeLensProvider *CodeLensOptions `json:"codeLensProvider,omitempty"`
108109
}
109110

110111
type CompletionOptions struct {
@@ -189,6 +190,20 @@ type CodeAction struct {
189190
Command Command `json:"command"`
190191
}
191192

193+
type CodeLensOptions struct {
194+
ResolveProvider *bool `json:"resolveProvider,omitempty"`
195+
}
196+
197+
type CodeLensParams struct {
198+
TextDocument TextDocumentIdentifier `json:"textDocument"`
199+
}
200+
201+
type CodeLens struct {
202+
Range Range `json:"range"`
203+
Command *Command `json:"command,omitempty"`
204+
Data *any `json:"data,omitempty"`
205+
}
206+
192207
type Command struct {
193208
Title string `json:"title"`
194209
Tooltip string `json:"tooltip"`

0 commit comments

Comments
 (0)
Please sign in to comment.