144 lines
3.2 KiB
Go
144 lines
3.2 KiB
Go
package cmd
|
|
|
|
import (
|
|
"bufio"
|
|
"encoding/base64"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"strings"
|
|
|
|
"github.com/dgraph-io/badger/v4"
|
|
"github.com/spf13/cobra"
|
|
)
|
|
|
|
var restoreCmd = &cobra.Command{
|
|
Use: "restore [DB]",
|
|
Short: "Restore key/value pairs from an NDJSON dump",
|
|
Aliases: []string{"import"},
|
|
Args: cobra.MaximumNArgs(1),
|
|
RunE: restore,
|
|
SilenceUsage: true,
|
|
}
|
|
|
|
func restore(cmd *cobra.Command, args []string) error {
|
|
store := &Store{}
|
|
dbName := "default"
|
|
if len(args) == 1 {
|
|
parsed, err := store.parseDB(args[0], false)
|
|
if err != nil {
|
|
return fmt.Errorf("cannot restore '%s': %v", args[0], err)
|
|
}
|
|
dbName = parsed
|
|
}
|
|
|
|
reader, closer, err := restoreInput(cmd)
|
|
if err != nil {
|
|
return fmt.Errorf("cannot restore '%s': %v", dbName, err)
|
|
}
|
|
if closer != nil {
|
|
defer closer.Close()
|
|
}
|
|
|
|
db, err := store.open(dbName)
|
|
if err != nil {
|
|
return fmt.Errorf("cannot restore '%s': %v", dbName, err)
|
|
}
|
|
defer db.Close()
|
|
|
|
scanner := bufio.NewScanner(reader)
|
|
buf := make([]byte, 1024*1024)
|
|
scanner.Buffer(buf, 8*1024*1024)
|
|
|
|
wb := db.NewWriteBatch()
|
|
defer wb.Cancel()
|
|
|
|
lineNo := 0
|
|
var restored int
|
|
|
|
for scanner.Scan() {
|
|
lineNo++
|
|
line := strings.TrimSpace(scanner.Text())
|
|
if line == "" {
|
|
continue
|
|
}
|
|
|
|
var entry dumpEntry
|
|
if err := json.Unmarshal([]byte(line), &entry); err != nil {
|
|
return fmt.Errorf("cannot restore '%s': line %d: %w", dbName, lineNo, err)
|
|
}
|
|
if entry.Key == "" {
|
|
return fmt.Errorf("cannot restore '%s': line %d: missing key", dbName, lineNo)
|
|
}
|
|
|
|
value, err := decodeEntryValue(entry)
|
|
if err != nil {
|
|
return fmt.Errorf("cannot restore '%s': line %d: %w", dbName, lineNo, err)
|
|
}
|
|
|
|
entryMeta := byte(0x0)
|
|
if entry.Secret {
|
|
entryMeta = metaSecret
|
|
}
|
|
|
|
writeEntry := badger.NewEntry([]byte(entry.Key), value).WithMeta(entryMeta)
|
|
if entry.ExpiresAt != nil {
|
|
if *entry.ExpiresAt < 0 {
|
|
return fmt.Errorf("cannot restore '%s': line %d: expires_at must be >= 0", dbName, lineNo)
|
|
}
|
|
writeEntry.ExpiresAt = uint64(*entry.ExpiresAt)
|
|
}
|
|
|
|
if err := wb.SetEntry(writeEntry); err != nil {
|
|
return fmt.Errorf("cannot restore '%s': line %d: %w", dbName, lineNo, err)
|
|
}
|
|
restored++
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
return err
|
|
}
|
|
|
|
if err := wb.Flush(); err != nil {
|
|
return fmt.Errorf("cannot restore '%s': %v", dbName, err)
|
|
}
|
|
|
|
fmt.Fprintf(cmd.ErrOrStderr(), "Restored %d entries into @%s\n", restored, dbName)
|
|
return nil
|
|
}
|
|
|
|
func restoreInput(cmd *cobra.Command) (io.Reader, io.Closer, error) {
|
|
filePath, err := cmd.Flags().GetString("file")
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
if strings.TrimSpace(filePath) == "" {
|
|
return cmd.InOrStdin(), nil, nil
|
|
}
|
|
f, err := os.Open(filePath)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
return f, f, nil
|
|
}
|
|
|
|
func decodeEntryValue(entry dumpEntry) ([]byte, error) {
|
|
switch entry.Encoding {
|
|
case "", "text":
|
|
return []byte(entry.Value), nil
|
|
case "base64":
|
|
b, err := base64.StdEncoding.DecodeString(entry.Value)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return b, nil
|
|
default:
|
|
return nil, fmt.Errorf("unsupported encoding %q", entry.Encoding)
|
|
}
|
|
}
|
|
|
|
func init() {
|
|
restoreCmd.Flags().StringP("file", "f", "", "Path to an NDJSON dump (defaults to stdin)")
|
|
rootCmd.AddCommand(restoreCmd)
|
|
}
|