github.com/hdt3213/godis@v1.2.9/database/tx_utils.go (about)

     1  package database
     2  
     3  import (
     4  	"github.com/hdt3213/godis/aof"
     5  	"github.com/hdt3213/godis/lib/utils"
     6  	"strconv"
     7  )
     8  
     9  func readFirstKey(args [][]byte) ([]string, []string) {
    10  	// assert len(args) > 0
    11  	key := string(args[0])
    12  	return nil, []string{key}
    13  }
    14  
    15  func writeFirstKey(args [][]byte) ([]string, []string) {
    16  	key := string(args[0])
    17  	return []string{key}, nil
    18  }
    19  
    20  func writeAllKeys(args [][]byte) ([]string, []string) {
    21  	keys := make([]string, len(args))
    22  	for i, v := range args {
    23  		keys[i] = string(v)
    24  	}
    25  	return keys, nil
    26  }
    27  
    28  func readAllKeys(args [][]byte) ([]string, []string) {
    29  	keys := make([]string, len(args))
    30  	for i, v := range args {
    31  		keys[i] = string(v)
    32  	}
    33  	return nil, keys
    34  }
    35  
    36  func noPrepare(args [][]byte) ([]string, []string) {
    37  	return nil, nil
    38  }
    39  
    40  func rollbackFirstKey(db *DB, args [][]byte) []CmdLine {
    41  	key := string(args[0])
    42  	return rollbackGivenKeys(db, key)
    43  }
    44  
    45  func rollbackGivenKeys(db *DB, keys ...string) []CmdLine {
    46  	var undoCmdLines [][][]byte
    47  	for _, key := range keys {
    48  		entity, ok := db.GetEntity(key)
    49  		if !ok {
    50  			undoCmdLines = append(undoCmdLines,
    51  				utils.ToCmdLine("DEL", key),
    52  			)
    53  		} else {
    54  			undoCmdLines = append(undoCmdLines,
    55  				utils.ToCmdLine("DEL", key), // clean existed first
    56  				aof.EntityToCmd(key, entity).Args,
    57  				toTTLCmd(db, key).Args,
    58  			)
    59  		}
    60  	}
    61  	return undoCmdLines
    62  }
    63  
    64  func rollbackHashFields(db *DB, key string, fields ...string) []CmdLine {
    65  	var undoCmdLines [][][]byte
    66  	dict, errReply := db.getAsDict(key)
    67  	if errReply != nil {
    68  		return nil
    69  	}
    70  	if dict == nil {
    71  		undoCmdLines = append(undoCmdLines,
    72  			utils.ToCmdLine("DEL", key),
    73  		)
    74  		return undoCmdLines
    75  	}
    76  	for _, field := range fields {
    77  		entity, ok := dict.Get(field)
    78  		if !ok {
    79  			undoCmdLines = append(undoCmdLines,
    80  				utils.ToCmdLine("HDEL", key, field),
    81  			)
    82  		} else {
    83  			value, _ := entity.([]byte)
    84  			undoCmdLines = append(undoCmdLines,
    85  				utils.ToCmdLine("HSET", key, field, string(value)),
    86  			)
    87  		}
    88  	}
    89  	return undoCmdLines
    90  }
    91  
    92  func prepareSetCalculate(args [][]byte) ([]string, []string) {
    93  	keys := make([]string, len(args))
    94  	for i, arg := range args {
    95  		keys[i] = string(arg)
    96  	}
    97  	return nil, keys
    98  }
    99  
   100  func prepareSetCalculateStore(args [][]byte) ([]string, []string) {
   101  	dest := string(args[0])
   102  	keys := make([]string, len(args)-1)
   103  	keyArgs := args[1:]
   104  	for i, arg := range keyArgs {
   105  		keys[i] = string(arg)
   106  	}
   107  	return []string{dest}, keys
   108  }
   109  
   110  func rollbackSetMembers(db *DB, key string, members ...string) []CmdLine {
   111  	var undoCmdLines [][][]byte
   112  	set, errReply := db.getAsSet(key)
   113  	if errReply != nil {
   114  		return nil
   115  	}
   116  	if set == nil {
   117  		undoCmdLines = append(undoCmdLines,
   118  			utils.ToCmdLine("DEL", key),
   119  		)
   120  		return undoCmdLines
   121  	}
   122  	for _, member := range members {
   123  		ok := set.Has(member)
   124  		if !ok {
   125  			undoCmdLines = append(undoCmdLines,
   126  				utils.ToCmdLine("SREM", key, member),
   127  			)
   128  		} else {
   129  			undoCmdLines = append(undoCmdLines,
   130  				utils.ToCmdLine("SADD", key, member),
   131  			)
   132  		}
   133  	}
   134  	return undoCmdLines
   135  }
   136  
   137  // undoSetChange rollbacks SADD and SREM command
   138  func undoSetChange(db *DB, args [][]byte) []CmdLine {
   139  	key := string(args[0])
   140  	memberArgs := args[1:]
   141  	members := make([]string, len(memberArgs))
   142  	for i, mem := range memberArgs {
   143  		members[i] = string(mem)
   144  	}
   145  	return rollbackSetMembers(db, key, members...)
   146  }
   147  
   148  func rollbackZSetFields(db *DB, key string, fields ...string) []CmdLine {
   149  	var undoCmdLines [][][]byte
   150  	zset, errReply := db.getAsSortedSet(key)
   151  	if errReply != nil {
   152  		return nil
   153  	}
   154  	if zset == nil {
   155  		undoCmdLines = append(undoCmdLines,
   156  			utils.ToCmdLine("DEL", key),
   157  		)
   158  		return undoCmdLines
   159  	}
   160  	for _, field := range fields {
   161  		elem, ok := zset.Get(field)
   162  		if !ok {
   163  			undoCmdLines = append(undoCmdLines,
   164  				utils.ToCmdLine("ZREM", key, field),
   165  			)
   166  		} else {
   167  			score := strconv.FormatFloat(elem.Score, 'f', -1, 64)
   168  			undoCmdLines = append(undoCmdLines,
   169  				utils.ToCmdLine("ZADD", key, score, field),
   170  			)
   171  		}
   172  	}
   173  	return undoCmdLines
   174  }