From 825159556917d4d5c551818100f3b557b617b549 Mon Sep 17 00:00:00 2001
From: Charlie Egan <charlie@styra.com>
Date: Thu, 9 Jan 2025 20:02:53 +0000
Subject: [PATCH] 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 <charlie@styra.com>
---
 .../completion/providers/regov1/regov1.rego    |  1 +
 .../providers/regov1/regov1_test.rego          | 18 +++++++++++++++---
 .../providers/test_utils/test_utils.rego       | 13 +++++++++++++
 internal/embeds/schemas/regal-ast.json         |  5 +++++
 internal/lsp/completions/providers/options.go  |  1 +
 internal/lsp/completions/providers/policy.go   |  1 +
 internal/lsp/server.go                         | 12 +++++-------
 internal/lsp/types/types.go                    |  3 +++
 8 files changed, 44 insertions(+), 10 deletions(-)

diff --git a/bundle/regal/lsp/completion/providers/regov1/regov1.rego b/bundle/regal/lsp/completion/providers/regov1/regov1.rego
index 0620f29f..15df96e1 100644
--- a/bundle/regal/lsp/completion/providers/regov1/regov1.rego
+++ b/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]
diff --git a/bundle/regal/lsp/completion/providers/regov1/regov1_test.rego b/bundle/regal/lsp/completion/providers/regov1/regov1_test.rego
index a48b8170..a9069fba 100644
--- a/bundle/regal/lsp/completion/providers/regov1/regov1_test.rego
+++ b/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()
 }
diff --git a/bundle/regal/lsp/completion/providers/test_utils/test_utils.rego b/bundle/regal/lsp/completion/providers/test_utils/test_utils.rego
index 64da237a..d0b9c712 100644
--- a/bundle/regal/lsp/completion/providers/test_utils/test_utils.rego
+++ b/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,
+	},
+}}
diff --git a/internal/embeds/schemas/regal-ast.json b/internal/embeds/schemas/regal-ast.json
index a878a9d0..98b52e85 100644
--- a/internal/embeds/schemas/regal-ast.json
+++ b/internal/embeds/schemas/regal-ast.json
@@ -413,6 +413,11 @@
             },
             "input_dot_json_path": {
               "type": "string"
+            },
+            "rego_version": {
+              "type": "integer",
+              "minimum": 0,
+              "maximum": 3
             }
           }
         }
diff --git a/internal/lsp/completions/providers/options.go b/internal/lsp/completions/providers/options.go
index b9b0ae75..8f20f634 100644
--- a/internal/lsp/completions/providers/options.go
+++ b/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
 }
diff --git a/internal/lsp/completions/providers/policy.go b/internal/lsp/completions/providers/policy.go
index 53134bfa..44a3b47c 100644
--- a/internal/lsp/completions/providers/policy.go
+++ b/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)
 
diff --git a/internal/lsp/server.go b/internal/lsp/server.go
index 58672c2e..94dc5a30 100644
--- a/internal/lsp/server.go
+++ b/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}
diff --git a/internal/lsp/types/types.go b/internal/lsp/types/types.go
index 9adf2865..5c67b77e 100644
--- a/internal/lsp/types/types.go
+++ b/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 {