Iss 69 - Implement GETEX (#101)

GETEX implemented. Fixed issue in SortedSet.GetRandom where it would sometimes return an empty value in one of its indexes - @osteensco

echovault/api_generic.go

@@ -69,6 +69,28 @@ type PExpireOptions ExpireOptions
 type ExpireAtOptions ExpireOptions
 type PExpireAtOptions ExpireOptions
 
+// GetExOptions modifies the behaviour of
+//
+// EX - Set the specified expire time, in seconds.
+//
+// PX - Set the specified expire time, in milliseconds.
+//
+// EXAT - Set the specified Unix time at which the key will expire, in seconds.
+//
+// PXAT - Set the specified Unix time at which the key will expire, in milliseconds.
+//
+// PERSIST - Remove the time to live associated with the key.
+//
+// UNIXTIME - Number of seconds or miliseconds from now
+type GetExOptions struct {
+	EX       bool
+	PX       bool
+	EXAT     bool
+	PXAT     bool
+	PERSIST  bool
+	UNIXTIME int
+}
+
 // Set creates or modifies the value at the given key.
 //
 // Parameters:
@@ -323,7 +345,7 @@ func (server *EchoVault) Expire(key string, seconds int, options ExpireOptions)
 //
 // `key` - string.
 //
-// `milliseconds` - int - number of seconds from now.
+// `milliseconds` - int - number of milliseconds from now.
 //
 // `options` - PExpireOptions
 //
@@ -578,3 +600,49 @@ func (server *EchoVault) GetDel(key string) (string, error) {
 	}
 	return internal.ParseStringResponse(b)
 }
+
+// GetEx retrieves the value of the provided key and optionally sets its expiration
+//
+// Parameters:
+//
+// `key` - string - the key whose value should be retrieved and expiry set.
+//
+// `opts` - GetExOptions.
+//
+// Returns: A string representing the value at the specified key. If the value does not exist, an empty string is returned.
+func (server *EchoVault) GetEx(key string, opts GetExOptions) (string, error) {
+
+	cmd := make([]string, 2)
+
+	cmd[0] = "GETEX"
+	cmd[1] = key
+
+	var command string
+
+	switch {
+	case opts.EX:
+		command = "EX"
+	case opts.PX:
+		command = "PX"
+	case opts.EXAT:
+		command = "EXAT"
+	case opts.PXAT:
+		command = "PXAT"
+	case opts.PERSIST:
+		command = "PERSIST"
+	default:
+	}
+
+	if command != "" {
+		cmd = append(cmd, command)
+	}
+	if opts.UNIXTIME != 0 {
+		cmd = append(cmd, strconv.Itoa(opts.UNIXTIME))
+	}
+
+	b, err := server.handleCommand(server.context, internal.EncodeCommand(cmd), nil, false, true)
+	if err != nil {
+		return "", err
+	}
+	return internal.ParseStringResponse(b)
+}

echovault/api_generic_test.go

@@ -1400,3 +1400,140 @@ func TestEchoVault_GETDEL(t *testing.T) {
 		})
 	}
 }
