Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
64 changes: 64 additions & 0 deletions crates/ember-core/src/keyspace/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>) {
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.
Expand Down
79 changes: 79 additions & 0 deletions crates/ember-core/src/keyspace/proto.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
#[cfg(feature = "protobuf")]
use super::*;
#[cfg(feature = "protobuf")]
use crate::schema::SchemaRegistry;

#[cfg(feature = "protobuf")]
impl Keyspace {
Expand Down Expand Up @@ -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<String>) {
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.
///
Expand Down
62 changes: 62 additions & 0 deletions crates/ember-core/src/schema.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, SchemaError> {
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
Expand Down Expand Up @@ -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<String, SchemaError> {
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`.
Expand Down
68 changes: 68 additions & 0 deletions crates/ember-core/src/shard/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>,
type_name: Option<String>,
},
/// 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<String>,
type_name: Option<String>,
field_path: String,
field_value: String,
},
}

impl ShardRequest {
Expand Down Expand Up @@ -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,
&reg,
);
ShardResponse::Scan {
cursor: next_cursor,
keys,
}
}
// these requests are intercepted in process_message, not handled here
ShardRequest::Snapshot
| ShardRequest::SerializeSnapshot
Expand Down
6 changes: 5 additions & 1 deletion crates/ember-protocol/src/command/attributes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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,
Expand Down
20 changes: 20 additions & 0 deletions crates/ember-protocol/src/command/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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<String>,
count: Option<usize>,
type_name: Option<String>,
},

/// 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<String>,
type_name: Option<String>,
count: Option<usize>,
},

// --- client commands ---
/// CLIENT ID. Returns the unique ID of the current connection.
ClientId,
Expand Down
Loading
Loading