Merge pull request #1983 from josephschorr/datastore-overwrite

Ensure that the bootstrap overwrite flag actually fully overwrites

pkg/cmd/datastore/datastore.go

@@ -207,7 +207,7 @@ func RegisterDatastoreFlagsWithPrefix(flagSet *pflag.FlagSet, prefix string, opt
 	flagSet.Float64Var(&opts.MaxRevisionStalenessPercent, flagName("datastore-revision-quantization-max-staleness-percent"), defaults.MaxRevisionStalenessPercent, "float percentage (where 1 = 100%) of the revision quantization interval where we may opt to select a stale revision for performance reasons. Defaults to 0.1 (representing 10%)")
 	flagSet.BoolVar(&opts.ReadOnly, flagName("datastore-readonly"), defaults.ReadOnly, "set the service to read-only mode")
 	flagSet.StringSliceVar(&opts.BootstrapFiles, flagName("datastore-bootstrap-files"), defaults.BootstrapFiles, "bootstrap data yaml files to load")
-	flagSet.BoolVar(&opts.BootstrapOverwrite, flagName("datastore-bootstrap-overwrite"), defaults.BootstrapOverwrite, "overwrite any existing data with bootstrap data")
+	flagSet.BoolVar(&opts.BootstrapOverwrite, flagName("datastore-bootstrap-overwrite"), defaults.BootstrapOverwrite, "overwrite any existing data with bootstrap data (this can be quite slow)")
 	flagSet.DurationVar(&opts.BootstrapTimeout, flagName("datastore-bootstrap-timeout"), defaults.BootstrapTimeout, "maximum duration before timeout for the bootstrap data to be written")
 	flagSet.BoolVar(&opts.RequestHedgingEnabled, flagName("datastore-request-hedging"), defaults.RequestHedgingEnabled, "enable request hedging")
 	flagSet.DurationVar(&opts.RequestHedgingInitialSlowValue, flagName("datastore-request-hedging-initial-slow-value"), defaults.RequestHedgingInitialSlowValue, "initial value to use for slow datastore requests, before statistics have been collected")
@@ -325,26 +325,33 @@ func NewDatastore(ctx context.Context, options ...ConfigOption) (datastore.Datas
 		if err != nil {
 			return nil, fmt.Errorf("unable to determine datastore state before applying bootstrap data: %w", err)
 		}
-		if opts.BootstrapOverwrite || len(nsDefs) == 0 {
-			log.Ctx(ctx).Info().Strs("files", opts.BootstrapFiles).Msg("initializing datastore from bootstrap files")
-
-			if len(opts.BootstrapFiles) > 0 {
-				_, _, err = validationfile.PopulateFromFiles(ctx, ds, opts.BootstrapFiles)
-				if err != nil {
-					return nil, fmt.Errorf("failed to load bootstrap files: %w", err)
-				}
-			}
 
-			if len(opts.BootstrapFileContents) > 0 {
-				_, _, err = validationfile.PopulateFromFilesContents(ctx, ds, opts.BootstrapFileContents)
-				if err != nil {
-					return nil, fmt.Errorf("failed to load bootstrap file contents: %w", err)
-				}
+		if opts.BootstrapOverwrite {
+			log.Ctx(ctx).Info().Msg("deleting existing data before applying bootstrap data (this may take a bit)")
+			if err := datastore.DeleteAllData(ctx, ds); err != nil {
+				return nil, fmt.Errorf("failed to delete existing data before applying bootstrap data: %w", err)
 			}
-			log.Ctx(ctx).Info().Strs("files", opts.BootstrapFiles).Msg("completed datastore initialization from bootstrap files")
-		} else {
+			log.Ctx(ctx).Info().Msg("deleted existing data before applying bootstrap data")
+		} else if len(nsDefs) > 0 {
 			return nil, errors.New("cannot apply bootstrap data: schema or tuples already exist in the datastore. Delete existing data or set the flag --datastore-bootstrap-overwrite=true")
 		}
+
+		log.Ctx(ctx).Info().Strs("files", opts.BootstrapFiles).Msg("initializing datastore from bootstrap files")
+
+		if len(opts.BootstrapFiles) > 0 {
+			_, _, err = validationfile.PopulateFromFiles(ctx, ds, opts.BootstrapFiles)
+			if err != nil {
+				return nil, fmt.Errorf("failed to load bootstrap files: %w", err)
+			}
+		}
+
+		if len(opts.BootstrapFileContents) > 0 {
+			_, _, err = validationfile.PopulateFromFilesContents(ctx, ds, opts.BootstrapFileContents)
+			if err != nil {
+				return nil, fmt.Errorf("failed to load bootstrap file contents: %w", err)
+			}
+		}
+		log.Ctx(ctx).Info().Strs("files", opts.BootstrapFiles).Msg("completed datastore initialization from bootstrap files")
 	}
 
 	if opts.RequestHedgingEnabled {

pkg/datastore/test/basic.go

@@ -5,6 +5,9 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/authzed/spicedb/internal/testfixtures"
+	"github.com/authzed/spicedb/pkg/datastore"
 )
 
 func UseAfterCloseTest(t *testing.T, tester DatastoreTester) {
@@ -22,3 +25,53 @@ func UseAfterCloseTest(t *testing.T, tester DatastoreTester) {
 	_, err = ds.HeadRevision(context.Background())
 	require.Error(err)
 }
+
+func DeleteAllDataTest(t *testing.T, tester DatastoreTester) {
+	rawDS, err := tester.New(0, veryLargeGCInterval, veryLargeGCWindow, 1)
+	require.NoError(t, err)
+
+	ds, revision := testfixtures.StandardDatastoreWithCaveatedData(rawDS, require.New(t))
+	ctx := context.Background()
+
+	// Ensure at least a few relationships and namespaces exist.
+	reader := ds.SnapshotReader(revision)
+	nsDefs, err := reader.ListAllNamespaces(ctx)
+	require.NoError(t, err)
+	require.NotEmpty(t, nsDefs, "no namespace definitions provided")
+
+	foundRels := false
+	for _, nsDef := range nsDefs {
+		iter, err := reader.QueryRelationships(ctx, datastore.RelationshipsFilter{OptionalResourceType: nsDef.Definition.Name})
+		require.NoError(t, err)
+		t.Cleanup(iter.Close)
+
+		if iter.Next() != nil {
+			foundRels = true
+		}
+
+		iter.Close()
+	}
+	require.True(t, foundRels, "no relationships provided")
+
+	// Delete all data.
+	err = datastore.DeleteAllData(ctx, ds)
+	require.NoError(t, err)
+
+	// Ensure there are no relationships or namespaces.
+	headRev, err := ds.HeadRevision(ctx)
+	require.NoError(t, err)
+
+	reader = ds.SnapshotReader(headRev)
+	afterNSDefs, err := reader.ListAllNamespaces(ctx)
+	require.NoError(t, err)
+	require.Empty(t, afterNSDefs, "namespace definitions still exist")
+
+	for _, nsDef := range nsDefs {
+		iter, err := reader.QueryRelationships(ctx, datastore.RelationshipsFilter{OptionalResourceType: nsDef.Definition.Name})
+		require.NoError(t, err)
+		t.Cleanup(iter.Close)
+
+		require.Nil(t, iter.Next(), "relationships still exist")
+		iter.Close()
+	}
+}

pkg/datastore/test/datastore.go

@@ -175,6 +175,7 @@ func AllWithExceptions(t *testing.T, tester DatastoreTester, except Categories)
 
 	t.Run("TestRelationshipCounters", func(t *testing.T) { RelationshipCountersTest(t, tester) })
 	t.Run("TestUpdateRelationshipCounter", func(t *testing.T) { UpdateRelationshipCounterTest(t, tester) })
+	t.Run("TestDeleteAllData", func(t *testing.T) { DeleteAllDataTest(t, tester) })
 }
 
 // All runs all generic datastore tests on a DatastoreTester.

pkg/datastore/util.go

@@ -1,5 +1,11 @@
 package datastore
 
+import (
+	"context"
+
+	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
+)
+
 // DefinitionsOf returns just the schema definitions found in the list of revisioned
 // definitions.
 func DefinitionsOf[T SchemaDefinition](revisionedDefinitions []RevisionedDefinition[T]) []T {
@@ -9,3 +15,49 @@ func DefinitionsOf[T SchemaDefinition](revisionedDefinitions []RevisionedDefinit
 	}
 	return definitions
 }
+
+// DeleteAllData deletes all data from the datastore. Should only be used when explicitly requested.
+// The data is transactionally deleted, which means it may time out.
+func DeleteAllData(ctx context.Context, ds Datastore) error {
+	_, err := ds.ReadWriteTx(ctx, func(ctx context.Context, rwt ReadWriteTransaction) error {
+		nsDefs, err := rwt.ListAllNamespaces(ctx)
+		if err != nil {
+			return err
+		}
+
+		// Delete all relationships.
+		namespaceNames := make([]string, 0, len(nsDefs))
+		for _, nsDef := range nsDefs {
+			_, err = rwt.DeleteRelationships(ctx, &v1.RelationshipFilter{
+				ResourceType: nsDef.Definition.Name,
+			})
+			if err != nil {
+				return err
+			}
+			namespaceNames = append(namespaceNames, nsDef.Definition.Name)
+		}
+
+		// Delete all caveats.
+		caveatDefs, err := rwt.ListAllCaveats(ctx)
+		if err != nil {
+			return err
+		}
+
+		caveatNames := make([]string, 0, len(caveatDefs))
+		for _, caveatDef := range caveatDefs {
+			caveatNames = append(caveatNames, caveatDef.Definition.Name)
+		}
+
+		if err := rwt.DeleteCaveats(ctx, caveatNames); err != nil {
+			return err
+		}
+
+		// Delete all namespaces.
+		if err := rwt.DeleteNamespaces(ctx, namespaceNames...); err != nil {
+			return err
+		}
+
+		return nil
+	})
+	return err
+}