chore: make development system filesystem-aware

Author: tstirrat15

internal/lsp/handlers.go

@@ -57,10 +57,11 @@ func (s *Server) computeDiagnostics(ctx context.Context, uri lsp.DocumentURI) ([
 			return &jsonrpc2.Error{Code: jsonrpc2.CodeInternalError, Message: "file not found"}
 		}
 
+		overlayFS := newLSPOverlayFS(uriToSourceDir(uri), files)
 		devCtx, devErrs, err := development.NewDevContext(ctx, &developerv1.RequestContext{
 			Schema:        file.contents,
 			Relationships: nil,
-		})
+		}, development.WithSourceFS(overlayFS))
 		if err != nil {
 			return err
 		}
@@ -197,7 +198,8 @@ func (s *Server) getCompiledContents(path lsp.DocumentURI, files *persistent.Map
 		return compiled, nil
 	}
 
-	justCompiled, derr, err := development.CompileSchema(file.contents)
+	overlayFS := newLSPOverlayFS(uriToSourceDir(path), files)
+	justCompiled, derr, err := development.CompileSchema(file.contents, development.WithSourceFS(overlayFS))
 	if err != nil {
 		return nil, err
 	}

internal/lsp/overlay.go

@@ -0,0 +1,76 @@
+package lsp
+
+import (
+	"io/fs"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"github.com/jzelinskie/persistent"
+	"github.com/sourcegraph/go-lsp"
+)
+
+// lspOverlayFS is an fs.FS that serves open editor files from memory
+// and falls back to disk for everything else. It is rooted at sourceDir.
+type lspOverlayFS struct {
+	base      fs.FS
+	sourceDir string
+	files     *persistent.Map[lsp.DocumentURI, trackedFile]
+}
+
+var _ fs.FS = &lspOverlayFS{}
+
+func newLSPOverlayFS(sourceDir string, files *persistent.Map[lsp.DocumentURI, trackedFile]) fs.FS {
+	return &lspOverlayFS{
+		base:      os.DirFS(sourceDir),
+		sourceDir: sourceDir,
+		files:     files,
+	}
+}
+
+func (o *lspOverlayFS) Open(name string) (fs.File, error) {
+	absPath := filepath.Join(o.sourceDir, filepath.FromSlash(name))
+	uri := lsp.DocumentURI("file://" + absPath)
+	if file, ok := o.files.Get(uri); ok {
+		return newMemFile(name, file.contents), nil
+	}
+	return o.base.Open(name)
+}
+
+// memFile implements fs.File for an in-memory string.
+type memFile struct {
+	name    string
+	content string
+	reader  *strings.Reader
+}
+
+var _ fs.File = &memFile{}
+
+func newMemFile(name, content string) *memFile {
+	return &memFile{name: name, content: content, reader: strings.NewReader(content)}
+}
+
+func (m *memFile) Read(b []byte) (int, error) { return m.reader.Read(b) }
+func (m *memFile) Close() error               { return nil }
+func (m *memFile) Stat() (fs.FileInfo, error) {
+	return memFileInfo{name: m.name, size: int64(len(m.content))}, nil
+}
+
+type memFileInfo struct {
+	name string
+	size int64
+}
+
+func (i memFileInfo) Name() string       { return filepath.Base(i.name) }
+func (i memFileInfo) Size() int64        { return i.size }
+func (i memFileInfo) Mode() fs.FileMode  { return 0o444 }
+func (i memFileInfo) ModTime() time.Time { return time.Time{} }
+func (i memFileInfo) IsDir() bool        { return false }
+func (i memFileInfo) Sys() any           { return nil }
+
+// uriToSourceDir extracts the containing directory from a file:// URI.
+func uriToSourceDir(uri lsp.DocumentURI) string {
+	path := strings.TrimPrefix(string(uri), "file://")
+	return filepath.Dir(path)
+}

pkg/development/devcontext.go

