diff --git a/services/database/backuptool/backuptool.go b/services/database/backuptool/backuptool.go index 482b7ed6cb360490c00b4a87febb8b4615dff5ad..704dff871dee40d4cfd1e08bcf89922cbd6844f0 100644 --- a/services/database/backuptool/backuptool.go +++ b/services/database/backuptool/backuptool.go @@ -4,8 +4,8 @@ import ( "bytes" "encoding/binary" "encoding/gob" - "errors" "flag" + "fmt" "io" "log" "os" @@ -19,6 +19,7 @@ var ( configFile = flag.String("config", "/etc/djrandom/server.conf", "Config file location") dump = flag.Bool("dump", false, "Dump the database") restore = flag.Bool("restore", false, "Restore a database dump") + verbose = flag.Bool("verbose", false, "Print statistics on the dump") ) // List of database buckets that we should dump. Other data can be @@ -47,12 +48,12 @@ func ReadRecord(in io.Reader) (*Record, error) { } data := make([]byte, sz) - n, err := in.Read(data) + n, err := io.ReadFull(in, data) if err != nil { return nil, err } if int64(n) != sz { - return nil, errors.New("short read") + return nil, fmt.Errorf("short read (got %d, want %d)", n, sz) } var record Record @@ -68,8 +69,8 @@ func WriteRecord(out io.Writer, record Record) error { if err := gob.NewEncoder(&buf).Encode(&record); err != nil { return err } - sz := int64(buf.Len()) enc := buf.Bytes() + sz := int64(len(enc)) binary.Write(out, binary.LittleEndian, sz) _, err := out.Write(enc) @@ -85,6 +86,7 @@ func scanner(db *db_client.DatabaseRpcClient, bucketCh chan string, outCh chan R for bucket := range bucketCh { // Call Scan() in batches, to avoid overloading the server. + total := 0 batchSize := 1000 startKey, endKey := "", "\xff" for startKey < endKey { @@ -109,6 +111,7 @@ func scanner(db *db_client.DatabaseRpcClient, bucketCh chan string, outCh chan R Data: data, } n++ + total++ startKey = key } // An empty result means we've reached the end @@ -119,6 +122,9 @@ func scanner(db *db_client.DatabaseRpcClient, bucketCh chan string, outCh chan R // Start the next scan from the next key. startKey = db_client.IncrementKey(startKey) } + if *verbose { + log.Printf("bucket %s: %d entries", bucket, total) + } } } @@ -198,6 +204,7 @@ func doRestore(db *db_client.DatabaseRpcClient) { } func main() { + log.SetFlags(0) flag.Parse() config.Parse(*configFile)