Merge pull request #2 from Peter-Sh/main

Support for SET options and SETEX command
cybergarage · May 9, 2024 · 1c4623c · 1c4623c
2 parents 0ac0af5 + ace2d17
commit 1c4623c
Show file tree

Hide file tree

Showing 5 changed files with 107 additions and 6 deletions.
diff --git a/doc/cmds/string.csv b/doc/cmds/string.csv
@@ -15,8 +15,8 @@ O,MGET,1.0.0,
 O,MSET,1.0.1,
 O,MSETNX,1.0.1,
 -,PSETEX,2.6.0,
-O,SET,1.0.0,Any options are not supported yet
--,SETEX,2.0.0,
+O,SET,1.0.0,
+O,SETEX,2.0.0,
 O,SETNX,2.0.0,
 -,SETRANGE,2.2.0,
 O,STRLEN,2.2.0,

diff --git a/redis/core_commander.go b/redis/core_commander.go
@@ -239,11 +239,26 @@ func (server *Server) registerCoreExecutors() {
 	})
 
 	server.RegisterExexutor("SET", func(conn *Conn, cmd string, args Arguments) (*Message, error) {
-		opt := newDefaultSetOption()
 		key, val, err := nextSetArguments(cmd, args)
 		if err != nil {
 			return nil, err
 		}
+
+		opt, err := nextSetOptionArguments(cmd, args)
+		if err != nil {
+			return nil, err
+		}
+		return server.userCommandHandler.Set(conn, key, val, opt)
+	})
+
+	server.RegisterExexutor("SETEX", func(conn *Conn, cmd string, args Arguments) (*Message, error) {
+		key, seconds, val, err := nextSetExArguments(cmd, args)
+		if err != nil {
+			return nil, err
+		}
+		opt := newDefaultSetOption()
+		opt.EX = time.Duration(seconds) * time.Second
+
 		return server.userCommandHandler.Set(conn, key, val, opt)
 	})
 

diff --git a/redis/errors.go b/redis/errors.go
@@ -32,6 +32,8 @@ const (
 	errorMissingCommandArgument = "%s: missing argument (%s) %w"
 	errorUnkownCommandArgument  = "%s: unknown argument (%s)"
 	errorInvalidCommandArgument = "%s: %w argument (%s - %s)"
+	errorUseOnlyOnce            = "%s may be used only once"
+	errorShouldBeGreaterThanInt = "%s should be greater than %d"
 )
 
 // NewErrNotSupported returns a new ErrNotSupported.

diff --git a/redis/handler_func.go b/redis/handler_func.go
@@ -16,6 +16,7 @@ package redis
 
 import (
 	"errors"
+	"fmt"
 	"regexp"
 	"strconv"
 	"strings"
@@ -113,6 +114,25 @@ func nextSetArguments(cmd string, args Arguments) (string, string, error) {
 	return key, val, err
 }
 
+func nextSetExArguments(cmd string, args Arguments) (string, int, string, error) {
+	key, err := args.NextString()
+	if err != nil {
+		return "", 0, "", newMissingArgumentError(cmd, "key", err)
+	}
+	seconds, err := args.NextInteger()
+	if err != nil {
+		return "", 0, "", newMissingArgumentError(cmd, "seconds", err)
+	}
+	if seconds < 1 {
+		return "", 0, "", newInvalidArgumentError(cmd, "seconds", fmt.Errorf(errorShouldBeGreaterThanInt, "argument", 0))
+	}
+	val, err := args.NextString()
+	if err != nil {
+		return "", 0, "", newMissingArgumentError(cmd, "value", err)
+	}
+	return key, seconds, val, err
+}
+
 func nextMGetArguments(cmd string, args Arguments) ([]string, error) {
 	return nextStringArrayArguments(cmd, "keys", args)
 }
@@ -121,6 +141,71 @@ func nextMSetArguments(cmd string, args Arguments) (map[string]string, error) {
 	return nextStringMapArguments(cmd, args)
 }
 
+func nextSetOptionArguments(cmd string, args Arguments) (SetOption, error) {
+	opt := newDefaultSetOption()
+	for {
+		argStr, err := args.NextString()
+		if err != nil {
+			if errors.Is(err, proto.ErrEOM) {
+				break
+			} else {
+				return opt, err
+			}
+		}
+		argStr = strings.ToUpper(argStr)
+		switch argStr {
+		case "NX":
+			if opt.NX || opt.XX {
+				return opt, newInvalidArgumentError(cmd, argStr, fmt.Errorf(errorUseOnlyOnce, "NX|XX"))
+			}
+			opt.NX = true
+		case "XX":
+			if opt.NX || opt.XX {
+				return opt, newInvalidArgumentError(cmd, argStr, fmt.Errorf(errorUseOnlyOnce, "NX|XX"))
+			}
+			opt.XX = true
+		case "EX", "PX", "EXAT", "PXAT":
+			if opt.EX > 0 || opt.PX > 0 || !opt.EXAT.IsZero() || !opt.PXAT.IsZero() {
+				return opt, newInvalidArgumentError(cmd, argStr, fmt.Errorf(errorUseOnlyOnce, "EX|PX|EXAT|PXAT"))
+			}
+			argInt, err := args.NextInteger()
+			if err != nil {
+				if errors.Is(err, proto.ErrEOM) {
+					return opt, newMissingArgumentError(cmd, argStr, err)
+				} else {
+					return opt, err
+				}
+			}
+			if argInt < 1 {
+				return opt, newInvalidArgumentError(cmd, argStr, fmt.Errorf(errorShouldBeGreaterThanInt, "expire", 0))
+			}
+			switch argStr {
+			case "EX":
+				opt.EX = time.Duration(argInt) * time.Second
+			case "PX":
+				opt.PX = time.Duration(argInt) * time.Millisecond
+			case "EXAT":
+				opt.EXAT = time.Unix(int64(argInt), 0)
+			case "PXAT":
+				opt.PXAT = time.UnixMilli(int64(argInt))
+			}
+		case "KEEPTTL":
+			if opt.KEEPTTL {
+				return opt, newInvalidArgumentError(cmd, argStr, fmt.Errorf(errorUseOnlyOnce, ""))
+			}
+			opt.KEEPTTL = true
+		case "GET":
+			if opt.GET {
+				return opt, newInvalidArgumentError(cmd, argStr, fmt.Errorf(errorUseOnlyOnce, ""))
+			}
+			opt.GET = true
+		default:
+			return opt, newUnkownArgumentError(cmd, argStr)
+		}
+	}
+	return opt, nil
+}
+
 // Hash argument fuctions
 
 func nextHashArgument(cmd string, args Arguments) (string, error) {

diff --git a/redis/options.go b/redis/options.go
@@ -102,14 +102,13 @@ func newScanTypeFromString(str string) (ScanType, error) {
 }
 
 func newDefaultSetOption() SetOption {
-	now := time.Now()
 	return SetOption{
 		NX:      false,
 		XX:      false,
 		EX:      0,
 		PX:      0,
-		EXAT:    now,
-		PXAT:    now,
+		EXAT:    time.Time{},
+		PXAT:    time.Time{},
 		KEEPTTL: false,
 		GET:     false,
 	}