mirror of
https://gitee.com/he3db/he3pg.git
synced 2024-11-29 18:58:35 +08:00
129 lines
3.0 KiB
Go
129 lines
3.0 KiB
Go
package cmd
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io/ioutil"
|
|
"strconv"
|
|
|
|
"github.com/aws/aws-sdk-go/aws"
|
|
"github.com/aws/aws-sdk-go/aws/credentials"
|
|
"github.com/aws/aws-sdk-go/aws/session"
|
|
"github.com/aws/aws-sdk-go/service/s3"
|
|
"github.com/spf13/cobra"
|
|
"github.com/tikv/client-go/txnkv"
|
|
)
|
|
|
|
var restoreCmd = &cobra.Command{
|
|
Use: "restore",
|
|
Short: "Restore He3DB",
|
|
Long: "Welcome to use hbr for He3DB restore",
|
|
Run: runRestore,
|
|
}
|
|
|
|
func init() {
|
|
rootCmd.AddCommand(restoreCmd)
|
|
}
|
|
|
|
func runRestore(cmd *cobra.Command, args []string) {
|
|
access_key, _ := cmd.Flags().GetString("access_key")
|
|
secret_key, _ := cmd.Flags().GetString("secret_key")
|
|
endpoint, _ := cmd.Flags().GetString("endpoint")
|
|
region, _ := cmd.Flags().GetString("region")
|
|
bucket, _ := cmd.Flags().GetString("bucket")
|
|
pd, _ := cmd.Flags().GetString("pd")
|
|
backup_name, _ := cmd.Flags().GetString("name")
|
|
|
|
if access_key == "" || secret_key == "" || endpoint == "" || region == "" || bucket == "" || pd == "" || backup_name == "" {
|
|
fmt.Printf("PARAMETER ERROR!\n")
|
|
return
|
|
}
|
|
|
|
client, err := txnkv.NewClient([]string{pd})
|
|
if err != nil {
|
|
fmt.Printf("Connect Tikv Error!\n%v\n", err)
|
|
return
|
|
}
|
|
defer func() {
|
|
client.Close()
|
|
}()
|
|
|
|
sess, err := session.NewSession(&aws.Config{
|
|
Region: aws.String(region),
|
|
Endpoint: aws.String(endpoint),
|
|
Credentials: credentials.NewStaticCredentials(access_key, secret_key, ""),
|
|
S3ForcePathStyle: aws.Bool(true),
|
|
})
|
|
if err != nil {
|
|
fmt.Printf("Connect S3 Error!\n%v\n", err)
|
|
return
|
|
}
|
|
s3_client := s3.New(sess)
|
|
|
|
count := 0
|
|
input := &s3.ListObjectsInput{
|
|
Bucket: aws.String(bucket),
|
|
Prefix: aws.String(backup_name),
|
|
}
|
|
for {
|
|
resp, err := s3_client.ListObjects(input)
|
|
if err != nil {
|
|
fmt.Printf("S3 ListObjects Error!\n%v\n", err)
|
|
return
|
|
}
|
|
|
|
for _, keys := range resp.Contents {
|
|
out, err := s3_client.GetObject(&s3.GetObjectInput{
|
|
Bucket: aws.String(bucket),
|
|
Key: aws.String(*keys.Key),
|
|
})
|
|
if err != nil {
|
|
fmt.Printf("S3 GetObject Error!\n%v\n", err)
|
|
return
|
|
}
|
|
defer out.Body.Close()
|
|
|
|
data, err := ioutil.ReadAll(out.Body)
|
|
if err != nil {
|
|
fmt.Printf("out.Body.Read!\n%v\n", err)
|
|
return
|
|
}
|
|
|
|
txn, err := client.Begin()
|
|
if err != nil {
|
|
fmt.Printf("Tikv Transaction Begin Error!\n%v\n", err)
|
|
return
|
|
}
|
|
|
|
fmt.Printf("filename:%s\n", (*keys.Key)[len(backup_name)+1:])
|
|
|
|
ret := make([]byte, (len(*keys.Key)-len(backup_name)-1)/2)
|
|
index := 0
|
|
for i := len(backup_name) + 1; i < len(*keys.Key); i += 2 {
|
|
value, _ := strconv.ParseUint((*keys.Key)[i:i+2], 16, 8)
|
|
ret[index] = byte(0xff & value)
|
|
index++
|
|
|
|
}
|
|
|
|
if err := txn.Set(ret, data); err != nil {
|
|
fmt.Printf("Tikv Set Error!\n%v\n", err)
|
|
return
|
|
}
|
|
|
|
if err := txn.Commit(context.TODO()); err != nil {
|
|
fmt.Printf("Tikv Transaction Commit Error!\n%v\n", err)
|
|
return
|
|
}
|
|
count++
|
|
}
|
|
if resp.NextMarker == nil {
|
|
fmt.Printf("Done!\n")
|
|
break
|
|
}
|
|
input.Marker = resp.NextMarker
|
|
}
|
|
fmt.Printf("N:%v\n", count)
|
|
fmt.Printf("Done!\n")
|
|
}
|