+
+func TestEchoVault_GETEX(t *testing.T) {
+	mockClock := clock.NewClock()
+	server := createEchoVault()
+
+	tests := []struct {
+		name        string
+		presetValue interface{}
+		getExOpts   GetExOptions
+		key         string
+		want        string
+		wantEx      int
+		wantErr     bool
+	}{
+		{
+			name:        "1. Return string from existing key, no expire options",
+			presetValue: "value1",
+			getExOpts:   GetExOptions{},
+			key:         "key1",
+			want:        "value1",
+			wantEx:      -1,
+			wantErr:     false,
+		},
+		{
+			name:        "2. Return empty string if the key does not exist",
+			presetValue: nil,
+			getExOpts:   GetExOptions{EX: true, UNIXTIME: int(mockClock.Now().Add(100 * time.Second).Unix())},
+			key:         "key2",
+			want:        "",
+			wantEx:      0,
+			wantErr:     false,
+		},
+		{
+			name:        "3. Return key set expiry with EX",
+			presetValue: "value3",
+			getExOpts:   GetExOptions{EX: true, UNIXTIME: 100},
+			key:         "key3",
+			want:        "value3",
+			wantEx:      100,
+			wantErr:     false,
+		},
+		{
+			name:        "4. Return key set expiry with PX",
+			presetValue: "value4",
+			getExOpts:   GetExOptions{PX: true, UNIXTIME: 100000},
+			key:         "key4",
+			want:        "value4",
+			wantEx:      100,
+			wantErr:     false,
+		},
+		{
+			name:        "5. Return key set expiry with EXAT",
+			presetValue: "value5",
+			getExOpts:   GetExOptions{EXAT: true, UNIXTIME: int(mockClock.Now().Add(100 * time.Second).Unix())},
+			key:         "key5",
+			want:        "value5",
+			wantEx:      100,
+			wantErr:     false,
+		},
+		{
+			name:        "6. Return key set expiry with PXAT",
+			presetValue: "value6",
+			getExOpts:   GetExOptions{PXAT: true, UNIXTIME: int(mockClock.Now().Add(100 * time.Second).UnixMilli())},
+			key:         "key6",
+			want:        "value6",
+			wantEx:      100,
+			wantErr:     false,
+		},
+		{
+			name:        "7. Return key passing PERSIST",
+			presetValue: "value7",
+			getExOpts:   GetExOptions{PERSIST: true},
+			key:         "key7",
+			want:        "value7",
+			wantEx:      -1,
+			wantErr:     false,
+		},
+		{
+			name:        "8. Return key passing PERSIST, and include a UNIXTIME",
+			presetValue: "value8",
+			getExOpts:   GetExOptions{PERSIST: true, UNIXTIME: int(mockClock.Now().Add(100 * time.Second).Unix())},
+			key:         "key8",
+			want:        "value8",
+			wantEx:      -1,
+			wantErr:     false,
+		},
+		{
+			name:        "9. Return key and attempt to set expiry with EX without providing UNIXTIME",
+			presetValue: "value9",
+			getExOpts:   GetExOptions{EX: true},
+			key:         "key9",
+			want:        "value9",
+			wantEx:      -1,
+			wantErr:     false,
+		},
+		{
+			name:        "10. Return key and attempt to set expiry with PXAT without providing UNIXTIME",
+			presetValue: "value10",
+			getExOpts:   GetExOptions{PXAT: true},
+			key:         "key10",
+			want:        "value10",
+			wantEx:      -1,
+			wantErr:     false,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.presetValue != nil {
+				err := presetValue(server, context.Background(), tt.key, tt.presetValue)
+				if err != nil {
+					t.Error(err)
+					return
+				}
+			}
+			//Check value received
+			got, err := server.GetEx(tt.key, tt.getExOpts)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("GETEX() GET error = %v, wantErr %v", err, tt.wantErr)
+				return
+			}
+			if got != tt.want {
+				t.Errorf("GETEX() GET - got = %v, want %v", got, tt.want)
+			}
+			//Check expiry was set
+			if tt.presetValue != nil {
+				actual, err := server.TTL(tt.key)
+				if (err != nil) != tt.wantErr {
+					t.Errorf("GETEX() EXPIRY error = %v, wantErr %v", err, tt.wantErr)
+					return
+				}
+				if actual != tt.wantEx {
+					t.Errorf("GETEX() EXPIRY - got = %v, want %v", actual, tt.wantEx)
+				}
+			}
+		})
+	}
+}

internal/modules/generic/commands.go

