From 786bebe58bf7e86e3cf75e4c400edf3b691382f8 Mon Sep 17 00:00:00 2001
From: Todd Gruben <todd@molecula.com>
Date: Fri, 8 Oct 2021 10:49:41 -0500
Subject: [PATCH] partial backup/restore

---
 cmd/backup.go        |  1 +
 ctl/backup.go        | 18 +++++++++++++++-
 ctl/restore.go       | 50 +++++++++++++++++++++++++++++++++++++-------
 testBackupRestore.sh | 22 +++++++++++++++++++
 4 files changed, 83 insertions(+), 8 deletions(-)

diff --git a/cmd/backup.go b/cmd/backup.go
index 68baeb014..d81029d7a 100644
--- a/cmd/backup.go
+++ b/cmd/backup.go
@@ -40,6 +40,7 @@ Backs up a FeatureBase server to a local, tar-formatted snapshot file.
 	flags.BoolVar(&cmd.NoSync, "no-sync", false, "disable file sync")
 	flags.IntVar(&cmd.Concurrency, "concurrency", cmd.Concurrency, "number of concurrent backup goroutines")
 	flags.StringVar(&cmd.Host, "host", "localhost:10101", "host:port of FeatureBase.")
+	flags.StringVar(&cmd.Index, "index", "", "index to backup, default backs up all indexes. ")
 	ctl.SetTLSConfig(flags, "", &cmd.TLS.CertificatePath, &cmd.TLS.CertificateKeyPath, &cmd.TLS.CACertPath, &cmd.TLS.SkipVerify, &cmd.TLS.EnableClientVerification)
 	return ccmd
 }
diff --git a/ctl/backup.go b/ctl/backup.go
index e68ce52d2..cdcfb1427 100644
--- a/ctl/backup.go
+++ b/ctl/backup.go
@@ -24,7 +24,7 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/molecula/featurebase/v2"
+	pilosa "github.com/molecula/featurebase/v2"
 	"github.com/molecula/featurebase/v2/http"
 	"github.com/molecula/featurebase/v2/server"
 	"github.com/molecula/featurebase/v2/topology"
@@ -38,6 +38,9 @@ type BackupCommand struct { // nolint: maligned
 	// Destination host and port.
 	Host string `json:"host"`
 
+	// Optional Index filter
+	Index string `json:"index"`
+
 	// Path to write the backup to.
 	OutputDir string
 
@@ -91,6 +94,19 @@ func (cmd *BackupCommand) Run(ctx context.Context) (err error) {
 	if err != nil {
 		return fmt.Errorf("getting schema: %w", err)
 	}
+	if cmd.Index != "" {
+		for _, idx := range indexes {
+			if idx.Name == cmd.Index {
+				indexes = make([]*pilosa.IndexInfo, 0)
+				indexes = append(indexes, idx)
+				break
+			}
+		}
+		if len(indexes) <= 0 {
+			return fmt.Errorf("Index not found to back up")
+		}
+	}
+
 	schema := &pilosa.Schema{Indexes: indexes}
 
 	// Ensure output directory doesn't exist; then create output directory.
diff --git a/ctl/restore.go b/ctl/restore.go
index 26d72c788..cf03b9eab 100644
--- a/ctl/restore.go
+++ b/ctl/restore.go
@@ -17,6 +17,7 @@ package ctl
 import (
 	"context"
 	"crypto/tls"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -26,9 +27,10 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/molecula/featurebase/v2"
+	pilosa "github.com/molecula/featurebase/v2"
 	"github.com/molecula/featurebase/v2/server"
 	"github.com/molecula/featurebase/v2/topology"
+	"github.com/molecula/featurebase/v2/vprint"
 	"golang.org/x/sync/errgroup"
 )
 
@@ -102,7 +104,6 @@ func (cmd *RestoreCommand) Run(ctx context.Context) (err error) {
 	} else if err := cmd.restoreIDAlloc(ctx, primary); err != nil {
 		return fmt.Errorf("cannot restore idalloc: %w", err)
 	}
-
 	if err := cmd.restoreShards(ctx); err != nil {
 		return fmt.Errorf("cannot restore shards: %w", err)
 	} else if err := cmd.restoreIndexTranslation(ctx); err != nil {
@@ -128,11 +129,46 @@ func (cmd *RestoreCommand) restoreSchema(ctx context.Context, primary *topology.
 	}
 	defer f.Close()
 
-	cmd.Logger().Printf("Load Schema")
-	url := primary.URI.Path("/schema")
-
-	var client http.Client
-	_, err = client.Post(url, "application/json", f)
+	existingSchema, err := cmd.client.Schema(ctx)
+	if len(existingSchema) == 0 {
+		cmd.Logger().Printf("Load Schema")
+		url := primary.URI.Path("/schema")
+		var client http.Client
+		_, err = client.Post(url, "application/json", f)
+	} else {
+		schema := &pilosa.Schema{}
+		if err := json.NewDecoder(f).Decode(schema); err != nil {
+			if err != nil {
+				return err
+			}
+		}
+		exists := func(indexName string) bool {
+			for _, i := range existingSchema {
+				if i.Name == indexName {
+					return true
+				}
+			}
+			return false
+		}
+		//NOTE SHOULD ONLY BE ONE
+		for _, index := range schema.Indexes {
+			if exists(index.Name) {
+				return errors.New(fmt.Sprintf("Index Exists %v", index.Name))
+			}
+			vprint.VV("Create INDEX %v", index.Name)
+			err = cmd.client.CreateIndex(ctx, index.Name, index.Options)
+			if err != nil {
+				return err
+			}
+			for _, field := range index.Fields {
+				vprint.VV("Create Field %v", field.Name)
+				err = cmd.client.CreateFieldWithOptions(ctx, index.Name, field.Name, field.Options)
+				if err != nil {
+					return err
+				}
+			}
+		}
+	}
 	return err
 }
 
diff --git a/testBackupRestore.sh b/testBackupRestore.sh
index 6d0dbc67f..7bcbe3894 100755
--- a/testBackupRestore.sh
+++ b/testBackupRestore.sh
@@ -33,3 +33,25 @@ else
 	echo "FAIL Single"
 	exit 1
 fi
+
+datagen --source texas_health -e 9999 --pilosa.index newsink --pilosa.batch-size 10000 --pilosa.hosts pilosa0:10101
+before=$(/featurebase chksum --host pilosa0:10101)
+/featurebase backup -o newbackupdir --host pilosa0:10101 --index newsink
+curl -X DELETE -s pilosa0:10101/index/newsink
+/featurebase restore -s newbackupdir --host pilosa0:10101
+after=$(/featurebase chksum --host pilosa0:10101)
+if [ "$before" = "$after" ]; then
+	echo "PASS Cluster Table"
+else
+	echo "FAIL Single Table"
+	exit 1
+fi
+/featurebase restore -s newbackupdir --host pilosax:10101
+single=$(/featurebase chksum --host pilosax:10101)
+if [ "$before" = "$single" ]; then
+	echo "PASS Single Table"
+	exit 0
+else
+	echo "FAIL Single Table"
+	exit 1
+fi