fix: sorted set remove request (#133)

momentohq · Mar 14, 2023 · 8370836 · 8370836
1 parent 4985173
commit 8370836
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 2 deletions.
diff --git a/momento/sorted_set_fetch.go b/momento/sorted_set_fetch.go
@@ -80,9 +80,18 @@ func (r *SortedSetFetchRequest) initGrpcRequest(scsDataClient) error {
 	}
 
 	switch numResults := r.NumberOfResults.(type) {
+	case *FetchAllElements:
+		grpcReq.NumResults = &pb.XSortedSetFetchRequest_All{}
 	case FetchAllElements:
+		grpcReq.NumResults = &pb.XSortedSetFetchRequest_All{}
 	case nil:
 		grpcReq.NumResults = &pb.XSortedSetFetchRequest_All{}
+	case *FetchLimitedElements:
+		grpcReq.NumResults = &pb.XSortedSetFetchRequest_Limit{
+			Limit: &pb.XSortedSetFetchRequest_XLimit{
+				Limit: numResults.Limit,
+			},
+		}
 	case FetchLimitedElements:
 		grpcReq.NumResults = &pb.XSortedSetFetchRequest_Limit{
 			Limit: &pb.XSortedSetFetchRequest_XLimit{

diff --git a/momento/sorted_set_remove.go b/momento/sorted_set_remove.go
@@ -42,7 +42,7 @@ type RemoveAllElements struct{}
 func (RemoveAllElements) isSortedSetRemoveNumElement() {}
 
 type RemoveSomeElements struct {
-	elementsToRemove []Bytes
+	Elements []Bytes
 }
 
 func (RemoveSomeElements) isSortedSetRemoveNumElement() {}
@@ -63,12 +63,20 @@ func (r *SortedSetRemoveRequest) initGrpcRequest(scsDataClient) error {
 	}
 
 	switch toRemove := r.ElementsToRemove.(type) {
+	case RemoveAllElements:
+		grpcReq.RemoveElements = &pb.XSortedSetRemoveRequest_All{}
 	case *RemoveAllElements:
 		grpcReq.RemoveElements = &pb.XSortedSetRemoveRequest_All{}
+	case RemoveSomeElements:
+		grpcReq.RemoveElements = &pb.XSortedSetRemoveRequest_Some{
+			Some: &pb.XSortedSetRemoveRequest_XSome{
+				ElementName: momentoBytesListToPrimitiveByteList(toRemove.Elements),
+			},
+		}
 	case *RemoveSomeElements:
 		grpcReq.RemoveElements = &pb.XSortedSetRemoveRequest_Some{
 			Some: &pb.XSortedSetRemoveRequest_XSome{
-				ElementName: momentoBytesListToPrimitiveByteList(toRemove.elementsToRemove),
+				ElementName: momentoBytesListToPrimitiveByteList(toRemove.Elements),
 			},
 		}
 	default: