Repository: polarsignals/frostdb
Pull request #500 (2023-08-24 19:27:56)

Cleanup db
Made clearing the storage a little more natural, and with fewer footguns. Now it uses `db.Close` to perform clear storage, which means that it's no longer valid to use this DB afterwards.

This is preferable to previously with `DropStorage` because the wal isn't meant to be live reloaded and that created a lot of footguns. So instead to acheive the same results one would clear the db storage, and reopen the database again. 
```
db.Close(WithClearStorage())
db = c.DB(....)
```


Patch:
diff --git a/db.go b/db.go
index 222d425..dddd37a 100644
--- a/db.go
+++ b/db.go
@@ -775,7 +775,23 @@ func (db *DB) recover(ctx context.Context, wal WAL) error {
 	return nil
 }
 
-func (db *DB) Close() error {
+type CloseOption func(*closeOptions)
+
+type closeOptions struct {
+	clearStorage bool
+}
+
+func WithClearStorage() CloseOption {
+	return func(o *closeOptions) {
+		o.clearStorage = true
+	}
+}
+
+func (db *DB) Close(options ...CloseOption) error {
+	opts := &closeOptions{}
+	for _, opt := range options {
+		opt(opts)
+	}
 	level.Info(db.logger).Log("msg", "closing DB")
 	shouldPersist := len(db.sinks) > 0 && !db.columnStore.manualBlockRotation
 	for _, table := range db.tables {
@@ -796,41 +812,8 @@ func (db *DB) Close() error {
 		return err
 	}
 
-	if shouldPersist {
-		// If we've successfully persisted all the table blocks we can remove
-		// the wal and snapshots. Move the two directories to a trash dir since
-		// it is an O(1) operation. The trash dir is cleaned up on startup in
-		// case we cannot delete the WAL/snapshots in time.
-		trashDir := db.trashDir()
-
-		if moveErr := func() error {
-			_ = os.Mkdir(trashDir, os.FileMode(0o755))
-			// Create a temporary directory in the trash dir to avoid clashing
-			// with other wal/snapshot dirs that might not have been removed
-			// previously.
-			tmpPath, err := os.MkdirTemp(trashDir, "")
-			if err != nil {
-				return err
-			}
-			if err := os.Rename(db.snapshotsDir(), filepath.Join(tmpPath, snapshotsPath)); err != nil && !os.IsNotExist(err) {
-				return err
-			}
-			if err := os.Rename(db.walDir(), filepath.Join(tmpPath, walPath)); err != nil && !os.IsNotExist(err) {
-				return err
-			}
-			return nil
-		}(); moveErr != nil {
-			// If we failed to move the wal/snapshots to the trash dir, fall
-			// back to attempting to remove them with RemoveAll.
-			if err := os.RemoveAll(db.snapshotsDir()); err != nil {
-				return fmt.Errorf("%v: %v", moveErr, err)
-			}
-			if err := os.RemoveAll(db.walDir()); err != nil {
-				return fmt.Errorf("%v: %v", moveErr, err)
-			}
-			return moveErr
-		}
-		if err := os.RemoveAll(trashDir); err != nil {
+	if shouldPersist || opts.clearStorage {
+		if err := db.dropStorage(); err != nil {
 			return err
 		}
 		level.Info(db.logger).Log("msg", "cleaned up wal & snapshots")
@@ -1115,7 +1098,36 @@ func validateName(name string) bool {
 	return !strings.Contains(name, "/")
 }
 
-// DropStorage drops all data in the storage.
-func (db *DB) DropStorage() error {
-	return os.RemoveAll(db.storagePath)
+// dropStorage removes snapshots and WAL data from the storage directory.
+func (db *DB) dropStorage() error {
+	trashDir := db.trashDir()
+
+	if moveErr := func() error {
+		_ = os.Mkdir(trashDir, os.FileMode(0o755))
+		// Create a temporary directory in the trash dir to avoid clashing
+		// with other wal/snapshot dirs that might not have been removed
+		// previously.
+		tmpPath, err := os.MkdirTemp(trashDir, "")
+		if err != nil {
+			return err
+		}
+		if err := os.Rename(db.snapshotsDir(), filepath.Join(tmpPath, snapshotsPath)); err != nil && !os.IsNotExist(err) {
+			return err
+		}
+		if err := os.Rename(db.walDir(), filepath.Join(tmpPath, walPath)); err != nil && !os.IsNotExist(err) {
+			return err
+		}
+		return nil
+	}(); moveErr != nil {
+		// If we failed to move the wal/snapshots to the trash dir, fall
+		// back to attempting to remove them with RemoveAll.
+		if err := os.RemoveAll(db.snapshotsDir()); err != nil {
+			return fmt.Errorf("%v: %v", moveErr, err)
+		}
+		if err := os.RemoveAll(db.walDir()); err != nil {
+			return fmt.Errorf("%v: %v", moveErr, err)
+		}
+		return moveErr
+	}
+	return os.RemoveAll(trashDir)
 }
diff --git a/dynparquet/example.go b/dynparquet/example.go
index 4301693..8648067 100644
--- a/dynparquet/example.go
+++ b/dynparquet/example.go
@@ -54,7 +54,47 @@ func (s Samples) ToBuffer(schema *Schema) (*Buffer, error) {
 	return pb, nil
 }
 
-func (s Samples) ToRecord(schema *arrow.Schema) (arrow.Record, error) {
+func (s Samples) ToRecord() (arrow.Record, error) {
+	fields := []arrow.Field{
+		{
+			Name: "example_type",
+			Type: &arrow.DictionaryType{
+				IndexType: &arrow.Int32Type{},
+				ValueType: &arrow.BinaryType{},
+			},
+		},
+	}
+	seen := map[string]struct{}{}
+	for _, smpl := range s {
+		for _, lbl := range smpl.Labels {
+			if _, ok := seen[lbl.Name]; ok {
+				continue
+			}
+			seen[lbl.Name] = struct{}{}
+			fields = append(fields,
+				arrow.Field{
+					Name:     "labels." + lbl.Name,
+					Nullable: true,
+					Type: &arrow.DictionaryType{
+						IndexType: &arrow.Int32Type{},
+						ValueType: &arrow.BinaryType{},
+					},
+				},
+			)
+		}
+	}
+	fields = append(fields, arrow.Field{
+		Name: "stacktrace",
+		Type: &arrow.DictionaryType{
+			IndexType: &arrow.Int32Type{},
+			ValueType: &arrow.BinaryType{},
+		},
+	},
+	)
+	fields = append(fields, arrow.Field{Name: "timestamp", Type: arrow.PrimitiveTypes.Int64})
+	fields = append(fields, arrow.Field{Name: "value", Type: arrow.PrimitiveTypes.Int64})
+	schema := arrow.NewSchema(fields, nil)
+
 	bld := array.NewRecordBuilder(memory.NewGoAllocator(), schema)
 	defer bld.Release()
 


Test patch:
diff --git a/compaction_test.go b/compaction_test.go
index 9880b91..5e1031f 100644
--- a/compaction_test.go
+++ b/compaction_test.go
@@ -12,8 +12,6 @@ import (
 
 	"github.com/polarsignals/frostdb/dynparquet"
 	"github.com/polarsignals/frostdb/parts"
-	"github.com/polarsignals/frostdb/pqarrow"
-	"github.com/polarsignals/frostdb/query/logicalplan"
 )
 
 // insertSamples is a helper function to insert a deterministic sample with a
@@ -53,16 +51,7 @@ func insertSampleRecords(ctx context.Context, t *testing.T, table *Table, timest
 		})
 	}
 
-	ps, err := table.Schema().GetDynamicParquetSchema(map[string][]string{
-		"labels": {"label1"},
-	})
-	require.NoError(t, err)
-	defer table.Schema().PutPooledParquetSchema(ps)
-
-	sc, err := pqarrow.ParquetSchemaToArrowSchema(ctx, ps.Schema, logicalplan.IterOptions{})
-	require.NoError(t, err)
-
-	ar, err := samples.ToRecord(sc)
+	ar, err := samples.ToRecord()
 	require.NoError(t, err)
 
 	tx, err := table.InsertRecord(ctx, ar)
diff --git a/db_test.go b/db_test.go
index c6234e3..973ba74 100644
--- a/db_test.go
+++ b/db_test.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/apache/arrow/go/v13/arrow"
 	"github.com/apache/arrow/go/v13/arrow/memory"
+	"github.com/go-kit/log/level"
 	"github.com/google/uuid"
 	"github.com/parquet-go/parquet-go"
 	"github.com/stretchr/testify/require"
@@ -20,7 +21,6 @@ import (
 
 	"github.com/polarsignals/frostdb/dynparquet"
 	schemapb "github.com/polarsignals/frostdb/gen/proto/go/frostdb/schema/v1alpha1"
-	"github.com/polarsignals/frostdb/pqarrow"
 	"github.com/polarsignals/frostdb/query"
 	"github.com/polarsignals/frostdb/query/logicalplan"
 )
@@ -138,17 +138,7 @@ func TestDBWithWAL(t *testing.T) {
 
 		switch isArrow {
 		case true:
-
-			ps, err := table.Schema().GetDynamicParquetSchema(map[string][]string{
-				"labels": {"label1", "label2", "label3", "label4"},
-			})
-			require.NoError(t, err)
-			defer table.Schema().PutPooledParquetSchema(ps)
-
-			sc, err := pqarrow.ParquetSchemaToArrowSchema(ctx, ps.Schema, logicalplan.IterOptions{})
-			require.NoError(t, err)
-
-			rec, err := samples.ToRecord(sc)
+			rec, err := samples.ToRecord()
 			require.NoError(t, err)
 
 			ctx := context.Background()
@@ -180,16 +170,7 @@ func TestDBWithWAL(t *testing.T) {
 
 		switch isArrow {
 		case true:
-			ps, err := table.Schema().GetDynamicParquetSchema(map[string][]string{
-				"labels": {"label1", "label2"},
-			})
-			require.NoError(t, err)
-			defer table.Schema().PutPooledParquetSchema(ps)
-
-			sc, err := pqarrow.ParquetSchemaToArrowSchema(ctx, ps.Schema, logicalplan.IterOptions{})
-			require.NoError(t, err)
-
-			rec, err := samples.ToRecord(sc)
+			rec, err := samples.ToRecord()
 			require.NoError(t, err)
 
 			ctx := context.Background()
@@ -220,16 +201,7 @@ func TestDBWithWAL(t *testing.T) {
 
 		switch isArrow {
 		case true:
-			ps, err := table.Schema().GetDynamicParquetSchema(map[string][]string{
-				"labels": {"label1", "label2", "label3"},
-			})
-			require.NoError(t, err)
-			defer table.Schema().PutPooledParquetSchema(ps)
-
-			sc, err := pqarrow.ParquetSchemaToArrowSchema(ctx, ps.Schema, logicalplan.IterOptions{})
-			require.NoError(t, err)
-
-			rec, err := samples.ToRecord(sc)
+			rec, err := samples.ToRecord()
 			require.NoError(t, err)
 
 			ctx := context.Background()
@@ -1317,16 +1289,7 @@ func Test_DB_TableWrite_ArrowRecord(t *testing.T) {
 		},
 	}
 
-	ps, err := table.Schema().GetDynamicParquetSchema(map[string][]string{
-		"labels": {"label1", "label2", "label3"},
-	})
-	require.NoError(t, err)
-	defer table.Schema().PutPooledParquetSchema(ps)
-
-	sc, err := pqarrow.ParquetSchemaToArrowSchema(ctx, ps.Schema, logicalplan.IterOptions{})
-	require.NoError(t, err)
-
-	r, err := samples.ToRecord(sc)
+	r, err := samples.ToRecord()
 	require.NoError(t, err)
 
 	_, err = table.InsertRecord(ctx, r)
@@ -1810,3 +1773,75 @@ func Test_DB_Limiter(t *testing.T) {
 		})
 	}
 }
+
+// DropStorage ensures that a database can continue on after drop storage is called.
+func Test_DB_DropStorage(t *testing.T) {
+	logger := newTestLogger(t)
+	mem := memory.NewCheckedAllocator(memory.DefaultAllocator)
+	defer mem.AssertSize(t, 0)
+
+	config := NewTableConfig(
+		dynparquet.SampleDefinition(),
+	)
+
+	dir := t.TempDir()
+
+	c, err := New(
+		WithLogger(logger),
+		WithWAL(),
+		WithStoragePath(dir),
+		WithActiveMemorySize(1024*1024),
+	)
+	defer func() {
+		_ = c.Close()
+	}()
+	require.NoError(t, err)
+	db, err := c.DB(context.Background(), "test")
+	require.NoError(t, err)
+	table, err := db.Table("test", config)
+	require.NoError(t, err)
+
+	samples := dynparquet.NewTestSamples()
+
+	ctx := context.Background()
+	for i := 0; i < 100; i++ {
+		r, err := samples.ToRecord()
+		require.NoError(t, err)
+		defer r.Release()
+		_, err = table.InsertRecord(ctx, r)
+		require.NoError(t, err)
+	}
+	countRows := func(expected int) {
+		rows := 0
+		engine := query.NewEngine(mem, db.TableProvider())
+		err = engine.ScanTable("test").
+			Execute(context.Background(), func(ctx context.Context, r arrow.Record) error {
+				rows += int(r.NumRows())
+				return nil
+			})
+		require.NoError(t, err)
+		require.Equal(t, expected, rows)
+	}
+	countRows(300)
+
+	level.Debug(logger).Log("msg", "dropping storage")
+	require.NoError(t, db.Close(WithClearStorage()))
+
+	// Open a new store against the dropped storage, and expect empty db
+	c, err = New(
+		WithLogger(logger),
+		WithWAL(),
+		WithStoragePath(dir),
+		WithActiveMemorySize(1024*1024),
+	)
+	defer func() {
+		_ = c.Close()
+	}()
+	require.NoError(t, err)
+	level.Debug(logger).Log("msg", "opening new db")
+	db, err = c.DB(context.Background(), "test")
+	require.NoError(t, err)
+	_, err = db.Table("test", config)
+	require.NoError(t, err)
+	countRows(0)
+}
diff --git a/table_test.go b/table_test.go
index 66d91a5..e3f6934 100644
--- a/table_test.go
+++ b/table_test.go
@@ -1172,19 +1172,10 @@ func Test_L0Query(t *testing.T) {
 		Value:     3,
 	}}
 
-	ps, err := table.Schema().GetDynamicParquetSchema(map[string][]string{
-		"labels": {"label1", "label2", "label3", "label4"},
-	})
+	r, err := samples.ToRecord()
 	require.NoError(t, err)
-	defer table.Schema().PutPooledParquetSchema(ps)
 
 	ctx := context.Background()
-	sc, err := pqarrow.ParquetSchemaToArrowSchema(ctx, ps.Schema, logicalplan.IterOptions{})
-	require.NoError(t, err)
-
-	r, err := samples.ToRecord(sc)
-	require.NoError(t, err)
-
 	_, err = table.InsertRecord(ctx, r)
 	require.NoError(t, err)
 
@@ -1427,19 +1418,10 @@ func Test_Table_Size(t *testing.T) {
 		samples := dynparquet.NewTestSamples()
 		switch isArrow {
 		case true:
-			ps, err := table.Schema().GetDynamicParquetSchema(map[string][]string{
-				"labels": {"node", "namespace", "container"},
-			})
+			rec, err := samples.ToRecord()
 			require.NoError(t, err)
-			defer table.Schema().PutPooledParquetSchema(ps)
 
 			ctx := context.Background()
-			sc, err := pqarrow.ParquetSchemaToArrowSchema(ctx, ps.Schema, logicalplan.IterOptions{})
-			require.NoError(t, err)
-
-			rec, err := samples.ToRecord(sc)
-			require.NoError(t, err)
-
 			_, err = table.InsertRecord(ctx, rec)
 			require.NoError(t, err)