@@ -52,15 +52,15 @@ type DevContext struct {
 
 // NewDevContext creates a new DevContext from the specified request context, parsing and populating
 // the datastore as needed.
-func NewDevContext(ctx context.Context, requestContext *devinterface.RequestContext) (*DevContext, *devinterface.DeveloperErrors, error) {
+func NewDevContext(ctx context.Context, requestContext *devinterface.RequestContext, opts ...CompileOption) (*DevContext, *devinterface.DeveloperErrors, error) {
 	ds, err := memdb.NewMemdbDatastore(0, 0*time.Second, memdb.DisableGC)
 	if err != nil {
 		return nil, nil, err
 	}
 	dl := datalayer.NewDataLayer(ds)
 	ctx = datalayer.ContextWithDataLayer(ctx, dl)
 
-	dctx, devErrs, nerr := newDevContextWithDataLayer(ctx, requestContext, dl)
+	dctx, devErrs, nerr := newDevContextWithDataLayer(ctx, requestContext, dl, opts...)
 	if nerr != nil || devErrs != nil {
 		// If any form of error occurred, immediately close the data layer
 		derr := dl.Close()
@@ -74,9 +74,9 @@ func NewDevContext(ctx context.Context, requestContext *devinterface.RequestCont
 	return dctx, nil, nil
 }
 
-func newDevContextWithDataLayer(ctx context.Context, requestContext *devinterface.RequestContext, dl datalayer.DataLayer) (*DevContext, *devinterface.DeveloperErrors, error) {
+func newDevContextWithDataLayer(ctx context.Context, requestContext *devinterface.RequestContext, dl datalayer.DataLayer, opts ...CompileOption) (*DevContext, *devinterface.DeveloperErrors, error) {
 	// Compile the schema and load its caveats and namespaces into the datastore.
-	compiled, devError, err := CompileSchema(requestContext.Schema)
+	compiled, devError, err := CompileSchema(requestContext.Schema, opts...)
 	if err != nil {
 		return nil, nil, err
 	}

pkg/development/schema.go

@@ -2,6 +2,7 @@ package development
 
 import (
 	"errors"
+	"io/fs"
 
 	"github.com/ccoveille/go-safecast/v2"
 
@@ -11,14 +12,37 @@ import (
 	"github.com/authzed/spicedb/pkg/schemadsl/input"
 )
 
+// CompileOption configures schema compilation.
+type CompileOption func(*compileConfig)
+
+type compileConfig struct {
+	fsys fs.FS
+}
+
+// WithSourceFS enables import resolution using the given filesystem.
+// The filesystem should be rooted at the directory containing the schema file.
+func WithSourceFS(fsys fs.FS) CompileOption {
+	return func(cfg *compileConfig) { cfg.fsys = fsys }
+}
+
 // CompileSchema compiles a schema into its caveat and namespace definition(s), returning a developer
 // error if the schema could not be compiled. The non-developer error is returned only if an
 // internal errors occurred.
-func CompileSchema(schema string) (*compiler.CompiledSchema, *devinterface.DeveloperError, error) {
+func CompileSchema(schema string, opts ...CompileOption) (*compiler.CompiledSchema, *devinterface.DeveloperError, error) {
+	cfg := &compileConfig{}
+	for _, o := range opts {
+		o(cfg)
+	}
+
+	var compilerOpts []compiler.Option
+	if cfg.fsys != nil {
+		compilerOpts = append(compilerOpts, compiler.SourceFS(cfg.fsys))
+	}
+
 	compiled, err := compiler.Compile(compiler.InputSchema{
 		Source:       input.Source("schema"),
 		SchemaString: schema,
-	}, compiler.AllowUnprefixedObjectType())
+	}, compiler.AllowUnprefixedObjectType(), compilerOpts...)
 
 	var contextError compiler.WithContextError
 	if errors.As(err, &contextError) {

pkg/schemadsl/compiler/translator.go

@@ -840,6 +840,9 @@ func translateImports(itctx importResolutionContext, root *dslNode) error {
 			if err := validateFilepath(importPath); err != nil {
 				return err
 			}
+			if itctx.sourceFS == nil {
+				return fmt.Errorf("import statement found but no source filesystem was configured for compilation")
+			}
 			filePath := filepath.Join(itctx.sourcePrefix, importPath)
 
 			newSourcePrefix := filepath.Dir(filePath)