diff --git a/server/server.go b/server/server.go index ae340f4..817c0ee 100644 --- a/server/server.go +++ b/server/server.go @@ -36,6 +36,7 @@ const ( // Sorted set commands CMD_ZADD = "ZADD" CMD_ZCARD = "ZCARD" + CMD_ZDIFF = "ZDIFF" CMD_ZREMOVE = "ZREMOVE" CMD_ZMEMBERS = "ZMEMBERS" ) @@ -148,6 +149,8 @@ func (s *Server) HandleCommand() { s.zAdd(cmd) case CMD_ZCARD: s.zCard(cmd) + case CMD_ZDIFF: + s.zDiff(cmd) case CMD_ZREMOVE: s.zRemove(cmd) case CMD_ZMEMBERS: diff --git a/server/zset.go b/server/zset.go index f9d23fa..700e6d5 100644 --- a/server/zset.go +++ b/server/zset.go @@ -52,3 +52,18 @@ func (s *Server) zRemove(cmd Command) { return } } + +func (s *Server) zDiff(cmd Command) { + if len(cmd.Args) < 2 { + cmd.error(ErrNotEnoughArgs) + return + } + elements, err := s.DB.ZDiff(cmd.Args[0], cmd.Args[1]) + if err != nil { + cmd.error(err) + return + } + for index, element := range elements { + cmd.write(fmt.Sprintf("%d) %s", index+1, element)) + } +} diff --git a/store/zset.go b/store/zset.go index b10e6e5..8a9290c 100644 --- a/store/zset.go +++ b/store/zset.go @@ -1,6 +1,10 @@ package store -import "github.com/Devansh3712/tandb/zset" +import ( + "fmt" + + "github.com/Devansh3712/tandb/zset" +) func (s *Store) ZAdd(set, key string) { s.Mutex.Lock() @@ -46,3 +50,20 @@ func (s *Store) ZRemove(set, key string) error { } return value.Remove(key) } + +func (s *Store) ZDiff(s1, s2 string) ([]string, error) { + s.Mutex.RLock() + defer s.Mutex.RUnlock() + + v1, ok := s.ZSets[s1] + if !ok { + return nil, fmt.Errorf("the set %s does not exist", s1) + } + v2, ok := s.ZSets[s2] + if !ok { + return nil, fmt.Errorf("the set %s does not exist", s2) + } + + diff := v1.Difference(v2) + return diff.Members(), nil +}