feat: add demo command

internal/middleware/pertoken/pertoken.go

@@ -29,15 +29,17 @@ const (
 type MiddlewareForTesting struct {
 	datastoreByToken *sync.Map
 	configFilePaths  []string
+	configContents   map[string][]byte
 	caveatTypeSet    *caveattypes.TypeSet
 }
 
 // NewMiddleware returns a new per-token datastore middleware that initializes each datastore with the data in the
 // config files.
-func NewMiddleware(configFilePaths []string, caveatTypeSet *caveattypes.TypeSet) *MiddlewareForTesting {
+func NewMiddleware(configFilePaths []string, configContents map[string][]byte, caveatTypeSet *caveattypes.TypeSet) *MiddlewareForTesting {
 	return &MiddlewareForTesting{
 		datastoreByToken: &sync.Map{},
 		configFilePaths:  configFilePaths,
+		configContents:   configContents,
 		caveatTypeSet:    caveatTypeSet,
 	}
 }
@@ -66,6 +68,13 @@ func (m *MiddlewareForTesting) getOrCreateDatastore(ctx context.Context) (datast
 		return nil, fmt.Errorf("failed to load config files: %w", err)
 	}
 
+	if len(m.configContents) > 0 {
+		_, _, err = validationfile.PopulateFromFilesContents(ctx, datalayer.NewDataLayer(ds), m.caveatTypeSet, m.configContents)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load config contents: %w", err)
+		}
+	}
+
 	// Squash the revisions so that the caller sees all the populated data.
 	ds.(squashable).SquashRevisionsForTesting()
 

internal/middleware/pertoken/pertoken_test.go

@@ -143,7 +143,7 @@ validation: null
 
 	// Create middleware
 	caveatTypeSet := caveattypes.MustNewStandardTypeSet()
-	middleware := NewMiddleware([]string{configFile}, caveatTypeSet.TypeSet)
+	middleware := NewMiddleware([]string{configFile}, nil, caveatTypeSet.TypeSet)
 
 	s := &perTokenMiddlewareTestSuite{
 		InterceptorTestSuite: &testpb.InterceptorTestSuite{
@@ -160,6 +160,46 @@ validation: null
 	suite.Run(t, s)
 }
 
+func TestPerTokenMiddlewareWithConfigContents(t *testing.T) {
+	configContent := `---
+schema: >-
+  definition example/user {}
+
+  definition example/project {
+      relation reader: example/user
+      permission read = reader
+  }
+relationships: >-
+  example/project:pied_piper#reader@example/user:tarben
+assertions:
+  assertTrue: []
+  assertFalse: []
+validation: null
+`
+
+	caveatTypeSet := caveattypes.MustNewStandardTypeSet()
+	middleware := NewMiddleware(nil, map[string][]byte{"generated-demo.yaml": []byte(configContent)}, caveatTypeSet.TypeSet)
+
+	ctx := metadata.AppendToOutgoingContext(context.Background(), "authorization", "bearer sometoken")
+	_, err := middleware.getOrCreateDatastore(ctx)
+	require.NoError(t, err)
+
+	ds, err := middleware.getOrCreateDatastore(ctx)
+	require.NoError(t, err)
+
+	rev, err := ds.HeadRevision(context.Background())
+	require.NoError(t, err)
+
+	it, err := ds.SnapshotReader(rev.Revision).QueryRelationships(context.Background(), datastore.RelationshipsFilter{
+		OptionalResourceType: "example/project",
+	})
+	require.NoError(t, err)
+
+	rels, err := datastore.IteratorToSlice(it)
+	require.NoError(t, err)
+	require.Len(t, rels, 1)
+}
+
 func (s *perTokenMiddlewareTestSuite) TestUnaryInterceptor_WithMissingToken() {
 	resp, err := s.Client.Ping(s.SimpleCtx(), &testpb.PingRequest{Value: "test"})
 	s.Require().NoError(err)

pkg/cmd/demo.go

@@ -0,0 +1,68 @@
+package cmd
+
+import (
+	"os/signal"
+	"syscall"
+
+	"github.com/spf13/cobra"
+
+	"github.com/authzed/spicedb/pkg/cmd/demo"
+	"github.com/authzed/spicedb/pkg/cmd/server"
+	"github.com/authzed/spicedb/pkg/cmd/termination"
+)
+
+func RegisterDemoFlags(cmd *cobra.Command, config *demo.Config) {
+	RegisterTestingFlags(cmd, config.TestServerConfig)
+	cmd.Flags().StringVar(&config.DataSet, "dataset", demo.DefaultDataSetName, "demo dataset to preload into each token-scoped datastore")
+}
+
+func NewDemoCommand(programName string, config *demo.Config) *cobra.Command {
+	return &cobra.Command{
+		Use:   "demo",
+		Short: "test server with an in-memory datastore and pre-loaded data",
+		Long: `An in-memory spicedb server with pre-loaded data which serves completely isolated datastores per client-supplied auth token used.
+
+Examples:
+  # Run the demo server with the basic dataset
+  spicedb demo --dataset basic
+
+  # Run the demo server with a custom dataset loaded from a file
+  spicedb demo --dataset custom --load-configs /path/to/config.yaml
+
+  # Available demo datasets include the following:
+  #   sample: user, groups, documents with relationships between them
+  #   basic: only relationships, no schema, for testing schema-less operations
+  #   docs: a dataset used for documentation examples, with a variety of schema and relationships
+  #   custom: no dataset provided by default, intended for use with user-provided datasets via --load-configs
+  #   entitlements: a dataset modeling a SaaS application with users, teams, and projects
+  #   github: a dataset modeling GitHub with users, teams, organizations, and repositories
+  #   user-defined: a dataset modeling a user-defined application with users, roles, and projects
+`,
+		PreRunE: server.DefaultPreRunE(programName),
+		RunE: termination.PublishError(func(cmd *cobra.Command, args []string) error {
+			bootstrapContents, err := config.BootstrapConfig()
+			if err != nil {
+				return err
+			}
+
+			if len(bootstrapContents) > 0 {
+				if config.TestServerConfig.LoadConfigsContents == nil {
+					config.TestServerConfig.LoadConfigsContents = map[string][]byte{}
+				}
+				for name, contents := range bootstrapContents {
+					config.TestServerConfig.LoadConfigsContents[name] = contents
+				}
+			}
+
+			srv, err := config.TestServerConfig.Complete(cmd.Context())
+			if err != nil {
+				return err
+			}
+
+			signalctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
+			defer stop()
+
+			return srv.Run(signalctx)
+		}),
+	}
+}