@@ -269,9 +269,12 @@ func handleExpire(params internal.HandlerFuncParams) ([]byte, error) {
 	if err != nil {
 		return nil, errors.New("expire time must be integer")
 	}
-	expireAt := params.GetClock().Now().Add(time.Duration(n) * time.Second)
+
+	var expireAt time.Time
 	if strings.ToLower(params.Command[0]) == "pexpire" {
 		expireAt = params.GetClock().Now().Add(time.Duration(n) * time.Millisecond)
+	} else {
+		expireAt = params.GetClock().Now().Add(time.Duration(n) * time.Second)
 	}
 
 	if !keyExists {
@@ -333,9 +336,12 @@ func handleExpireAt(params internal.HandlerFuncParams) ([]byte, error) {
 	if err != nil {
 		return nil, errors.New("expire time must be integer")
 	}
-	expireAt := time.Unix(n, 0)
+
+	var expireAt time.Time
 	if strings.ToLower(params.Command[0]) == "pexpireat" {
 		expireAt = time.UnixMilli(n)
+	} else {
+		expireAt = time.Unix(n, 0)
 	}
 
 	if !keyExists {
@@ -712,6 +718,73 @@ func handleGetdel(params internal.HandlerFuncParams) ([]byte, error) {
 	return []byte(fmt.Sprintf("+%v\r\n", value)), nil
 }
 
+func handleGetex(params internal.HandlerFuncParams) ([]byte, error) {
+	keys, err := getExKeyFunc(params.Command)
+	if err != nil {
+		return nil, err
+	}
+
+	key := keys.ReadKeys[0]
+	keyExists := params.KeysExist(params.Context, []string{key})[key]
+
+	if !keyExists {
+		return []byte("$-1\r\n"), nil
+	}
+
+	value := params.GetValues(params.Context, []string{key})[key]
+
+	exkey := keys.WriteKeys[0]
+
+	cmdLen := len(params.Command)
+
+	// Handle no expire options provided
+	if cmdLen == 2 {
+		return []byte(fmt.Sprintf("+%v\r\n", value)), nil
+	}
+
+	// Handle persist
+	exCommand := strings.ToUpper(params.Command[2])
+	// If time is provided with PERSIST it is effectively ignored
+	if exCommand == "persist" {
+		// getValues will update key access so no need here
+		params.SetExpiry(params.Context, exkey, time.Time{}, false)
+		return []byte(fmt.Sprintf("+%v\r\n", value)), nil
+	}
+
+	// Handle exipre command passed but no time provided
+	if cmdLen == 3 {
+		return []byte(fmt.Sprintf("+%v\r\n", value)), nil
+	}
+
+	// Extract time
+	exTimeString := params.Command[3]
+	n, err := strconv.ParseInt(exTimeString, 10, 64)
+	if err != nil {
+		return []byte("$-1\r\n"), errors.New("expire time must be integer")
+	}
+
+	var expireAt time.Time
+	switch exCommand {
+	case "EX":
+		expireAt = params.GetClock().Now().Add(time.Duration(n) * time.Second)
+	case "PX":
+		expireAt = params.GetClock().Now().Add(time.Duration(n) * time.Millisecond)
+	case "EXAT":
+		expireAt = time.Unix(n, 0)
+	case "PXAT":
+		expireAt = time.UnixMilli(n)
+	case "PERSIST":
+		expireAt = time.Time{}
+	default:
+		return nil, fmt.Errorf("unknown option %s -- '%v'", strings.ToUpper(exCommand), params.Command)
+	}
+
+	params.SetExpiry(params.Context, exkey, expireAt, false)
+
+	return []byte(fmt.Sprintf("+%v\r\n", value)), nil
+
+}
+
 func Commands() []internal.Command {
 	return []internal.Command{
 		{
@@ -1004,5 +1077,14 @@ Delete all the keys in the currently selected database. This command is always s
 			KeyExtractionFunc: getDelKeyFunc,
 			HandlerFunc:       handleGetdel,
 		},
+		{
+			Command:           "getex",
+			Module:            constants.GenericModule,
+			Categories:        []string{constants.WriteCategory, constants.FastCategory},
+			Description:       "(GETEX key [EX seconds | PX milliseconds | EXAT unix-time-seconds | PXAT unix-time-milliseconds | PERSIST]) Get the value of key and optionally set its expiration. GETEX is similar to [GET], but is a write command with additional options.",
+			Sync:              true,
+			KeyExtractionFunc: getExKeyFunc,
+			HandlerFunc:       handleGetex,
+		},
 	}
 }

internal/modules/generic/commands_test.go

@@ -2930,4 +2930,221 @@ func Test_Generic(t *testing.T) {
 		}
 	})
 
+	t.Run("Test_HandleGETEX", func(t *testing.T) {
+		t.Parallel()
+		conn, err := internal.GetConnection("localhost", port)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		defer func() {
+			_ = conn.Close()
+		}()
+		client := resp.NewConn(conn)
+
+		tests := []struct {
+			name             string
+			command          []string
+			presetValues     map[string]KeyData
+			expectedResponse string
+			expectedValues   map[string]KeyData
+			expectedError    error
+		}{
+			{
+				name:    "1. Get key and set new expire by seconds",
+				command: []string{"GETEX", "GetExKey1", "EX", "100"},
+				presetValues: map[string]KeyData{
+					"GetExKey1": {Value: "value1", ExpireAt: time.Time{}},
+				},
+				expectedResponse: "value1",
+				expectedValues: map[string]KeyData{
+					"GetExKey1": {Value: "value1", ExpireAt: mockClock.Now().Add(100 * time.Second)},
+				},
+				expectedError: nil,
+			},
+			{
+				name:    "2. Get key and set new expire by milliseconds",
+				command: []string{"GETEX", "GetExKey2", "PX", "1000"},
+				presetValues: map[string]KeyData{
+					"GetExKey2": {Value: "value2", ExpireAt: time.Time{}},
+				},
+				expectedResponse: "value2",
+				expectedValues: map[string]KeyData{
+					"GetExKey2": {Value: "value2", ExpireAt: mockClock.Now().Add(1000 * time.Millisecond)},
+				},
+				expectedError: nil,
+			},
+			{
+				name:    "3. Get key and set new expire at by seconds",
+				command: []string{"GETEX", "GetExKey3", "EXAT", fmt.Sprintf("%d", mockClock.Now().Add(100*time.Second).Unix())},
+				presetValues: map[string]KeyData{
+					"GetExKey3": {Value: "value3", ExpireAt: time.Time{}},
+				},
+				expectedResponse: "value3",
+				expectedValues: map[string]KeyData{
+					"GetExKey3": {Value: "value3", ExpireAt: mockClock.Now().Add(100 * time.Second)},
+				},
+				expectedError: nil,
+			},
+			{
+				name:    "4. Get key and set new expire at by milliseconds",
+				command: []string{"GETEX", "GetExKey4", "PXAT", fmt.Sprintf("%d", mockClock.Now().Add(1000*time.Millisecond).UnixMilli())},
+				presetValues: map[string]KeyData{
+					"GetExKey4": {Value: "value4", ExpireAt: time.Time{}},
+				},
+				expectedResponse: "value4",
+				expectedValues: map[string]KeyData{
+					"GetExKey4": {Value: "value4", ExpireAt: mockClock.Now().Add(1000 * time.Millisecond)},
+				},
+				expectedError: nil,
+			},
+			{
+				name:    "5. Get key and persist",
+				command: []string{"GETEX", "GetExKey5", "PERSIST"},
+				presetValues: map[string]KeyData{
+					"GetExKey5": {Value: "value5", ExpireAt: time.Time{}},
+				},
+				expectedResponse: "value5",
+				expectedValues: map[string]KeyData{
+					"GetExKey5": {Value: "value5", ExpireAt: time.Time{}},
+				},
+				expectedError: nil,
+			},
+			{
+				name:    "6. Get key when no expire options are passed",
+				command: []string{"GETEX", "GetExKey6"},
+				presetValues: map[string]KeyData{
+					"GetExKey6": {Value: "value6", ExpireAt: time.Time{}},
+				},
+				expectedResponse: "value6",
+				expectedValues: map[string]KeyData{
+					"GetExKey6": {Value: "value6", ExpireAt: time.Time{}},
+				},
+				expectedError: nil,
+			},
+			{
+				name:             "7. Return empty string when key doesn't exist",
+				command:          []string{"GETEX", "GetExKey7", "PXAT", "1000"},
+				presetValues:     nil,
+				expectedResponse: "",
+				expectedValues:   nil,
+				expectedError:    nil,
+			},
+			{
+				name:    "8. Get key and don't set expiration when time not provided",
+				command: []string{"GETEX", "GetExKey8", "PXAT"},
+				presetValues: map[string]KeyData{
+					"GetExKey8": {Value: "value8", ExpireAt: time.Time{}},
+				},
+				expectedResponse: "value8",
+				expectedValues: map[string]KeyData{
+					"GetExKey8": {Value: "value8", ExpireAt: time.Time{}},
+				},
+				expectedError: nil,
+			},
+			{
+				name:    "9. Return error when expire time is not a valid integer",
+				command: []string{"GETEX", "GetExKey9", "EX", "notAnInt"},
+				presetValues: map[string]KeyData{
+					"GetExKey9": {Value: "value9", ExpireAt: time.Time{}},
+				},
+				expectedResponse: "",
+				expectedValues:   nil,
+				expectedError:    errors.New("expire time must be integer"),
+			},
+			{
+				name:             "10. Command too short",
+				command:          []string{"GETEX"},
+				presetValues:     nil,
+				expectedResponse: "",
+				expectedValues:   nil,
+				expectedError:    errors.New(constants.WrongArgsResponse),
+			},
+			{
+				name:             "11. Command too long",
+				command:          []string{"GETEX", "GetExKey11", "EX", "1000", "PERSIST"},
+				presetValues:     nil,
+				expectedResponse: "",
+				expectedValues:   nil,
+				expectedError:    errors.New(constants.WrongArgsResponse),
+			},
+		}
+
+		for _, test := range tests {
+			t.Run(test.name, func(t *testing.T) {
+				if test.presetValues != nil {
+					for k, v := range test.presetValues {
+						command := []resp.Value{resp.StringValue("SET"), resp.StringValue(k), resp.StringValue(v.Value.(string))}
+						if !v.ExpireAt.Equal(time.Time{}) {
+							command = append(command, []resp.Value{
+								resp.StringValue("PX"),
+								resp.StringValue(fmt.Sprintf("%d", v.ExpireAt.Sub(mockClock.Now()).Milliseconds())),
+							}...)
+						}
+						if err = client.WriteArray(command); err != nil {
+							t.Error(err)
+						}
+						res, _, err := client.ReadValue()
+						if err != nil {
+							t.Error(err)
+						}
+						if !strings.EqualFold(res.String(), "ok") {
+							t.Errorf("expected preset response to be OK, got %s", res.String())
+						}
+					}
+				}
+
+				command := make([]resp.Value, len(test.command))
+				for i, c := range test.command {
+					command[i] = resp.StringValue(c)
+				}
+
+				if err = client.WriteArray(command); err != nil {
+					t.Error(err)
+				}
+
+				res, _, err := client.ReadValue()
+				if err != nil {
+					t.Error(err)
+				}
+
+				if test.expectedError != nil {
+					if res.Error() == nil || !strings.Contains(res.Error().Error(), test.expectedError.Error()) {
+						t.Errorf("expected error \"%s\", got \"%v\"", test.expectedError.Error(), res.Error())
+					}
+					return
+				}
+
+				if res.String() != test.expectedResponse {
+					t.Errorf("expected response %s, got %s", test.expectedResponse, res.String())
+				}
+
+				if test.expectedValues == nil {
+					return
+				}
+
+				for key, expected := range test.expectedValues {
+					// Compare the expiry of the key with what's expected
+					if err = client.WriteArray([]resp.Value{resp.StringValue("PTTL"), resp.StringValue(key)}); err != nil {
+						t.Error(err)
+					}
+
+					res, _, err = client.ReadValue()
+					if err != nil {
+						t.Error(err)
+					}
+					if expected.ExpireAt.Equal(time.Time{}) {
+						if res.Integer() != -1 {
+							t.Error("expected key to be persisted, it was not.")
+						}
+						continue
+					}
+					if res.Integer() != int(expected.ExpireAt.Sub(mockClock.Now()).Milliseconds()) {
+						t.Errorf("expected expiry %d, got %d", expected.ExpireAt.Sub(mockClock.Now()).Milliseconds(), res.Integer())
+					}
+				}
+			})
+		}
+	})
+
 }

internal/modules/generic/key_funcs.go

@@ -212,3 +212,14 @@ func getDelKeyFunc(cmd []string) (internal.KeyExtractionFuncResult, error) {
 		WriteKeys: cmd[1:],
 	}, nil
 }
