Restructure spicedb commands

cmd/spicedb/main.go

@@ -16,6 +16,10 @@ import (
 
 	log "github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/pkg/cmd"
+	"github.com/authzed/spicedb/pkg/cmd/cockroachdb"
+	"github.com/authzed/spicedb/pkg/cmd/memory"
+	"github.com/authzed/spicedb/pkg/cmd/mysql"
+	"github.com/authzed/spicedb/pkg/cmd/postgres"
 	cmdutil "github.com/authzed/spicedb/pkg/cmd/server"
 	"github.com/authzed/spicedb/pkg/cmd/testserver"
 	_ "github.com/authzed/spicedb/pkg/runtime"
@@ -56,6 +60,7 @@ func main() {
 	if err != nil {
 		log.Fatal().Err(err).Msg("failed to register datastore command")
 	}
+	datastoreCmd.Hidden = true
 
 	cmd.RegisterDatastoreRootFlags(datastoreCmd)
 	rootCmd.AddCommand(datastoreCmd)
@@ -74,14 +79,44 @@ func main() {
 	cmd.RegisterMigrateFlags(migrateCmd)
 	rootCmd.AddCommand(migrateCmd)
 
+	// Add datastore commands
+	rootCmd.AddGroup(&cobra.Group{
+		ID:    "datastores",
+		Title: "Datastores:",
+	})
+	pgCmd, err := postgres.NewPostgresCommand(rootCmd.Use)
+	if err != nil {
+		log.Fatal().Err(err).Msg("failed to register serve flags")
+	}
+	rootCmd.AddCommand(pgCmd)
+	crdbCmd, err := cockroachdb.NewCommand(rootCmd.Use)
+	if err != nil {
+		log.Fatal().Err(err).Msg("failed to register serve flags")
+	}
+	rootCmd.AddCommand(crdbCmd)
+	memCmd, err := memory.NewCommand(rootCmd.Use)
+	if err != nil {
+		log.Fatal().Err(err).Msg("failed to register serve flags")
+	}
+	rootCmd.AddCommand(memCmd)
+	myCmd := mysql.NewCommand(rootCmd.Use)
+	rootCmd.AddCommand(myCmd)
+
 	// Add server commands
 	serverConfig := cmdutil.NewConfigWithOptionsAndDefaults()
 	serveCmd := cmd.NewServeCommand(rootCmd.Use, serverConfig)
 	if err := cmd.RegisterServeFlags(serveCmd, serverConfig); err != nil {
 		log.Fatal().Err(err).Msg("failed to register server flags")
 	}
+	serveCmd.Hidden = true
 	rootCmd.AddCommand(serveCmd)
 
+	// Add developer tools
+	rootCmd.AddGroup(&cobra.Group{
+		ID:    "devtools",
+		Title: "Developer Tools:",
+	})
+
 	devtoolsCmd := cmd.NewDevtoolsCommand(rootCmd.Use)
 	cmd.RegisterDevtoolsFlags(devtoolsCmd)
 	rootCmd.AddCommand(devtoolsCmd)
@@ -117,6 +152,11 @@ func main() {
 	})
 
 	if err := rootCmd.Execute(); err != nil {
+		// Ensure that logging has been set-up before printing an error.
+		if preRunErr := rootCmd.PersistentPreRunE(rootCmd, nil); preRunErr != nil {
+			panic(preRunErr)
+		}
+
 		if !errors.Is(err, errParsing) {
 			log.Err(err).Msg("terminated with errors")
 		}

go.mod

@@ -91,7 +91,6 @@ require (
 	go.opentelemetry.io/otel v1.28.0
 	go.opentelemetry.io/otel/sdk v1.28.0
 	go.opentelemetry.io/otel/trace v1.28.0
-	go.uber.org/atomic v1.11.0
 	go.uber.org/goleak v1.3.0
 	golang.org/x/exp v0.0.0-20240506185415-9bf2ced13842
 	golang.org/x/mod v0.19.0
@@ -381,6 +380,7 @@ require (
 	go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracehttp v1.22.0 // indirect
 	go.opentelemetry.io/otel/metric v1.28.0 // indirect
 	go.opentelemetry.io/proto/otlp v1.0.0 // indirect
+	go.uber.org/atomic v1.11.0 // indirect
 	go.uber.org/automaxprocs v1.5.3 // indirect
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.26.0 // indirect

internal/graph/checkingresourcestream.go

@@ -5,8 +5,8 @@ import (
 	"context"
 	"slices"
 	"sync"
+	"sync/atomic"
 
-	"go.uber.org/atomic"
 	"golang.org/x/exp/maps"
 
 	"github.com/authzed/spicedb/internal/dispatch"

pkg/closer/closer.go

@@ -0,0 +1,46 @@
+package closer
+
+import (
+	"io"
+
+	"github.com/hashicorp/go-multierror"
+)
+
+type Stack struct {
+	closers []func() error
+}
+
+func (c *Stack) AddWithError(closer func() error) {
+	c.closers = append(c.closers, closer)
+}
+
+func (c *Stack) AddCloser(closer io.Closer) {
+	if closer != nil {
+		c.closers = append(c.closers, closer.Close)
+	}
+}
+
+func (c *Stack) AddWithoutError(closer func()) {
+	c.closers = append(c.closers, func() error {
+		closer()
+		return nil
+	})
+}
+
+func (c *Stack) Close() error {
+	var err error
+	// closer in reverse order how it's expected in deferred funcs
+	for i := len(c.closers) - 1; i >= 0; i-- {
+		if closerErr := c.closers[i](); closerErr != nil {
+			err = multierror.Append(err, closerErr)
+		}
+	}
+	return err
+}
+
+func (c *Stack) CloseIfError(err error) error {
+	if err != nil {
+		return c.Close()
+	}
+	return nil
+}

pkg/cmd/cockroachdb/cockroachdb.go

@@ -0,0 +1,137 @@
+package cockroachdb
+
+import (
+	"fmt"
+	"time"
+
+	"github.com/go-logr/zerologr"
+	"github.com/jzelinskie/cobrautil/v2"
+	"github.com/jzelinskie/cobrautil/v2/cobraotel"
+	"github.com/spf13/cobra"
+
+	"github.com/authzed/spicedb/internal/datastore/crdb/migrations"
+	"github.com/authzed/spicedb/internal/logging"
+	pkgcmd "github.com/authzed/spicedb/pkg/cmd"
+	"github.com/authzed/spicedb/pkg/cmd/server"
+	"github.com/authzed/spicedb/pkg/migrate"
+	"github.com/authzed/spicedb/pkg/releases"
+	"github.com/authzed/spicedb/pkg/runtime"
+)
+
+func NewCommand(programName string) (*cobra.Command, error) {
+	crdbCmd := &cobra.Command{
+		Use:     "cockroachdb",
+		Aliases: []string{"cockroach", "crdb"},
+		Short:   "Perform operations on data stored in CockroachDB",
+		GroupID: "datastores",
+		Hidden:  false,
+	}
+	migrationsCmd := NewMigrationCommand(programName)
+	crdbCmd.AddCommand(migrationsCmd)
+
+	cfg := &server.Config{}
+	cfg.DatastoreConfig.Engine = "cockroachdb"
+	cfg.NamespaceCacheConfig = pkgcmd.NamespaceCacheConfig
+	cfg.ClusterDispatchCacheConfig = server.CacheConfig{}
+	cfg.DispatchCacheConfig = server.CacheConfig{}
+
+	serveCmd := &cobra.Command{
+		Use:     "serve-grpc",
+		Short:   "Serve the SpiceDB gRPC API services",
+		Example: pkgcmd.ServeExample(programName),
+		PreRunE: cobrautil.CommandStack(
+			cobraotel.New("spicedb", cobraotel.WithLogger(zerologr.New(&logging.Logger))).RunE(),
+			releases.CheckAndLogRunE(),
+			runtime.RunE(),
+		),
+		RunE: pkgcmd.ServeGRPCRunE(cfg),
+	}
+
+	nfs := cobrautil.NewNamedFlagSets(serveCmd)
+	if err := pkgcmd.RegisterCRDBDatastoreFlags(serveCmd, nfs.FlagSet("CockroachDB Datastore"), cfg); err != nil {
+		return nil, err
+	}
+
+	postRegisterFn, err := pkgcmd.RegisterCommonServeFlags(programName, serveCmd, nfs, cfg, true)
+	if err != nil {
+		return nil, err
+	}
+
+	// Flags must be registered to the command after flags are set.
+	nfs.AddFlagSets(serveCmd)
+	if err := postRegisterFn(); err != nil {
+		return nil, err
+	}
+
+	crdbCmd.AddCommand(serveCmd)
+
+	return crdbCmd, nil
+}
+
+func NewMigrationCommand(programName string) *cobra.Command {
+	migrationsCmd := &cobra.Command{
+		Use:   "migrations",
+		Short: "Perform migrations and schema changes",
+	}
+
+	headCmd := &cobra.Command{
+		Use:   "head",
+		Short: "Print the latest migration",
+		Args:  cobra.ExactArgs(0),
+		RunE: func(cmd *cobra.Command, args []string) error {
+			head, err := migrations.CRDBMigrations.HeadRevision()
+			if err != nil {
+				return fmt.Errorf("unable to compute head revision: %w", err)
+			}
+			_, err = fmt.Println(head)
+			return err
+		},
+	}
+	migrationsCmd.AddCommand(headCmd)
+
+	execCmd := &cobra.Command{
+		Use:   "exec <target migration>",
+		Short: "Execute all migrations up to and including the provided migration",
+		Args:  cobra.ExactArgs(1),
+		RunE:  ExecMigrationRunE,
+	}
+	migrationsCmd.AddCommand(execCmd)
+	RegisterMigrationExecFlags(execCmd)
+
+	return migrationsCmd
+}
+
+func RegisterMigrationExecFlags(cmd *cobra.Command) {
+	cmd.Flags().String("crdb-uri", "postgres://roach:password@localhost:5432/spicedb", "connection string in URI format")
+	cmd.Flags().Uint64("backfill-batch-size", 1000, "batch size used when backfilling data")
+	cmd.Flags().Duration("timeout", 1*time.Hour, "maximum execution duration for an individual migration")
+}
+
+func ExecMigrationRunE(cmd *cobra.Command, args []string) error {
+	ctx := cmd.Context()
+
+	revision := args[0]
+	if revision == "head" {
+		head, err := migrations.CRDBMigrations.HeadRevision()
+		if err != nil {
+			return fmt.Errorf("unable to compute head revision: %w", err)
+		}
+		revision = head
+	}
+
+	logging.Ctx(ctx).Info().Str("target", revision).Msg("executing migrations")
+
+	migrationDriver, err := migrations.NewCRDBDriver(cobrautil.MustGetStringExpanded(cmd, "crdb-uri"))
+	if err != nil {
+		return fmt.Errorf("unable to create cockroachdb migration driver: %w", err)
+	}
+
+	return migrate.RunMigration(
+		cmd.Context(),
+		migrationDriver,
+		migrations.CRDBMigrations,
+		revision,
+		cobrautil.MustGetDuration(cmd, "timeout"),
+		cobrautil.MustGetUint64(cmd, "backfill-batch-size"),
+	)
+}

pkg/cmd/devtools.go

@@ -49,6 +49,7 @@ func NewDevtoolsCommand(programName string) *cobra.Command {
 		Use:     "serve-devtools",
 		Short:   "runs the developer tools service",
 		Long:    "Serves the authzed.api.v0.DeveloperService which is used for development tooling such as the Authzed Playground",
+		GroupID: "devtools",
 		PreRunE: server.DefaultPreRunE(programName),
 		RunE:    termination.PublishError(runfunc),
 		Args:    cobra.ExactArgs(0),

pkg/cmd/lsp.go

@@ -4,13 +4,8 @@ import (
 	"context"
 	"time"
 
-	"github.com/go-logr/zerologr"
-	"github.com/jzelinskie/cobrautil/v2"
-	"github.com/jzelinskie/cobrautil/v2/cobrazerolog"
-	"github.com/rs/zerolog"
 	"github.com/spf13/cobra"
 
-	"github.com/authzed/spicedb/internal/logging"
 	"github.com/authzed/spicedb/internal/lsp"
 	"github.com/authzed/spicedb/pkg/cmd/termination"
 	"github.com/authzed/spicedb/pkg/releases"
@@ -37,17 +32,10 @@ func RegisterLSPFlags(cmd *cobra.Command, config *LSPConfig) error {
 
 func NewLSPCommand(programName string, config *LSPConfig) *cobra.Command {
 	return &cobra.Command{
-		Use:   "lsp",
-		Short: "serve language server protocol",
-		PreRunE: cobrautil.CommandStack(
-			cobrautil.SyncViperDotEnvPreRunE(programName, "spicedb.env", zerologr.New(&logging.Logger)),
-			cobrazerolog.New(
-				cobrazerolog.WithTarget(func(logger zerolog.Logger) {
-					logging.SetGlobalLogger(logger)
-				}),
-			).RunE(),
-			releases.CheckAndLogRunE(),
-		),
+		Use:     "lsp",
+		Short:   "serve language server protocol",
+		GroupID: "devtools",
+		PreRunE: releases.CheckAndLogRunE(),
 		RunE: termination.PublishError(func(cmd *cobra.Command, args []string) error {
 			srv, err := config.Complete(cmd.Context())
 			if err != nil {

pkg/cmd/memory/memory.go

@@ -0,0 +1,60 @@
+package memory
+
+import (
+	"github.com/go-logr/zerologr"
+	"github.com/jzelinskie/cobrautil/v2"
+	"github.com/jzelinskie/cobrautil/v2/cobraotel"
+	"github.com/spf13/cobra"
+
+	"github.com/authzed/spicedb/internal/logging"
+	pkgcmd "github.com/authzed/spicedb/pkg/cmd"
+	"github.com/authzed/spicedb/pkg/cmd/server"
+	"github.com/authzed/spicedb/pkg/releases"
+	"github.com/authzed/spicedb/pkg/runtime"
+)
+
+const (
+	exampleWithoutTLS = "memory serve-grpc --grpc-preshared-key secretKeyHere"
+	exampleWithTLS    = "memory serve-grpc --grpc-preshared-key secretKeyHere --grpc-tls-cert path/to/cert --grpc-tls-key path/to/key"
+)
+
+func NewCommand(programName string) (*cobra.Command, error) {
+	memCmd := &cobra.Command{
+		Use:     "memory",
+		Aliases: []string{"mem"},
+		Short:   "Perform operations on data stored in non-persistent memory",
+		GroupID: "datastores",
+	}
+
+	cfg := &server.Config{}
+	cfg.DatastoreConfig.Engine = "memory"
+	cfg.NamespaceCacheConfig = pkgcmd.NamespaceCacheConfig
+	cfg.ClusterDispatchCacheConfig = server.CacheConfig{}
+	cfg.DispatchCacheConfig = server.CacheConfig{}
+
+	serveCmd := &cobra.Command{
+		Use:     "serve-grpc",
+		Short:   "Serve the SpiceDB gRPC API services",
+		Example: pkgcmd.ServeExample(programName, exampleWithoutTLS, exampleWithTLS),
+		PreRunE: cobrautil.CommandStack(
+			cobraotel.New("spicedb", cobraotel.WithLogger(zerologr.New(&logging.Logger))).RunE(),
+			releases.CheckAndLogRunE(),
+			runtime.RunE(),
+		),
+		RunE: pkgcmd.ServeGRPCRunE(cfg),
+	}
+	nfs := cobrautil.NewNamedFlagSets(serveCmd)
+	postRegisterFn, err := pkgcmd.RegisterCommonServeFlags(programName, serveCmd, nfs, cfg, false)
+	if err != nil {
+		return nil, err
+	}
+
+	// Flags must be registered to the command after flags are set.
+	nfs.AddFlagSets(serveCmd)
+	if err := postRegisterFn(); err != nil {
+		return nil, err
+	}
+	memCmd.AddCommand(serveCmd)
+
+	return memCmd, nil
+}

pkg/cmd/migrate.go

@@ -118,6 +118,7 @@ func migrateRun(cmd *cobra.Command, args []string) error {
 	return fmt.Errorf("cannot migrate datastore engine type: %s", datastoreEngine)
 }
 
+// TODO(jzelinskie): deprecated, replace with migrate.RunMigration()
 func runMigration[D migrate.Driver[C, T], C any, T any](
 	ctx context.Context,
 	driver D,