From 97a10e4fdb9f280ba858c67ccab251e90b6b4933 Mon Sep 17 00:00:00 2001
From: Kelvin Mwinuka <kelvin_mwinuka@epam.com>
Date: Sat, 9 Mar 2024 23:28:46 +0800
Subject: [PATCH] Implemented EXPIRTE and PEXPIRE command handler

---
 src/modules/generic/commands.go  | 68 +++++++++++++++++++++++++++++++-
 src/modules/generic/key_funcs.go |  4 +-
 2 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/src/modules/generic/commands.go b/src/modules/generic/commands.go
index 17b9d50..deea0ba 100644
--- a/src/modules/generic/commands.go
+++ b/src/modules/generic/commands.go
@@ -7,6 +7,7 @@ import (
 	"github.com/echovault/echovault/src/utils"
 	"log"
 	"net"
+	"strconv"
 	"strings"
 	"time"
 )
@@ -319,8 +320,71 @@ func handleTTL(ctx context.Context, cmd []string, server utils.Server, _ *net.Co
 }
 
 func handleExpire(ctx context.Context, cmd []string, server utils.Server, _ *net.Conn) ([]byte, error) {
-	// Handle EXPIRE and PEXPIRE
-	return nil, errors.New("command not implemented yet")
+	keys, err := expireKeyFunc(cmd)
+	if err != nil {
+		return nil, err
+	}
+
+	key := keys[0]
+
+	// Extract time
+	n, err := strconv.ParseInt(cmd[2], 10, 64)
+	if err != nil {
+		return nil, errors.New("expire time must be integer")
+	}
+	expireAt := time.Now().Add(time.Duration(n) * time.Second)
+	if strings.ToLower(cmd[0]) == "pexpire" {
+		expireAt = time.Now().Add(time.Duration(n) * time.Millisecond)
+	}
+
+	if !server.KeyExists(key) {
+		return []byte(":0\r\n"), nil
+	}
+
+	if _, err = server.KeyLock(ctx, key); err != nil {
+		return nil, err
+	}
+	defer server.KeyUnlock(key)
+
+	if len(cmd) == 3 {
+		server.SetExpiry(ctx, key, expireAt, true)
+		return []byte(":1\r\n"), nil
+	}
+
+	currentExpireAt := server.GetExpiry(ctx, key)
+
+	switch strings.ToLower(cmd[3]) {
+	case "nx":
+		if currentExpireAt != (time.Time{}) {
+			return []byte(":0\r\n"), nil
+		}
+		server.SetExpiry(ctx, key, expireAt, false)
+	case "xx":
+		if currentExpireAt == (time.Time{}) {
+			return []byte(":0\r\n"), nil
+		}
+		server.SetExpiry(ctx, key, expireAt, false)
+	case "gt":
+		if currentExpireAt == (time.Time{}) {
+			return []byte(":0\r\n"), nil
+		}
+		if expireAt.Before(currentExpireAt) {
+			return []byte(":0\r\n"), nil
+		}
+		server.SetExpiry(ctx, key, expireAt, false)
+	case "lt":
+		if currentExpireAt != (time.Time{}) {
+			if currentExpireAt.Before(expireAt) {
+				return []byte(":0\r\n"), nil
+			}
+			server.SetExpiry(ctx, key, expireAt, false)
+		}
+		server.SetExpiry(ctx, key, expireAt, false)
+	default:
+		return nil, fmt.Errorf("unknown option %s", strings.ToUpper(cmd[0]))
+	}
+
+	return []byte(":1\r\n"), nil
 }
 
 func handleExpireAt(ctx context.Context, cmd []string, server utils.Server, _ *net.Conn) ([]byte, error) {
diff --git a/src/modules/generic/key_funcs.go b/src/modules/generic/key_funcs.go
index 13a436b..1a39f9e 100644
--- a/src/modules/generic/key_funcs.go
+++ b/src/modules/generic/key_funcs.go
@@ -68,14 +68,14 @@ func ttlKeyFunc(cmd []string) ([]string, error) {
 }
 
 func expireKeyFunc(cmd []string) ([]string, error) {
-	if len(cmd) != 4 {
+	if len(cmd) < 3 || len(cmd) > 4 {
 		return nil, errors.New(utils.WrongArgsResponse)
 	}
 	return []string{cmd[1]}, nil
 }
 
 func expireAtKeyFunc(cmd []string) ([]string, error) {
-	if len(cmd) != 4 {
+	if len(cmd) < 3 || len(cmd) > 4 {
 		return nil, errors.New(utils.WrongArgsResponse)
 	}
 	return []string{cmd[1]}, nil