lsp/templating: gracefully handle unknown root

internal/lsp/server.go

@@ -356,10 +356,10 @@ func (l *LanguageServer) StartDiagnosticsWorker(ctx context.Context) {
 			case <-ctx.Done():
 				return
 			case job := <-workspaceLintRuns:
-				// if there are no files in the cache, then there is no need to
-				// run the aggregate report. This can happen if the server is
-				// very slow to start up.
-				if len(l.cache.GetAllFiles()) == 0 {
+				// if there are no parsed modules in the cache, then there is
+				// no need to run the aggregate report. This can happen if the
+				// server is very slow to start up.
+				if len(l.cache.GetAllModules()) == 0 {
 					continue
 				}
 
@@ -986,6 +986,12 @@ func (l *LanguageServer) StartTemplateWorker(ctx context.Context) {
 		case <-ctx.Done():
 			return
 		case job := <-l.templateFileJobs:
+			// disable the templating feature for files in the workspace root.
+			if filepath.Dir(uri.ToPath(l.clientIdentifier, job.URI)) ==
+				uri.ToPath(l.clientIdentifier, l.workspaceRootURI) {
+				continue
+			}
+
 			// determine the new contents for the file, if permitted
 			newContents, err := l.templateContentsForFile(job.URI)
 			if err != nil {
@@ -1084,6 +1090,13 @@ func (l *LanguageServer) StartWebServer(ctx context.Context) {
 }
 
 func (l *LanguageServer) templateContentsForFile(fileURI string) (string, error) {
+	// this function should not be called with files in the root, but if it is,
+	// then it is an error to prevent unwanted behavior.
+	if filepath.Dir(uri.ToPath(l.clientIdentifier, fileURI)) ==
+		uri.ToPath(l.clientIdentifier, l.workspaceRootURI) {
+		return "", errors.New("this function does not template files in the workspace root")
+	}
+
 	content, ok := l.cache.GetFileContents(fileURI)
 	if !ok {
 		return "", fmt.Errorf("failed to get file contents for URI %q", fileURI)
@@ -1109,6 +1122,16 @@ func (l *LanguageServer) templateContentsForFile(fileURI string) (string, error)
 		return "", fmt.Errorf("failed to get potential roots during templating of new file: %w", err)
 	}
 
+	// handle the case where the root is unknown by providing the server's root
+	// dir as a defacto root. This allows templating of files when there is no
+	// known root, but the package could be determined based on the file path
+	// relative to the server's workspace root
+	if len(roots) == 1 && roots[0] == dir {
+		roots = []string{uri.ToPath(l.clientIdentifier, l.workspaceRootURI)}
+	}
+
+	roots = append(roots, uri.ToPath(l.clientIdentifier, l.workspaceRootURI))
+
 	longestPrefixRoot := ""
 
 	for _, root := range roots {
@@ -2024,6 +2047,12 @@ func (l *LanguageServer) handleTextDocumentFormatting(
 		oldContent, ok = l.cache.GetFileContents(params.TextDocument.URI)
 	}
 
+	// disable the templating feature for files in the workspace root.
+	if filepath.Dir(uri.ToPath(l.clientIdentifier, params.TextDocument.URI)) ==
+		uri.ToPath(l.clientIdentifier, l.workspaceRootURI) {
+		return []types.TextEdit{}, nil
+	}
+
 	// if the file is empty, then the formatters will fail, so we template
 	// instead
 	if oldContent == "" {

internal/lsp/server_formatting_test.go

@@ -3,6 +3,7 @@ package lsp
 import (
 	"context"
 	"encoding/json"
+	"path/filepath"
 	"testing"
 
 	"github.com/sourcegraph/jsonrpc2"
@@ -30,7 +31,7 @@ func TestFormatting(t *testing.T) {
 		t.Fatalf("failed to create and init language server: %s", err)
 	}
 
-	mainRegoURI := fileURIScheme + tempDir + mainRegoFileName
+	mainRegoURI := fileURIScheme + filepath.Join(tempDir, "main/main.rego")
 
 	// Simple as possible — opa fmt should just remove a newline
 	content := `package main

internal/lsp/server_template_test.go

@@ -41,14 +41,6 @@ func TestTemplateContentsForFile(t *testing.T) {
 			},
 			ExpectedError: "file on disk already has contents",
 		},
-		"empty file is templated as main when no root": {
-			FileKey:           "foo/bar.rego",
-			CacheFileContents: "",
-			DiskContents: map[string]string{
-				"foo/bar.rego": "",
-			},
-			ExpectedContents: "package main\n\nimport rego.v1\n",
-		},
 		"empty file is templated based on root": {
 			FileKey:           "foo/bar.rego",
 			CacheFileContents: "",
@@ -100,7 +92,6 @@ func TestTemplateContentsForFile(t *testing.T) {
 
 			ctx := context.Background()
 
-			// create a new language server
 			s := NewLanguageServer(ctx, &LanguageServerOptions{ErrorLog: newTestLogger(t)})
 			s.workspaceRootURI = uri.FromPath(clients.IdentifierGeneric, td)
 
@@ -124,6 +115,73 @@ func TestTemplateContentsForFile(t *testing.T) {
 	}
 }
 
+func TestTemplateContentsForFileInWorkspaceRoot(t *testing.T) {
+	t.Parallel()
+
+	td := t.TempDir()
+
+	err := os.MkdirAll(filepath.Join(td, ".regal"), 0o755)
+	if err != nil {
+		t.Fatalf("failed to create directory %s: %s", filepath.Join(td, ".regal"), err)
+	}
+
+	err = os.WriteFile(filepath.Join(td, ".regal/config.yaml"), []byte{}, 0o600)
+	if err != nil {
+		t.Fatalf("failed to create file %s: %s", filepath.Join(td, ".regal"), err)
+	}
+
+	ctx := context.Background()
+
+	s := NewLanguageServer(ctx, &LanguageServerOptions{ErrorLog: newTestLogger(t)})
+	s.workspaceRootURI = uri.FromPath(clients.IdentifierGeneric, td)
+
+	fileURI := uri.FromPath(clients.IdentifierGeneric, filepath.Join(td, "foo.rego"))
+
+	s.cache.SetFileContents(fileURI, "")
+
+	_, err = s.templateContentsForFile(fileURI)
+	if err == nil {
+		t.Fatalf("expected error")
+	}
+
+	if !strings.Contains(err.Error(), "this function does not template files in the workspace root") {
+		t.Fatalf("expected error about root templating, got %s", err.Error())
+	}
+}
+
+func TestTemplateContentsForFileWithUnknownRoot(t *testing.T) {
+	t.Parallel()
+
+	td := t.TempDir()
+
+	ctx := context.Background()
+
+	s := NewLanguageServer(ctx, &LanguageServerOptions{ErrorLog: newTestLogger(t)})
+	s.workspaceRootURI = uri.FromPath(clients.IdentifierGeneric, td)
+
+	err := os.MkdirAll(filepath.Join(td, "foo"), 0o755)
+	if err != nil {
+		t.Fatalf("failed to create directory %s: %s", filepath.Join(td, "foo"), err)
+	}
+
+	fileURI := uri.FromPath(clients.IdentifierGeneric, filepath.Join(td, "foo/bar.rego"))
+
+	s.cache.SetFileContents(fileURI, "")
+
+	newContents, err := s.templateContentsForFile(fileURI)
+	if err != nil {
+		t.Fatalf("unexpected error: %s", err)
+	}
+
+	exp := `package foo
+
+import rego.v1
+`
+	if exp != newContents {
+		t.Errorf("unexpected content: %s, want %s", newContents, exp)
+	}
+}
+
 func TestNewFileTemplating(t *testing.T) {
 	t.Parallel()