Implementation of Copy command (#141)

* Added COPY command - @zenc0derr 

---------

Co-authored-by: Tejesh Kumar S <zenc0derr>
Co-authored-by: Kelvin Clement Mwinuka <[email protected]>

internal/modules/generic/commands.go

@@ -595,6 +595,7 @@ func handleIncrByFloat(params internal.HandlerFuncParams) ([]byte, error) {
 	response := fmt.Sprintf("$%d\r\n%g\r\n", len(fmt.Sprintf("%g", newValue)), newValue)
 	return []byte(response), nil
 }
+
 func handleDecrBy(params internal.HandlerFuncParams) ([]byte, error) {
 	// Extract key from command
 	keys, err := decrByKeyFunc(params.Command)
@@ -870,6 +871,50 @@ func handleObjIdleTime(params internal.HandlerFuncParams) ([]byte, error) {
 	return []byte(fmt.Sprintf("+%v\r\n", idletime)), nil
 }
 
+func handleCopy(params internal.HandlerFuncParams) ([]byte, error) {
+	keys, err := copyKeyFunc(params.Command)
+	if err != nil {
+		return nil, err
+	}
+
+	options, err := getCopyCommandOptions(params.Command[3:], CopyOptions{})
+	if err != nil {
+		return nil, err
+	}
+	sourceKey := keys.ReadKeys[0]
+	destinationKey := keys.WriteKeys[0]
+	sourceKeyExists := params.KeysExist(params.Context, []string{sourceKey})[sourceKey]
+
+	if !sourceKeyExists {
+		return []byte(":0\r\n"), nil
+	}
+
+	if !options.replace {
+		destinationKeyExists := params.KeysExist(params.Context, []string{destinationKey})[destinationKey]
+
+		if destinationKeyExists {
+			return []byte(":0\r\n"), nil
+		}
+	}
+
+	value := params.GetValues(params.Context, []string{sourceKey})[sourceKey]
+
+	ctx := context.WithoutCancel(params.Context)
+
+	if options.database != "" {
+		database, _ := strconv.Atoi(options.database)
+		ctx = context.WithValue(ctx, "Database", database)
+	}
+
+	if err = params.SetValues(ctx, map[string]interface{}{
+		destinationKey: value,
+	}); err != nil {
+		return nil, err
+	}
+
+	return []byte(":1\r\n"), nil
+}
+
 func handleMove(params internal.HandlerFuncParams) ([]byte, error) {
 	keys, err := moveKeyFunc(params.Command)
 	if err != nil {
@@ -1255,6 +1300,18 @@ The command is only available when the maxmemory-policy configuration directive
 			KeyExtractionFunc: objIdleTimeKeyFunc,
 			HandlerFunc:       handleObjIdleTime,
 		},
+		{
+			Command:    "copy",
+			Module:     constants.GenericModule,
+			Categories: []string{constants.KeyspaceCategory, constants.WriteCategory, constants.SlowCategory},
+			Description: `(COPY source destination [DB destination-db] [REPLACE]) 
+Copies the value stored at the source key to the destination key.
+The command returns zero when the destination key already exists. 
+The REPLACE option removes the destination key before copying the value to it.`,
+			Sync:              false,
+			KeyExtractionFunc: copyKeyFunc,
+			HandlerFunc:       handleCopy,
+		},
 		{
 			Command:           "move",
 			Module:            constants.GenericModule,

internal/modules/generic/commands_test.go

@@ -3333,6 +3333,168 @@ func Test_Generic(t *testing.T) {
 		}
 	})
 
+	t.Run("Test_HandleCOPY", 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
+			sourceKeyPresetValue interface{}
+			sourcekey            string
+			destKeyPresetValue   interface{}
+			destinationKey       string
+			database             string
+			replace              bool
+			expectedValue        string
+			expectedResponse     string
+		}{
+			{
+				name:                 "1. Copy Value into non existing key",
+				sourceKeyPresetValue: "value1",
+				sourcekey:            "skey1",
+				destKeyPresetValue:   nil,
+				destinationKey:       "dkey1",
+				database:             "0",
+				replace:              false,
+				expectedValue:        "value1",
+				expectedResponse:     "1",
+			},
+			{
+				name:                 "2. Copy Value into existing key without replace option",
+				sourceKeyPresetValue: "value2",
+				sourcekey:            "skey2",
+				destKeyPresetValue:   "dValue2",
+				destinationKey:       "dkey2",
+				database:             "0",
+				replace:              false,
+				expectedValue:        "dValue2",
+				expectedResponse:     "0",
+			},
+			{
+				name:                 "3. Copy Value into existing key with replace option",
+				sourceKeyPresetValue: "value3",
+				sourcekey:            "skey3",
+				destKeyPresetValue:   "dValue3",
+				destinationKey:       "dkey3",
+				database:             "0",
+				replace:              true,
+				expectedValue:        "value3",
+				expectedResponse:     "1",
+			},
+			{
+				name:                 "4. Copy Value into different database",
+				sourceKeyPresetValue: "value4",
+				sourcekey:            "skey4",
+				destKeyPresetValue:   nil,
+				destinationKey:       "dkey4",
+				database:             "1",
+				replace:              true,
+				expectedValue:        "value4",
+				expectedResponse:     "1",
+			},
+		}
+
+		for _, tt := range tests {
+			t.Run(tt.name, func(t *testing.T) {
+				if tt.sourceKeyPresetValue != nil {
+					cmd := []resp.Value{resp.StringValue("Set"), resp.StringValue(tt.sourcekey), resp.StringValue(tt.sourceKeyPresetValue.(string))}
+
+					err := client.WriteArray(cmd)
+					if err != nil {
+						t.Error(err)
+					}
+
+					rd, _, err := client.ReadValue()
+					if err != nil {
+						t.Error(err)
+					}
+
+					if !strings.EqualFold(rd.String(), "ok") {
+						t.Errorf("expected preset response to be \"OK\", got %s", rd.String())
+					}
+				}
+
+				if tt.destKeyPresetValue != nil {
+					cmd := []resp.Value{resp.StringValue("Set"), resp.StringValue(tt.destinationKey), resp.StringValue(tt.destKeyPresetValue.(string))}
+
+					err := client.WriteArray(cmd)
+					if err != nil {
+						t.Error(err)
+					}
+
+					rd, _, err := client.ReadValue()
+					if err != nil {
+						t.Error(err)
+					}
+
+					if !strings.EqualFold(rd.String(), "ok") {
+						t.Errorf("expected preset response to be \"OK\", got %s", rd.String())
+					}
+				}
+
+				command := []resp.Value{resp.StringValue("COPY"), resp.StringValue(tt.sourcekey), resp.StringValue(tt.destinationKey)}
+
+				if tt.database != "0" {
+					command = append(command, resp.StringValue("DB"), resp.StringValue(tt.database))
+				}
+
+				if tt.replace {
+					command = append(command, resp.StringValue("REPLACE"))
+				}
+
+				err := client.WriteArray(command)
+				if err != nil {
+					t.Error(err)
+				}
+
+				rd, _, err := client.ReadValue()
+				if err != nil {
+					t.Error(err)
+				}
+				if !strings.EqualFold(rd.String(), tt.expectedResponse) {
+					t.Errorf("expected response to be %s, but got %s", tt.expectedResponse, rd.String())
+				}
+
+				if tt.database != "0" {
+					selectCommand := []resp.Value{resp.StringValue("SELECT"), resp.StringValue(tt.database)}
+
+					err := client.WriteArray(selectCommand)
+					if err != nil {
+						t.Error(err)
+					}
+					_, _, err = client.ReadValue()
+					if err != nil {
+						t.Error(err)
+					}
+				}
+
+				getCommand := []resp.Value{resp.StringValue("GET"), resp.StringValue(tt.destinationKey)}
+
+				err = client.WriteArray(getCommand)
+				if err != nil {
+					t.Error(err)
+				}
+
+				rd, _, err = client.ReadValue()
+				if err != nil {
+					t.Error(err)
+				}
+				if !strings.EqualFold(rd.String(), tt.expectedValue) {
+					t.Errorf("expected value in destinaton key to be %s, but got %s", tt.expectedValue, rd.String())
+				}
+			})
+		}
+
+	})
+
 	t.Run("Test_HandleMOVE", func(t *testing.T) {
 		t.Parallel()
 
@@ -3474,7 +3636,7 @@ func Test_Generic(t *testing.T) {
 }
 
 // Certain commands will need to be tested in a server with an eviction policy.
-// This is for testing against an LFU evictiona policy.
+// This is for testing against an LFU eviction policy.
 func Test_LFU_Generic(t *testing.T) {
 	// mockClock := clock.NewClock()
 	port, err := internal.GetFreePort()

internal/modules/generic/key_funcs.go

@@ -268,6 +268,18 @@ func objIdleTimeKeyFunc(cmd []string) (internal.KeyExtractionFuncResult, error)
 	}, nil
 }
 
+func copyKeyFunc(cmd []string) (internal.KeyExtractionFuncResult, error) {
+	if len(cmd) < 3 && len(cmd)>6{
+		return internal.KeyExtractionFuncResult{}, errors.New(constants.WrongArgsResponse)
+	}
+
+	return internal.KeyExtractionFuncResult{
+		Channels:  make([]string, 0),
+		ReadKeys:  cmd[1:2],
+		WriteKeys: cmd[2:3],
+	}, nil
+}
+
 func moveKeyFunc(cmd []string) (internal.KeyExtractionFuncResult, error) {
 	if len(cmd) != 3 {
 		return internal.KeyExtractionFuncResult{}, errors.New(constants.WrongArgsResponse)

internal/modules/generic/utils.go

@@ -29,6 +29,11 @@ type SetOptions struct {
 	expireAt interface{} // Exact expireAt time un unix milliseconds
 }
 
+type CopyOptions struct {
+	database string
+	replace bool
+}
+
 func getSetCommandOptions(clock clock.Clock, cmd []string, options SetOptions) (SetOptions, error) {
 	if len(cmd) == 0 {
 		return options, nil
@@ -116,3 +121,32 @@ func getSetCommandOptions(clock clock.Clock, cmd []string, options SetOptions) (
 		return SetOptions{}, fmt.Errorf("unknown option %s for set command", strings.ToUpper(cmd[0]))
 	}
 }
+
+func getCopyCommandOptions(cmd []string, options CopyOptions) (CopyOptions, error) {
+	if len(cmd) == 0 {
+		return options, nil
+	}
+
+	switch strings.ToLower(cmd[0]){
+	case "replace":
+		options.replace = true
+		return getCopyCommandOptions(cmd[1:], options)
+
+	case "db":
+		if len(cmd) < 2 {
+			return CopyOptions{}, errors.New("syntax error")
+		}
+
+		_, err := strconv.Atoi(cmd[1])
+		if err != nil {
+			return CopyOptions{}, errors.New("value is not an integer or out of range")
+		}
+
+		options.database = cmd [1]
+		return getCopyCommandOptions(cmd[2:], options)
+
+
+	default:
+		return CopyOptions{}, fmt.Errorf("unknown option %s for copy command", strings.ToUpper(cmd[0]))
+	}
+}

sugardb/api_generic.go

@@ -137,6 +137,16 @@ type GetExOption interface {
 
 func (x GetExOpt) isGetExOpt() GetExOpt { return x }
 
+// COPYOptions is a struct wrapper for all optional parameters of the Copy command.
+//
+// `Database` - string - Logical database index
+//
+// `Replace` - bool - Whether to replace the destination key if it exists
+type COPYOptions struct {
+	Database string
+	Replace  bool
+}
+
 // Set creates or modifies the value at the given key.
 //
 // Parameters:
@@ -719,6 +729,33 @@ func (server *SugarDB) Type(key string) (string, error) {
 	return internal.ParseStringResponse(b)
 }
 
+// Copy copies a value of a source key to destination key.
+//
+// Parameters:
+//
+// `source` - string - the source key from which data is to be copied
+//
+// `destination` - string - the destination key where data should be copied
+//
+// Returns: 1 if the copy is successful. 0 if the copy is unsuccessful
+func (server *SugarDB) Copy(sourceKey, destinationKey string, options COPYOptions) (int, error) {
+	cmd := []string{"COPY", sourceKey, destinationKey}
+
+	if options.Database != "" {
+		cmd = append(cmd, "db", options.Database)
+	}
+
+	if options.Replace {
+		cmd = append(cmd, "replace")
+	}
+
+	b, err := server.handleCommand(server.context, internal.EncodeCommand(cmd), nil, false, true)
+	if err != nil {
+		return 0, err
+	}
+	return internal.ParseIntegerResponse(b)
+}
+
 // Move key from currently selected database to specified destination database and return 1.
 // When key already exists in the destination database, or it does not exist in the source database, it does nothing and returns 0.
 //