44 "context"
55 "errors"
66 "fmt"
7+ "os"
8+ "path/filepath"
9+ "regexp"
710 "strings"
811
912 "github.com/open-policy-agent/opa/v1/ast"
@@ -13,10 +16,15 @@ import (
1316 "github.com/open-policy-agent/opa/v1/topdown/print"
1417
1518 rbundle "github.com/open-policy-agent/regal/bundle"
19+ rio "github.com/open-policy-agent/regal/internal/io"
20+ rrego "github.com/open-policy-agent/regal/internal/lsp/rego"
1621 rquery "github.com/open-policy-agent/regal/internal/lsp/rego/query"
22+ "github.com/open-policy-agent/regal/internal/lsp/types"
1723 "github.com/open-policy-agent/regal/internal/lsp/uri"
24+ rparse "github.com/open-policy-agent/regal/internal/parse"
1825 "github.com/open-policy-agent/regal/internal/util"
1926 "github.com/open-policy-agent/regal/pkg/config"
27+ "github.com/open-policy-agent/regal/pkg/roast/encoding"
2028 "github.com/open-policy-agent/regal/pkg/roast/transform"
2129
2230 _ "github.com/open-policy-agent/regal/pkg/builtins"
2937 Roots : & []string {"workspace" }, // no data in this bundle so no roots are used, however, roots must be set
3038 Metadata : map [string ]any {"name" : "workspace" },
3139 }
40+ regalEvalUseAsInputComment = regexp .MustCompile (`^\s*regal eval:\s*use-as-input` )
3241)
3342
3443type EvalResult struct {
@@ -45,6 +54,131 @@ type PrintHook struct {
4554 FileNameBase string
4655}
4756
57+ func (l * LanguageServer ) handleEvalCommand (ctx context.Context , args types.CommandArgs ) error {
58+ if args .Target == "" || args .Query == "" {
59+ l .log .Message ("expected command target and query, got target %q, query %q" , args .Target , args .Query )
60+
61+ return nil
62+ }
63+
64+ contents , module , ok := l .cache .GetContentAndModule (args .Target )
65+ if ! ok {
66+ l .log .Message ("failed to get content or module for file %q" , args .Target )
67+
68+ return nil
69+ }
70+
71+ var pq * rquery.Prepared
72+
73+ pq , err := l .queryCache .GetOrSet (ctx , l .regoStore , rquery .RuleHeadLocations )
74+ if err != nil {
75+ l .log .Message ("failed to prepare query %s" , rquery .RuleHeadLocations , err )
76+
77+ return nil
78+ }
79+
80+ var allRuleHeadLocations rrego.RuleHeads
81+
82+ allRuleHeadLocations , err = rrego .AllRuleHeadLocations (
83+ ctx , pq , filepath .Base (uri .ToPath (args .Target )), contents , module ,
84+ )
85+ if err != nil {
86+ l .log .Message ("failed to get rule head locations: %s" , err )
87+
88+ return nil
89+ }
90+
91+ // if there are none, then it's a package evaluation
92+ ruleHeadLocations := allRuleHeadLocations [args .Query ]
93+
94+ var inputMap map [string ]any
95+
96+ var inputPath string
97+
98+ // When the first comment in the file is `regal eval: use-as-input`, the AST of that module is
99+ // used as the input rather than the contents of input.json/yaml. This is a development feature for
100+ // working on rules (built-in or custom), allowing querying the AST of the module directly.
101+ if len (module .Comments ) > 0 && regalEvalUseAsInputComment .Match (module .Comments [0 ].Text ) {
102+ //nolint:staticcheck
103+ inputMap , err = rparse .PrepareAST (l .toRelativePath (args .Target ), contents , module )
104+ if err != nil {
105+ l .log .Message ("failed to prepare module: %s" , err )
106+
107+ return nil
108+ }
109+ } else {
110+ // Normal mode — try to find the input.json/yaml file in the workspace and use as input
111+ // NOTE that we don't break on missing input, as some rules don't depend on that, and should
112+ // still be evaluable. We may consider returning some notice to the user though.
113+ inputPath , inputMap = rio .FindInput (uri .ToPath (args .Target ), l .workspacePath ())
114+
115+ ruleName := strings .TrimPrefix (args .Query , module .Package .Path .String ()+ "." )
116+
117+ if inputPath == "" && ! l .supressInputPrompt {
118+ created , err := l .handleInputSkeletonPrompt (ctx , args .Target , ruleName , args .Row )
119+ // Bubbling up an error here if input.json creation fails for any reason.
120+ if err != nil {
121+ return err
122+ }
123+
124+ if created {
125+ return nil
126+ }
127+ }
128+ }
129+
130+ var result EvalResult
131+
132+ if result , err = l .EvalInWorkspace (ctx , args .Query , inputMap ); err != nil {
133+ return fmt .Errorf ("failed to evaluate workspace path: %w" , err )
134+ }
135+
136+ target := "package"
137+ if len (ruleHeadLocations ) > 0 {
138+ target = strings .TrimPrefix (args .Query , module .Package .Path .String ()+ "." )
139+ }
140+
141+ if l .featureFlags .InlineEvaluationProvider && l .getClient ().SupportsEvalCodelensDisplayInline () {
142+ responseParams := map [string ]any {
143+ "result" : result ,
144+ "line" : args .Row ,
145+ "target" : target ,
146+ // only used when the target is 'package'
147+ "package" : strings .TrimPrefix (module .Package .Path .String (), "data." ),
148+ // only used when the target is a rule
149+ "rule_head_locations" : ruleHeadLocations ,
150+ }
151+
152+ responseResult := map [string ]any {}
153+
154+ // Use a timeout context for RPC to ensure it completes during graceful shutdown
155+ rpcCtx , rpcCancel := context .WithTimeout (context .Background (), rpcTimeout )
156+
157+ //nolint:contextcheck
158+ if err = l .conn .Call (rpcCtx , "regal/showEvalResult" , responseParams , & responseResult ); err != nil {
159+ l .log .Message ("regal/showEvalResult failed: %v" , err .Error ())
160+ }
161+
162+ rpcCancel ()
163+ } else {
164+ output := filepath .Join (l .workspacePath (), "output.json" )
165+
166+ var f * os.File
167+ if f , err = os .OpenFile (output , os .O_CREATE | os .O_TRUNC | os .O_WRONLY , 0o755 ); err == nil {
168+ value := result .Value
169+ if result .IsUndefined {
170+ value = emptyStringAnyMap // undefined displays as an empty object
171+ }
172+
173+ err = encoding .NewIndentEncoder (f , "" , " " ).Encode (value )
174+
175+ rio .CloseIgnore (f )
176+ }
177+ }
178+
179+ return err
180+ }
181+
48182func (l * LanguageServer ) Eval (
49183 ctx context.Context , query string , input map [string ]any , printHook print.Hook ,
50184) (rego.ResultSet , error ) {
0 commit comments