lsp/completions: No rego.v1 import when in v1 file (#1327)

Also, use new version cache to determine version for formatting

Signed-off-by: Charlie Egan <[email protected]>

bundle/regal/lsp/completion/providers/regov1/regov1.rego

@@ -11,6 +11,7 @@ import data.regal.lsp.completion.location
 # description: completion suggestion for rego.v1
 items contains item if {
 	not strings.any_prefix_match(input.regal.file.lines, "import rego.v1")
+	not input.regal.context.rego_version == 3 # the rego.v1 import is not used in v1 rego
 
 	position := location.to_position(input.regal.context.location)
 	line := input.regal.file.lines[position.line]

bundle/regal/lsp/completion/providers/regov1/regov1_test.rego

@@ -8,7 +8,7 @@ test_regov1_completion_on_typing if {
 	policy := `package policy
 
 import r`
-	items := provider.items with input as util.input_with_location(policy, {"row": 3, "col": 9})
+	items := provider.items with input as util.input_with_location_and_version(policy, {"row": 3, "col": 9}, 0)
 	items == {{
 		"label": "rego.v1",
 		"kind": 9,
@@ -33,7 +33,7 @@ test_regov1_completion_on_invoked if {
 	policy := `package policy
 
 import `
-	items := provider.items with input as util.input_with_location(policy, {"row": 3, "col": 8})
+	items := provider.items with input as util.input_with_location_and_version(policy, {"row": 3, "col": 8}, 0)
 	items == {{
 		"label": "rego.v1",
 		"kind": 9,
@@ -60,6 +60,18 @@ test_no_regov1_completion_if_already_imported if {
 import rego.v1
 
 import r`
-	items := provider.items with input as util.input_with_location(policy, {"row": 5, "col": 9})
+	items := provider.items with input as util.input_with_location_and_version(policy, {"row": 5, "col": 9}, 0)
+	items == set()
+}
+
+test_no_regov1_completion_if_v1_file if {
+	policy := `package policy
+
+import r`
+	items := provider.items with input as util.input_with_location_and_version(
+		policy,
+		{"row": 3, "col": 9},
+		3, # RegoV1
+	)
 	items == set()
 }

bundle/regal/lsp/completion/providers/test_utils/test_utils.rego

@@ -28,3 +28,16 @@ input_with_location(policy, location) := {"regal": {
 	},
 	"context": {"location": location},
 }}
+
+# METADATA
+# description: same as input_with_location but with option to set rego_version too
+input_with_location_and_version(policy, location, rego_version) := {"regal": {
+	"file": {
+		"name": "p.rego",
+		"lines": split(policy, "\n"),
+	},
+	"context": {
+		"location": location,
+		"rego_version": rego_version,
+	},
+}}

internal/embeds/schemas/regal-ast.json

@@ -413,6 +413,11 @@
             },
             "input_dot_json_path": {
               "type": "string"
+            },
+            "rego_version": {
+              "type": "integer",
+              "minimum": 0,
+              "maximum": 3
             }
           }
         }

internal/lsp/completions/providers/options.go

@@ -12,4 +12,5 @@ type Options struct {
 	Builtins         map[string]*ast.Builtin
 	RootURI          string
 	ClientIdentifier clients.Identifier
+	RegoVersion      ast.RegoVersion
 }

internal/lsp/completions/providers/policy.go

@@ -67,6 +67,7 @@ func (p *Policy) Run(
 	inputContext["client_identifier"] = opts.ClientIdentifier
 	inputContext["workspace_root"] = uri.ToPath(opts.ClientIdentifier, opts.RootURI)
 	inputContext["path_separator"] = rio.PathSeparator
+	inputContext["rego_version"] = opts.RegoVersion
 
 	workspacePath := uri.ToPath(opts.ClientIdentifier, opts.RootURI)
 

internal/lsp/server.go

@@ -1764,6 +1764,7 @@ func (l *LanguageServer) handleTextDocumentCompletion(
 		ClientIdentifier: l.clientIdentifier,
 		RootURI:          l.workspaceRootURI,
 		Builtins:         l.builtinsForCurrentCapabilities(),
+		RegoVersion:      l.determineVersionForFile(params.TextDocument.URI),
 	})
 	if err != nil {
 		return nil, fmt.Errorf("failed to find completions: %w", err)
@@ -2137,8 +2138,6 @@ func (l *LanguageServer) handleTextDocumentFormatting(
 			return []types.TextEdit{}, nil
 		}
 
-		// TODO: We'll want to use version here too, to determine if "import rego.v1" should be used
-
 		newContent, err := l.templateContentsForFile(params.TextDocument.URI)
 		if err != nil {
 			return nil, fmt.Errorf("failed to template contents as a templating fallback: %w", err)
@@ -2168,13 +2167,12 @@ func (l *LanguageServer) handleTextDocumentFormatting(
 
 	switch formatter {
 	case "opa-fmt", "opa-fmt-rego-v1":
-		opts := format.Opts{}
+		opts := format.Opts{
+			RegoVersion: l.determineVersionForFile(params.TextDocument.URI),
+		}
+
 		if formatter == "opa-fmt-rego-v1" {
 			opts.RegoVersion = ast.RegoV0CompatV1
-		} else {
-			if module, ok := l.cache.GetModule(params.TextDocument.URI); ok {
-				opts.RegoVersion = module.RegoVersion()
-			}
 		}
 
 		f := &fixes.Fmt{OPAFmtOpts: opts}

internal/lsp/types/types.go

@@ -1,6 +1,8 @@
 package types
 
 import (
+	"github.com/open-policy-agent/opa/v1/ast"
+
 	"github.com/styrainc/regal/internal/lsp/types/completion"
 	"github.com/styrainc/regal/internal/lsp/types/symbols"
 )
@@ -129,6 +131,7 @@ type CompletionParams struct {
 	TextDocument TextDocumentIdentifier `json:"textDocument"`
 	Context      CompletionContext      `json:"context"`
 	Position     Position               `json:"position"`
+	RegoVersion  ast.RegoVersion        `json:"regoVersion"`
 }
 
 type CompletionContext struct {