added setrange and getrange functions

2025-10-07 17:51:10 +08:00 · 2021-06-28 21:20:20 +03:00
parent 3e4531348c
commit 7b0df5cda3
2 changed files with 357 additions and 1 deletions
--- a/string.go
+++ b/string.go
@@ -486,7 +486,7 @@ func execStrLen(db *DB, args [][]byte) redis.Reply {
 	return reply.MakeIntReply(int64(len(bytes)))
 }

-// execAppend sets string value and time to live to the given key
+// execAppend sets string value to the given key
 func execAppend(db *DB, args [][]byte) redis.Reply {
 	key := string(args[0])
 	bytes, err := db.getAsString(key)
@@ -501,6 +501,85 @@ func execAppend(db *DB, args [][]byte) redis.Reply {
 	return reply.MakeIntReply(int64(len(bytes)))
 }

+// execSetRange overwrites part of the string stored at key, starting at the specified offset.
+// If the offset is larger than the current length of the string at key, the string is padded with zero-bytes.
+func execSetRange(db *DB, args [][]byte) redis.Reply {
+	key := string(args[0])
+	offset, errNative := strconv.ParseInt(string(args[1]), 10, 64)
+	if errNative != nil {
+		return reply.MakeErrReply(errNative.Error())
+	}
+	value := args[2]
+	bytes, err := db.getAsString(key)
+	if err != nil {
+		return err
+	}
+	bytesLen := int64(len(bytes))
+	if bytesLen < offset {
+		diff := offset - bytesLen
+		diffArray := make([]byte, diff)
+		bytes = append(bytes, diffArray...)
+		bytesLen = int64(len(bytes))
+	}
+	for i := 0; i < len(value); i++ {
+		idx := offset + int64(i)
+		if idx >= bytesLen {
+			bytes = append(bytes, value[i])
+		} else {
+			bytes[idx] = value[i]
+		}
+	}
+	db.PutEntity(key, &DataEntity{
+		Data: bytes,
+	})
+	db.AddAof(makeAofCmd("setRange", args))
+	return reply.MakeIntReply(int64(len(bytes)))
+}
+
+func execGetRange(db *DB, args [][]byte) redis.Reply {
+	key := string(args[0])
+	startIdx, errNative := strconv.ParseInt(string(args[1]), 10, 64)
+	if errNative != nil {
+		return reply.MakeErrReply(errNative.Error())
+	}
+	endIdx, errNative := strconv.ParseInt(string(args[2]), 10, 64)
+	if errNative != nil {
+		return reply.MakeErrReply(errNative.Error())
+	}
+
+	bytes, err := db.getAsString(key)
+	if err != nil {
+		return err
+	}
+
+	if bytes == nil {
+		return reply.MakeNullBulkReply()
+	}
+
+	bytesLen := int64(len(bytes))
+	if startIdx < -1*bytesLen {
+		return &reply.NullBulkReply{}
+	} else if startIdx < 0 {
+		startIdx = bytesLen + startIdx
+	} else if startIdx >= bytesLen {
+		return &reply.NullBulkReply{}
+	}
+	if endIdx < -1*bytesLen {
+		return &reply.NullBulkReply{}
+	} else if endIdx < 0 {
+		endIdx = bytesLen + endIdx + 1
+	} else if endIdx < bytesLen {
+		endIdx = endIdx + 1
+	} else {
+		endIdx = bytesLen
+	}
+	if startIdx > endIdx {
+		return reply.MakeNullBulkReply()
+	}
+
+	return reply.MakeBulkReply(bytes[startIdx:endIdx])
+}
+
 func init() {
 	RegisterCommand("Set", execSet, writeFirstKey, rollbackFirstKey, -3)
 	RegisterCommand("SetNx", execSetNX, writeFirstKey, rollbackFirstKey, 3)
@@ -518,4 +597,6 @@ func init() {
 	RegisterCommand("DecrBy", execDecrBy, writeFirstKey, rollbackFirstKey, 3)
 	RegisterCommand("StrLen", execStrLen, readFirstKey, nil, 2)
 	RegisterCommand("Append", execAppend, writeFirstKey, rollbackFirstKey, 3)
+	RegisterCommand("SetRange", execSetRange, writeFirstKey, rollbackFirstKey, 4)
+	RegisterCommand("GetRange", execGetRange, readFirstKey, nil, 4)
 }
--- a/string_test.go
+++ b/string_test.go
@@ -372,3 +372,278 @@ func TestAppend_KeyNotExist(t *testing.T) {
 	asserts.AssertIntReply(t, val, len(key))
 }

+func TestSetRange_StringExist(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	key2 := utils.RandString(3)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("SetRange", key, fmt.Sprint(0), key2))
+	val, ok := actual.(*reply.IntReply)
+	if !ok {
+		t.Errorf("expect int bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	result := len(key2 + key[3:])
+	asserts.AssertIntReply(t, val, result)
+}
+
+func TestSetRange_StringExist_OffsetOutOfLen(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	key2 := utils.RandString(3)
+	emptyByteLen := 5
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("SetRange", key, fmt.Sprint(len(key)+emptyByteLen), key2))
+	val, ok := actual.(*reply.IntReply)
+	if !ok {
+		t.Errorf("expect int bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	result := len(key + string(make([]byte, emptyByteLen)) + key2)
+	asserts.AssertIntReply(t, val, result)
+}
+
+func TestSetRange_StringNotExist(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("SetRange", key, fmt.Sprint(0), key))
+	val, ok := actual.(*reply.IntReply)
+	if !ok {
+		t.Errorf("expect int bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+	asserts.AssertIntReply(t, val, len(key))
+}
+
+func TestGetRange_StringExist(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(0), fmt.Sprint(len(key))))
+	val, ok := actual.(*reply.BulkReply)
+	if !ok {
+		t.Errorf("expect bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertBulkReply(t, val, key)
+}
+
+func TestGetRange_RangeLargeThenDataLen(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(0), fmt.Sprint(len(key)+2)))
+	val, ok := actual.(*reply.BulkReply)
+	if !ok {
+		t.Errorf("expect bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertBulkReply(t, val, key)
+}
+
+func TestGetRange_StringNotExist(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(0), fmt.Sprint(len(key))))
+	val, ok := actual.(*reply.NullBulkReply)
+	if !ok {
+		t.Errorf("expect nil bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertNullBulk(t, val)
+}
+
+func TestGetRange_StringExist_GetPartial(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(0), fmt.Sprint(len(key)/2)))
+	val, ok := actual.(*reply.BulkReply)
+	if !ok {
+		t.Errorf("expect bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertBulkReply(t, val, key[:(len(key)/2)+1])
+}
+
+func TestGetRange_StringExist_EndIdxOutOfRange(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	emptyByteLen := 2
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(0), fmt.Sprint(len(key)+emptyByteLen)))
+	val, ok := actual.(*reply.BulkReply)
+	if !ok {
+		t.Errorf("expect bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertBulkReply(t, val, key)
+}
+
+func TestGetRange_StringExist_StartIdxEndIdxAreSame(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	emptyByteLen := 2
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(len(key)+emptyByteLen), fmt.Sprint(len(key)+emptyByteLen)))
+	val, ok := actual.(*reply.NullBulkReply)
+	if !ok {
+		t.Errorf("expect nil bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertNullBulk(t, val)
+}
+
+func TestGetRange_StringExist_StartIdxGreaterThanEndIdx(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(len(key)+1), fmt.Sprint(len(key))))
+	val, ok := actual.(*reply.NullBulkReply)
+	if !ok {
+		t.Errorf("expect nil bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertNullBulk(t, val)
+}
+
+func TestGetRange_StringExist_StartIdxEndIdxAreNegative(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(-1*len(key)), fmt.Sprint(-1)))
+	val, ok := actual.(*reply.BulkReply)
+	if !ok {
+		t.Errorf("expect bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertBulkReply(t, val, key)
+}
+
+func TestGetRange_StringExist_StartIdxNegative(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(-1*len(key)), fmt.Sprint(len(key)/2)))
+	val, ok := actual.(*reply.BulkReply)
+	if !ok {
+		t.Errorf("expect bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertBulkReply(t, val, key[0:(len(key)/2)+1])
+}
+
+func TestGetRange_StringExist_EndIdxNegative(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(0), fmt.Sprint(-len(key)/2)))
+	val, ok := actual.(*reply.BulkReply)
+	if !ok {
+		t.Errorf("expect bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertBulkReply(t, val, key[0:(len(key)/2)+1])
+}
+
+func TestGetRange_StringExist_StartIsOutOfRange(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(-len(key)-3), fmt.Sprint(len(key))))
+	val, ok := actual.(*reply.NullBulkReply)
+	if !ok {
+		t.Errorf("expect bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertNullBulk(t, val)
+}
+
+func TestGetRange_StringExist_EndIdxIsOutOfRange(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(0), fmt.Sprint(-len(key)-3)))
+	val, ok := actual.(*reply.NullBulkReply)
+	if !ok {
+		t.Errorf("expect bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertNullBulk(t, val)
+}
+
+func TestGetRange_StringExist_StartIdxGreaterThanDataLen(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(len(key)+1), fmt.Sprint(0)))
+	val, ok := actual.(*reply.NullBulkReply)
+	if !ok {
+		t.Errorf("expect bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	asserts.AssertNullBulk(t, val)
+}
+
+func TestGetRange_StringExist_StartIdxIncorrectFormat(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+	incorrectValue := "incorrect"
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, incorrectValue, fmt.Sprint(0)))
+	val, ok := actual.(*reply.StandardErrReply)
+	if !ok {
+		t.Errorf("expect standart bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	errorMsg := fmt.Sprintf("strconv.ParseInt: parsing \"%s\": invalid syntax", incorrectValue)
+	asserts.AssertErrReply(t, val, errorMsg)
+}
+
+func TestGetRange_StringExist_EndIdxIncorrectFormat(t *testing.T) {
+	testDB.Flush()
+	key := utils.RandString(10)
+	testDB.Exec(nil, utils.ToCmdLine2("SET", key, key))
+	incorrectValue := "incorrect"
+
+	actual := testDB.Exec(nil, utils.ToCmdLine("GetRange", key, fmt.Sprint(0), incorrectValue))
+	val, ok := actual.(*reply.StandardErrReply)
+	if !ok {
+		t.Errorf("expect standart bulk reply, get: %s", string(actual.ToBytes()))
+		return
+	}
+
+	errorMsg := fmt.Sprintf("strconv.ParseInt: parsing \"%s\": invalid syntax", incorrectValue)
+	asserts.AssertErrReply(t, val, errorMsg)
+}