feat: db reset skip seed flag (#2665)

* wip * feat: db reset --skip-seed flag Add a flag option to the command Closes #2534 * chore: apply suggestions from code review Co-authored-by: Han Qiao <[email protected]> * feat: add db.seed enabled config params * chore: remove unused type * Update db.go * Apply suggestions from code review --------- Co-authored-by: Han Qiao <[email protected]>
supabase · Sep 11, 2024 · e60bc21 · e60bc21
1 parent 6003aae
commit e60bc21
Show file tree

Hide file tree

Showing 5 changed files with 68 additions and 0 deletions.
diff --git a/cmd/db.go b/cmd/db.go
@@ -183,10 +183,15 @@ var (
 		},
 	}
 
+	noSeed bool
+
 	dbResetCmd = &cobra.Command{
 		Use:   "reset",
 		Short: "Resets the local database to current migrations",
 		RunE: func(cmd *cobra.Command, args []string) error {
+			if noSeed {
+				utils.Config.Db.Seed.Enabled = false
+			}
 			return reset.Run(cmd.Context(), migrationVersion, flags.DbConfig, afero.NewOsFs())
 		},
 	}
@@ -304,6 +309,7 @@ func init() {
 	resetFlags.String("db-url", "", "Resets the database specified by the connection string (must be percent-encoded).")
 	resetFlags.Bool("linked", false, "Resets the linked project with local migrations.")
 	resetFlags.Bool("local", true, "Resets the local database with local migrations.")
+	resetFlags.BoolVar(&noSeed, "no-seed", false, "Skip running the seed script after reset.")
 	dbResetCmd.MarkFlagsMutuallyExclusive("db-url", "linked", "local")
 	resetFlags.StringVar(&migrationVersion, "version", "", "Reset up to the specified version.")
 	dbCmd.AddCommand(dbResetCmd)

diff --git a/internal/db/reset/reset_test.go b/internal/db/reset/reset_test.go
@@ -357,6 +357,33 @@ func TestResetRemote(t *testing.T) {
 		assert.NoError(t, err)
 	})
 
+	t.Run("resets remote database with seed config disabled", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		path := filepath.Join(utils.MigrationsDir, "0_schema.sql")
+		require.NoError(t, afero.WriteFile(fsys, path, nil, 0644))
+		seedPath := filepath.Join(utils.SeedDataPath)
+		// Will raise an error when seeding
+		require.NoError(t, afero.WriteFile(fsys, seedPath, []byte("INSERT INTO test_table;"), 0644))
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		conn.Query(migration.ListSchemas, escapedSchemas).
+			Reply("SELECT 1", []interface{}{"private"}).
+			Query("DROP SCHEMA IF EXISTS private CASCADE").
+			Reply("DROP SCHEMA").
+			Query(migration.DropObjects).
+			Reply("INSERT 0")
+		helper.MockMigrationHistory(conn).
+			Query(migration.INSERT_MIGRATION_VERSION, "0", "schema", nil).
+			Reply("INSERT 0 1")
+		utils.Config.Db.Seed.Enabled = false
+		// Run test
+		err := resetRemote(context.Background(), "", dbConfig, fsys, conn.Intercept)
+		// No error should be raised since we're skipping the seed
+		assert.NoError(t, err)
+	})
+
 	t.Run("throws error on connect failure", func(t *testing.T) {
 		// Setup in-memory fs
 		fsys := afero.NewMemMapFs()

diff --git a/internal/migration/apply/apply.go b/internal/migration/apply/apply.go
@@ -20,6 +20,9 @@ func MigrateAndSeed(ctx context.Context, version string, conn *pgx.Conn, fsys af
 	if err := migration.ApplyMigrations(ctx, migrations, conn, afero.NewIOFS(fsys)); err != nil {
 		return err
 	}
+	if !utils.Config.Db.Seed.Enabled {
+		return nil
+	}
 	return SeedDatabase(ctx, conn, fsys)
 }
 

diff --git a/internal/migration/apply/apply_test.go b/internal/migration/apply/apply_test.go
@@ -38,6 +38,30 @@ func TestMigrateDatabase(t *testing.T) {
 		assert.NoError(t, err)
 	})
 
+	t.Run("skip seeding when seed config is disabled", func(t *testing.T) {
+		// Setup in-memory fs
+		fsys := afero.NewMemMapFs()
+		path := filepath.Join(utils.MigrationsDir, "0_test.sql")
+		sql := "create schema public"
+		require.NoError(t, afero.WriteFile(fsys, path, []byte(sql), 0644))
+		seedPath := filepath.Join(utils.SeedDataPath)
+		// This will raise an error when seeding
+		require.NoError(t, afero.WriteFile(fsys, seedPath, []byte("INSERT INTO test_table;"), 0644))
+		// Setup mock postgres
+		conn := pgtest.NewConn()
+		defer conn.Close(t)
+		helper.MockMigrationHistory(conn).
+			Query(sql).
+			Reply("CREATE SCHEMA").
+			Query(migration.INSERT_MIGRATION_VERSION, "0", "test", []string{sql}).
+			Reply("INSERT 0 1")
+		utils.Config.Db.Seed.Enabled = false
+		// Run test
+		err := MigrateAndSeed(context.Background(), "", conn.MockClient(t), fsys)
+		// No error should be returned since seeding is skipped
+		assert.NoError(t, err)
+	})
+
 	t.Run("ignores empty local directory", func(t *testing.T) {
 		assert.NoError(t, MigrateAndSeed(context.Background(), "", nil, afero.NewMemMapFs()))
 	})

diff --git a/pkg/config/config.go b/pkg/config/config.go
@@ -160,6 +160,11 @@ type (
 		Password     string `toml:"-"`
 		RootKey      string `toml:"-" mapstructure:"root_key"`
 		Pooler       pooler `toml:"pooler"`
+		Seed         seed   `toml:"seed"`
+	}
+
+	seed struct {
+		Enabled bool `toml:"enabled"`
 	}
 
 	pooler struct {
@@ -457,6 +462,9 @@ func NewConfig(editors ...ConfigEditor) config {
 				EncryptionKey: "12345678901234567890123456789032",
 				SecretKeyBase: "EAx3IQ/wRG1v47ZD4NE4/9RzBI8Jmil3x0yhcW4V2NHBP6c2iPIzwjofi2Ep4HIG",
 			},
+			Seed: seed{
+				Enabled: true,
+			},
 		},
 		Realtime: realtime{
 			Image:           realtimeImage,