Go: implement SUnionStore (#2805)

Go: implement SUnionStore --------- Signed-off-by: Chloe Yip <[email protected]> Co-authored-by: James Xin <[email protected]> Co-authored-by: Yury-Fridlyand <[email protected]>
valkey-io · Dec 13, 2024 · 8b049d4 · 8b049d4
1 parent 3995396
commit 8b049d4
Show file tree

Hide file tree

Showing 4 changed files with 144 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,6 @@
 #### Changes
 
+* Go: Add SUNIONSTORE command ([#2805](https://github.com/valkey-io/valkey-glide/pull/2805)
 * Java: bump `netty` version ([#2795](https://github.com/valkey-io/valkey-glide/pull/2795))
 * Java: Bump protobuf (protoc) version ([#2796](https://github.com/valkey-io/valkey-glide/pull/2796), [#2800](https://github.com/valkey-io/valkey-glide/pull/2800))
 * Go: Add `SInterStore` ([#2779](https://github.com/valkey-io/valkey-glide/issues/2779))

diff --git a/go/api/base_client.go b/go/api/base_client.go
@@ -534,6 +534,15 @@ func (client *baseClient) SRem(key string, members []string) (Result[int64], err
 	return handleLongResponse(result)
 }
 
+func (client *baseClient) SUnionStore(destination string, keys []string) (Result[int64], error) {
+	result, err := client.executeCommand(C.SUnionStore, append([]string{destination}, keys...))
+	if err != nil {
+		return CreateNilInt64Result(), err
+	}
+
+	return handleLongResponse(result)
+}
+
 func (client *baseClient) SMembers(key string) (map[Result[string]]struct{}, error) {
 	result, err := client.executeCommand(C.SMembers, []string{key})
 	if err != nil {

diff --git a/go/api/set_commands.go b/go/api/set_commands.go
@@ -324,4 +324,27 @@ type SetCommands interface {
 	//
 	// [valkey.io]: https://valkey.io/commands/smismember/
 	SMIsMember(key string, members []string) ([]Result[bool], error)
+
+	// SUnionStore stores the members of the union of all given sets specified by `keys` into a new set at `destination`.
+	//
+	// Note: When in cluster mode, `destination` and all `keys` must map to the same hash slot.
+	//
+	// See [valkey.io] for details.
+	//
+	// Parameters:
+	//	 destination - The key of the destination set.
+	//   keys - The keys from which to retrieve the set members.
+	//
+	// Return value:
+	//   The number of elements in the resulting set.
+	//
+	// Example:
+	//   result, err := client.SUnionStore("my_set", []string{"set1", "set2"})
+	//   if err != nil {
+	//       fmt.Println(result.Value())
+	//   }
+	//   // Output: 2 - Two elements were stored at "my_set", and those elements are the union of "set1" and "set2".
+	//
+	// [valkey.io]: https://valkey.io/commands/sunionstore/
+	SUnionStore(destination string, keys []string) (Result[int64], error)
 }
diff --git a/go/integTest/shared_commands_test.go b/go/integTest/shared_commands_test.go
@@ -4,6 +4,7 @@ package integTest
 
 import (
 	"math"
+	"reflect"
 	"time"
 
 	"github.com/google/uuid"
@@ -1340,6 +1341,116 @@ func (suite *GlideTestSuite) TestSRem_WithExistingKeyAndDifferentMembers() {
 	})
 }
 
+func (suite *GlideTestSuite) TestSUnionStore() {
+	suite.runWithDefaultClients(func(client api.BaseClient) {
+		key1 := "{key}-1-" + uuid.NewString()
+		key2 := "{key}-2-" + uuid.NewString()
+		key3 := "{key}-3-" + uuid.NewString()
+		key4 := "{key}-4-" + uuid.NewString()
+		stringKey := "{key}-5-" + uuid.NewString()
+		nonExistingKey := "{key}-6-" + uuid.NewString()
+
+		memberArray1 := []string{"a", "b", "c"}
+		memberArray2 := []string{"c", "d", "e"}
+		memberArray3 := []string{"e", "f", "g"}
+		expected1 := map[api.Result[string]]struct{}{
+			api.CreateStringResult("a"): {},
+			api.CreateStringResult("b"): {},
+			api.CreateStringResult("c"): {},
+			api.CreateStringResult("d"): {},
+			api.CreateStringResult("e"): {},
+		}
+		expected2 := map[api.Result[string]]struct{}{
+			api.CreateStringResult("a"): {},
+			api.CreateStringResult("b"): {},
+			api.CreateStringResult("c"): {},
+			api.CreateStringResult("d"): {},
+			api.CreateStringResult("e"): {},
+			api.CreateStringResult("f"): {},
+			api.CreateStringResult("g"): {},
+		}
+		t := suite.T()
+
+		res1, err := client.SAdd(key1, memberArray1)
+		assert.NoError(t, err)
+		assert.Equal(t, int64(3), res1.Value())
+
+		res2, err := client.SAdd(key2, memberArray2)
+		assert.NoError(t, err)
+		assert.Equal(t, int64(3), res2.Value())
+
+		res3, err := client.SAdd(key3, memberArray3)
+		assert.NoError(t, err)
+		assert.Equal(t, int64(3), res3.Value())
+
+		// store union in new key
+		res4, err := client.SUnionStore(key4, []string{key1, key2})
+		assert.NoError(t, err)
+		assert.Equal(t, int64(5), res4.Value())
+
+		res5, err := client.SMembers(key4)
+		assert.NoError(t, err)
+		assert.Len(t, res5, 5)
+		assert.True(t, reflect.DeepEqual(res5, expected1))
+
+		// overwrite existing set
+		res6, err := client.SUnionStore(key1, []string{key4, key2})
+		assert.NoError(t, err)
+		assert.Equal(t, int64(5), res6.Value())
+
+		res7, err := client.SMembers(key1)
+		assert.NoError(t, err)
+		assert.Len(t, res7, 5)
+		assert.True(t, reflect.DeepEqual(res7, expected1))
+
+		// overwrite one of the source keys
+		res8, err := client.SUnionStore(key2, []string{key4, key2})
+		assert.NoError(t, err)
+		assert.Equal(t, int64(5), res8.Value())
+
+		res9, err := client.SMembers(key2)
+		assert.NoError(t, err)
+		assert.Len(t, res9, 5)
+		assert.True(t, reflect.DeepEqual(res9, expected1))
+
+		// union with non-existing key
+		res10, err := client.SUnionStore(key2, []string{nonExistingKey})
+		assert.NoError(t, err)
+		assert.Equal(t, int64(0), res10.Value())
+
+		// check that the key is now empty
+		members1, err := client.SMembers(key2)
+		assert.NoError(t, err)
+		assert.Empty(t, members1)
+
+		// invalid argument - key list must not be empty
+		res11, err := client.SUnionStore(key4, []string{})
+		assert.Equal(suite.T(), int64(0), res11.Value())
+		assert.NotNil(suite.T(), err)
+		assert.IsType(suite.T(), &api.RequestError{}, err)
+
+		// non-set key
+		_, err = client.Set(stringKey, "value")
+		assert.NoError(t, err)
+
+		res12, err := client.SUnionStore(key4, []string{stringKey, key1})
+		assert.Equal(suite.T(), int64(0), res12.Value())
+		assert.NotNil(suite.T(), err)
+		assert.IsType(suite.T(), &api.RequestError{}, err)
+
+		// overwrite destination when destination is not a set
+		res13, err := client.SUnionStore(stringKey, []string{key1, key3})
+		assert.NoError(t, err)
+		assert.Equal(t, int64(7), res13.Value())
+
+		// check that the key is now empty
+		res14, err := client.SMembers(stringKey)
+		assert.NoError(t, err)
+		assert.Len(t, res14, 7)
+		assert.True(t, reflect.DeepEqual(res14, expected2))
+	})
+}
+
 func (suite *GlideTestSuite) TestSMembers() {
 	suite.runWithDefaultClients(func(client api.BaseClient) {
 		key := uuid.NewString()