github.com/hdt3213/godis@v1.2.9/cluster/copy.go (about)

     1  package cluster
     2  
     3  import (
     4  	"github.com/hdt3213/godis/interface/redis"
     5  	"github.com/hdt3213/godis/lib/utils"
     6  	"github.com/hdt3213/godis/redis/protocol"
     7  	"strconv"
     8  	"strings"
     9  )
    10  
    11  const copyToAnotherDBErr = "ERR Copying to another database is not allowed in cluster mode"
    12  const noReplace = "NoReplace"
    13  const useReplace = "UseReplace"
    14  
    15  // Copy copies the value stored at the source key to the destination key.
    16  // the origin and the destination must within the same node.
    17  func Copy(cluster *Cluster, c redis.Connection, args [][]byte) redis.Reply {
    18  	if len(args) < 3 {
    19  		return protocol.MakeErrReply("ERR wrong number of arguments for 'copy' command")
    20  	}
    21  	srcKey := string(args[1])
    22  	destKey := string(args[2])
    23  	srcNode := cluster.peerPicker.PickNode(srcKey)
    24  	destNode := cluster.peerPicker.PickNode(destKey)
    25  	replaceFlag := noReplace
    26  	if len(args) > 3 {
    27  		for i := 3; i < len(args); i++ {
    28  			arg := strings.ToLower(string(args[i]))
    29  			if arg == "db" {
    30  				return protocol.MakeErrReply(copyToAnotherDBErr)
    31  			} else if arg == "replace" {
    32  				replaceFlag = useReplace
    33  			} else {
    34  				return protocol.MakeSyntaxErrReply()
    35  			}
    36  		}
    37  	}
    38  
    39  	if srcNode == destNode {
    40  		return cluster.relay(srcNode, c, args)
    41  	}
    42  	groupMap := map[string][]string{
    43  		srcNode:  {srcKey},
    44  		destNode: {destKey},
    45  	}
    46  
    47  	txID := cluster.idGenerator.NextID()
    48  	txIDStr := strconv.FormatInt(txID, 10)
    49  	// prepare Copy from
    50  	srcPrepareResp := cluster.relayPrepare(srcNode, c, makeArgs("Prepare", txIDStr, "CopyFrom", srcKey))
    51  	if protocol.IsErrorReply(srcPrepareResp) {
    52  		// rollback src node
    53  		requestRollback(cluster, c, txID, map[string][]string{srcNode: {srcKey}})
    54  		return srcPrepareResp
    55  	}
    56  	srcPrepareMBR, ok := srcPrepareResp.(*protocol.MultiBulkReply)
    57  	if !ok || len(srcPrepareMBR.Args) < 2 {
    58  		requestRollback(cluster, c, txID, map[string][]string{srcNode: {srcKey}})
    59  		return protocol.MakeErrReply("ERR invalid prepare response")
    60  	}
    61  	// prepare Copy to
    62  	destPrepareResp := cluster.relayPrepare(destNode, c, utils.ToCmdLine3("Prepare", []byte(txIDStr),
    63  		[]byte("CopyTo"), []byte(destKey), srcPrepareMBR.Args[0], srcPrepareMBR.Args[1], []byte(replaceFlag)))
    64  	if protocol.IsErrorReply(destPrepareResp) {
    65  		// rollback src node
    66  		requestRollback(cluster, c, txID, groupMap)
    67  		return destPrepareResp
    68  	}
    69  	if _, errReply := requestCommit(cluster, c, txID, groupMap); errReply != nil {
    70  		requestRollback(cluster, c, txID, groupMap)
    71  		return errReply
    72  	}
    73  	return protocol.MakeIntReply(1)
    74  }
    75  
    76  // prepareCopyFrom is prepare-function for CopyFrom, see prepareFuncMap
    77  func prepareCopyFrom(cluster *Cluster, conn redis.Connection, cmdLine CmdLine) redis.Reply {
    78  	if len(cmdLine) != 2 {
    79  		return protocol.MakeArgNumErrReply("CopyFrom")
    80  	}
    81  	key := string(cmdLine[1])
    82  	existResp := cluster.db.ExecWithLock(conn, utils.ToCmdLine("Exists", key))
    83  	if protocol.IsErrorReply(existResp) {
    84  		return existResp
    85  	}
    86  	existIntResp := existResp.(*protocol.IntReply)
    87  	if existIntResp.Code == 0 {
    88  		return protocol.MakeErrReply("ERR no such key")
    89  	}
    90  	return cluster.db.ExecWithLock(conn, utils.ToCmdLine2("DumpKey", key))
    91  }
    92  
    93  func prepareCopyTo(cluster *Cluster, conn redis.Connection, cmdLine CmdLine) redis.Reply {
    94  	if len(cmdLine) != 5 {
    95  		return protocol.MakeArgNumErrReply("CopyTo")
    96  	}
    97  	key := string(cmdLine[1])
    98  	replaceFlag := string(cmdLine[4])
    99  	existResp := cluster.db.ExecWithLock(conn, utils.ToCmdLine("Exists", key))
   100  	if protocol.IsErrorReply(existResp) {
   101  		return existResp
   102  	}
   103  	existIntResp := existResp.(*protocol.IntReply)
   104  	if existIntResp.Code == 1 {
   105  		if replaceFlag == noReplace {
   106  			return protocol.MakeErrReply(keyExistsErr)
   107  		}
   108  	}
   109  	return protocol.MakeOkReply()
   110  }
   111  
   112  func init() {
   113  	registerPrepareFunc("CopyFrom", prepareCopyFrom)
   114  	registerPrepareFunc("CopyTo", prepareCopyTo)
   115  }