+
+func getExKeyFunc(cmd []string) (internal.KeyExtractionFuncResult, error) {
+	if len(cmd) < 2 || len(cmd) > 4 {
+		return internal.KeyExtractionFuncResult{}, errors.New(constants.WrongArgsResponse)
+	}
+	return internal.KeyExtractionFuncResult{
+		Channels:  make([]string, 0),
+		ReadKeys:  cmd[1:2],
+		WriteKeys: cmd[1:2],
+	}, nil
+}

internal/modules/sorted_set/commands_test.go

@@ -16,17 +16,18 @@ package sorted_set_test
 
 import (
 	"errors"
+	"math"
+	"slices"
+	"strconv"
+	"strings"
+	"testing"
+
 	"github.com/echovault/echovault/echovault"
 	"github.com/echovault/echovault/internal"
 	"github.com/echovault/echovault/internal/config"
 	"github.com/echovault/echovault/internal/constants"
 	"github.com/echovault/echovault/internal/modules/sorted_set"
 	"github.com/tidwall/resp"
-	"math"
-	"slices"
-	"strconv"
-	"strings"
-	"testing"
 )
 
 func Test_SortedSet(t *testing.T) {
@@ -2376,7 +2377,7 @@ func Test_SortedSet(t *testing.T) {
 				expectedError: nil,
 			},
 			{
-				name:          "2. Return error when the source key is not a sorted set.",
+				name:          "3. Return error when the source key is not a sorted set.",
 				key:           "ZrandMemberKey3",
 				presetValue:   "Default value",
 				command:       []string{"ZRANDMEMBER", "ZrandMemberKey3"},

internal/modules/sorted_set/sorted_set.go

@@ -86,15 +86,15 @@ func (set *SortedSet) GetRandom(count int) []MemberParam {
 		}
 	} else {
 		// If count is positive only allow unique values
+
 		for i := 0; i < internal.AbsInt(count); {
 			n = rand.Intn(len(members))
 			if !slices.ContainsFunc(res, func(m MemberParam) bool {
 				return m.Value == members[n].Value
 			}) {
 				res = append(res, members[n])
-				slices.DeleteFunc(members, func(m MemberParam) bool {
-					return m.Value == members[n].Value
-				})
+				members[n] = members[len(members)-1]
+				members = members[:len(members)-1]
 				i++
 			}
 		}