diff --git a/clients/ember-go/ember.go b/clients/ember-go/ember.go index 1a3994ff..a1018fe7 100644 --- a/clients/ember-go/ember.go +++ b/clients/ember-go/ember.go @@ -632,3 +632,235 @@ func (c *Client) PubSubNumSub(ctx context.Context, channels ...string) (map[stri } return result, nil } + +// Expiretime returns the absolute unix expiry timestamp in seconds (-1 = no expiry, -2 = missing). +func (c *Client) Expiretime(ctx context.Context, key string) (int64, error) { + resp, err := c.rpc.Expiretime(c.ctx(ctx), &pb.ExpiretimeRequest{Key: key}) + if err != nil { + return 0, err + } + return resp.Value, nil +} + +// Pexpiretime returns the absolute unix expiry timestamp in milliseconds (-1 = no expiry, -2 = missing). +func (c *Client) Pexpiretime(ctx context.Context, key string) (int64, error) { + resp, err := c.rpc.Pexpiretime(c.ctx(ctx), &pb.PexpiretimeRequest{Key: key}) + if err != nil { + return 0, err + } + return resp.Value, nil +} + +// Expireat sets the expiry at an absolute unix timestamp (seconds). Returns true if set. +func (c *Client) Expireat(ctx context.Context, key string, timestamp uint64) (bool, error) { + resp, err := c.rpc.Expireat(c.ctx(ctx), &pb.ExpireatRequest{Key: key, Timestamp: timestamp}) + if err != nil { + return false, err + } + return resp.Value, nil +} + +// Pexpireat sets the expiry at an absolute unix timestamp (milliseconds). Returns true if set. +func (c *Client) Pexpireat(ctx context.Context, key string, timestampMs uint64) (bool, error) { + resp, err := c.rpc.Pexpireat(c.ctx(ctx), &pb.PexpireatRequest{Key: key, TimestampMs: timestampMs}) + if err != nil { + return false, err + } + return resp.Value, nil +} + +// Getset sets key to value and returns the old value. Returns nil if the key didn't exist. +func (c *Client) Getset(ctx context.Context, key string, value []byte) ([]byte, error) { + resp, err := c.rpc.Getset(c.ctx(ctx), &pb.GetsetRequest{Key: key, Value: value}) + if err != nil { + return nil, err + } + return resp.Value, nil +} + +// Msetnx sets multiple key-value pairs only if none exist. Returns true if all were set. +func (c *Client) Msetnx(ctx context.Context, pairs map[string][]byte) (bool, error) { + kvs := make([]*pb.KeyValue, 0, len(pairs)) + for k, v := range pairs { + kvs = append(kvs, &pb.KeyValue{Key: k, Value: v}) + } + resp, err := c.rpc.Msetnx(c.ctx(ctx), &pb.MsetnxRequest{Pairs: kvs}) + if err != nil { + return false, err + } + return resp.Value, nil +} + +// Getbit returns the bit at offset in the string stored at key. +func (c *Client) Getbit(ctx context.Context, key string, offset uint64) (int64, error) { + resp, err := c.rpc.Getbit(c.ctx(ctx), &pb.GetbitRequest{Key: key, Offset: offset}) + if err != nil { + return 0, err + } + return resp.Value, nil +} + +// Setbit sets or clears the bit at offset. Returns the original bit value. +func (c *Client) Setbit(ctx context.Context, key string, offset uint64, value uint32) (int64, error) { + resp, err := c.rpc.Setbit(c.ctx(ctx), &pb.SetbitRequest{Key: key, Offset: offset, Value: value}) + if err != nil { + return 0, err + } + return resp.Value, nil +} + +// Bitcount counts set bits in the string at key. Pass nil range for the whole string. +func (c *Client) Bitcount(ctx context.Context, key string, start, end int64, unit string, hasRange bool) (int64, error) { + resp, err := c.rpc.Bitcount(c.ctx(ctx), &pb.BitcountRequest{ + Key: key, + HasRange: hasRange, + Start: start, + End: end, + Unit: unit, + }) + if err != nil { + return 0, err + } + return resp.Value, nil +} + +// Bitpos finds the first set or clear bit in the string at key. +func (c *Client) Bitpos(ctx context.Context, key string, bit uint32, start, end int64, unit string, hasRange bool) (int64, error) { + resp, err := c.rpc.Bitpos(c.ctx(ctx), &pb.BitposRequest{ + Key: key, + Bit: bit, + HasRange: hasRange, + Start: start, + End: end, + Unit: unit, + }) + if err != nil { + return 0, err + } + return resp.Value, nil +} + +// Bitop performs a bitwise operation between strings. op is "AND", "OR", "XOR", or "NOT". +func (c *Client) Bitop(ctx context.Context, op, dest string, keys []string) (int64, error) { + resp, err := c.rpc.Bitop(c.ctx(ctx), &pb.BitopRequest{Op: op, Dest: dest, Keys: keys}) + if err != nil { + return 0, err + } + return resp.Value, nil +} + +// Smove atomically moves member from source to destination. Returns true if moved. +func (c *Client) Smove(ctx context.Context, source, destination, member string) (bool, error) { + resp, err := c.rpc.Smove(c.ctx(ctx), &pb.SmoveRequest{ + Source: source, + Destination: destination, + Member: member, + }) + if err != nil { + return false, err + } + return resp.Value, nil +} + +// Sintercard returns the cardinality of the set intersection. limit=0 means no limit. +func (c *Client) Sintercard(ctx context.Context, limit uint64, keys ...string) (int64, error) { + resp, err := c.rpc.Sintercard(c.ctx(ctx), &pb.SintercardRequest{Keys: keys, Limit: limit}) + if err != nil { + return 0, err + } + return resp.Value, nil +} + +// LmpopResult is the result of an LMPOP command. +type LmpopResult struct { + Key string + Elements [][]byte +} + +// Lmpop pops elements from the first non-empty list. left=true for LEFT, false for RIGHT. +// Returns nil if all lists are empty. +func (c *Client) Lmpop(ctx context.Context, left bool, count uint32, keys ...string) (*LmpopResult, error) { + resp, err := c.rpc.Lmpop(c.ctx(ctx), &pb.LmpopRequest{ + Keys: keys, + Left: left, + Count: count, + }) + if err != nil { + return nil, err + } + if !resp.Found { + return nil, nil + } + return &LmpopResult{Key: resp.Key, Elements: resp.Elements}, nil +} + +// ZmpopResult is the result of a ZMPOP command. +type ZmpopResult struct { + Key string + Members []ScoreMemberResult +} + +// ScoreMemberResult is a member-score pair from ZMPOP. +type ScoreMemberResult struct { + Member string + Score float64 +} + +// Zmpop pops elements from the first non-empty sorted set. min=true for MIN, false for MAX. +// Returns nil if all sorted sets are empty. +func (c *Client) Zmpop(ctx context.Context, min bool, count uint32, keys ...string) (*ZmpopResult, error) { + resp, err := c.rpc.Zmpop(c.ctx(ctx), &pb.ZmpopRequest{ + Keys: keys, + Min: min, + Count: count, + }) + if err != nil { + return nil, err + } + if !resp.Found { + return nil, nil + } + members := make([]ScoreMemberResult, len(resp.Members)) + for i, m := range resp.Members { + members[i] = ScoreMemberResult{Member: m.Member, Score: m.Score} + } + return &ZmpopResult{Key: resp.Key, Members: members}, nil +} + +// Hrandfield returns random fields from the hash at key. +// count=nil returns a single field; positive = distinct fields, negative = allow repeats. +func (c *Client) Hrandfield(ctx context.Context, key string, count *int32, withValues bool) ([][]byte, error) { + req := &pb.HrandfieldRequest{Key: key, WithValues: withValues} + if count != nil { + req.HasCount = true + req.Count = *count + } + resp, err := c.rpc.Hrandfield(c.ctx(ctx), req) + if err != nil { + return nil, err + } + result := make([][]byte, len(resp.Values)) + for i, v := range resp.Values { + result[i] = v + } + return result, nil +} + +// Zrandmember returns random members from the sorted set at key. +// count=nil returns a single member; positive = distinct members, negative = allow repeats. +func (c *Client) Zrandmember(ctx context.Context, key string, count *int32, withScores bool) ([][]byte, error) { + req := &pb.ZrandmemberRequest{Key: key, WithScores: withScores} + if count != nil { + req.HasCount = true + req.Count = *count + } + resp, err := c.rpc.Zrandmember(c.ctx(ctx), req) + if err != nil { + return nil, err + } + result := make([][]byte, len(resp.Values)) + for i, v := range resp.Values { + result[i] = v + } + return result, nil +} diff --git a/clients/ember-py/ember/client.py b/clients/ember-py/ember/client.py index f84d65d1..0a4c601c 100644 --- a/clients/ember-py/ember/client.py +++ b/clients/ember-py/ember/client.py @@ -505,3 +505,164 @@ def pubsub_numpat(self) -> int: metadata=self._metadata(), ) return resp.value + + # --- expiry --- + + def expiretime(self, key: str) -> int: + """Get the absolute unix expiry timestamp in seconds (-1 = no expiry, -2 = missing).""" + resp = self._stub.Expiretime( + ember_pb2.ExpiretimeRequest(key=key), + metadata=self._metadata(), + ) + return resp.value + + def pexpiretime(self, key: str) -> int: + """Get the absolute unix expiry timestamp in milliseconds.""" + resp = self._stub.Pexpiretime( + ember_pb2.PexpiretimeRequest(key=key), + metadata=self._metadata(), + ) + return resp.value + + def expireat(self, key: str, timestamp: int) -> bool: + """Set expiry at an absolute unix timestamp (seconds).""" + resp = self._stub.Expireat( + ember_pb2.ExpireatRequest(key=key, timestamp=timestamp), + metadata=self._metadata(), + ) + return resp.value + + def pexpireat(self, key: str, timestamp_ms: int) -> bool: + """Set expiry at an absolute unix timestamp (milliseconds).""" + resp = self._stub.Pexpireat( + ember_pb2.PexpireatRequest(key=key, timestamp_ms=timestamp_ms), + metadata=self._metadata(), + ) + return resp.value + + # --- strings --- + + def getset(self, key: str, value: bytes) -> bytes | None: + """Set key to value and return the old value, or None if the key didn't exist.""" + resp = self._stub.Getset( + ember_pb2.GetsetRequest(key=key, value=value), + metadata=self._metadata(), + ) + return resp.value if resp.HasField("value") else None + + def msetnx(self, pairs: dict[str, bytes]) -> bool: + """Set multiple keys only if none exist. Returns True if all were set.""" + kvs = [ember_pb2.KeyValue(key=k, value=v) for k, v in pairs.items()] + resp = self._stub.Msetnx( + ember_pb2.MsetnxRequest(pairs=kvs), + metadata=self._metadata(), + ) + return resp.value + + # --- bitmaps --- + + def getbit(self, key: str, offset: int) -> int: + """Return the bit at offset in the string stored at key.""" + resp = self._stub.Getbit( + ember_pb2.GetbitRequest(key=key, offset=offset), + metadata=self._metadata(), + ) + return resp.value + + def setbit(self, key: str, offset: int, value: int) -> int: + """Set or clear the bit at offset. Returns the original bit value.""" + resp = self._stub.Setbit( + ember_pb2.SetbitRequest(key=key, offset=offset, value=value), + metadata=self._metadata(), + ) + return resp.value + + def bitcount(self, key: str, start: int = 0, end: int = -1, unit: str = "BYTE", has_range: bool = False) -> int: + """Count set bits in the string. Set has_range=True to use start/end/unit.""" + resp = self._stub.Bitcount( + ember_pb2.BitcountRequest( + key=key, has_range=has_range, start=start, end=end, unit=unit + ), + metadata=self._metadata(), + ) + return resp.value + + def bitpos(self, key: str, bit: int, start: int = 0, end: int = -1, unit: str = "BYTE", has_range: bool = False) -> int: + """Find the first set or clear bit. Set has_range=True to use start/end/unit.""" + resp = self._stub.Bitpos( + ember_pb2.BitposRequest( + key=key, bit=bit, has_range=has_range, start=start, end=end, unit=unit + ), + metadata=self._metadata(), + ) + return resp.value + + def bitop(self, op: str, dest: str, *keys: str) -> int: + """Perform a bitwise operation. op is 'AND', 'OR', 'XOR', or 'NOT'.""" + resp = self._stub.Bitop( + ember_pb2.BitopRequest(op=op, dest=dest, keys=list(keys)), + metadata=self._metadata(), + ) + return resp.value + + # --- sets --- + + def smove(self, source: str, destination: str, member: str) -> bool: + """Atomically move member from source to destination.""" + resp = self._stub.Smove( + ember_pb2.SmoveRequest(source=source, destination=destination, member=member), + metadata=self._metadata(), + ) + return resp.value + + def sintercard(self, *keys: str, limit: int = 0) -> int: + """Return the cardinality of the set intersection. limit=0 means no limit.""" + resp = self._stub.Sintercard( + ember_pb2.SintercardRequest(keys=list(keys), limit=limit), + metadata=self._metadata(), + ) + return resp.value + + # --- lists --- + + def lmpop(self, *keys: str, left: bool = True, count: int = 1) -> tuple[str, list[bytes]] | None: + """Pop elements from the first non-empty list. Returns (key, elements) or None.""" + resp = self._stub.Lmpop( + ember_pb2.LmpopRequest(keys=list(keys), left=left, count=count), + metadata=self._metadata(), + ) + if not resp.found: + return None + return resp.key, list(resp.elements) + + # --- hash --- + + def hrandfield(self, key: str, count: int | None = None, with_values: bool = False) -> list[bytes]: + """Get random field(s) from the hash. Returns a flat list (interleaved with values if with_values=True).""" + req = ember_pb2.HrandfieldRequest(key=key, with_values=with_values) + if count is not None: + req.has_count = True + req.count = count + resp = self._stub.Hrandfield(req, metadata=self._metadata()) + return list(resp.values) + + # --- sorted sets --- + + def zmpop(self, *keys: str, min: bool = True, count: int = 1) -> tuple[str, list[tuple[str, float]]] | None: + """Pop elements from the first non-empty sorted set. Returns (key, [(member, score)]) or None.""" + resp = self._stub.Zmpop( + ember_pb2.ZmpopRequest(keys=list(keys), min=min, count=count), + metadata=self._metadata(), + ) + if not resp.found: + return None + return resp.key, [(m.member, m.score) for m in resp.members] + + def zrandmember(self, key: str, count: int | None = None, with_scores: bool = False) -> list[bytes]: + """Get random member(s) from the sorted set. Returns a flat list (interleaved with scores if with_scores=True).""" + req = ember_pb2.ZrandmemberRequest(key=key, with_scores=with_scores) + if count is not None: + req.has_count = True + req.count = count + resp = self._stub.Zrandmember(req, metadata=self._metadata()) + return list(resp.values) diff --git a/clients/ember-ts/src/client.ts b/clients/ember-ts/src/client.ts index afb17733..e46707ee 100644 --- a/clients/ember-ts/src/client.ts +++ b/clients/ember-ts/src/client.ts @@ -1407,4 +1407,236 @@ export class EmberClient { async slowLogReset(): Promise { await this.call('slowLogReset', {}); } + + // --------------------------------------------------------------------------- + // expiry (extended) + // --------------------------------------------------------------------------- + + /** + * Returns the absolute unix expiry timestamp in seconds. + * Returns -1 if no expiry, -2 if the key does not exist. + */ + async expiretime(key: string): Promise { + const res = await this.call<{ value: number }>('expiretime', { key }); + return res.value; + } + + /** + * Returns the absolute unix expiry timestamp in milliseconds. + * Returns -1 if no expiry, -2 if the key does not exist. + */ + async pexpiretime(key: string): Promise { + const res = await this.call<{ value: number }>('pexpiretime', { key }); + return res.value; + } + + /** + * Sets the expiry to an absolute unix timestamp (seconds). + * Returns `true` if the timeout was set. + */ + async expireat(key: string, timestamp: number): Promise { + const res = await this.call<{ value: boolean }>('expireat', { key, timestamp }); + return res.value; + } + + /** + * Sets the expiry to an absolute unix timestamp (milliseconds). + * Returns `true` if the timeout was set. + */ + async pexpireat(key: string, timestampMs: number): Promise { + const res = await this.call<{ value: boolean }>('pexpireat', { key, timestampMs }); + return res.value; + } + + // --------------------------------------------------------------------------- + // strings (extended) + // --------------------------------------------------------------------------- + + /** + * Atomically sets `key` to `value` and returns the old value. + * Returns `null` if the key did not exist. + */ + async getset(key: string, value: Buffer | string): Promise { + const res = await this.call<{ value?: Buffer }>('getset', { + key, + value: Buffer.isBuffer(value) ? value : Buffer.from(value), + }); + return res.value ?? null; + } + + /** + * Sets multiple keys only if none of them exist. + * Returns `true` if all keys were set, `false` if any existed. + */ + async msetnx(pairs: Array<{ key: string; value: Buffer | string }>): Promise { + const res = await this.call<{ value: boolean }>('msetnx', { + pairs: pairs.map(p => ({ + key: p.key, + value: Buffer.isBuffer(p.value) ? p.value : Buffer.from(p.value), + })), + }); + return res.value; + } + + // --------------------------------------------------------------------------- + // bitmaps + // --------------------------------------------------------------------------- + + /** + * Returns the bit at `offset` in the string stored at `key`. + */ + async getbit(key: string, offset: number): Promise { + const res = await this.call<{ value: number }>('getbit', { key, offset }); + return res.value; + } + + /** + * Sets or clears the bit at `offset`. Returns the original bit value. + */ + async setbit(key: string, offset: number, value: 0 | 1): Promise { + const res = await this.call<{ value: number }>('setbit', { key, offset, value }); + return res.value; + } + + /** + * Counts the number of set bits in the string at `key`. + * Pass `range` as `{ start, end, unit }` to count within a range. + */ + async bitcount(key: string, range?: { start: number; end: number; unit?: 'BYTE' | 'BIT' }): Promise { + const req: Record = { key }; + if (range) { + req.hasRange = true; + req.start = range.start; + req.end = range.end; + req.unit = range.unit ?? 'BYTE'; + } + const res = await this.call<{ value: number }>('bitcount', req); + return res.value; + } + + /** + * Finds the first set (`bit=1`) or clear (`bit=0`) bit in the string at `key`. + */ + async bitpos(key: string, bit: 0 | 1, range?: { start: number; end: number; unit?: 'BYTE' | 'BIT' }): Promise { + const req: Record = { key, bit }; + if (range) { + req.hasRange = true; + req.start = range.start; + req.end = range.end; + req.unit = range.unit ?? 'BYTE'; + } + const res = await this.call<{ value: number }>('bitpos', req); + return res.value; + } + + /** + * Performs a bitwise operation between strings. + * `op` is `"AND"`, `"OR"`, `"XOR"`, or `"NOT"`. + * Result is stored at `dest`. Returns the length of the resulting string. + */ + async bitop(op: 'AND' | 'OR' | 'XOR' | 'NOT', dest: string, keys: string[]): Promise { + const res = await this.call<{ value: number }>('bitop', { op, dest, keys }); + return res.value; + } + + // --------------------------------------------------------------------------- + // sets (extended) + // --------------------------------------------------------------------------- + + /** + * Atomically moves `member` from `source` to `destination`. + * Returns `true` if the move succeeded. + */ + async smove(source: string, destination: string, member: string): Promise { + const res = await this.call<{ value: boolean }>('smove', { source, destination, member }); + return res.value; + } + + /** + * Returns the cardinality of the intersection of `keys`. + * `limit=0` means no limit. + */ + async sintercard(keys: string[], limit = 0): Promise { + const res = await this.call<{ value: number }>('sintercard', { keys, limit }); + return res.value; + } + + // --------------------------------------------------------------------------- + // lists (extended) + // --------------------------------------------------------------------------- + + /** + * Pops up to `count` elements from the first non-empty list in `keys`. + * `left=true` pops from the head, `left=false` from the tail. + * Returns `null` if all lists are empty. + */ + async lmpop( + keys: string[], + left: boolean, + count = 1, + ): Promise<{ key: string; elements: Buffer[] } | null> { + const res = await this.call<{ found: boolean; key: string; elements: Buffer[] }>('lmpop', { + keys, + left, + count, + }); + if (!res.found) return null; + return { key: res.key, elements: res.elements }; + } + + // --------------------------------------------------------------------------- + // hash (extended) + // --------------------------------------------------------------------------- + + /** + * Returns random field(s) from the hash at `key`. + * `count=undefined` returns a single field. + * `withValues=true` returns interleaved field-value pairs. + */ + async hrandfield(key: string, count?: number, withValues = false): Promise { + const req: Record = { key, withValues }; + if (count != null) { + req.hasCount = true; + req.count = count; + } + const res = await this.call<{ values: Buffer[] }>('hrandfield', req); + return res.values ?? []; + } + + // --------------------------------------------------------------------------- + // sorted sets (extended) + // --------------------------------------------------------------------------- + + /** + * Pops up to `count` elements from the first non-empty sorted set in `keys`. + * `min=true` pops minimum-score members. Returns `null` if all sorted sets are empty. + */ + async zmpop( + keys: string[], + min: boolean, + count = 1, + ): Promise<{ key: string; members: Array<{ member: string; score: number }> } | null> { + const res = await this.call<{ + found: boolean; + key: string; + members: Array<{ member: string; score: number }>; + }>('zmpop', { keys, min, count }); + if (!res.found) return null; + return { key: res.key, members: res.members }; + } + + /** + * Returns random member(s) from the sorted set at `key`. + * `count=undefined` returns a single member. + * `withScores=true` returns interleaved member-score pairs. + */ + async zrandmember(key: string, count?: number, withScores = false): Promise { + const req: Record = { key, withScores }; + if (count != null) { + req.hasCount = true; + req.count = count; + } + const res = await this.call<{ values: Buffer[] }>('zrandmember', req); + return res.values ?? []; + } } diff --git a/crates/ember-cli/src/commands.rs b/crates/ember-cli/src/commands.rs index fa87a3b0..556e032e 100644 --- a/crates/ember-cli/src/commands.rs +++ b/crates/ember-cli/src/commands.rs @@ -85,6 +85,12 @@ pub static COMMANDS: &[CommandInfo] = &[ group: "string", summary: "get a substring of the string stored at a key", }, + CommandInfo { + name: "GETSET", + args: "key value", + group: "string", + summary: "set a key's value and return its old value", + }, CommandInfo { name: "INCR", args: "key", @@ -115,6 +121,12 @@ pub static COMMANDS: &[CommandInfo] = &[ group: "string", summary: "set multiple keys to multiple values", }, + CommandInfo { + name: "MSETNX", + args: "key value [key value ...]", + group: "string", + summary: "set multiple keys only if none of them exist", + }, CommandInfo { name: "PSETEX", args: "key milliseconds value", @@ -157,6 +169,37 @@ pub static COMMANDS: &[CommandInfo] = &[ group: "string", summary: "get a substring of the string stored at a key (alias for GETRANGE)", }, + // --- bitmap --- + CommandInfo { + name: "BITCOUNT", + args: "key [start end [BYTE|BIT]]", + group: "bitmap", + summary: "count set bits in a string", + }, + CommandInfo { + name: "BITOP", + args: "AND|OR|XOR|NOT destkey key [key ...]", + group: "bitmap", + summary: "perform bitwise operations between strings", + }, + CommandInfo { + name: "BITPOS", + args: "key bit [start [end [BYTE|BIT]]]", + group: "bitmap", + summary: "find first set or clear bit in a string", + }, + CommandInfo { + name: "GETBIT", + args: "key offset", + group: "bitmap", + summary: "return the bit value at offset in the string", + }, + CommandInfo { + name: "SETBIT", + args: "key offset value", + group: "bitmap", + summary: "set or clear the bit at offset in the string", + }, // --- generic --- CommandInfo { name: "COPY", @@ -182,6 +225,18 @@ pub static COMMANDS: &[CommandInfo] = &[ group: "generic", summary: "set a key's time to live in seconds", }, + CommandInfo { + name: "EXPIREAT", + args: "key timestamp", + group: "generic", + summary: "set expiry at an absolute unix timestamp (seconds)", + }, + CommandInfo { + name: "EXPIRETIME", + args: "key", + group: "generic", + summary: "get the absolute unix expiry timestamp of a key in seconds", + }, CommandInfo { name: "KEYS", args: "pattern", @@ -206,6 +261,18 @@ pub static COMMANDS: &[CommandInfo] = &[ group: "generic", summary: "set a key's time to live in milliseconds", }, + CommandInfo { + name: "PEXPIREAT", + args: "key timestamp-ms", + group: "generic", + summary: "set expiry at an absolute unix timestamp (milliseconds)", + }, + CommandInfo { + name: "PEXPIRETIME", + args: "key", + group: "generic", + summary: "get the absolute unix expiry timestamp of a key in milliseconds", + }, CommandInfo { name: "PTTL", args: "key", @@ -291,9 +358,15 @@ pub static COMMANDS: &[CommandInfo] = &[ group: "list", summary: "get the length of a list", }, + CommandInfo { + name: "LMPOP", + args: "numkeys key [key ...] LEFT|RIGHT [COUNT n]", + group: "list", + summary: "pop elements from the first non-empty list", + }, CommandInfo { name: "LPOP", - args: "key", + args: "key [count]", group: "list", summary: "remove and return the first element of a list", }, @@ -335,7 +408,7 @@ pub static COMMANDS: &[CommandInfo] = &[ }, CommandInfo { name: "RPOP", - args: "key", + args: "key [count]", group: "list", summary: "remove and return the last element of a list", }, @@ -394,6 +467,12 @@ pub static COMMANDS: &[CommandInfo] = &[ group: "hash", summary: "get the values of multiple hash fields", }, + CommandInfo { + name: "HRANDFIELD", + args: "key [count [WITHVALUES]]", + group: "hash", + summary: "get one or more random fields from a hash", + }, CommandInfo { name: "HSCAN", args: "key cursor [MATCH pattern] [COUNT count]", @@ -425,6 +504,12 @@ pub static COMMANDS: &[CommandInfo] = &[ group: "set", summary: "get the number of members in a set", }, + CommandInfo { + name: "SINTERCARD", + args: "numkeys key [key ...] [LIMIT count]", + group: "set", + summary: "intersect multiple sets and return the cardinality", + }, CommandInfo { name: "SISMEMBER", args: "key member", @@ -437,6 +522,12 @@ pub static COMMANDS: &[CommandInfo] = &[ group: "set", summary: "get all members in a set", }, + CommandInfo { + name: "SMOVE", + args: "source destination member", + group: "set", + summary: "move a member from one set to another", + }, CommandInfo { name: "SREM", args: "key member [member ...]", @@ -462,6 +553,18 @@ pub static COMMANDS: &[CommandInfo] = &[ group: "sorted_set", summary: "get the number of members in a sorted set", }, + CommandInfo { + name: "ZMPOP", + args: "numkeys key [key ...] MIN|MAX [COUNT n]", + group: "sorted_set", + summary: "pop elements from the first non-empty sorted set", + }, + CommandInfo { + name: "ZRANDMEMBER", + args: "key [count [WITHSCORES]]", + group: "sorted_set", + summary: "get one or more random members from a sorted set", + }, CommandInfo { name: "ZRANGE", args: "key start stop [WITHSCORES]", diff --git a/crates/ember-client/src/commands.rs b/crates/ember-client/src/commands.rs index 03a2ad42..80ab6c9f 100644 --- a/crates/ember-client/src/commands.rs +++ b/crates/ember-client/src/commands.rs @@ -1297,6 +1297,389 @@ impl Client { Ok(Subscriber::new(self)) } + // --- expiry commands --- + + /// Returns the absolute unix expiry timestamp in seconds, or -1 if no expiry, -2 if key missing. + pub async fn expiretime(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"EXPIRETIME", key.as_bytes())).await?; + integer(frame) + } + + /// Returns the absolute unix expiry timestamp in milliseconds, or -1 if no expiry, -2 if key missing. + pub async fn pexpiretime(&mut self, key: &str) -> Result { + let frame = self.send_frame(cmd2(b"PEXPIRETIME", key.as_bytes())).await?; + integer(frame) + } + + /// Sets expiry at an absolute unix timestamp (seconds). Returns `true` if the timeout was set. + pub async fn expireat(&mut self, key: &str, timestamp: u64) -> Result { + let ts = timestamp.to_string(); + let frame = self + .send_frame(cmd3(b"EXPIREAT", key.as_bytes(), ts.as_bytes())) + .await?; + bool_flag(frame) + } + + /// Sets expiry at an absolute unix timestamp (milliseconds). Returns `true` if the timeout was set. + pub async fn pexpireat(&mut self, key: &str, timestamp_ms: u64) -> Result { + let ts = timestamp_ms.to_string(); + let frame = self + .send_frame(cmd3(b"PEXPIREAT", key.as_bytes(), ts.as_bytes())) + .await?; + bool_flag(frame) + } + + // --- string commands (extended) --- + + /// Sets `key` to `value`, returning the old value. Returns `None` if the key didn't exist. + pub async fn getset( + &mut self, + key: &str, + value: impl AsRef<[u8]>, + ) -> Result, ClientError> { + let frame = self + .send_frame(cmd3(b"GETSET", key.as_bytes(), value.as_ref())) + .await?; + optional_bytes(frame) + } + + /// Sets multiple key-value pairs only if none of the keys already exist. + /// Returns `true` if all keys were set, `false` if any key existed. + pub async fn msetnx>( + &mut self, + pairs: &[(&str, V)], + ) -> Result { + let mut parts = Vec::with_capacity(1 + pairs.len() * 2); + parts.push(Frame::Bulk(Bytes::from_static(b"MSETNX"))); + for (k, v) in pairs { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(v.as_ref()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + bool_flag(frame) + } + + // --- bitmap commands --- + + /// Returns the bit value at `offset` in the string stored at `key`. + pub async fn getbit(&mut self, key: &str, offset: u64) -> Result { + let off = offset.to_string(); + let frame = self + .send_frame(cmd3(b"GETBIT", key.as_bytes(), off.as_bytes())) + .await?; + integer(frame) + } + + /// Sets or clears the bit at `offset`. Returns the original bit value. + pub async fn setbit(&mut self, key: &str, offset: u64, value: u8) -> Result { + let off = offset.to_string(); + let val = value.to_string(); + let frame = self + .send_frame(cmd4(b"SETBIT", key.as_bytes(), off.as_bytes(), val.as_bytes())) + .await?; + integer(frame) + } + + /// Counts set bits in the string at `key`. + /// + /// `range`: optional `(start, end, unit)` where unit is `"BYTE"` or `"BIT"`. + pub async fn bitcount( + &mut self, + key: &str, + range: Option<(i64, i64, &str)>, + ) -> Result { + let frame = if let Some((start, end, unit)) = range { + let s = start.to_string(); + let e = end.to_string(); + let mut parts = Vec::with_capacity(5); + parts.push(Frame::Bulk(Bytes::from_static(b"BITCOUNT"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(s.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(e.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(unit.as_bytes()))); + self.send_frame(Frame::Array(parts)).await? + } else { + self.send_frame(cmd2(b"BITCOUNT", key.as_bytes())).await? + }; + integer(frame) + } + + /// Finds the first set (`bit=1`) or clear (`bit=0`) bit in the string at `key`. + /// + /// `range`: optional `(start, end, unit)` where unit is `"BYTE"` or `"BIT"`. + pub async fn bitpos( + &mut self, + key: &str, + bit: u8, + range: Option<(i64, i64, &str)>, + ) -> Result { + let b = bit.to_string(); + let frame = if let Some((start, end, unit)) = range { + let s = start.to_string(); + let e = end.to_string(); + let mut parts = Vec::with_capacity(6); + parts.push(Frame::Bulk(Bytes::from_static(b"BITPOS"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(key.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(b.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(s.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(e.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(unit.as_bytes()))); + self.send_frame(Frame::Array(parts)).await? + } else { + self.send_frame(cmd3(b"BITPOS", key.as_bytes(), b.as_bytes())) + .await? + }; + integer(frame) + } + + /// Performs a bitwise operation between strings. + /// + /// `op` is `"AND"`, `"OR"`, `"XOR"`, or `"NOT"`. Result is stored at `dest`. + /// Returns the length of the resulting string. + pub async fn bitop( + &mut self, + op: &str, + dest: &str, + keys: &[&str], + ) -> Result { + let mut parts = Vec::with_capacity(3 + keys.len()); + parts.push(Frame::Bulk(Bytes::from_static(b"BITOP"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(op.as_bytes()))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(dest.as_bytes()))); + for k in keys { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + integer(frame) + } + + // --- set commands (extended) --- + + /// Atomically moves `member` from `src` to `dst`. Returns `true` if the move succeeded. + pub async fn smove( + &mut self, + src: &str, + dst: &str, + member: &str, + ) -> Result { + let frame = self + .send_frame(cmd4(b"SMOVE", src.as_bytes(), dst.as_bytes(), member.as_bytes())) + .await?; + bool_flag(frame) + } + + /// Returns the cardinality of the intersection of multiple sets. `limit=0` means no limit. + pub async fn sintercard(&mut self, keys: &[&str], limit: usize) -> Result { + let n = keys.len().to_string(); + let lim = limit.to_string(); + let mut parts = Vec::with_capacity(3 + keys.len() + 2); + parts.push(Frame::Bulk(Bytes::from_static(b"SINTERCARD"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(n.as_bytes()))); + for k in keys { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + } + if limit > 0 { + parts.push(Frame::Bulk(Bytes::from_static(b"LIMIT"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(lim.as_bytes()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + integer(frame) + } + + // --- list commands (extended) --- + + /// Pops up to `count` elements from the first non-empty list in `keys`. + /// + /// `left=true` pops from the head, `left=false` from the tail. + /// Returns `None` if all lists are empty, or `Some((key, elements))`. + pub async fn lmpop( + &mut self, + keys: &[&str], + left: bool, + count: usize, + ) -> Result)>, ClientError> { + let n = keys.len().to_string(); + let dir = if left { b"LEFT" as &[u8] } else { b"RIGHT" }; + let cnt = count.to_string(); + let mut parts = Vec::with_capacity(4 + keys.len() + 2); + parts.push(Frame::Bulk(Bytes::from_static(b"LMPOP"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(n.as_bytes()))); + for k in keys { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + } + parts.push(Frame::Bulk(Bytes::copy_from_slice(dir))); + if count > 0 { + parts.push(Frame::Bulk(Bytes::from_static(b"COUNT"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(cnt.as_bytes()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + match frame { + Frame::Null => Ok(None), + Frame::Array(mut elems) if elems.len() == 2 => { + let key = match elems.remove(0) { + Frame::Bulk(b) => String::from_utf8(b.to_vec()) + .map_err(|_| ClientError::Protocol("key is not valid UTF-8".into()))?, + other => { + return Err(ClientError::Protocol(format!( + "expected bulk key in LMPOP response, got {other:?}" + ))) + } + }; + let elements = bytes_vec(elems.remove(0))?; + Ok(Some((key, elements))) + } + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "unexpected LMPOP response: {other:?}" + ))), + } + } + + // --- hash commands (extended) --- + + /// Returns random field(s) from the hash at `key`. + /// + /// `count=None` returns a single field name as a one-element Vec. + /// Positive count returns that many distinct fields; negative allows repeats. + pub async fn hrandfield( + &mut self, + key: &str, + count: Option, + ) -> Result, ClientError> { + let frame = if let Some(n) = count { + let s = n.to_string(); + self.send_frame(cmd3(b"HRANDFIELD", key.as_bytes(), s.as_bytes())) + .await? + } else { + self.send_frame(cmd2(b"HRANDFIELD", key.as_bytes())).await? + }; + match frame { + Frame::Bulk(b) => Ok(vec![b]), + Frame::Null => Ok(vec![]), + Frame::Array(_) => bytes_vec(frame), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "unexpected HRANDFIELD response: {other:?}" + ))), + } + } + + /// Returns random field-value pairs from the hash at `key`. + /// + /// `count` controls how many pairs to return (positive = distinct, negative = allow repeats). + pub async fn hrandfield_withvalues( + &mut self, + key: &str, + count: i64, + ) -> Result, ClientError> { + let s = count.to_string(); + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"HRANDFIELD")), + Frame::Bulk(Bytes::copy_from_slice(key.as_bytes())), + Frame::Bulk(Bytes::copy_from_slice(s.as_bytes())), + Frame::Bulk(Bytes::from_static(b"WITHVALUES")), + ])) + .await?; + pairs(frame) + } + + // --- sorted set commands (extended) --- + + /// Pops up to `count` elements from the first non-empty sorted set in `keys`. + /// + /// `min=true` pops minimum-score members, `min=false` pops maximum-score members. + /// Returns `None` if all sorted sets are empty, or `Some((key, members))`. + pub async fn zmpop( + &mut self, + keys: &[&str], + min: bool, + count: usize, + ) -> Result)>, ClientError> { + let n = keys.len().to_string(); + let dir = if min { b"MIN" as &[u8] } else { b"MAX" }; + let cnt = count.to_string(); + let mut parts = Vec::with_capacity(4 + keys.len() + 2); + parts.push(Frame::Bulk(Bytes::from_static(b"ZMPOP"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(n.as_bytes()))); + for k in keys { + parts.push(Frame::Bulk(Bytes::copy_from_slice(k.as_bytes()))); + } + parts.push(Frame::Bulk(Bytes::copy_from_slice(dir))); + if count > 0 { + parts.push(Frame::Bulk(Bytes::from_static(b"COUNT"))); + parts.push(Frame::Bulk(Bytes::copy_from_slice(cnt.as_bytes()))); + } + let frame = self.send_frame(Frame::Array(parts)).await?; + match frame { + Frame::Null => Ok(None), + Frame::Array(mut elems) if elems.len() == 2 => { + let key = match elems.remove(0) { + Frame::Bulk(b) => String::from_utf8(b.to_vec()) + .map_err(|_| ClientError::Protocol("key is not valid UTF-8".into()))?, + other => { + return Err(ClientError::Protocol(format!( + "expected bulk key in ZMPOP response, got {other:?}" + ))) + } + }; + let members = scored_members(elems.remove(0))?; + Ok(Some((key, members))) + } + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "unexpected ZMPOP response: {other:?}" + ))), + } + } + + /// Returns random member(s) from the sorted set at `key`. + /// + /// `count=None` returns a single member name. Positive count returns distinct members; + /// negative allows repeats. + pub async fn zrandmember( + &mut self, + key: &str, + count: Option, + ) -> Result, ClientError> { + let frame = if let Some(n) = count { + let s = n.to_string(); + self.send_frame(cmd3(b"ZRANDMEMBER", key.as_bytes(), s.as_bytes())) + .await? + } else { + self.send_frame(cmd2(b"ZRANDMEMBER", key.as_bytes())).await? + }; + match frame { + Frame::Bulk(b) => Ok(vec![b]), + Frame::Null => Ok(vec![]), + Frame::Array(_) => bytes_vec(frame), + Frame::Error(e) => Err(ClientError::Server(e)), + other => Err(ClientError::Protocol(format!( + "unexpected ZRANDMEMBER response: {other:?}" + ))), + } + } + + /// Returns random member-score pairs from the sorted set at `key`. + /// + /// `count` controls how many pairs to return (positive = distinct, negative = allow repeats). + pub async fn zrandmember_withscores( + &mut self, + key: &str, + count: i64, + ) -> Result, ClientError> { + let s = count.to_string(); + let frame = self + .send_frame(Frame::Array(vec![ + Frame::Bulk(Bytes::from_static(b"ZRANDMEMBER")), + Frame::Bulk(Bytes::copy_from_slice(key.as_bytes())), + Frame::Bulk(Bytes::copy_from_slice(s.as_bytes())), + Frame::Bulk(Bytes::from_static(b"WITHSCORES")), + ])) + .await?; + scored_members(frame) + } + // --- pipeline --- /// Executes all commands queued in `pipeline` as a single batch. diff --git a/crates/ember-server/src/grpc.rs b/crates/ember-server/src/grpc.rs index 99c7c40c..2d8852e6 100644 --- a/crates/ember-server/src/grpc.rs +++ b/crates/ember-server/src/grpc.rs @@ -11,7 +11,7 @@ use std::time::{Duration, Instant}; use bytes::Bytes; use ember_core::{Engine, ShardRequest, ShardResponse, TtlResult, Value}; -use ember_protocol::command::ScoreBound; +use ember_protocol::command::{BitOpKind, BitRange, BitRangeUnit, ScoreBound}; use subtle::ConstantTimeEq; use tokio_stream::wrappers::ReceiverStream; use tonic::service::interceptor::InterceptedService; @@ -3273,6 +3273,628 @@ impl EmberCache for EmberService { })) } + // ----------------------------------------------------------------------- + // keys (new) + // ----------------------------------------------------------------------- + + async fn expiretime( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + let resp = self + .route( + &req.key, + ShardRequest::Expiretime { key: req.key.clone() }, + ) + .await?; + self.record_command(start, "EXPIRETIME"); + + match resp { + ShardResponse::Integer(n) => Ok(Response::new(IntResponse { value: n })), + other => Err(unexpected_response(&other)), + } + } + + async fn pexpiretime( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + let resp = self + .route( + &req.key, + ShardRequest::Pexpiretime { key: req.key.clone() }, + ) + .await?; + self.record_command(start, "PEXPIRETIME"); + + match resp { + ShardResponse::Integer(n) => Ok(Response::new(IntResponse { value: n })), + other => Err(unexpected_response(&other)), + } + } + + async fn expireat( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + let resp = self + .route( + &req.key, + ShardRequest::Expireat { + key: req.key.clone(), + timestamp: req.timestamp, + }, + ) + .await?; + self.record_command(start, "EXPIREAT"); + + match resp { + ShardResponse::Bool(v) => Ok(Response::new(BoolResponse { value: v })), + other => Err(unexpected_response(&other)), + } + } + + async fn pexpireat( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + let resp = self + .route( + &req.key, + ShardRequest::Pexpireat { + key: req.key.clone(), + timestamp_ms: req.timestamp_ms, + }, + ) + .await?; + self.record_command(start, "PEXPIREAT"); + + match resp { + ShardResponse::Bool(v) => Ok(Response::new(BoolResponse { value: v })), + other => Err(unexpected_response(&other)), + } + } + + // ----------------------------------------------------------------------- + // strings (new) + // ----------------------------------------------------------------------- + + async fn getset( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + let resp = self + .route( + &req.key, + ShardRequest::GetSet { + key: req.key.clone(), + value: req.value.into(), + }, + ) + .await?; + self.record_command(start, "GETSET"); + + match resp { + ShardResponse::Value(opt) => Ok(Response::new(GetResponse { + value: opt.and_then(|v| match v { + Value::String(b) => Some(b.to_vec()), + _ => None, + }), + })), + ShardResponse::WrongType => Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )), + other => Err(unexpected_response(&other)), + } + } + + async fn msetnx( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + if req.pairs.is_empty() { + return Err(Status::invalid_argument("at least one pair required")); + } + for p in &req.pairs { + validate_key(&p.key, &self.ctx.limits)?; + } + let pairs: Vec<(String, Bytes)> = req + .pairs + .into_iter() + .map(|kv| (kv.key, kv.value.into())) + .collect(); + let first_key = pairs[0].0.clone(); + let resp = self + .route(&first_key, ShardRequest::MSetNx { pairs }) + .await?; + self.record_command(start, "MSETNX"); + + match resp { + ShardResponse::Bool(v) => Ok(Response::new(BoolResponse { value: v })), + other => Err(unexpected_response(&other)), + } + } + + // ----------------------------------------------------------------------- + // bitmaps + // ----------------------------------------------------------------------- + + async fn getbit( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + let resp = self + .route( + &req.key, + ShardRequest::GetBit { + key: req.key.clone(), + offset: req.offset, + }, + ) + .await?; + self.record_command(start, "GETBIT"); + + match resp { + ShardResponse::Integer(n) => Ok(Response::new(IntResponse { value: n })), + ShardResponse::WrongType => Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )), + other => Err(unexpected_response(&other)), + } + } + + async fn setbit( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + if req.value > 1 { + return Err(Status::invalid_argument("bit value must be 0 or 1")); + } + let resp = self + .route( + &req.key, + ShardRequest::SetBit { + key: req.key.clone(), + offset: req.offset, + value: req.value as u8, + }, + ) + .await?; + self.record_command(start, "SETBIT"); + + match resp { + ShardResponse::Integer(n) => Ok(Response::new(IntResponse { value: n })), + ShardResponse::WrongType => Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )), + other => Err(unexpected_response(&other)), + } + } + + async fn bitcount( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + let range = if req.has_range { + let unit = if req.unit == "BIT" { + BitRangeUnit::Bit + } else { + BitRangeUnit::Byte + }; + Some(BitRange { + start: req.start, + end: req.end, + unit, + }) + } else { + None + }; + let resp = self + .route( + &req.key, + ShardRequest::BitCount { + key: req.key.clone(), + range, + }, + ) + .await?; + self.record_command(start, "BITCOUNT"); + + match resp { + ShardResponse::Integer(n) => Ok(Response::new(IntResponse { value: n })), + ShardResponse::WrongType => Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )), + other => Err(unexpected_response(&other)), + } + } + + async fn bitpos( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + if req.bit > 1 { + return Err(Status::invalid_argument("bit argument must be 0 or 1")); + } + let range = if req.has_range { + let unit = if req.unit == "BIT" { + BitRangeUnit::Bit + } else { + BitRangeUnit::Byte + }; + Some(BitRange { + start: req.start, + end: req.end, + unit, + }) + } else { + None + }; + let resp = self + .route( + &req.key, + ShardRequest::BitPos { + key: req.key.clone(), + bit: req.bit as u8, + range, + }, + ) + .await?; + self.record_command(start, "BITPOS"); + + match resp { + ShardResponse::Integer(n) => Ok(Response::new(IntResponse { value: n })), + ShardResponse::WrongType => Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )), + other => Err(unexpected_response(&other)), + } + } + + async fn bitop( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + if req.keys.is_empty() { + return Err(Status::invalid_argument("at least one source key required")); + } + validate_key(&req.dest, &self.ctx.limits)?; + for k in &req.keys { + validate_key(k, &self.ctx.limits)?; + } + let op = match req.op.to_uppercase().as_str() { + "AND" => BitOpKind::And, + "OR" => BitOpKind::Or, + "XOR" => BitOpKind::Xor, + "NOT" => BitOpKind::Not, + other => { + return Err(Status::invalid_argument(format!( + "unsupported BITOP operation: {other}" + ))); + } + }; + let resp = self + .route( + &req.dest, + ShardRequest::BitOp { + op, + dest: req.dest.clone(), + keys: req.keys, + }, + ) + .await?; + self.record_command(start, "BITOP"); + + match resp { + ShardResponse::Integer(n) => Ok(Response::new(IntResponse { value: n })), + ShardResponse::WrongType => Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )), + other => Err(unexpected_response(&other)), + } + } + + // ----------------------------------------------------------------------- + // sets (new) + // ----------------------------------------------------------------------- + + async fn smove( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.source, &self.ctx.limits)?; + validate_key(&req.destination, &self.ctx.limits)?; + let resp = self + .route( + &req.source, + ShardRequest::SMove { + source: req.source.clone(), + destination: req.destination, + member: req.member, + }, + ) + .await?; + self.record_command(start, "SMOVE"); + + match resp { + ShardResponse::Bool(v) => Ok(Response::new(BoolResponse { value: v })), + ShardResponse::WrongType => Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )), + other => Err(unexpected_response(&other)), + } + } + + async fn sintercard( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + let keys = req.keys; + if keys.is_empty() { + return Err(Status::invalid_argument("at least one key required")); + } + for k in &keys { + validate_key(k, &self.ctx.limits)?; + } + let limit = req.limit as usize; + let first_key = keys[0].clone(); + let resp = self + .route( + &first_key, + ShardRequest::SInterCard { keys, limit }, + ) + .await?; + self.record_command(start, "SINTERCARD"); + + match resp { + ShardResponse::Integer(n) => Ok(Response::new(IntResponse { value: n })), + ShardResponse::WrongType => Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )), + other => Err(unexpected_response(&other)), + } + } + + // ----------------------------------------------------------------------- + // lists (new) + // ----------------------------------------------------------------------- + + async fn lmpop( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + if req.keys.is_empty() { + return Err(Status::invalid_argument("at least one key required")); + } + for k in &req.keys { + validate_key(k, &self.ctx.limits)?; + } + let count = req.count.max(1) as usize; + + for key in &req.keys { + let resp = self + .route( + key, + ShardRequest::LmpopSingle { + key: key.clone(), + left: req.left, + count, + }, + ) + .await?; + match resp { + ShardResponse::Array(items) if !items.is_empty() => { + self.record_command(start, "LMPOP"); + return Ok(Response::new(LmpopResponse { + found: true, + key: key.clone(), + elements: items.into_iter().map(|b| b.to_vec()).collect(), + })); + } + ShardResponse::Array(_) | ShardResponse::Value(None) => continue, + ShardResponse::WrongType => { + return Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )); + } + other => return Err(unexpected_response(&other)), + } + } + + self.record_command(start, "LMPOP"); + Ok(Response::new(LmpopResponse { + found: false, + key: String::new(), + elements: vec![], + })) + } + + // ----------------------------------------------------------------------- + // sorted sets (new) + // ----------------------------------------------------------------------- + + async fn zmpop( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + if req.keys.is_empty() { + return Err(Status::invalid_argument("at least one key required")); + } + for k in &req.keys { + validate_key(k, &self.ctx.limits)?; + } + let count = req.count.max(1) as usize; + + for key in &req.keys { + let resp = self + .route( + key, + ShardRequest::ZmpopSingle { + key: key.clone(), + min: req.min, + count, + }, + ) + .await?; + match resp { + ShardResponse::ZPopResult(members) if !members.is_empty() => { + self.record_command(start, "ZMPOP"); + return Ok(Response::new(ZmpopResponse { + found: true, + key: key.clone(), + members: members + .into_iter() + .map(|(m, s)| ScoreMember { score: s, member: m }) + .collect(), + })); + } + ShardResponse::ZPopResult(_) | ShardResponse::Value(None) => continue, + ShardResponse::WrongType => { + return Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )); + } + other => return Err(unexpected_response(&other)), + } + } + + self.record_command(start, "ZMPOP"); + Ok(Response::new(ZmpopResponse { + found: false, + key: String::new(), + members: vec![], + })) + } + + // ----------------------------------------------------------------------- + // hashes (new) + // ----------------------------------------------------------------------- + + async fn hrandfield( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + let count = if req.has_count { Some(req.count as i64) } else { None }; + let with_values = req.with_values; + let resp = self + .route( + &req.key, + ShardRequest::HRandField { + key: req.key.clone(), + count, + with_values, + }, + ) + .await?; + self.record_command(start, "HRANDFIELD"); + + match resp { + ShardResponse::HRandFieldResult(pairs) => { + let values = if with_values { + let mut v = Vec::with_capacity(pairs.len() * 2); + for (field, val) in pairs { + v.push(field.into_bytes()); + v.push(val.map(|b| b.to_vec()).unwrap_or_default()); + } + v + } else { + pairs.into_iter().map(|(f, _)| f.into_bytes()).collect() + }; + Ok(Response::new(ArrayResponse { values })) + } + ShardResponse::WrongType => Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )), + other => Err(unexpected_response(&other)), + } + } + + async fn zrandmember( + &self, + request: Request, + ) -> Result, Status> { + let start = Instant::now(); + let req = request.into_inner(); + validate_key(&req.key, &self.ctx.limits)?; + let count = if req.has_count { Some(req.count as i64) } else { None }; + let with_scores = req.with_scores; + let resp = self + .route( + &req.key, + ShardRequest::ZRandMember { + key: req.key.clone(), + count, + with_scores, + }, + ) + .await?; + self.record_command(start, "ZRANDMEMBER"); + + match resp { + ShardResponse::ZRandMemberResult(pairs) => { + let values = if with_scores { + let mut v = Vec::with_capacity(pairs.len() * 2); + for (member, score) in pairs { + v.push(member.into_bytes()); + if let Some(s) = score { + v.push(format!("{s}").into_bytes()); + } + } + v + } else { + pairs.into_iter().map(|(m, _)| m.into_bytes()).collect() + }; + Ok(Response::new(ArrayResponse { values })) + } + ShardResponse::WrongType => Err(Status::invalid_argument( + "WRONGTYPE Operation against a key holding the wrong kind of value", + )), + other => Err(unexpected_response(&other)), + } + } + // ----------------------------------------------------------------------- // pipeline (bidirectional streaming) // ----------------------------------------------------------------------- @@ -3494,6 +4116,35 @@ async fn handle_pipeline_command( // extended server Time => time => TimeResp, LastSave => last_save => IntVal, + + // new keys + Expiretime => expiretime => IntVal, + Pexpiretime => pexpiretime => IntVal, + Expireat => expireat => BoolVal, + Pexpireat => pexpireat => BoolVal, + + // new strings + Getset => getset => Get, + Msetnx => msetnx => BoolVal, + + // bitmaps + Getbit => getbit => IntVal, + Setbit => setbit => IntVal, + Bitcount => bitcount => IntVal, + Bitpos => bitpos => IntVal, + Bitop => bitop => IntVal, + + // new sets + Smove => smove => BoolVal, + Sintercard => sintercard => IntVal, + + // new lists + Lmpop => lmpop => Lmpop, + + // new sorted sets + Zmpop => zmpop => Zmpop, + Hrandfield => hrandfield => Array, + Zrandmember => zrandmember => Array, }) } diff --git a/proto/ember/v1/ember.proto b/proto/ember/v1/ember.proto index cd0a7a35..e86c0dbc 100644 --- a/proto/ember/v1/ember.proto +++ b/proto/ember/v1/ember.proto @@ -174,6 +174,47 @@ service EmberCache { rpc Time(TimeRequest) returns (TimeResponse); rpc LastSave(LastSaveRequest) returns (IntResponse); + // --- keys (new) --- + + rpc Expiretime(ExpiretimeRequest) returns (IntResponse); + rpc Pexpiretime(PexpiretimeRequest) returns (IntResponse); + rpc Expireat(ExpireatRequest) returns (BoolResponse); + rpc Pexpireat(PexpireatRequest) returns (BoolResponse); + + // --- strings (new) --- + + rpc Getset(GetsetRequest) returns (GetResponse); + rpc Msetnx(MsetnxRequest) returns (BoolResponse); + + // --- bitmaps --- + + rpc Getbit(GetbitRequest) returns (IntResponse); + rpc Setbit(SetbitRequest) returns (IntResponse); + rpc Bitcount(BitcountRequest) returns (IntResponse); + rpc Bitpos(BitposRequest) returns (IntResponse); + rpc Bitop(BitopRequest) returns (IntResponse); + + // --- sets (new) --- + + rpc Smove(SmoveRequest) returns (BoolResponse); + rpc Sintercard(SintercardRequest) returns (IntResponse); + + // --- lists (new) --- + + rpc Lmpop(LmpopRequest) returns (LmpopResponse); + + // --- sorted sets (new) --- + + rpc Zmpop(ZmpopRequest) returns (ZmpopResponse); + + // --- hash (new) --- + + rpc Hrandfield(HrandfieldRequest) returns (ArrayResponse); + + // --- sorted set random --- + + rpc Zrandmember(ZrandmemberRequest) returns (ArrayResponse); + // --- streaming --- // bidirectional streaming for batch operations, matching RESP3 pipelining. @@ -320,6 +361,51 @@ message SetRangeRequest { bytes value = 3; } +message GetsetRequest { + string key = 1; + bytes value = 2; +} + +message MsetnxRequest { + repeated KeyValue pairs = 1; +} + +message GetbitRequest { + string key = 1; + uint64 offset = 2; +} + +message SetbitRequest { + string key = 1; + uint64 offset = 2; + uint32 value = 3; +} + +message BitcountRequest { + string key = 1; + bool has_range = 2; + int64 start = 3; + int64 end = 4; + // "BYTE" or "BIT" (only used when has_range is true) + string unit = 5; +} + +message BitposRequest { + string key = 1; + uint32 bit = 2; + bool has_range = 3; + int64 start = 4; + int64 end = 5; + string unit = 6; +} + +message BitopRequest { + // "AND", "OR", "XOR", or "NOT" + string op = 1; + string dest = 2; + repeated string keys = 3; +} + // --------------------------------------------------------------------------- // keys // --------------------------------------------------------------------------- @@ -400,6 +486,24 @@ message TouchRequest { repeated string keys = 1; } +message ExpiretimeRequest { + string key = 1; +} + +message PexpiretimeRequest { + string key = 1; +} + +message ExpireatRequest { + string key = 1; + uint64 timestamp = 2; +} + +message PexpireatRequest { + string key = 1; + uint64 timestamp_ms = 2; +} + // --------------------------------------------------------------------------- // lists // --------------------------------------------------------------------------- @@ -484,6 +588,19 @@ message LMoveRequest { bool dst_left = 4; } +message LmpopRequest { + repeated string keys = 1; + // true = LEFT, false = RIGHT + bool left = 2; + uint32 count = 3; +} + +message LmpopResponse { + bool found = 1; + string key = 2; + repeated bytes elements = 3; +} + // --------------------------------------------------------------------------- // hashes // --------------------------------------------------------------------------- @@ -561,6 +678,13 @@ message HScanResponse { repeated FieldValue fields = 2; } +message HrandfieldRequest { + string key = 1; + bool has_count = 2; + int32 count = 3; + bool with_values = 4; +} + // --------------------------------------------------------------------------- // sets // --------------------------------------------------------------------------- @@ -631,6 +755,18 @@ message SMisMemberRequest { repeated string members = 2; } +message SmoveRequest { + string source = 1; + string destination = 2; + string member = 3; +} + +message SintercardRequest { + repeated string keys = 1; + // 0 means no limit + uint64 limit = 2; +} + message SScanRequest { string key = 1; uint64 cursor = 2; @@ -770,6 +906,26 @@ message ZUnionRequest { bool with_scores = 2; } +message ZmpopRequest { + repeated string keys = 1; + // true = MIN, false = MAX + bool min = 2; + uint32 count = 3; +} + +message ZmpopResponse { + bool found = 1; + string key = 2; + repeated ScoreMember members = 3; +} + +message ZrandmemberRequest { + string key = 1; + bool has_count = 2; + int32 count = 3; + bool with_scores = 4; +} + message ZScanRequest { string key = 1; uint64 cursor = 2; @@ -1124,6 +1280,39 @@ message PipelineRequest { // extended server TimeRequest time = 110; LastSaveRequest last_save = 111; + + // new keys + ExpiretimeRequest expiretime = 112; + PexpiretimeRequest pexpiretime = 113; + ExpireatRequest expireat = 114; + PexpireatRequest pexpireat = 115; + + // new strings + GetsetRequest getset = 116; + MsetnxRequest msetnx = 117; + + // bitmaps + GetbitRequest getbit = 118; + SetbitRequest setbit = 119; + BitcountRequest bitcount = 120; + BitposRequest bitpos = 121; + BitopRequest bitop = 122; + + // new sets + SmoveRequest smove = 123; + SintercardRequest sintercard = 124; + + // new lists + LmpopRequest lmpop = 125; + + // new sorted sets + ZmpopRequest zmpop = 126; + + // new hash + HrandfieldRequest hrandfield = 127; + + // sorted set random + ZrandmemberRequest zrandmember = 128; } } @@ -1163,6 +1352,8 @@ message PipelineResponse { ZScanResponse zscan = 32; SScanResponse sscan = 33; TimeResponse time_resp = 34; + LmpopResponse lmpop = 35; + ZmpopResponse zmpop = 36; } }