Add badger.compression to Dgraph restore (#6987)

hypermodeinc · Nov 30, 2020 · 784b245 · 784b245
1 parent 8e5e538
commit 784b245
Show file tree

Hide file tree

Showing 7 changed files with 21 additions and 11 deletions.
diff --git a/dgraph/cmd/bulk/run.go b/dgraph/cmd/bulk/run.go
@@ -158,10 +158,6 @@ func run() {
 	if opt.Version {
 		os.Exit(0)
 	}
-	if opt.BadgerCompressionLevel < 0 {
-		fmt.Printf("Invalid compression level: %d. It should be non-negative",
-			opt.BadgerCompressionLevel)
-	}
 
 	totalCache := int64(Bulk.Conf.GetInt("badger.cache-mb"))
 	x.AssertTruef(totalCache >= 0, "ERROR: Cache size must be non-negative")

diff --git a/ee/backup/run.go b/ee/backup/run.go
@@ -39,6 +39,7 @@ var ExportBackup x.SubCommand
 
 var opt struct {
 	backupId    string
+	compression string
 	location    string
 	pdir        string
 	zero        string
@@ -111,6 +112,10 @@ $ dgraph restore -p . -l /var/backups/dgraph -z localhost:5080
 	}
 
 	flag := Restore.Cmd.Flags()
+	flag.StringVar(&opt.compression, "badger.compression", "snappy",
+		"[none, zstd:level, snappy] Specifies the compression algorithm and the compression"+
+			"level (if applicable) for the postings directory. none would disable compression,"+
+			" while zstd:1 would set zstd compression at level 1.")
 	flag.StringVarP(&opt.location, "location", "l", "",
 		"Sets the source location URI (required).")
 	flag.StringVarP(&opt.pdir, "postings", "p", "",
@@ -215,8 +220,10 @@ func runRestoreCmd() error {
 		zc = pb.NewZeroClient(zero)
 	}
 
+	ctype, clevel := x.ParseCompression(opt.compression)
+
 	start = time.Now()
-	result := worker.RunRestore(opt.pdir, opt.location, opt.backupId, opt.key)
+	result := worker.RunRestore(opt.pdir, opt.location, opt.backupId, opt.key, ctype, clevel)
 	if result.Err != nil {
 		return result.Err
 	}

diff --git a/systest/backup/encryption/backup_test.go b/systest/backup/encryption/backup_test.go
@@ -28,6 +28,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/dgraph-io/badger/v2/options"
 	"github.com/dgraph-io/dgo/v200/protos/api"
 	"github.com/dgraph-io/dgraph/ee/enc"
 	"github.com/dgraph-io/dgraph/testutil"
@@ -354,7 +355,7 @@ func runFailingRestore(t *testing.T, backupLocation, lastDir string, commitTs ui
 	require.NotNil(t, k)
 	require.NoError(t, err)
 
-	result := worker.RunRestore("./data/restore", backupLocation, lastDir, k)
+	result := worker.RunRestore("./data/restore", backupLocation, lastDir, k, options.Snappy, 0)
 	require.Error(t, result.Err)
 	require.Contains(t, result.Err.Error(), "expected a BackupNum value of 1")
 }

diff --git a/systest/backup/filesystem/backup_test.go b/systest/backup/filesystem/backup_test.go
@@ -30,6 +30,7 @@ import (
 
 	"google.golang.org/grpc/credentials"
 
+	"github.com/dgraph-io/badger/v2/options"
 	"github.com/dgraph-io/dgo/v200"
 	"github.com/dgraph-io/dgo/v200/protos/api"
 	"github.com/stretchr/testify/require"
@@ -356,7 +357,7 @@ func runRestore(t *testing.T, backupLocation, lastDir string, commitTs uint64) m
 	require.NoError(t, os.RemoveAll(restoreDir))
 
 	t.Logf("--- Restoring from: %q", backupLocation)
-	result := worker.RunRestore("./data/restore", backupLocation, lastDir, x.SensitiveByteSlice(nil))
+	result := worker.RunRestore("./data/restore", backupLocation, lastDir, x.SensitiveByteSlice(nil), options.Snappy, 0)
 	require.NoError(t, result.Err)
 
 	for i, pdir := range []string{"p1", "p2", "p3"} {
@@ -379,7 +380,7 @@ func runFailingRestore(t *testing.T, backupLocation, lastDir string, commitTs ui
 	// calling restore.
 	require.NoError(t, os.RemoveAll(restoreDir))
 
-	result := worker.RunRestore("./data/restore", backupLocation, lastDir, x.SensitiveByteSlice(nil))
+	result := worker.RunRestore("./data/restore", backupLocation, lastDir, x.SensitiveByteSlice(nil), options.Snappy, 0)
 	require.Error(t, result.Err)
 	require.Contains(t, result.Err.Error(), "expected a BackupNum value of 1")
 }

diff --git a/systest/backup/minio-large/backup_test.go b/systest/backup/minio-large/backup_test.go
@@ -28,6 +28,7 @@ import (
 
 	"google.golang.org/grpc/credentials"
 
+	"github.com/dgraph-io/badger/v2/options"
 	"github.com/dgraph-io/dgo/v200"
 	"github.com/dgraph-io/dgo/v200/protos/api"
 	minio "github.com/minio/minio-go/v6"
@@ -166,7 +167,7 @@ func runRestore(t *testing.T, backupLocation, lastDir string, commitTs uint64) m
 	require.NoError(t, os.MkdirAll(restoreDir, os.ModePerm))
 
 	t.Logf("--- Restoring from: %q", backupLocation)
-	result := worker.RunRestore("./data/restore", backupLocation, lastDir, x.SensitiveByteSlice(nil))
+	result := worker.RunRestore("./data/restore", backupLocation, lastDir, x.SensitiveByteSlice(nil), options.Snappy, 0)
 	require.NoError(t, result.Err)
 
 	restored1, err := testutil.GetPredicateValues("./data/restore/p1", "name1", commitTs)

diff --git a/systest/backup/minio/backup_test.go b/systest/backup/minio/backup_test.go
@@ -30,6 +30,7 @@ import (
 
 	"google.golang.org/grpc/credentials"
 
+	"github.com/dgraph-io/badger/v2/options"
 	"github.com/dgraph-io/dgo/v200"
 	"github.com/dgraph-io/dgo/v200/protos/api"
 	minio "github.com/minio/minio-go/v6"
@@ -360,7 +361,7 @@ func runFailingRestore(t *testing.T, backupLocation, lastDir string, commitTs ui
 	// calling restore.
 	require.NoError(t, os.RemoveAll(restoreDir))
 
-	result := worker.RunRestore("./data/restore", backupLocation, lastDir, x.SensitiveByteSlice(nil))
+	result := worker.RunRestore("./data/restore", backupLocation, lastDir, x.SensitiveByteSlice(nil), options.Snappy, 0)
 	require.Error(t, result.Err)
 	require.Contains(t, result.Err.Error(), "expected a BackupNum value of 1")
 }

diff --git a/worker/restore.go b/worker/restore.go
@@ -24,6 +24,7 @@ import (
 	"path/filepath"
 
 	"github.com/dgraph-io/badger/v2"
+	"github.com/dgraph-io/badger/v2/options"
 	bpb "github.com/dgraph-io/badger/v2/pb"
 	"github.com/pkg/errors"
 
@@ -35,7 +36,7 @@ import (
 )
 
 // RunRestore calls badger.Load and tries to load data into a new DB.
-func RunRestore(pdir, location, backupId string, key x.SensitiveByteSlice) LoadResult {
+func RunRestore(pdir, location, backupId string, key x.SensitiveByteSlice, ctype options.CompressionType, clevel int) LoadResult {
 	// Create the pdir if it doesn't exist.
 	if err := os.MkdirAll(pdir, 0700); err != nil {
 		return LoadResult{0, 0, err}
@@ -64,6 +65,8 @@ func RunRestore(pdir, location, backupId string, key x.SensitiveByteSlice) LoadR
 			// The badger DB should be opened only after creating the backup
 			// file reader and verifying the encryption in the backup file.
 			db, err := badger.OpenManaged(badger.DefaultOptions(dir).
+				WithCompression(ctype).
+				WithZSTDCompressionLevel(clevel).
 				WithSyncWrites(false).
 				WithValueThreshold(1 << 10).
 				WithBlockCacheSize(100 * (1 << 20)).