From 059de689e3a58de6ddbcb9b6390fbf87a2b0c02c Mon Sep 17 00:00:00 2001 From: Yuvraj Gosain Date: Wed, 23 Oct 2024 00:40:57 +0530 Subject: [PATCH] Migrating GEOADD and GEODIST command --- integration_tests/commands/http/geo_tests.go | 86 +++++++++ integration_tests/commands/resp/geo_tests.go | 86 +++++++++ .../commands/websocket/geo_tests.go | 87 +++++++++ internal/eval/commands.go | 22 ++- internal/eval/eval.go | 133 ------------- internal/eval/eval_test.go | 129 +++++++++---- internal/eval/store_eval.go | 175 ++++++++++++++++++ internal/server/cmd_meta.go | 12 ++ internal/worker/cmd_meta.go | 9 + 9 files changed, 559 insertions(+), 180 deletions(-) create mode 100644 integration_tests/commands/http/geo_tests.go create mode 100644 integration_tests/commands/resp/geo_tests.go create mode 100644 integration_tests/commands/websocket/geo_tests.go diff --git a/integration_tests/commands/http/geo_tests.go b/integration_tests/commands/http/geo_tests.go new file mode 100644 index 000000000..9357544d2 --- /dev/null +++ b/integration_tests/commands/http/geo_tests.go @@ -0,0 +1,86 @@ +package http + +import ( + "testing" + + "gotest.tools/v3/assert" +) + +func TestGeoAdd(t *testing.T) { + exec := NewHTTPCommandExecutor() + + testCases := []struct { + name string + commands []HTTPCommand + expected []interface{} + }{ + { + name: "GEOADD with wrong number of arguments", + commands: []HTTPCommand{ + {Command: "GEOADD", Body: map[string]interface{}{"key": "mygeo", "values": []interface{}{"1.2", "2.4"}}}, + }, + expected: []interface{}{"ERR wrong number of arguments for 'geoadd' command"}, + }, + { + name: "GEOADD Commands with new member and updating it", + commands: []HTTPCommand{ + {Command: "GEOADD", Body: map[string]interface{}{"key": "mygeo", "values": []interface{}{"1.2", "2.4", "NJ"}}}, + {Command: "GEOADD", Body: map[string]interface{}{"key": "mygeo", "values": []interface{}{"1.24", "2.48", "NJ"}}}, + }, + expected: []interface{}{float64(1), float64(0)}, + }, + { + name: "GEOADD Adding both XX and NX options together", + commands: []HTTPCommand{ + {Command: "GEOADD", Body: map[string]interface{}{"key": "mygeo", "values": []interface{}{"XX", "NX", "1.2", "2.4", "NJ"}}}, + }, + expected: []interface{}{"ERR XX and NX options at the same time are not compatible"}, + }, + { + name: "GEOADD Invalid Longitude", + commands: []HTTPCommand{ + {Command: "GEOADD", Body: map[string]interface{}{"key": "mygeo", "values": []interface{}{"181", "2.4", "MT"}}}, + }, + expected: []interface{}{"ERR invalid longitude"}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.commands { + result, _ := exec.FireCommand(cmd) + assert.Equal(t, tc.expected[i], result, "Value mismatch for cmd %s", cmd) + } + }) + } +} + +func TestGeoDist(t *testing.T) { + exec := NewHTTPCommandExecutor() + + testCases := []struct { + name string + commands []HTTPCommand + expected []interface{} + }{ + { + name: "GEODIST b/w existing points", + commands: []HTTPCommand{ + {Command: "GEOADD", Body: map[string]interface{}{"key": "points", "values": []interface{}{"13.361389", "38.115556", "Palermo"}}}, + {Command: "GEOADD", Body: map[string]interface{}{"key": "points", "values": []interface{}{"15.087269", "37.502669", "Catania"}}}, + {Command: "GEODIST", Body: map[string]interface{}{"key": "points", "values": []interface{}{"Palermo", "Catania"}}}, + {Command: "GEODIST", Body: map[string]interface{}{"key": "points", "values": []interface{}{"Palermo", "Catania", "km"}}}, + }, + expected: []interface{}{float64(1), float64(1), float64(166274.144), float64(166.2741)}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.commands { + result, _ := exec.FireCommand(cmd) + assert.Equal(t, tc.expected[i], result, "Value mismatch for cmd %s", cmd) + } + }) + } +} diff --git a/integration_tests/commands/resp/geo_tests.go b/integration_tests/commands/resp/geo_tests.go new file mode 100644 index 000000000..068bfef04 --- /dev/null +++ b/integration_tests/commands/resp/geo_tests.go @@ -0,0 +1,86 @@ +package resp + +import ( + "testing" + "time" + + "gotest.tools/v3/assert" +) + +func TestGeoAdd(t *testing.T) { + conn := getLocalConnection() + defer conn.Close() + + testCases := []struct { + name string + cmds []string + expect []interface{} + }{ + { + name: "GeoAdd With Wrong Number of Arguments", + cmds: []string{"GEOADD mygeo 1 2"}, + expect: []interface{}{"ERR wrong number of arguments for 'geoadd' command"}, + }, + { + name: "GeoAdd With Adding New Member And Updating it", + cmds: []string{"GEOADD mygeo 1.21 1.44 NJ", "GEOADD mygeo 1.22 1.54 NJ"}, + expect: []interface{}{int64(1), int64(0)}, + }, + { + name: "GeoAdd With Adding New Member And Updating it with NX", + cmds: []string{"GEOADD mygeo 1.21 1.44 MD", "GEOADD mygeo 1.22 1.54 MD"}, + expect: []interface{}{int64(1), int64(0)}, + }, + { + name: "GEOADD with both NX and XX options", + cmds: []string{"GEOADD mygeo NX XX 1.21 1.44 MD"}, + expect: []interface{}{"ERR XX and NX options at the same time are not compatible"}, + }, + { + name: "GEOADD invalid longitude", + cmds: []string{"GEOADD mygeo 181.0 1.44 MD"}, + expect: []interface{}{"ERR invalid longitude"}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.cmds { + result := FireCommand(conn, cmd) + assert.Equal(t, tc.expect[i], result, "Value mismatch for cmd %s", cmd) + } + }) + } +} + +func TestGeoDist(t *testing.T) { + conn := getLocalConnection() + defer conn.Close() + + testCases := []struct { + name string + cmds []string + expect []interface{} + delays []time.Duration + }{ + { + name: "GEODIST b/w existing points", + cmds: []string{ + "GEOADD points 13.361389 38.115556 Palermo", + "GEOADD points 15.087269 37.502669 Catania", + "GEODIST points Palermo Catania", + "GEODIST points Palermo Catania km", + }, + expect: []interface{}{int64(1), int64(1), "166274.144", "166.2741"}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.cmds { + result := FireCommand(conn, cmd) + assert.Equal(t, tc.expect[i], result, "Value mismatch for cmd %s", cmd) + } + }) + } +} diff --git a/integration_tests/commands/websocket/geo_tests.go b/integration_tests/commands/websocket/geo_tests.go new file mode 100644 index 000000000..e2c5d240b --- /dev/null +++ b/integration_tests/commands/websocket/geo_tests.go @@ -0,0 +1,87 @@ +package websocket + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestGeoAdd(t *testing.T) { + exec := NewWebsocketCommandExecutor() + conn := exec.ConnectToServer() + + testCases := []struct { + name string + cmds []string + expect []interface{} + }{ + { + name: "GeoAdd With Wrong Number of Arguments", + cmds: []string{"GEOADD mygeo 1 2"}, + expect: []interface{}{"ERR wrong number of arguments for 'geoadd' command"}, + }, + { + name: "GeoAdd With Adding New Member And Updating it", + cmds: []string{"GEOADD mygeo 1.21 1.44 NJ", "GEOADD mygeo 1.22 1.54 NJ"}, + expect: []interface{}{float64(1), float64(0)}, + }, + { + name: "GeoAdd With Adding New Member And Updating it with NX", + cmds: []string{"GEOADD mygeo NX 1.21 1.44 MD", "GEOADD mygeo 1.22 1.54 MD"}, + expect: []interface{}{float64(1), float64(0)}, + }, + { + name: "GEOADD with both NX and XX options", + cmds: []string{"GEOADD mygeo NX XX 1.21 1.44 DEL"}, + expect: []interface{}{"ERR XX and NX options at the same time are not compatible"}, + }, + { + name: "GEOADD invalid longitude", + cmds: []string{"GEOADD mygeo 181.0 1.44 MD"}, + expect: []interface{}{"ERR invalid longitude"}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.cmds { + result, err := exec.FireCommandAndReadResponse(conn, cmd) + assert.Nil(t, err) + assert.Equal(t, tc.expect[i], result, "Value mismatch for cmd %s", cmd) + } + }) + } +} + +func TestGeoDist(t *testing.T) { + exec := NewWebsocketCommandExecutor() + conn := exec.ConnectToServer() + defer conn.Close() + + testCases := []struct { + name string + cmds []string + expect []interface{} + }{ + { + name: "GEODIST b/w existing points", + cmds: []string{ + "GEOADD points 13.361389 38.115556 Palermo", + "GEOADD points 15.087269 37.502669 Catania", + "GEODIST points Palermo Catania", + "GEODIST points Palermo Catania km", + }, + expect: []interface{}{float64(1), float64(1), float64(166274.144), float64(166.2741)}, + }, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + for i, cmd := range tc.cmds { + result, err := exec.FireCommandAndReadResponse(conn, cmd) + assert.Nil(t, err) + assert.Equal(t, tc.expect[i], result, "Value mismatch for cmd %s", cmd) + } + }) + } +} diff --git a/internal/eval/commands.go b/internal/eval/commands.go index d07c405bc..d4dc01259 100644 --- a/internal/eval/commands.go +++ b/internal/eval/commands.go @@ -1119,18 +1119,20 @@ var ( NewEval: evalHINCRBYFLOAT, } geoAddCmdMeta = DiceCmdMeta{ - Name: "GEOADD", - Info: `Adds one or more members to a geospatial index. The key is created if it doesn't exist.`, - Arity: -5, - Eval: evalGEOADD, - KeySpecs: KeySpecs{BeginIndex: 1}, + Name: "GEOADD", + Info: `Adds one or more members to a geospatial index. The key is created if it doesn't exist.`, + Arity: -5, + IsMigrated: true, + NewEval: evalGEOADD, + KeySpecs: KeySpecs{BeginIndex: 1}, } geoDistCmdMeta = DiceCmdMeta{ - Name: "GEODIST", - Info: `Returns the distance between two members in the geospatial index.`, - Arity: -4, - Eval: evalGEODIST, - KeySpecs: KeySpecs{BeginIndex: 1}, + Name: "GEODIST", + Info: `Returns the distance between two members in the geospatial index.`, + Arity: -4, + IsMigrated: true, + NewEval: evalGEODIST, + KeySpecs: KeySpecs{BeginIndex: 1}, } jsonstrappendCmdMeta = DiceCmdMeta{ Name: "JSON.STRAPPEND", diff --git a/internal/eval/eval.go b/internal/eval/eval.go index 774ade1de..ee7c42968 100644 --- a/internal/eval/eval.go +++ b/internal/eval/eval.go @@ -5,7 +5,6 @@ import ( "errors" "fmt" "log/slog" - "math" "math/bits" "regexp" "sort" @@ -15,8 +14,6 @@ import ( "unicode" "unsafe" - "github.com/dicedb/dice/internal/eval/geo" - "github.com/dicedb/dice/internal/eval/sortedset" "github.com/dicedb/dice/internal/object" "github.com/rs/xid" @@ -4122,136 +4119,6 @@ func evalBITFIELDRO(args []string, store *dstore.Store) []byte { return bitfieldEvalGeneric(args, store, true) } -func evalGEOADD(args []string, store *dstore.Store) []byte { - if len(args) < 4 { - return diceerrors.NewErrArity("GEOADD") - } - - key := args[0] - var nx, xx bool - startIdx := 1 - - // Parse options - for startIdx < len(args) { - option := strings.ToUpper(args[startIdx]) - if option == "NX" { - nx = true - startIdx++ - } else if option == "XX" { - xx = true - startIdx++ - } else { - break - } - } - - // Check if we have the correct number of arguments after parsing options - if (len(args)-startIdx)%3 != 0 { - return diceerrors.NewErrArity("GEOADD") - } - - if xx && nx { - return diceerrors.NewErrWithMessage("ERR XX and NX options at the same time are not compatible") - } - - // Get or create sorted set - obj := store.Get(key) - var ss *sortedset.Set - if obj != nil { - var err []byte - ss, err = sortedset.FromObject(obj) - if err != nil { - return err - } - } else { - ss = sortedset.New() - } - - added := 0 - for i := startIdx; i < len(args); i += 3 { - longitude, err := strconv.ParseFloat(args[i], 64) - if err != nil || math.IsNaN(longitude) || longitude < -180 || longitude > 180 { - return diceerrors.NewErrWithMessage("ERR invalid longitude") - } - - latitude, err := strconv.ParseFloat(args[i+1], 64) - if err != nil || math.IsNaN(latitude) || latitude < -85.05112878 || latitude > 85.05112878 { - return diceerrors.NewErrWithMessage("ERR invalid latitude") - } - - member := args[i+2] - _, exists := ss.Get(member) - - // Handle XX option: Only update existing elements - if xx && !exists { - continue - } - - // Handle NX option: Only add new elements - if nx && exists { - continue - } - - hash := geo.EncodeHash(latitude, longitude) - - wasInserted := ss.Upsert(hash, member) - if wasInserted { - added++ - } - } - - obj = store.NewObj(ss, -1, object.ObjTypeSortedSet, object.ObjEncodingBTree) - store.Put(key, obj) - - return clientio.Encode(added, false) -} - -func evalGEODIST(args []string, store *dstore.Store) []byte { - if len(args) < 3 || len(args) > 4 { - return diceerrors.NewErrArity("GEODIST") - } - - key := args[0] - member1 := args[1] - member2 := args[2] - unit := "m" - if len(args) == 4 { - unit = strings.ToLower(args[3]) - } - - // Get the sorted set - obj := store.Get(key) - if obj == nil { - return clientio.RespNIL - } - - ss, err := sortedset.FromObject(obj) - if err != nil { - return err - } - - // Get the scores (geohashes) for both members - score1, ok := ss.Get(member1) - if !ok { - return clientio.RespNIL - } - score2, ok := ss.Get(member2) - if !ok { - return clientio.RespNIL - } - - lat1, lon1 := geo.DecodeHash(score1) - lat2, lon2 := geo.DecodeHash(score2) - - distance := geo.GetDistance(lon1, lat1, lon2, lat2) - - result, err := geo.ConvertDistance(distance, unit) - if err != nil { - return err - } - - return clientio.Encode(utils.RoundToDecimals(result, 4), false) -} // evalJSONSTRAPPEND appends a string value to the JSON string value at the specified path // in the JSON object saved at the key in arguments. diff --git a/internal/eval/eval_test.go b/internal/eval/eval_test.go index 5a23a9309..0f8905f68 100644 --- a/internal/eval/eval_test.go +++ b/internal/eval/eval_test.go @@ -6420,84 +6420,130 @@ func testEvalBitFieldRO(t *testing.T, store *dstore.Store) { func testEvalGEOADD(t *testing.T, store *dstore.Store) { tests := map[string]evalTestCase{ "GEOADD with wrong number of arguments": { - input: []string{"mygeo", "1", "2"}, - output: diceerrors.NewErrArity("GEOADD"), + input: []string{"mygeo", "1", "2"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongArgumentCount("GEOADD"), + }, }, "GEOADD with non-numeric longitude": { - input: []string{"mygeo", "long", "40.7128", "NewYork"}, - output: diceerrors.NewErrWithMessage("ERR invalid longitude"), + input: []string{"mygeo", "long", "40.7128", "NewYork"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("invalid longitude"), + }, }, "GEOADD with non-numeric latitude": { - input: []string{"mygeo", "-74.0060", "lat", "NewYork"}, - output: diceerrors.NewErrWithMessage("ERR invalid latitude"), + input: []string{"mygeo", "-74.0060", "lat", "NewYork"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("invalid latitude"), + }, }, "GEOADD new member to non-existing key": { - setup: func() {}, - input: []string{"mygeo", "-74.0060", "40.7128", "NewYork"}, - output: clientio.Encode(int64(1), false), + setup: func() {}, + input: []string{"mygeo", "-74.0060", "40.7128", "NewYork"}, + migratedOutput: EvalResponse{ + Result: 1, + Error: nil, + }, }, "GEOADD existing member with updated coordinates": { setup: func() { evalGEOADD([]string{"mygeo", "-74.0060", "40.7128", "NewYork"}, store) }, - input: []string{"mygeo", "-73.9352", "40.7304", "NewYork"}, - output: clientio.Encode(int64(0), false), + input: []string{"mygeo", "-73.9352", "40.7304", "NewYork"}, + migratedOutput: EvalResponse{ + Result: 0, + Error: nil, + }, }, "GEOADD multiple members": { setup: func() { evalGEOADD([]string{"mygeo", "-74.0060", "40.7128", "NewYork"}, store) }, - input: []string{"mygeo", "-118.2437", "34.0522", "LosAngeles", "-87.6298", "41.8781", "Chicago"}, - output: clientio.Encode(int64(2), false), + input: []string{"mygeo", "-118.2437", "34.0522", "LosAngeles", "-87.6298", "41.8781", "Chicago"}, + migratedOutput: EvalResponse{ + Result: 2, + Error: nil, + }, }, "GEOADD with NX option (new member)": { - input: []string{"mygeo", "NX", "-122.4194", "37.7749", "SanFrancisco"}, - output: clientio.Encode(int64(1), false), + input: []string{"mygeo", "NX", "-122.4194", "37.7749", "SanFrancisco"}, + migratedOutput: EvalResponse{ + Result: 1, + Error: nil, + }, }, "GEOADD with NX option (existing member)": { setup: func() { evalGEOADD([]string{"mygeo", "-74.0060", "40.7128", "NewYork"}, store) }, - input: []string{"mygeo", "NX", "-73.9352", "40.7304", "NewYork"}, - output: clientio.Encode(int64(0), false), + input: []string{"mygeo", "NX", "-73.9352", "40.7304", "NewYork"}, + migratedOutput: EvalResponse{ + Result: 0, + Error: nil, + }, }, "GEOADD with XX option (new member)": { - input: []string{"mygeo", "XX", "-71.0589", "42.3601", "Boston"}, - output: clientio.Encode(int64(0), false), + input: []string{"mygeo", "XX", "-71.0589", "42.3601", "Boston"}, + migratedOutput: EvalResponse{ + Result: 0, + Error: nil, + }, }, "GEOADD with XX option (existing member)": { setup: func() { evalGEOADD([]string{"mygeo", "-74.0060", "40.7128", "NewYork"}, store) }, - input: []string{"mygeo", "XX", "-73.9352", "40.7304", "NewYork"}, - output: clientio.Encode(int64(0), false), + input: []string{"mygeo", "XX", "-73.9352", "40.7304", "NewYork"}, + migratedOutput: EvalResponse{ + Result: 0, + Error: nil, + }, }, "GEOADD with both NX and XX options": { input: []string{"mygeo", "NX", "XX", "-74.0060", "40.7128", "NewYork"}, output: diceerrors.NewErrWithMessage("ERR XX and NX options at the same time are not compatible"), + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("XX and NX options at the same time are not compatible"), + }, }, "GEOADD with invalid option": { - input: []string{"mygeo", "INVALID", "-74.0060", "40.7128", "NewYork"}, - output: diceerrors.NewErrArity("GEOADD"), + input: []string{"mygeo", "INVALID", "-74.0060", "40.7128", "NewYork"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongArgumentCount("GEOADD"), + }, }, "GEOADD to a key of wrong type": { setup: func() { store.Put("mygeo", store.NewObj("string_value", -1, object.ObjTypeString, object.ObjEncodingRaw)) }, - input: []string{"mygeo", "-74.0060", "40.7128", "NewYork"}, - output: []byte("-ERR Existing key has wrong Dice type\r\n"), + input: []string{"mygeo", "-74.0060", "40.7128", "NewYork"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongTypeOperation, + }, }, "GEOADD with longitude out of range": { - input: []string{"mygeo", "181.0", "40.7128", "Invalid"}, - output: diceerrors.NewErrWithMessage("ERR invalid longitude"), + input: []string{"mygeo", "181.0", "40.7128", "Invalid"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("invalid longitude"), + }, }, "GEOADD with latitude out of range": { - input: []string{"mygeo", "-74.0060", "91.0", "Invalid"}, - output: diceerrors.NewErrWithMessage("ERR invalid latitude"), + input: []string{"mygeo", "-74.0060", "91.0", "Invalid"}, + migratedOutput: EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("invalid latitude"), + }, }, } - runEvalTests(t, tests, evalGEOADD, store) + runMigratedEvalTests(t, tests, evalGEOADD, store) } func testEvalGEODIST(t *testing.T, store *dstore.Store) { @@ -6507,28 +6553,37 @@ func testEvalGEODIST(t *testing.T, store *dstore.Store) { evalGEOADD([]string{"points", "13.361389", "38.115556", "Palermo"}, store) evalGEOADD([]string{"points", "15.087269", "37.502669", "Catania"}, store) }, - input: []string{"points", "Palermo", "Catania"}, - output: clientio.Encode(float64(166274.1440), false), // Example value + input: []string{"points", "Palermo", "Catania"}, + migratedOutput: EvalResponse{ + Result: float64(166274.1440), + Error: nil, + }, }, "GEODIST with units (km)": { setup: func() { evalGEOADD([]string{"points", "13.361389", "38.115556", "Palermo"}, store) evalGEOADD([]string{"points", "15.087269", "37.502669", "Catania"}, store) }, - input: []string{"points", "Palermo", "Catania", "km"}, - output: clientio.Encode(float64(166.2741), false), // Example value + input: []string{"points", "Palermo", "Catania", "km"}, + migratedOutput: EvalResponse{ + Result: float64(166.2741), + Error: nil, + }, }, "GEODIST to same point": { setup: func() { evalGEOADD([]string{"points", "13.361389", "38.115556", "Palermo"}, store) }, - input: []string{"points", "Palermo", "Palermo"}, - output: clientio.Encode(float64(0.0000), false), // Expecting distance 0 formatted to 4 decimals + input: []string{"points", "Palermo", "Palermo"}, + migratedOutput: EvalResponse{ + Result: float64(0.0000), + Error: nil, + }, }, // Add other test cases here... } - runEvalTests(t, tests, evalGEODIST, store) + runMigratedEvalTests(t, tests, evalGEODIST, store) } func testEvalSINTER(t *testing.T, store *dstore.Store) { diff --git a/internal/eval/store_eval.go b/internal/eval/store_eval.go index c1bafd570..fc2da3bec 100644 --- a/internal/eval/store_eval.go +++ b/internal/eval/store_eval.go @@ -10,6 +10,7 @@ import ( "github.com/bytedance/sonic" "github.com/dicedb/dice/internal/clientio" diceerrors "github.com/dicedb/dice/internal/errors" + "github.com/dicedb/dice/internal/eval/geo" "github.com/dicedb/dice/internal/eval/sortedset" "github.com/dicedb/dice/internal/object" "github.com/dicedb/dice/internal/server/utils" @@ -1643,3 +1644,177 @@ func evalZPOPMIN(args []string, store *dstore.Store) *EvalResponse { Error: nil, } } + +func evalGEOADD(args []string, store *dstore.Store) *EvalResponse { + if len(args) < 4 { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongArgumentCount("GEOADD"), + } + } + + key := args[0] + var nx, xx bool + startIdx := 1 + + // Parse options + for startIdx < len(args) { + option := strings.ToUpper(args[startIdx]) + if option == "NX" { + nx = true + startIdx++ + } else if option == "XX" { + xx = true + startIdx++ + } else { + break + } + } + + // Check if we have the correct number of arguments after parsing options + if (len(args)-startIdx)%3 != 0 { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongArgumentCount("GEOADD"), + } + } + + if xx && nx { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("XX and NX options at the same time are not compatible"), + } + } + + // Get or create sorted set + obj := store.Get(key) + var ss *sortedset.Set + if obj != nil { + var err []byte + ss, err = sortedset.FromObject(obj) + if err != nil { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongTypeOperation, + } + } + } else { + ss = sortedset.New() + } + + added := 0 + for i := startIdx; i < len(args); i += 3 { + longitude, err := strconv.ParseFloat(args[i], 64) + if err != nil || math.IsNaN(longitude) || longitude < -180 || longitude > 180 { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("invalid longitude"), + } + } + + latitude, err := strconv.ParseFloat(args[i+1], 64) + if err != nil || math.IsNaN(latitude) || latitude < -85.05112878 || latitude > 85.05112878 { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrGeneral("invalid latitude"), + } + } + + member := args[i+2] + _, exists := ss.Get(member) + + // Handle XX option: Only update existing elements + if xx && !exists { + continue + } + + // Handle NX option: Only add new elements + if nx && exists { + continue + } + + hash := geo.EncodeHash(latitude, longitude) + + wasInserted := ss.Upsert(hash, member) + if wasInserted { + added++ + } + } + + obj = store.NewObj(ss, -1, object.ObjTypeSortedSet, object.ObjEncodingBTree) + store.Put(key, obj) + + return &EvalResponse{ + Result: added, + Error: nil, + } +} + +func evalGEODIST(args []string, store *dstore.Store) *EvalResponse { + if len(args) < 3 || len(args) > 4 { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongArgumentCount("GEODIST"), + } + } + + key := args[0] + member1 := args[1] + member2 := args[2] + unit := "m" + if len(args) == 4 { + unit = strings.ToLower(args[3]) + } + + // Get the sorted set + obj := store.Get(key) + if obj == nil { + return &EvalResponse{ + Result: clientio.NIL, + Error: nil, + } + } + + ss, err := sortedset.FromObject(obj) + if err != nil { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongTypeOperation, + } + } + + // Get the scores (geohashes) for both members + score1, ok := ss.Get(member1) + if !ok { + return &EvalResponse{ + Result: nil, + Error: nil, + } + } + score2, ok := ss.Get(member2) + if !ok { + return &EvalResponse{ + Result: nil, + Error: nil, + } + } + + lat1, lon1 := geo.DecodeHash(score1) + lat2, lon2 := geo.DecodeHash(score2) + + distance := geo.GetDistance(lon1, lat1, lon2, lat2) + + result, err := geo.ConvertDistance(distance, unit) + + if err != nil { + return &EvalResponse{ + Result: nil, + Error: diceerrors.ErrWrongTypeOperation, + } + } + + return &EvalResponse{ + Result: utils.RoundToDecimals(result, 4), + Error: nil, + } +} diff --git a/internal/server/cmd_meta.go b/internal/server/cmd_meta.go index 8b5af5227..0bc47a9df 100644 --- a/internal/server/cmd_meta.go +++ b/internal/server/cmd_meta.go @@ -147,6 +147,16 @@ var ( CmdType: SingleShard, } + geoaddCmdMeta = CmdsMeta{ + Cmd: "GEOADD", + CmdType: SingleShard, + } + + geodistCmdMeta = CmdsMeta{ + Cmd: "GEODIST", + CmdType: SingleShard, + } + // Metadata for multishard commands would go here. // These commands require both breakup and gather logic. @@ -197,5 +207,7 @@ func init() { WorkerCmdsMeta["HINCRBY"] = hincrbyCmdMeta WorkerCmdsMeta["HINCRBYFLOAT"] = hincrbyfloatCmdMeta WorkerCmdsMeta["HRANDFIELD"] = hrandfieldCmdMeta + WorkerCmdsMeta["GEOADD"] = geoaddCmdMeta + WorkerCmdsMeta["GEODIST"] = geodistCmdMeta // Additional commands (multishard, custom) can be added here as needed. } diff --git a/internal/worker/cmd_meta.go b/internal/worker/cmd_meta.go index 9be05a82b..aa98c1480 100644 --- a/internal/worker/cmd_meta.go +++ b/internal/worker/cmd_meta.go @@ -84,6 +84,8 @@ const ( CmdHRandField = "HRANDFIELD" CmdGetRange = "GETRANGE" CmdAppend = "APPEND" + CmdGeoAdd = "GEOADD" + CmdGeoDist = "GEODIST" ) type CmdMeta struct { @@ -162,6 +164,13 @@ var CommandsMeta = map[string]CmdMeta{ CmdType: SingleShard, }, + CmdGeoAdd: { + CmdType: SingleShard, + }, + CmdGeoDist: { + CmdType: SingleShard, + }, + // Multi-shard commands. CmdRename: { CmdType: MultiShard,