feat: Implement use import (#2919)

Co-authored-by: Maria Ines Parnisari <maria.ines.parnisari@authzed.com>

Author: tstirrat15

CHANGELOG.md

@@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 ### Changed
 - Begin deprecation of library "github.com/dlmiddlecote/sqlstats" (https://github.com/authzed/spicedb/pull/2904).
   NOTE: in a future release, MySQL metrics will change.
+- Add support for imports and partials to the schemadsl package that drives the LSP and development server (https://github.com/authzed/spicedb/pull/2919).
 
 ### Fixed
 - enforce graceful shutdown on serve and serve-testing (https://github.com/authzed/spicedb/pull/2888)

internal/services/shared/errors.go

@@ -135,7 +135,7 @@ func RewriteError(ctx context.Context, err error, config *ConfigForErrors) error
 	return rerr
 }
 
-func rewriteError(ctx context.Context, err error, config *ConfigForErrors) error {
+func rewriteError(ctx context.Context, err error, _ *ConfigForErrors) error {
 	// Check if the error can be directly used.
 	if _, ok := status.FromError(err); ok {
 		return err

internal/services/v1/schema.go

@@ -127,6 +127,9 @@ func (ss *schemaServer) WriteSchema(ctx context.Context, in *v1.WriteSchemaReque
 	}
 
 	opts = append(opts, compiler.CaveatTypeSet(ss.caveatTypeSet))
+	// Imports don't make sense in a schema written directly to SpiceDB;
+	// the user must first compile them with `zed`
+	opts = append(opts, compiler.DisallowImportFlag())
 
 	compiled, err := compiler.Compile(compiler.InputSchema{
 		Source:       input.Source("schema"),

internal/services/v1/schema_test.go

@@ -61,6 +61,27 @@ func TestSchemaWriteInvalidNamespace(t *testing.T) {
 	grpcutil.RequireStatus(t, codes.FailedPrecondition, err)
 }
 
+// NOTE: imports must be handled by precompilation;
+// a write of a schema with an import statement is an error.
+func TestSchemaWriteImportsDisallowed(t *testing.T) {
+	conn, cleanup, _, _ := testserver.NewTestServer(require.New(t), 0, memdb.DisableGC, true, tf.EmptyDatastore)
+	t.Cleanup(cleanup)
+	client := v1.NewSchemaServiceClient(conn)
+
+	_, err := client.WriteSchema(t.Context(), &v1.WriteSchemaRequest{
+		Schema: `
+		use import
+
+		import "foo/bar/baz.zed"
+		
+		definition document {
+			relation viewer: user
+		}
+	`,
+	})
+	grpcutil.RequireStatus(t, codes.InvalidArgument, err)
+}
+
 func TestSchemaWriteAndReadBack(t *testing.T) {
 	conn, cleanup, _, _ := testserver.NewTestServer(require.New(t), 0, memdb.DisableGC, true, tf.EmptyDatastore)
 	t.Cleanup(cleanup)

pkg/composableschemadsl/compiler/compiler.go

@@ -117,7 +117,7 @@ type ObjectPrefixOption func(*config)
 // Compile compilers the input schema into a set of namespace definition protos.
 func Compile(schema InputSchema, prefix ObjectPrefixOption, opts ...Option) (*CompiledSchema, error) {
 	cfg := &config{
-		allowedFlags: mapz.NewSet[string](expirationFlag, selfFlag),
+		allowedFlags: mapz.NewSet(expirationFlag, selfFlag),
 	}
 
 	prefix(cfg) // required option

pkg/schemadsl/compiler/compiler.go

@@ -3,6 +3,8 @@ package compiler
 import (
 	"errors"
 	"fmt"
+	"io/fs"
+	"os"
 
 	"google.golang.org/protobuf/proto"
 
@@ -56,6 +58,10 @@ type config struct {
 	objectTypePrefix *string
 	allowedFlags     *mapz.Set[string]
 	caveatTypeSet    *caveattypes.TypeSet
+
+	// In an import context, this is the FS containing
+	// the importing schema (as opposed to imported schemas)
+	sourceFS fs.FS
 }
 
 func SkipValidation() Option { return func(cfg *config) { cfg.skipValidation = true } }
@@ -76,29 +82,50 @@ func CaveatTypeSet(cts *caveattypes.TypeSet) Option {
 	return func(cfg *config) { cfg.caveatTypeSet = cts }
 }
 
+// Config that supplies the root source folder for compilation. Required
+// for relative import syntax to work properly.
+func SourceFolder(sourceFolder string) Option {
+	return func(cfg *config) { cfg.sourceFS = os.DirFS(sourceFolder) }
+}
+
+// Config that supplies the fs.FS for compilation as an alternative to
+// SourceFolder.
+func SourceFS(fsys fs.FS) Option {
+	return func(cfg *config) { cfg.sourceFS = fsys }
+}
+
 const (
 	expirationFlag   = "expiration"
 	selfFlag         = "self"
 	typeCheckingFlag = "typechecking"
 	partialFlag      = "partial"
+	importFlag       = "import"
 )
 
-var allowedFlags = mapz.NewSet(expirationFlag, selfFlag, typeCheckingFlag, partialFlag)
+func allowedFlags() *mapz.Set[string] {
+	return mapz.NewSet(expirationFlag, selfFlag, typeCheckingFlag, partialFlag, importFlag)
+}
 
 func DisallowExpirationFlag() Option {
 	return func(cfg *config) {
 		cfg.allowedFlags.Delete(expirationFlag)
 	}
 }
 
+func DisallowImportFlag() Option {
+	return func(cfg *config) {
+		cfg.allowedFlags.Delete(importFlag)
+	}
+}
+
 type Option func(*config)
 
 type ObjectPrefixOption func(*config)
 
 // Compile compilers the input schema into a set of namespace definition protos.
 func Compile(schema InputSchema, prefix ObjectPrefixOption, opts ...Option) (*CompiledSchema, error) {
 	cfg := &config{
-		allowedFlags: allowedFlags,
+		allowedFlags: allowedFlags(),
 	}
 
 	prefix(cfg) // required option
@@ -107,14 +134,36 @@ func Compile(schema InputSchema, prefix ObjectPrefixOption, opts ...Option) (*Co
 		fn(cfg)
 	}
 
-	mapper := newPositionMapper(schema)
-	root := parser.Parse(createAstNode, schema.Source, schema.SchemaString).(*dslNode)
-	errs := root.FindAll(dslshape.NodeTypeError)
-	if len(errs) > 0 {
-		err := errorNodeToError(errs[0], mapper)
+	root, mapper, err := parseSchema(schema)
+	if err != nil {
 		return nil, err
 	}
 
+	present, err := validateImportPresence(cfg.allowedFlags.Has(importFlag), root)
+	if err != nil {
+		// This condition should basically always be satisfied (we trigger errors off of the node),
+		// but we're defensive here in case the implementation changes.
+		var withNodeError withNodeError
+		if errors.As(err, &withNodeError) {
+			return nil, toContextError(withNodeError.Error(), withNodeError.errorSourceCode, withNodeError.node, mapper)
+		}
+		return nil, err
+	}
+
+	if present {
+		// NOTE: import translation is done separately so that partial references
+		// and definitions defined in separate files can correctly resolve.
+		err = translateImports(importResolutionContext{
+			globallyVisitedFiles: mapz.NewSet[string](),
+			locallyVisitedFiles:  mapz.NewSet[string](),
+			sourceFS:             cfg.sourceFS,
+			mapper:               mapper,
+		}, root)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	initialCompiledPartials := make(map[string][]*core.Relation)
 	caveatTypeSet := caveattypes.TypeSetOrDefault(cfg.caveatTypeSet)
 	compiled, err := translate(&translationContext{
@@ -141,6 +190,34 @@ func Compile(schema InputSchema, prefix ObjectPrefixOption, opts ...Option) (*Co
 	return compiled, nil
 }
 
+func parseSchema(schema InputSchema) (*dslNode, input.PositionMapper, error) {
+	mapper := newPositionMapper(schema)
+	root := parser.Parse(createAstNode, schema.Source, schema.SchemaString).(*dslNode)
+	errs := root.FindAll(dslshape.NodeTypeError)
+	if len(errs) > 0 {
+		err := errorNodeToError(errs[0], mapper)
+		return nil, nil, err
+	}
+	return root, mapper, nil
+}
+
+// validateImportPresence validates whether a given AST is valid based on whether
+// imports are allowed in the context. if they're present and disallowed it returns
+// a validation error; otherwise it returns the presence.
+func validateImportPresence(allowed bool, root *dslNode) (present bool, err error) {
+	present = false
+	for _, topLevelNode := range root.GetChildren() {
+		// Process import nodes; ignore the others
+		if topLevelNode.GetType() == dslshape.NodeTypeImport {
+			if !allowed {
+				return false, topLevelNode.Errorf("import statements are not allowed in this context")
+			}
+			present = true
+		}
+	}
+	return present, nil
+}
+
 func errorNodeToError(node *dslNode, mapper input.PositionMapper) error {
 	if node.GetType() != dslshape.NodeTypeError {
 		return errors.New("given none error node")

pkg/schemadsl/compiler/compiler_test.go

@@ -1327,19 +1327,6 @@ func TestCompile(t *testing.T) {
 				),
 			},
 		},
-		{
-			"duplicate use of expiration pragmas",
-			withTenantPrefix,
-			`
-			use expiration
-			use expiration
-
-			definition simple {
-				relation viewer: user with expiration
-			}`,
-			`found duplicate use flag`,
-			[]SchemaDefinition{},
-		},
 		{
 			"self is allowed when used in arrow and interpreted as relation name",
 			withTenantPrefix,
@@ -1424,19 +1411,6 @@ func TestCompile(t *testing.T) {
 			"Expected end of statement or definition, found: TokenTypeRightArrow",
 			[]SchemaDefinition{},
 		},
-		{
-			"duplicate use of self pragmas errors",
-			withTenantPrefix,
-			`
-			use self
-			use self
-
-			definition expressioned {
-				permission foos = ((arel->brel) + self) - drel
-			}`,
-			`found duplicate use flag`,
-			[]SchemaDefinition{},
-		},
 		{
 			"relation with expiration trait and caveat",
 			withTenantPrefix,

pkg/schemadsl/compiler/importer-test/README.md

@@ -0,0 +1,5 @@
+# Test Structure
+
+Every folder should have a `root.zed`. This is the target for compilation.
+
+Every folder will have an `expected.zed`, which is the output of the compilation process.

pkg/schemadsl/compiler/importer-test/circular-import/expected.zed

@@ -0,0 +1,8 @@
+definition user {}
+
+definition persona {}
+
+definition resource {
+	relation viewer: user
+	permission view = viewer
+}

pkg/schemadsl/compiler/importer-test/circular-import/root.zed

@@ -0,0 +1,8 @@
+use import
+
+import "subjects.zed"
+
+definition resource {
+  relation viewer: user
+  permission view = viewer
+}