diff --git a/crates/ember-core/src/keyspace/mod.rs b/crates/ember-core/src/keyspace/mod.rs index 21da495c..f92cf348 100644 --- a/crates/ember-core/src/keyspace/mod.rs +++ b/crates/ember-core/src/keyspace/mod.rs @@ -1242,6 +1242,70 @@ impl Keyspace { self.entries.is_empty() } + /// Scans proto keys starting from a cursor position. + /// + /// Returns only keys holding `Value::Proto` values. If `type_name` is + /// provided, further restricts to keys whose message type matches exactly. + /// Pattern matching follows the same glob rules as `scan_keys`. + #[cfg(feature = "protobuf")] + pub fn scan_proto_keys( + &self, + cursor: u64, + count: usize, + pattern: Option<&str>, + type_name: Option<&str>, + ) -> (u64, Vec) { + let mut keys = Vec::with_capacity(count); + let mut position = 0u64; + let target_count = if count == 0 { 10 } else { count }; + let compiled = pattern.map(GlobPattern::new); + + for (key, entry) in self.entries.iter() { + if entry.is_expired() { + continue; + } + + if position < cursor { + position += 1; + continue; + } + + // only proto entries + let entry_type = match &entry.value { + Value::Proto { type_name: t, .. } => t.as_str(), + _ => { + position += 1; + continue; + } + }; + + // optional type filter + if let Some(wanted) = type_name { + if entry_type != wanted { + position += 1; + continue; + } + } + + // optional key pattern + if let Some(ref pat) = compiled { + if !pat.matches(key) { + position += 1; + continue; + } + } + + keys.push(String::from(&**key)); + position += 1; + + if keys.len() >= target_count { + return (position, keys); + } + } + + (0, keys) + } + /// Scans keys starting from a cursor position. /// /// Returns the next cursor (0 if scan complete) and a batch of keys. diff --git a/crates/ember-core/src/keyspace/proto.rs b/crates/ember-core/src/keyspace/proto.rs index 8eadeb2f..c32573af 100644 --- a/crates/ember-core/src/keyspace/proto.rs +++ b/crates/ember-core/src/keyspace/proto.rs @@ -1,5 +1,7 @@ #[cfg(feature = "protobuf")] use super::*; +#[cfg(feature = "protobuf")] +use crate::schema::SchemaRegistry; #[cfg(feature = "protobuf")] impl Keyspace { @@ -84,6 +86,83 @@ impl Keyspace { } } + /// Scans all proto keys, returning those where the given field equals the + /// given value. Walks the keyspace using the same position-cursor logic as + /// `scan_proto_keys`. Skips keys where field decoding fails (e.g. wrong + /// type, nested/repeated field) rather than returning an error. + /// + /// `field_value` is compared against the field's string representation: + /// booleans as `"true"/"false"`, integers and floats as their decimal + /// string, strings verbatim. + pub fn scan_proto_find( + &self, + cursor: u64, + count: usize, + pattern: Option<&str>, + type_name: Option<&str>, + field_path: &str, + field_value: &str, + registry: &SchemaRegistry, + ) -> (u64, Vec) { + let mut keys = Vec::with_capacity(count); + let mut position = 0u64; + let target_count = if count == 0 { 10 } else { count }; + let compiled = pattern.map(GlobPattern::new); + + for (key, entry) in self.entries.iter() { + if entry.is_expired() { + continue; + } + + if position < cursor { + position += 1; + continue; + } + + let (entry_type, data) = match &entry.value { + Value::Proto { type_name: t, data } => (t.as_str(), data.as_ref()), + _ => { + position += 1; + continue; + } + }; + + // optional type filter + if let Some(wanted) = type_name { + if entry_type != wanted { + position += 1; + continue; + } + } + + // optional key pattern + if let Some(ref pat) = compiled { + if !pat.matches(key) { + position += 1; + continue; + } + } + + // field value comparison — skip on any error (wrong type, + // non-scalar field, field not found, decode error, etc.) + let matches = registry + .get_field_str(entry_type, data, field_path) + .map(|v| v == field_value) + .unwrap_or(false); + + if matches { + keys.push(String::from(&**key)); + } + position += 1; + + if keys.len() >= target_count { + return (position, keys); + } + } + + (0, keys) + } + /// Returns the protobuf message type name for a key, or `None` if /// the key doesn't exist. /// diff --git a/crates/ember-core/src/schema.rs b/crates/ember-core/src/schema.rs index 639c45b7..493b30da 100644 --- a/crates/ember-core/src/schema.rs +++ b/crates/ember-core/src/schema.rs @@ -255,6 +255,25 @@ impl SchemaRegistry { value_to_frame(&value, &field_desc) } + /// Reads a scalar field from an encoded protobuf message and returns it + /// as a `String` suitable for equality comparison in `PROTO.FIND`. + /// + /// Booleans are represented as `"true"` / `"false"`, integers as their + /// decimal string, floats via Rust's default `Display`, and strings + /// verbatim. Returns an error for complex types (message, list, map). + pub fn get_field_str( + &self, + type_name: &str, + data: &[u8], + field_path: &str, + ) -> Result { + let descriptor = self.find_message(type_name)?; + let msg = DynamicMessage::decode(descriptor, data) + .map_err(|e| SchemaError::ValidationFailed(e.to_string()))?; + let (value, field_desc) = resolve_field_path(&msg, field_path)?; + value_to_comparison_string(&value, &field_desc) + } + /// Updates a single scalar field in an encoded protobuf message. /// /// Decodes the message, walks the dot-separated `field_path`, parses @@ -446,6 +465,49 @@ fn value_to_frame( } } +/// Converts a prost_reflect scalar value to a `String` for comparison in +/// `PROTO.FIND`. Booleans become `"true"` / `"false"`, integers and floats +/// use their decimal representation, strings are returned verbatim. +/// Returns an error for complex types. +fn value_to_comparison_string( + value: &prost_reflect::Value, + field_desc: &FieldDescriptor, +) -> Result { + if field_desc.is_list() || field_desc.is_map() { + return Err(SchemaError::ValidationFailed( + "PROTO.FIND does not support repeated/map fields".into(), + )); + } + match value { + prost_reflect::Value::String(s) => Ok(s.clone()), + prost_reflect::Value::Bytes(b) => Ok(String::from_utf8_lossy(b).into_owned()), + prost_reflect::Value::I32(n) => Ok(n.to_string()), + prost_reflect::Value::I64(n) => Ok(n.to_string()), + prost_reflect::Value::U32(n) => Ok(n.to_string()), + prost_reflect::Value::U64(n) => Ok(n.to_string()), + prost_reflect::Value::F32(n) => Ok(n.to_string()), + prost_reflect::Value::F64(n) => Ok(n.to_string()), + prost_reflect::Value::Bool(b) => Ok(if *b { "true" } else { "false" }.into()), + prost_reflect::Value::EnumNumber(n) => { + if let Kind::Enum(enum_desc) = field_desc.kind() { + if let Some(val) = enum_desc.get_value(*n) { + return Ok(val.name().to_owned()); + } + } + Ok(n.to_string()) + } + prost_reflect::Value::Message(_) => Err(SchemaError::ValidationFailed( + "PROTO.FIND does not support nested message fields".into(), + )), + prost_reflect::Value::List(_) => Err(SchemaError::ValidationFailed( + "PROTO.FIND does not support repeated fields".into(), + )), + prost_reflect::Value::Map(_) => Err(SchemaError::ValidationFailed( + "PROTO.FIND does not support map fields".into(), + )), + } +} + /// Walks a dot-separated field path through a mutable `DynamicMessage`, /// returning the parent message (mutably), the leaf field name, and its /// descriptor. Used by `set_field` and `clear_field`. diff --git a/crates/ember-core/src/shard/mod.rs b/crates/ember-core/src/shard/mod.rs index 430fd584..7c28f1c0 100644 --- a/crates/ember-core/src/shard/mod.rs +++ b/crates/ember-core/src/shard/mod.rs @@ -778,6 +778,25 @@ pub enum ShardRequest { key: String, field_path: String, }, + /// Cursor-based scan over proto keys, optionally filtered by type name. + #[cfg(feature = "protobuf")] + ProtoScan { + cursor: u64, + count: usize, + pattern: Option, + type_name: Option, + }, + /// Cursor-based scan over proto keys, returning those where the given + /// field equals the given value. + #[cfg(feature = "protobuf")] + ProtoFind { + cursor: u64, + count: usize, + pattern: Option, + type_name: Option, + field_path: String, + field_value: String, + }, } impl ShardRequest { @@ -2488,6 +2507,55 @@ fn dispatch( }) }) } + #[cfg(feature = "protobuf")] + ShardRequest::ProtoScan { + cursor, + count, + pattern, + type_name, + } => { + let (next_cursor, keys) = ks.scan_proto_keys( + *cursor, + *count, + pattern.as_deref(), + type_name.as_deref(), + ); + ShardResponse::Scan { + cursor: next_cursor, + keys, + } + } + #[cfg(feature = "protobuf")] + ShardRequest::ProtoFind { + cursor, + count, + pattern, + type_name, + field_path, + field_value, + } => { + let registry = match schema_registry { + Some(r) => r, + None => return ShardResponse::Err("protobuf support is not enabled".into()), + }; + let reg = match registry.read() { + Ok(r) => r, + Err(_) => return ShardResponse::Err("schema registry lock poisoned".into()), + }; + let (next_cursor, keys) = ks.scan_proto_find( + *cursor, + *count, + pattern.as_deref(), + type_name.as_deref(), + field_path, + field_value, + ®, + ); + ShardResponse::Scan { + cursor: next_cursor, + keys, + } + } // these requests are intercepted in process_message, not handled here ShardRequest::Snapshot | ShardRequest::SerializeSnapshot diff --git a/crates/ember-protocol/src/command/attributes.rs b/crates/ember-protocol/src/command/attributes.rs index 995cc318..9e2d0af8 100644 --- a/crates/ember-protocol/src/command/attributes.rs +++ b/crates/ember-protocol/src/command/attributes.rs @@ -229,6 +229,8 @@ impl Command { Command::ProtoGetField { .. } => "proto.getfield", Command::ProtoSetField { .. } => "proto.setfield", Command::ProtoDelField { .. } => "proto.delfield", + Command::ProtoScan { .. } => "proto.scan", + Command::ProtoFind { .. } => "proto.find", // acl Command::AclWhoAmI => "acl", @@ -589,7 +591,9 @@ impl Command { | Command::ProtoType { .. } | Command::ProtoSchemas | Command::ProtoDescribe { .. } - | Command::ProtoGetField { .. } => READ | STRING | SLOW, + | Command::ProtoGetField { .. } + | Command::ProtoScan { .. } + | Command::ProtoFind { .. } => READ | STRING | SLOW, // ACL commands Command::AclWhoAmI => CONNECTION | FAST, diff --git a/crates/ember-protocol/src/command/mod.rs b/crates/ember-protocol/src/command/mod.rs index 5e6ea4ef..8c1ecbd6 100644 --- a/crates/ember-protocol/src/command/mod.rs +++ b/crates/ember-protocol/src/command/mod.rs @@ -790,6 +790,26 @@ pub enum Command { /// PROTO.DELFIELD `key` `field_path`. Clears a field to its default value. ProtoDelField { key: String, field_path: String }, + /// PROTO.SCAN `cursor` \[MATCH pattern\] \[COUNT count\] \[TYPE typename\]. + /// Cursor-based scan of proto keys, optionally filtered by message type. + ProtoScan { + cursor: u64, + pattern: Option, + count: Option, + type_name: Option, + }, + + /// PROTO.FIND `cursor` `field_path` `value` \[MATCH pattern\] \[TYPE typename\] \[COUNT count\]. + /// Scans proto keys, returning those where the given field equals the given value. + ProtoFind { + cursor: u64, + field_path: String, + field_value: String, + pattern: Option, + type_name: Option, + count: Option, + }, + // --- client commands --- /// CLIENT ID. Returns the unique ID of the current connection. ClientId, diff --git a/crates/ember-protocol/src/command/parse.rs b/crates/ember-protocol/src/command/parse.rs index 80822e3b..50a88de8 100644 --- a/crates/ember-protocol/src/command/parse.rs +++ b/crates/ember-protocol/src/command/parse.rs @@ -233,6 +233,8 @@ impl Command { "PROTO.GETFIELD" => parse_proto_getfield(&frames[1..]), "PROTO.SETFIELD" => parse_proto_setfield(&frames[1..]), "PROTO.DELFIELD" => parse_proto_delfield(&frames[1..]), + "PROTO.SCAN" => parse_proto_scan(&frames[1..]), + "PROTO.FIND" => parse_proto_find(&frames[1..]), "TIME" => parse_no_args("TIME", &frames[1..], Command::Time), "LASTSAVE" => parse_no_args("LASTSAVE", &frames[1..], Command::LastSave), "ROLE" => parse_no_args("ROLE", &frames[1..], Command::Role), @@ -3385,6 +3387,131 @@ fn parse_proto_delfield(args: &[Frame]) -> Result { Ok(Command::ProtoDelField { key, field_path }) } +fn parse_proto_scan(args: &[Frame]) -> Result { + if args.is_empty() { + return Err(wrong_arity("PROTO.SCAN")); + } + let cursor = parse_u64(&args[0], "PROTO.SCAN")?; + let mut pattern = None; + let mut count = None; + let mut type_name = None; + let mut idx = 1; + + while idx < args.len() { + let mut kw = [0u8; MAX_KEYWORD_LEN]; + let flag = uppercase_arg(&args[idx], &mut kw)?; + match flag { + "MATCH" => { + idx += 1; + if idx >= args.len() { + return Err(wrong_arity("PROTO.SCAN")); + } + pattern = Some(extract_string(&args[idx])?); + idx += 1; + } + "COUNT" => { + idx += 1; + if idx >= args.len() { + return Err(wrong_arity("PROTO.SCAN")); + } + let n = parse_u64(&args[idx], "PROTO.SCAN")?; + if n > MAX_SCAN_COUNT { + return Err(ProtocolError::InvalidCommandFrame(format!( + "PROTO.SCAN COUNT {n} exceeds max {MAX_SCAN_COUNT}" + ))); + } + count = Some(n as usize); + idx += 1; + } + "TYPE" => { + idx += 1; + if idx >= args.len() { + return Err(wrong_arity("PROTO.SCAN")); + } + type_name = Some(extract_string(&args[idx])?); + idx += 1; + } + _ => { + return Err(ProtocolError::InvalidCommandFrame(format!( + "unsupported PROTO.SCAN option '{flag}'" + ))); + } + } + } + + Ok(Command::ProtoScan { + cursor, + pattern, + count, + type_name, + }) +} + +fn parse_proto_find(args: &[Frame]) -> Result { + // minimum: cursor field_path value + if args.len() < 3 { + return Err(wrong_arity("PROTO.FIND")); + } + let cursor = parse_u64(&args[0], "PROTO.FIND")?; + let field_path = extract_string(&args[1])?; + let field_value = extract_string(&args[2])?; + let mut pattern = None; + let mut type_name = None; + let mut count = None; + let mut idx = 3; + + while idx < args.len() { + let mut kw = [0u8; MAX_KEYWORD_LEN]; + let flag = uppercase_arg(&args[idx], &mut kw)?; + match flag { + "MATCH" => { + idx += 1; + if idx >= args.len() { + return Err(wrong_arity("PROTO.FIND")); + } + pattern = Some(extract_string(&args[idx])?); + idx += 1; + } + "TYPE" => { + idx += 1; + if idx >= args.len() { + return Err(wrong_arity("PROTO.FIND")); + } + type_name = Some(extract_string(&args[idx])?); + idx += 1; + } + "COUNT" => { + idx += 1; + if idx >= args.len() { + return Err(wrong_arity("PROTO.FIND")); + } + let n = parse_u64(&args[idx], "PROTO.FIND")?; + if n > MAX_SCAN_COUNT { + return Err(ProtocolError::InvalidCommandFrame(format!( + "PROTO.FIND COUNT {n} exceeds max {MAX_SCAN_COUNT}" + ))); + } + count = Some(n as usize); + idx += 1; + } + _ => { + return Err(ProtocolError::InvalidCommandFrame(format!( + "unsupported PROTO.FIND option '{flag}'" + ))); + } + } + } + + Ok(Command::ProtoFind { + cursor, + field_path, + field_value, + pattern, + type_name, + count, + }) +} + fn parse_auth(args: &[Frame]) -> Result { match args.len() { 1 => { diff --git a/crates/ember-server/src/concurrent_handler.rs b/crates/ember-server/src/concurrent_handler.rs index f4903927..1a310d8f 100644 --- a/crates/ember-server/src/concurrent_handler.rs +++ b/crates/ember-server/src/concurrent_handler.rs @@ -1119,6 +1119,152 @@ async fn execute_concurrent( } } + #[cfg(feature = "protobuf")] + Command::ProtoScan { + cursor, + pattern, + count, + type_name, + } => { + if _engine.schema_registry().is_none() { + return Frame::Error("ERR protobuf support is not enabled".into()); + } + let shard_count = _engine.shard_count(); + let count = count.unwrap_or(10); + let (shard_id, position) = if cursor == 0 { + (0usize, 0u64) + } else { + let sid = (cursor >> 48) as usize; + let pos = cursor & 0xFFFF_FFFF_FFFF; + if sid >= shard_count { + return Frame::Array(vec![ + Frame::Bulk(Bytes::from("0")), + Frame::Array(vec![]), + ]); + } + (sid, pos) + }; + let mut all_keys = Vec::new(); + let mut current_shard = shard_id; + let mut current_pos = position; + while all_keys.len() < count && current_shard < shard_count { + let req = ember_core::ShardRequest::ProtoScan { + cursor: current_pos, + count: count.saturating_sub(all_keys.len()), + pattern: pattern.clone(), + type_name: type_name.clone(), + }; + match _engine.send_to_shard(current_shard, req).await { + Ok(ember_core::ShardResponse::Scan { + cursor: next_pos, + keys, + }) => { + all_keys.extend(keys); + if next_pos == 0 { + current_shard += 1; + current_pos = 0; + } else { + current_pos = next_pos; + break; + } + } + Ok(other) => { + return Frame::Error(format!("ERR unexpected shard response: {other:?}")) + } + Err(e) => return Frame::Error(format!("ERR {e}")), + } + } + let next_cursor = if current_shard >= shard_count { + 0 + } else { + ((current_shard as u64) << 48) | current_pos + }; + Frame::Array(vec![ + Frame::Bulk(Bytes::from(next_cursor.to_string())), + Frame::Array( + all_keys + .into_iter() + .map(|k| Frame::Bulk(Bytes::from(k))) + .collect(), + ), + ]) + } + + #[cfg(feature = "protobuf")] + Command::ProtoFind { + cursor, + field_path, + field_value, + pattern, + type_name, + count, + } => { + if _engine.schema_registry().is_none() { + return Frame::Error("ERR protobuf support is not enabled".into()); + } + let shard_count = _engine.shard_count(); + let count = count.unwrap_or(10); + let (shard_id, position) = if cursor == 0 { + (0usize, 0u64) + } else { + let sid = (cursor >> 48) as usize; + let pos = cursor & 0xFFFF_FFFF_FFFF; + if sid >= shard_count { + return Frame::Array(vec![ + Frame::Bulk(Bytes::from("0")), + Frame::Array(vec![]), + ]); + } + (sid, pos) + }; + let mut all_keys = Vec::new(); + let mut current_shard = shard_id; + let mut current_pos = position; + while all_keys.len() < count && current_shard < shard_count { + let req = ember_core::ShardRequest::ProtoFind { + cursor: current_pos, + count: count.saturating_sub(all_keys.len()), + pattern: pattern.clone(), + type_name: type_name.clone(), + field_path: field_path.clone(), + field_value: field_value.clone(), + }; + match _engine.send_to_shard(current_shard, req).await { + Ok(ember_core::ShardResponse::Scan { + cursor: next_pos, + keys, + }) => { + all_keys.extend(keys); + if next_pos == 0 { + current_shard += 1; + current_pos = 0; + } else { + current_pos = next_pos; + break; + } + } + Ok(other) => { + return Frame::Error(format!("ERR unexpected shard response: {other:?}")) + } + Err(e) => return Frame::Error(format!("ERR {e}")), + } + } + let next_cursor = if current_shard >= shard_count { + 0 + } else { + ((current_shard as u64) << 48) | current_pos + }; + Frame::Array(vec![ + Frame::Bulk(Bytes::from(next_cursor.to_string())), + Frame::Array( + all_keys + .into_iter() + .map(|k| Frame::Bulk(Bytes::from(k))) + .collect(), + ), + ]) + } + #[cfg(not(feature = "protobuf"))] Command::ProtoRegister { .. } | Command::ProtoSet { .. } @@ -1128,7 +1274,9 @@ async fn execute_concurrent( | Command::ProtoDescribe { .. } | Command::ProtoGetField { .. } | Command::ProtoSetField { .. } - | Command::ProtoDelField { .. } => { + | Command::ProtoDelField { .. } + | Command::ProtoScan { .. } + | Command::ProtoFind { .. } => { Frame::Error("ERR unknown command (protobuf support not compiled)".into()) } diff --git a/crates/ember-server/src/connection/exec/protobuf.rs b/crates/ember-server/src/connection/exec/protobuf.rs index 87329b50..0b556850 100644 --- a/crates/ember-server/src/connection/exec/protobuf.rs +++ b/crates/ember-server/src/connection/exec/protobuf.rs @@ -270,6 +270,165 @@ pub(in crate::connection) async fn proto_del_field( } } +#[cfg(feature = "protobuf")] +pub(in crate::connection) async fn proto_scan( + cursor: u64, + pattern: Option, + count: Option, + type_name: Option, + cx: &ExecCtx<'_>, +) -> Frame { + if cx.engine.schema_registry().is_none() { + return Frame::Error("ERR protobuf support is not enabled".into()); + } + + // cursor encoding: (shard_id << 48) | position_within_shard — same as SCAN + let shard_count = cx.engine.shard_count(); + let count = count.unwrap_or(10); + + let (shard_id, position) = if cursor == 0 { + (0usize, 0u64) + } else { + let shard_id = (cursor >> 48) as usize; + let position = cursor & 0xFFFF_FFFF_FFFF; + if shard_id >= shard_count { + return Frame::Array(vec![Frame::Bulk(Bytes::from("0")), Frame::Array(vec![])]); + } + (shard_id, position) + }; + + let mut all_keys = Vec::new(); + let mut current_shard = shard_id; + let mut current_pos = position; + + while all_keys.len() < count && current_shard < shard_count { + let req = ShardRequest::ProtoScan { + cursor: current_pos, + count: count.saturating_sub(all_keys.len()), + pattern: pattern.clone(), + type_name: type_name.clone(), + }; + match cx.engine.send_to_shard(current_shard, req).await { + Ok(ShardResponse::Scan { + cursor: next_pos, + keys, + }) => { + all_keys.extend(keys); + if next_pos == 0 { + current_shard += 1; + current_pos = 0; + } else { + current_pos = next_pos; + break; + } + } + Ok(other) => { + return Frame::Error(format!("ERR unexpected shard response: {other:?}")); + } + Err(e) => { + return Frame::Error(format!("ERR {e}")); + } + } + } + + let next_cursor = if current_shard >= shard_count { + 0 + } else { + ((current_shard as u64) << 48) | current_pos + }; + + let cursor_str = next_cursor.to_string(); + let keys_frames: Vec = all_keys + .into_iter() + .map(|k| Frame::Bulk(Bytes::from(k))) + .collect(); + Frame::Array(vec![ + Frame::Bulk(Bytes::from(cursor_str)), + Frame::Array(keys_frames), + ]) +} + +#[cfg(feature = "protobuf")] +pub(in crate::connection) async fn proto_find( + cursor: u64, + field_path: String, + field_value: String, + pattern: Option, + type_name: Option, + count: Option, + cx: &ExecCtx<'_>, +) -> Frame { + if cx.engine.schema_registry().is_none() { + return Frame::Error("ERR protobuf support is not enabled".into()); + } + + let shard_count = cx.engine.shard_count(); + let count = count.unwrap_or(10); + + let (shard_id, position) = if cursor == 0 { + (0usize, 0u64) + } else { + let shard_id = (cursor >> 48) as usize; + let position = cursor & 0xFFFF_FFFF_FFFF; + if shard_id >= shard_count { + return Frame::Array(vec![Frame::Bulk(Bytes::from("0")), Frame::Array(vec![])]); + } + (shard_id, position) + }; + + let mut all_keys = Vec::new(); + let mut current_shard = shard_id; + let mut current_pos = position; + + while all_keys.len() < count && current_shard < shard_count { + let req = ShardRequest::ProtoFind { + cursor: current_pos, + count: count.saturating_sub(all_keys.len()), + pattern: pattern.clone(), + type_name: type_name.clone(), + field_path: field_path.clone(), + field_value: field_value.clone(), + }; + match cx.engine.send_to_shard(current_shard, req).await { + Ok(ShardResponse::Scan { + cursor: next_pos, + keys, + }) => { + all_keys.extend(keys); + if next_pos == 0 { + current_shard += 1; + current_pos = 0; + } else { + current_pos = next_pos; + break; + } + } + Ok(other) => { + return Frame::Error(format!("ERR unexpected shard response: {other:?}")); + } + Err(e) => { + return Frame::Error(format!("ERR {e}")); + } + } + } + + let next_cursor = if current_shard >= shard_count { + 0 + } else { + ((current_shard as u64) << 48) | current_pos + }; + + let cursor_str = next_cursor.to_string(); + let keys_frames: Vec = all_keys + .into_iter() + .map(|k| Frame::Bulk(Bytes::from(k))) + .collect(); + Frame::Array(vec![ + Frame::Bulk(Bytes::from(cursor_str)), + Frame::Array(keys_frames), + ]) +} + /// Returns an error when protobuf support is not compiled in. #[cfg(not(feature = "protobuf"))] pub(in crate::connection) fn not_compiled() -> Frame { diff --git a/crates/ember-server/src/connection/execute.rs b/crates/ember-server/src/connection/execute.rs index 31be09eb..cb4507af 100644 --- a/crates/ember-server/src/connection/execute.rs +++ b/crates/ember-server/src/connection/execute.rs @@ -558,6 +558,25 @@ pub(super) async fn execute( Command::ProtoDelField { key, field_path } => { exec::protobuf::proto_del_field(key, field_path, &cx).await } + #[cfg(feature = "protobuf")] + Command::ProtoScan { + cursor, + pattern, + count, + type_name, + } => exec::protobuf::proto_scan(cursor, pattern, count, type_name, &cx).await, + #[cfg(feature = "protobuf")] + Command::ProtoFind { + cursor, + field_path, + field_value, + pattern, + type_name, + count, + } => { + exec::protobuf::proto_find(cursor, field_path, field_value, pattern, type_name, count, &cx) + .await + } #[cfg(not(feature = "protobuf"))] Command::ProtoRegister { .. } | Command::ProtoSet { .. } @@ -567,7 +586,9 @@ pub(super) async fn execute( | Command::ProtoDescribe { .. } | Command::ProtoGetField { .. } | Command::ProtoSetField { .. } - | Command::ProtoDelField { .. } => exec::protobuf::not_compiled(), + | Command::ProtoDelField { .. } + | Command::ProtoScan { .. } + | Command::ProtoFind { .. } => exec::protobuf::not_compiled(), Command::Quit => Frame::Simple("OK".into()), Command::Asking => Frame::Simple("OK".into()), diff --git a/tests/integration/src/proto.rs b/tests/integration/src/proto.rs index 6c1222dd..b3d8779e 100644 --- a/tests/integration/src/proto.rs +++ b/tests/integration/src/proto.rs @@ -1173,3 +1173,451 @@ async fn concurrent_schema_recovery_after_restart() { drop(data_dir); } + +// ---- PROTO.SCAN tests ---- + +/// Helper: decode a PROTO.SCAN / PROTO.FIND response into (next_cursor, keys). +fn decode_scan_response(frame: Frame) -> (u64, Vec) { + match frame { + Frame::Array(items) if items.len() == 2 => { + let cursor_str = match &items[0] { + Frame::Bulk(b) => std::str::from_utf8(b).unwrap().to_owned(), + Frame::Simple(s) => s.clone(), + other => panic!("expected cursor bulk, got {other:?}"), + }; + let cursor: u64 = cursor_str.parse().expect("cursor is u64"); + let keys = match &items[1] { + Frame::Array(ks) => ks + .iter() + .map(|k| match k { + Frame::Bulk(b) => std::str::from_utf8(b).unwrap().to_owned(), + other => panic!("expected key bulk, got {other:?}"), + }) + .collect(), + other => panic!("expected key array, got {other:?}"), + }; + (cursor, keys) + } + other => panic!("expected [cursor, [keys]], got {other:?}"), + } +} + +/// PROTO.SCAN with no arguments returns all proto keys across all pages. +#[tokio::test] +async fn proto_scan_all_keys() { + let server = start_proto_server(false); + let mut c = server.connect().await; + + let desc = make_multi_field_descriptor(); + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &desc]).await; + + // store 5 proto keys + for i in 1..=5u32 { + let data = encode_profile(&desc, &format!("user{i}"), i as i32, true); + c.cmd_raw(&[b"PROTO.SET", format!("user:{i}").as_bytes(), b"test.Profile", &data]) + .await; + } + + // also store a non-proto key to ensure it's excluded + c.cmd(&["SET", "plain:1", "hello"]).await; + + // collect all keys via cursor iteration + let mut all_keys = Vec::new(); + let mut cursor = 0u64; + loop { + let resp = c.cmd(&["PROTO.SCAN", &cursor.to_string()]).await; + let (next, keys) = decode_scan_response(resp); + all_keys.extend(keys); + if next == 0 { + break; + } + cursor = next; + } + + all_keys.sort(); + assert_eq!(all_keys.len(), 5); + for i in 1..=5u32 { + assert!(all_keys.contains(&format!("user:{i}"))); + } + // plain key must not appear + assert!(!all_keys.contains(&"plain:1".to_owned())); +} + +/// TYPE filter restricts PROTO.SCAN to keys of a specific message type. +#[tokio::test] +async fn proto_scan_type_filter() { + let server = start_proto_server(false); + let mut c = server.connect().await; + + // register two different message types + let profile_desc = make_multi_field_descriptor(); + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &profile_desc]) + .await; + + let user_desc = make_descriptor("users", "User", "username"); + c.cmd_raw(&[b"PROTO.REGISTER", b"users", &user_desc]).await; + + // store one of each + let pdata = encode_profile(&profile_desc, "alice", 30, true); + c.cmd_raw(&[b"PROTO.SET", b"profile:1", b"test.Profile", &pdata]) + .await; + + let udata = encode_message(&user_desc, "users.User", "username", "alice"); + c.cmd_raw(&[b"PROTO.SET", b"user:1", b"users.User", &udata]) + .await; + + // scan with TYPE=test.Profile — should only return profile:1 + let resp = c + .cmd(&["PROTO.SCAN", "0", "TYPE", "test.Profile"]) + .await; + let (_, keys) = decode_scan_response(resp); + assert_eq!(keys.len(), 1); + assert_eq!(keys[0], "profile:1"); + + // scan with TYPE=users.User — should only return user:1 + let resp = c.cmd(&["PROTO.SCAN", "0", "TYPE", "users.User"]).await; + let (_, keys) = decode_scan_response(resp); + assert_eq!(keys.len(), 1); + assert_eq!(keys[0], "user:1"); +} + +/// MATCH pattern narrows PROTO.SCAN results to matching key names. +#[tokio::test] +async fn proto_scan_match_pattern() { + let server = start_proto_server(false); + let mut c = server.connect().await; + + let desc = make_multi_field_descriptor(); + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &desc]).await; + + for i in 1..=3u32 { + let data = encode_profile(&desc, "x", i as i32, false); + c.cmd_raw(&[b"PROTO.SET", format!("profile:{i}").as_bytes(), b"test.Profile", &data]) + .await; + let data = encode_profile(&desc, "y", i as i32, false); + c.cmd_raw(&[b"PROTO.SET", format!("other:{i}").as_bytes(), b"test.Profile", &data]) + .await; + } + + let mut matched = Vec::new(); + let mut cursor = 0u64; + loop { + let resp = c + .cmd(&["PROTO.SCAN", &cursor.to_string(), "MATCH", "profile:*"]) + .await; + let (next, keys) = decode_scan_response(resp); + matched.extend(keys); + if next == 0 { + break; + } + cursor = next; + } + + assert_eq!(matched.len(), 3); + for k in &matched { + assert!(k.starts_with("profile:"), "unexpected key {k}"); + } +} + +/// Cursor iteration is stable — adding keys mid-scan doesn't cause duplicates or panics. +#[tokio::test] +async fn proto_scan_cursor_consistency() { + let server = start_proto_server(false); + let mut c = server.connect().await; + + let desc = make_multi_field_descriptor(); + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &desc]).await; + + for i in 1..=10u32 { + let data = encode_profile(&desc, "x", i as i32, true); + c.cmd_raw(&[b"PROTO.SET", format!("p:{i}").as_bytes(), b"test.Profile", &data]) + .await; + } + + // first page with COUNT 3 + let resp = c.cmd(&["PROTO.SCAN", "0", "COUNT", "3"]).await; + let (cursor, first_page) = decode_scan_response(resp); + assert!(!first_page.is_empty()); + + // add more keys while iterating + for i in 11..=15u32 { + let data = encode_profile(&desc, "y", i as i32, false); + c.cmd_raw(&[b"PROTO.SET", format!("p:{i}").as_bytes(), b"test.Profile", &data]) + .await; + } + + // continue iterating — must not panic or crash + if cursor != 0 { + let resp = c.cmd(&["PROTO.SCAN", &cursor.to_string(), "COUNT", "3"]).await; + let (_, _) = decode_scan_response(resp); + } +} + +// ---- PROTO.FIND tests ---- + +/// PROTO.FIND locates keys by scalar field value. +#[tokio::test] +async fn proto_find_scalar_match() { + let server = start_proto_server(false); + let mut c = server.connect().await; + + let desc = make_multi_field_descriptor(); + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &desc]).await; + + // store three profiles with different active values + let active_data = encode_profile(&desc, "alice", 25, true); + c.cmd_raw(&[b"PROTO.SET", b"profile:alice", b"test.Profile", &active_data]) + .await; + + let inactive_data = encode_profile(&desc, "bob", 30, false); + c.cmd_raw(&[b"PROTO.SET", b"profile:bob", b"test.Profile", &inactive_data]) + .await; + + let active2_data = encode_profile(&desc, "carol", 22, true); + c.cmd_raw(&[b"PROTO.SET", b"profile:carol", b"test.Profile", &active2_data]) + .await; + + // find by bool field + let mut found = Vec::new(); + let mut cursor = 0u64; + loop { + let resp = c + .cmd(&["PROTO.FIND", &cursor.to_string(), "active", "true"]) + .await; + let (next, keys) = decode_scan_response(resp); + found.extend(keys); + if next == 0 { + break; + } + cursor = next; + } + found.sort(); + assert_eq!(found.len(), 2); + assert!(found.contains(&"profile:alice".to_owned())); + assert!(found.contains(&"profile:carol".to_owned())); + + // find by int field + let resp = c + .cmd(&["PROTO.FIND", "0", "age", "30"]) + .await; + let (_, keys) = decode_scan_response(resp); + assert_eq!(keys.len(), 1); + assert_eq!(keys[0], "profile:bob"); + + // find by string field + let resp = c + .cmd(&["PROTO.FIND", "0", "name", "alice"]) + .await; + let (_, keys) = decode_scan_response(resp); + assert_eq!(keys.len(), 1); + assert_eq!(keys[0], "profile:alice"); +} + +/// PROTO.FIND with a dot-separated path searches nested message fields. +#[tokio::test] +async fn proto_find_nested_path() { + use prost_reflect::prost_types::{DescriptorProto, FieldDescriptorProto, FileDescriptorProto, FileDescriptorSet}; + + // build a descriptor with a nested Address.city field + let fds = FileDescriptorSet { + file: vec![FileDescriptorProto { + name: Some("nested.proto".into()), + package: Some("nested".into()), + message_type: vec![ + DescriptorProto { + name: Some("Address".into()), + field: vec![FieldDescriptorProto { + name: Some("city".into()), + number: Some(1), + r#type: Some(9), // TYPE_STRING + label: Some(1), + ..Default::default() + }], + ..Default::default() + }, + DescriptorProto { + name: Some("Person".into()), + field: vec![ + FieldDescriptorProto { + name: Some("name".into()), + number: Some(1), + r#type: Some(9), // TYPE_STRING + label: Some(1), + ..Default::default() + }, + FieldDescriptorProto { + name: Some("address".into()), + number: Some(2), + r#type: Some(11), // TYPE_MESSAGE + label: Some(1), + type_name: Some(".nested.Address".into()), + ..Default::default() + }, + ], + ..Default::default() + }, + ], + ..Default::default() + }], + }; + let mut desc_bytes = Vec::new(); + fds.encode(&mut desc_bytes).expect("encode descriptor"); + + let pool = DescriptorPool::decode(desc_bytes.as_slice()).expect("decode pool"); + let person_desc = pool.get_message_by_name("nested.Person").expect("find message"); + + let encode_person = |name: &str, city: &str| { + let addr_desc = pool.get_message_by_name("nested.Address").expect("find address"); + let mut addr = DynamicMessage::new(addr_desc); + addr.set_field_by_name("city", prost_reflect::Value::String(city.into())); + + let mut person = DynamicMessage::new(person_desc.clone()); + person.set_field_by_name("name", prost_reflect::Value::String(name.into())); + person.set_field_by_name("address", prost_reflect::Value::Message(addr)); + let mut buf = Vec::new(); + person.encode(&mut buf).expect("encode person"); + buf + }; + + let server = start_proto_server(false); + let mut c = server.connect().await; + + c.cmd_raw(&[b"PROTO.REGISTER", b"nested_schema", &desc_bytes]) + .await; + + let alice_data = encode_person("alice", "Seattle"); + c.cmd_raw(&[b"PROTO.SET", b"person:alice", b"nested.Person", &alice_data]) + .await; + + let bob_data = encode_person("bob", "Portland"); + c.cmd_raw(&[b"PROTO.SET", b"person:bob", b"nested.Person", &bob_data]) + .await; + + let carol_data = encode_person("carol", "Seattle"); + c.cmd_raw(&[b"PROTO.SET", b"person:carol", b"nested.Person", &carol_data]) + .await; + + // find by nested address.city + let mut found = Vec::new(); + let mut cursor = 0u64; + loop { + let resp = c + .cmd(&["PROTO.FIND", &cursor.to_string(), "address.city", "Seattle"]) + .await; + let (next, keys) = decode_scan_response(resp); + found.extend(keys); + if next == 0 { + break; + } + cursor = next; + } + found.sort(); + assert_eq!(found.len(), 2); + assert!(found.contains(&"person:alice".to_owned())); + assert!(found.contains(&"person:carol".to_owned())); +} + +/// PROTO.FIND combined with TYPE filter only inspects keys of the given type. +#[tokio::test] +async fn proto_find_type_filter() { + let server = start_proto_server(false); + let mut c = server.connect().await; + + let profile_desc = make_multi_field_descriptor(); + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &profile_desc]) + .await; + + let user_desc = make_descriptor("users", "User", "name"); + c.cmd_raw(&[b"PROTO.REGISTER", b"users", &user_desc]).await; + + // profiles with age=25 + let p1 = encode_profile(&profile_desc, "alice", 25, true); + c.cmd_raw(&[b"PROTO.SET", b"profile:1", b"test.Profile", &p1]) + .await; + + // user key with field "name" (not "age") + let u1 = encode_message(&user_desc, "users.User", "name", "alice"); + c.cmd_raw(&[b"PROTO.SET", b"user:1", b"users.User", &u1]) + .await; + + // PROTO.FIND age=25 TYPE=test.Profile — should find profile:1 but not crash on user:1 + let resp = c + .cmd(&["PROTO.FIND", "0", "age", "25", "TYPE", "test.Profile"]) + .await; + let (_, keys) = decode_scan_response(resp); + assert_eq!(keys.len(), 1); + assert_eq!(keys[0], "profile:1"); +} + +/// PROTO.FIND returns cursor 0 and empty array when no keys match. +#[tokio::test] +async fn proto_find_no_match() { + let server = start_proto_server(false); + let mut c = server.connect().await; + + let desc = make_multi_field_descriptor(); + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &desc]).await; + + let data = encode_profile(&desc, "alice", 25, false); + c.cmd_raw(&[b"PROTO.SET", b"profile:1", b"test.Profile", &data]) + .await; + + let resp = c + .cmd(&["PROTO.FIND", "0", "active", "true"]) + .await; + let (cursor, keys) = decode_scan_response(resp); + assert_eq!(cursor, 0); + assert!(keys.is_empty()); +} + +/// PROTO.FIND with COUNT paginates correctly — all pages together return the full match set. +#[tokio::test] +async fn proto_find_count_pagination() { + let server = start_proto_server(false); + let mut c = server.connect().await; + + let desc = make_multi_field_descriptor(); + c.cmd_raw(&[b"PROTO.REGISTER", b"profiles", &desc]).await; + + // 6 active profiles + for i in 1..=6u32 { + let data = encode_profile(&desc, &format!("user{i}"), i as i32, true); + c.cmd_raw(&[ + b"PROTO.SET", + format!("profile:{i}").as_bytes(), + b"test.Profile", + &data, + ]) + .await; + } + // 2 inactive profiles + for i in 7..=8u32 { + let data = encode_profile(&desc, &format!("user{i}"), i as i32, false); + c.cmd_raw(&[ + b"PROTO.SET", + format!("profile:{i}").as_bytes(), + b"test.Profile", + &data, + ]) + .await; + } + + let mut all_found = Vec::new(); + let mut cursor = 0u64; + loop { + let resp = c + .cmd(&["PROTO.FIND", &cursor.to_string(), "active", "true", "COUNT", "2"]) + .await; + let (next, keys) = decode_scan_response(resp); + all_found.extend(keys); + if next == 0 { + break; + } + cursor = next; + } + + assert_eq!(all_found.len(), 6); + for i in 1..=6u32 { + assert!(all_found.contains(&format!("profile:{i}"))); + } +}