use std::str;
use std::sync::Arc;
use webc_cmd::{Cmd, CmdHandler, ConnectionContext};
use wedb_embed::{RangeScoreSpec, RespValue, Result, Set, WeDb, ZAdd};
pub struct StandaloneHandler {
pub db: Arc<WeDb>,
}
impl StandaloneHandler {
pub fn new(db: Arc<WeDb>) -> Self {
Self { db }
}
}
impl CmdHandler for StandaloneHandler {
async fn handle(&self, ctx: &mut ConnectionContext, cmd: Cmd) -> RespValue {
handle_cmd_with_ctx(&self.db, ctx, cmd).await
}
}
pub async fn handle_cmd(db: &Arc<WeDb>, cmd: Cmd) -> RespValue {
let mut ctx = ConnectionContext::default();
handle_cmd_with_ctx(db, &mut ctx, cmd).await
}
pub async fn handle_cmd_with_ctx(
db: &Arc<WeDb>,
ctx: &mut ConnectionContext,
cmd: Cmd,
) -> RespValue {
match dispatch(db, ctx, cmd).await {
Ok(val) => val,
Err(e) => RespValue::error(format!("ERR {e}")),
}
}
async fn dispatch(db: &Arc<WeDb>, ctx: &mut ConnectionContext, cmd: Cmd) -> Result<RespValue> {
match cmd {
Cmd::Ping(msg) => match msg {
Some(m) => Ok(RespValue::Blob(m.into_bytes())),
None => Ok(RespValue::pong()),
},
Cmd::Echo(msg) => Ok(RespValue::Blob(msg.into_bytes())),
Cmd::Select(db_num) => {
ctx.set_db(db_num);
Ok(RespValue::ok())
}
Cmd::Auth { password, .. } => {
let token_key = format!("_meta:ns_token:{password}");
if let Some(ns_bytes) = db.meta_ks.get(token_key.as_bytes())? {
let ns_str = String::from_utf8_lossy(&ns_bytes).into_owned();
ctx.set_namespace(ns_str);
ctx.become_user();
} else {
ctx.set_namespace("default");
ctx.become_admin();
}
Ok(RespValue::ok())
}
Cmd::NamespaceAdd(ns, token) => {
wedb_embed::is_namespace_legal(&ns)?;
if ns == "default" || ns == "__namespace" {
return Err(wedb_embed::Error::invalid_data(
"ERR forbidden to add the default namespace",
));
}
if token.is_empty() {
return Err(wedb_embed::Error::invalid_data("ERR token cannot be empty"));
}
let name_key = format!("_meta:ns_name:{ns}");
let token_key = format!("_meta:ns_token:{token}");
if let Some(existing_token_bytes) = db.meta_ks.get(name_key.as_bytes())? {
let existing_token = String::from_utf8_lossy(&existing_token_bytes);
if existing_token == token {
return Ok(RespValue::ok());
}
return Err(wedb_embed::Error::invalid_data(
"ERR the namespace already exists",
));
}
if db.meta_ks.contains_key(token_key.as_bytes())? {
return Err(wedb_embed::Error::invalid_data(
"ERR the token already exists",
));
}
let mut batch = db.db.batch();
batch.insert(&db.meta_ks, token_key.as_bytes(), ns.as_bytes());
batch.insert(&db.meta_ks, name_key.as_bytes(), token.as_bytes());
batch.commit()?;
Ok(RespValue::ok())
}
Cmd::NamespaceSet(ns, token) => {
wedb_embed::is_namespace_legal(&ns)?;
if ns == "default" || ns == "__namespace" {
return Err(wedb_embed::Error::invalid_data(
"ERR forbidden to add the default namespace",
));
}
if token.is_empty() {
return Err(wedb_embed::Error::invalid_data("ERR token cannot be empty"));
}
let name_key = format!("_meta:ns_name:{ns}");
let token_key = format!("_meta:ns_token:{token}");
if let Some(existing_ns_bytes) = db.meta_ks.get(token_key.as_bytes())? {
let existing_ns = String::from_utf8_lossy(&existing_ns_bytes);
if existing_ns != ns {
return Err(wedb_embed::Error::invalid_data(
"ERR the token already exists",
));
}
}
let mut batch = db.db.batch();
if let Some(old_token_bytes) = db.meta_ks.get(name_key.as_bytes())? {
let old_token = String::from_utf8_lossy(&old_token_bytes);
if old_token != token {
batch.remove(
&db.meta_ks,
format!("_meta:ns_token:{old_token}").as_bytes(),
);
}
}
batch.insert(&db.meta_ks, token_key.as_bytes(), ns.as_bytes());
batch.insert(&db.meta_ks, name_key.as_bytes(), token.as_bytes());
batch.commit()?;
Ok(RespValue::ok())
}
Cmd::NamespaceDel(ns) => {
if ns == "default" || ns == "__namespace" {
return Err(wedb_embed::Error::invalid_data(
"ERR forbidden to delete the default namespace",
));
}
let name_key = format!("_meta:ns_name:{ns}");
let token_bytes = match db.meta_ks.get(name_key.as_bytes())? {
Some(t) => t,
None => {
return Err(wedb_embed::Error::invalid_data(
"ERR the namespace was not found",
));
}
};
let token = String::from_utf8_lossy(&token_bytes);
let mut batch = db.db.batch();
batch.remove(&db.meta_ks, format!("_meta:ns_token:{token}").as_bytes());
batch.remove(&db.meta_ks, name_key.as_bytes());
let ns_prefix = format!("\x00ns:{ns}:").into_bytes();
for item in db.data_ks.prefix(&ns_prefix) {
let k = item.key()?;
batch.remove(&db.data_ks, k);
}
for item in db.meta_ks.prefix(&ns_prefix) {
let k = item.key()?;
batch.remove(&db.meta_ks, k);
}
batch.commit()?;
Ok(RespValue::ok())
}
Cmd::NamespaceGet(ns) => {
if ns == "*" {
let prefix = b"_meta:ns_name:";
let mut list = Vec::new();
for item in db.meta_ks.prefix(prefix) {
let (k, v) = item.into_inner()?;
let name = k[prefix.len()..].to_vec();
list.push(RespValue::Blob(name));
list.push(RespValue::Blob(v.to_vec()));
}
list.push(RespValue::Blob(b"default".to_vec()));
list.push(RespValue::Blob(b"".to_vec()));
Ok(RespValue::Arr(list))
} else if ns == "default" || ns == "__namespace" {
Ok(RespValue::Blob(Vec::new()))
} else {
let name_key = format!("_meta:ns_name:{ns}");
match db.meta_ks.get(name_key.as_bytes())? {
Some(v) => Ok(RespValue::Blob(v.to_vec())),
None => Ok(RespValue::Null),
}
}
}
Cmd::NamespaceCurrent => Ok(RespValue::Blob(ctx.namespace.as_bytes().to_vec())),
Cmd::Command => Ok(RespValue::Arr(vec![RespValue::ok()])),
Cmd::Time => {
let now = ts_::sec();
Ok(RespValue::Arr(vec![
RespValue::Blob(format!("{now}").into_bytes()),
RespValue::Blob(b"0".to_vec()),
]))
}
Cmd::Get(key) => match db.get(key.as_bytes())? {
Some(v) => Ok(RespValue::Blob(v)),
None => Ok(RespValue::Null),
},
Cmd::Set {
key,
value,
ex,
px,
exat,
pxat,
nx,
xx,
keepttl: _,
get,
} => {
let mut opts = Vec::new();
if nx {
opts.push(Set::Nx);
}
if xx {
opts.push(Set::Xx);
}
if get {
opts.push(Set::Get);
}
if let Some(sec) = ex {
opts.push(Set::Ex(sec));
} else if let Some(ms) = px {
opts.push(Set::Px(ms));
} else if let Some(ts) = exat {
opts.push(Set::ExAt(ts));
} else if let Some(ts_ms) = pxat {
opts.push(Set::PxAt(ts_ms));
}
let prev = db.set(key.as_bytes(), &value, &opts)?;
if get {
match prev {
Some(v) => Ok(RespValue::Blob(v)),
None => Ok(RespValue::Null),
}
} else {
Ok(RespValue::ok())
}
}
Cmd::MGet(keys) => {
let mut res = Vec::with_capacity(keys.len());
for k in keys {
match db.get(k.as_bytes())? {
Some(v) => res.push(RespValue::Blob(v)),
None => res.push(RespValue::Null),
}
}
Ok(RespValue::Arr(res))
}
Cmd::MSet(pairs) => {
let p: Vec<(&[u8], &[u8])> = pairs
.iter()
.map(|(k, v)| (k.as_bytes(), v.as_slice()))
.collect();
db.mset(&p)?;
Ok(RespValue::ok())
}
Cmd::Incr(key) => {
let n = db.incrby(key.as_bytes(), 1)?;
Ok(RespValue::Int(n))
}
Cmd::Decr(key) => {
let n = db.decrby(key.as_bytes(), 1)?;
Ok(RespValue::Int(n))
}
Cmd::IncrBy(key, delta) => {
let n = db.incrby(key.as_bytes(), delta)?;
Ok(RespValue::Int(n))
}
Cmd::DecrBy(key, delta) => {
let n = db.decrby(key.as_bytes(), delta)?;
Ok(RespValue::Int(n))
}
Cmd::IncrByFloat(key, delta) => {
let f = db.incrbyfloat(key.as_bytes(), delta)?;
Ok(RespValue::Blob(format!("{f}").into_bytes()))
}
Cmd::IncrEx {
key,
by_float,
by_int,
saturate,
lbound,
ubound,
..
} => {
if let Some(delta_f) = by_float {
let cur = db
.get(key.as_bytes())?
.and_then(|v| String::from_utf8(v).ok())
.and_then(|s| s.parse::<f64>().ok())
.unwrap_or(0.0);
let mut target = cur + delta_f;
if let Some(lb) = lbound
&& target < lb
&& saturate
{
target = lb;
}
if let Some(ub) = ubound
&& target > ub
&& saturate
{
target = ub;
}
let actual_delta = target - cur;
let target_str = format!("{target}");
db.set(key.as_bytes(), target_str.as_bytes(), &[])?;
Ok(RespValue::Arr(vec![
RespValue::Blob(target_str.into_bytes()),
RespValue::Blob(format!("{actual_delta}").into_bytes()),
]))
} else {
let delta = by_int.unwrap_or(1);
let cur = db
.get(key.as_bytes())?
.and_then(|v| String::from_utf8(v).ok())
.and_then(|s| s.parse::<i64>().ok())
.unwrap_or(0);
let mut target = cur.saturating_add(delta);
if let Some(lb) = lbound
&& target < lb as i64
&& saturate
{
target = lb as i64;
}
if let Some(ub) = ubound
&& target > ub as i64
&& saturate
{
target = ub as i64;
}
let actual_delta = target.saturating_sub(cur);
let target_str = format!("{target}");
db.set(key.as_bytes(), target_str.as_bytes(), &[])?;
Ok(RespValue::Arr(vec![
RespValue::Int(target),
RespValue::Int(actual_delta),
]))
}
}
Cmd::BitField { key, ops } => {
let results = db.bitfield(key.as_bytes(), &ops)?;
let arr = results
.into_iter()
.map(|v| match v {
Some(wedb_embed::BitfieldValue::Signed(n)) => RespValue::Int(n),
Some(wedb_embed::BitfieldValue::Unsigned(n)) => RespValue::Int(n as i64),
None => RespValue::Null,
})
.collect();
Ok(RespValue::Arr(arr))
}
Cmd::BitFieldRo { key, ops } => {
let results = db.bitfield_read_only(key.as_bytes(), &ops)?;
let arr = results
.into_iter()
.map(|v| match v {
Some(wedb_embed::BitfieldValue::Signed(n)) => RespValue::Int(n),
Some(wedb_embed::BitfieldValue::Unsigned(n)) => RespValue::Int(n as i64),
None => RespValue::Null,
})
.collect();
Ok(RespValue::Arr(arr))
}
Cmd::StrLen(key) => {
let len = db.strlen(key.as_bytes())?;
Ok(RespValue::Int(len as i64))
}
Cmd::Append(key, val) => {
let len = db.append(key.as_bytes(), &val)?;
Ok(RespValue::Int(len as i64))
}
Cmd::GetRange(key, start, end) => {
let bytes = db.getrange(key.as_bytes(), start, end)?;
Ok(RespValue::Blob(bytes))
}
Cmd::SetRange(key, offset, val) => {
let len = db.setrange(key.as_bytes(), offset, &val)?;
Ok(RespValue::Int(len as i64))
}
Cmd::SetBit(key, offset, val) => {
let old = db.setbit(key.as_bytes(), offset as u64, val)?;
Ok(RespValue::Int(old as i64))
}
Cmd::GetBit(key, offset) => {
let bit = db.getbit(key.as_bytes(), offset as u64)?;
Ok(RespValue::Int(bit as i64))
}
Cmd::BitCount {
key, start, end, ..
} => {
let count = db.bitcount(key.as_bytes(), start, end)?;
Ok(RespValue::Int(count as i64))
}
Cmd::BitPos {
key,
bit,
start,
end,
} => {
let pos = db.bitpos(key.as_bytes(), bit, start, end)?;
Ok(RespValue::Int(pos))
}
Cmd::Del(keys) => {
let count = db.del(&keys)?;
Ok(RespValue::Int(count as i64))
}
Cmd::Exists(keys) => {
let count = db.exists(&keys)?;
Ok(RespValue::Int(count as i64))
}
Cmd::FlushAll | Cmd::FlushDb => {
db.flushall()?;
Ok(RespValue::ok())
}
Cmd::HSet(key, pairs) => {
let p: Vec<(&[u8], &[u8])> = pairs
.iter()
.map(|(k, v)| (k.as_bytes(), v.as_slice()))
.collect();
db.hmset(key.as_bytes(), &p)?;
Ok(RespValue::Int(pairs.len() as i64))
}
Cmd::HGet(key, field) => match db.hget(key.as_bytes(), field.as_bytes())? {
Some(v) => Ok(RespValue::Blob(v)),
None => Ok(RespValue::Null),
},
Cmd::HMGet(key, fields) => {
let f: Vec<&[u8]> = fields.iter().map(|f| f.as_bytes()).collect();
let vals = db.hmget(key.as_bytes(), &f)?;
Ok(RespValue::Arr(
vals.into_iter()
.map(|v| match v {
Some(b) => RespValue::Blob(b),
None => RespValue::Null,
})
.collect(),
))
}
Cmd::HGetAll(key) => {
let pairs = db.hgetall(key.as_bytes())?;
let mut res = Vec::with_capacity(pairs.len() * 2);
for (k, v) in pairs {
res.push(RespValue::Blob(k));
res.push(RespValue::Blob(v));
}
Ok(RespValue::Arr(res))
}
Cmd::HDel(key, fields) => {
let f: Vec<&[u8]> = fields.iter().map(|f| f.as_bytes()).collect();
let count = db.hdel(key.as_bytes(), &f)?;
Ok(RespValue::Int(count as i64))
}
Cmd::HGetDel { key, fields } => {
let mut results = Vec::with_capacity(fields.len());
for f in &fields {
let v = db.hget(key.as_bytes(), f.as_bytes())?;
match v {
Some(val) => {
db.hdel(key.as_bytes(), &[f.as_bytes()])?;
results.push(RespValue::Blob(val));
}
None => results.push(RespValue::Null),
}
}
if fields.len() == 1 {
Ok(results.into_iter().next().unwrap_or(RespValue::Null))
} else {
Ok(RespValue::Arr(results))
}
}
Cmd::Sort {
key,
offset,
count,
desc,
alpha,
..
}
| Cmd::SortRo {
key,
offset,
count,
desc,
alpha,
..
} => {
let mut elements: Vec<Vec<u8>> = db.smembers(key.as_bytes()).unwrap_or_default();
if elements.is_empty() {
elements = db.lrange(key.as_bytes(), 0, -1).unwrap_or_default();
}
if alpha {
elements.sort();
} else {
elements.sort_by(|a, b| {
let fa = str::from_utf8(a)
.ok()
.and_then(|s| s.parse::<f64>().ok())
.unwrap_or(0.0);
let fb = str::from_utf8(b)
.ok()
.and_then(|s| s.parse::<f64>().ok())
.unwrap_or(0.0);
fa.total_cmp(&fb)
});
}
if desc {
elements.reverse();
}
let sliced: Vec<RespValue> = elements
.into_iter()
.skip(offset)
.take(count.unwrap_or(usize::MAX))
.map(RespValue::Blob)
.collect();
Ok(RespValue::Arr(sliced))
}
Cmd::HLen(key) => {
let len = db.hlen(key.as_bytes())?;
Ok(RespValue::Int(len as i64))
}
Cmd::HKeys(key) => {
let keys = db.hkeys(key.as_bytes())?;
Ok(RespValue::Arr(
keys.into_iter().map(RespValue::Blob).collect(),
))
}
Cmd::HVals(key) => {
let vals = db.hvals(key.as_bytes())?;
Ok(RespValue::Arr(
vals.into_iter().map(RespValue::Blob).collect(),
))
}
Cmd::HExists(key, field) => {
let exists = db.hexists(key.as_bytes(), field.as_bytes())?;
Ok(RespValue::Int(if exists { 1 } else { 0 }))
}
Cmd::HIncrBy(key, field, delta) => {
let n = db.hincrby(key.as_bytes(), field.as_bytes(), delta)?;
Ok(RespValue::Int(n))
}
Cmd::LPush(key, elements) => {
let e: Vec<&[u8]> = elements.iter().map(|e| e.as_slice()).collect();
let len = db.lpush(key.as_bytes(), &e)?;
Ok(RespValue::Int(len as i64))
}
Cmd::RPush(key, elements) => {
let e: Vec<&[u8]> = elements.iter().map(|e| e.as_slice()).collect();
let len = db.rpush(key.as_bytes(), &e)?;
Ok(RespValue::Int(len as i64))
}
Cmd::LPop(key, count) => {
let c = count.unwrap_or(1);
let popped = db.lpop(key.as_bytes(), c)?;
if count.is_none() {
match popped.into_iter().next() {
Some(v) => Ok(RespValue::Blob(v)),
None => Ok(RespValue::Null),
}
} else {
Ok(RespValue::Arr(
popped.into_iter().map(RespValue::Blob).collect(),
))
}
}
Cmd::RPop(key, count) => {
let c = count.unwrap_or(1);
let popped = db.rpop(key.as_bytes(), c)?;
if count.is_none() {
match popped.into_iter().next() {
Some(v) => Ok(RespValue::Blob(v)),
None => Ok(RespValue::Null),
}
} else {
Ok(RespValue::Arr(
popped.into_iter().map(RespValue::Blob).collect(),
))
}
}
Cmd::LLen(key) => {
let len = db.llen(key.as_bytes())?;
Ok(RespValue::Int(len as i64))
}
Cmd::LRange(key, start, stop) => {
let items = db.lrange(key.as_bytes(), start, stop)?;
Ok(RespValue::Arr(
items.into_iter().map(RespValue::Blob).collect(),
))
}
Cmd::LIndex(key, idx) => match db.lindex(key.as_bytes(), idx)? {
Some(v) => Ok(RespValue::Blob(v)),
None => Ok(RespValue::Null),
},
Cmd::SAdd(key, members) => {
let m: Vec<&[u8]> = members.iter().map(|m| m.as_slice()).collect();
let count = db.sadd(key.as_bytes(), &m)?;
Ok(RespValue::Int(count as i64))
}
Cmd::SRem(key, members) => {
let m: Vec<&[u8]> = members.iter().map(|m| m.as_slice()).collect();
let count = db.srem(key.as_bytes(), &m)?;
Ok(RespValue::Int(count as i64))
}
Cmd::SCard(key) => {
let count = db.scard(key.as_bytes())?;
Ok(RespValue::Int(count as i64))
}
Cmd::SMembers(key) => {
let members = db.smembers(key.as_bytes())?;
Ok(RespValue::Arr(
members.into_iter().map(RespValue::Blob).collect(),
))
}
Cmd::SIsMember(key, member) => {
let exists = db.sismember(key.as_bytes(), &member)?;
Ok(RespValue::Int(if exists { 1 } else { 0 }))
}
Cmd::SMIsMember(key, members) => {
let m: Vec<&[u8]> = members.iter().map(|m| m.as_slice()).collect();
let exists = db.smismember(key.as_bytes(), &m)?;
Ok(RespValue::Arr(
exists
.into_iter()
.map(|b| RespValue::Int(if b { 1 } else { 0 }))
.collect(),
))
}
Cmd::SPop(key, count) => match count {
Some(cnt) => {
let popped = db.spop(key.as_bytes(), cnt)?;
Ok(RespValue::Arr(
popped.into_iter().map(RespValue::Blob).collect(),
))
}
None => {
let popped = db.spop(key.as_bytes(), 1)?;
match popped.into_iter().next() {
Some(val) => Ok(RespValue::Blob(val)),
None => Ok(RespValue::Null),
}
}
},
Cmd::SRandMember(key, count) => match count {
Some(cnt) => {
let res = db.srandmember(key.as_bytes(), cnt)?;
Ok(RespValue::Arr(
res.into_iter().map(RespValue::Blob).collect(),
))
}
None => {
let res = db.srandmember(key.as_bytes(), 1)?;
match res.into_iter().next() {
Some(val) => Ok(RespValue::Blob(val)),
None => Ok(RespValue::Null),
}
}
},
Cmd::SMove { src, dst, member } => {
let moved = db.smove(src.as_bytes(), dst.as_bytes(), &member)?;
Ok(RespValue::Int(if moved { 1 } else { 0 }))
}
Cmd::SDiff(keys) => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
let res = db.sdiff(&k)?;
Ok(RespValue::Arr(
res.into_iter().map(RespValue::Blob).collect(),
))
}
Cmd::SUnion(keys) => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
let res = db.sunion(&k)?;
Ok(RespValue::Arr(
res.into_iter().map(RespValue::Blob).collect(),
))
}
Cmd::SInter(keys) => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
let res = db.sinter(&k)?;
Ok(RespValue::Arr(
res.into_iter().map(RespValue::Blob).collect(),
))
}
Cmd::SInterCard { keys, limit } => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
let count = db.sintercard(&k, limit)?;
Ok(RespValue::Int(count as i64))
}
Cmd::SDiffCard { keys, limit } => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
let count = db.sdiffcard(&k, limit)?;
Ok(RespValue::Int(count as i64))
}
Cmd::SUnionCard { keys, limit } => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
let count = db.sunioncard(&k, limit)?;
Ok(RespValue::Int(count as i64))
}
Cmd::SDiffStore(dst, keys) => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
let count = db.sdiffstore(dst.as_bytes(), &k)?;
Ok(RespValue::Int(count as i64))
}
Cmd::SUnionStore(dst, keys) => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
let count = db.sunionstore(dst.as_bytes(), &k)?;
Ok(RespValue::Int(count as i64))
}
Cmd::SInterStore(dst, keys) => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
let count = db.sinterstore(dst.as_bytes(), &k)?;
Ok(RespValue::Int(count as i64))
}
Cmd::SScan {
key,
cursor,
pattern,
count,
} => {
let pat_bytes = pattern.as_deref().map(|p| p.as_bytes());
let (next_cursor, items) = db.sscan(key.as_bytes(), cursor, pat_bytes, count)?;
Ok(RespValue::Arr(vec![
RespValue::Blob(format!("{next_cursor}").into_bytes()),
RespValue::Arr(items.into_iter().map(RespValue::Blob).collect()),
]))
}
Cmd::ZAdd {
key,
members,
nx,
xx,
gt,
lt,
ch,
incr,
} => {
let m: Vec<(f64, &[u8])> = members.iter().map(|(s, m)| (*s, m.as_slice())).collect();
let mut flags = Vec::new();
if nx {
flags.push(ZAdd::Nx);
}
if xx {
flags.push(ZAdd::Xx);
}
if gt {
flags.push(ZAdd::Gt);
}
if lt {
flags.push(ZAdd::Lt);
}
if ch {
flags.push(ZAdd::Ch);
}
if incr {
flags.push(ZAdd::Incr);
}
let count = db.zadd(key.as_bytes(), &m, &flags)?;
Ok(RespValue::Int(count as i64))
}
Cmd::ZScore(key, member) => match db.zscore(key.as_bytes(), &member)? {
Some(s) => Ok(RespValue::Blob(format!("{s}").into_bytes())),
None => Ok(RespValue::Null),
},
Cmd::ZCard(key) => {
let count = db.zcard(key.as_bytes())?;
Ok(RespValue::Int(count as i64))
}
Cmd::ZRem(key, members) => {
let m: Vec<&[u8]> = members.iter().map(|m| m.as_slice()).collect();
let count = db.zrem(key.as_bytes(), &m)?;
Ok(RespValue::Int(count as i64))
}
Cmd::ZIncrBy(key, delta, member) => {
let new_score = db.zincrby(key.as_bytes(), delta, &member)?;
Ok(RespValue::Blob(format!("{new_score}").into_bytes()))
}
Cmd::ZRank { key, member, .. } => match db.zrank(key.as_bytes(), &member)? {
Some(rank) => Ok(RespValue::Int(rank as i64)),
None => Ok(RespValue::Null),
},
Cmd::ZRevRank { key, member, .. } => match db.zrevrank(key.as_bytes(), &member)? {
Some(rank) => Ok(RespValue::Int(rank as i64)),
None => Ok(RespValue::Null),
},
Cmd::ZRange {
key,
min,
max,
by_score,
by_lex: _,
rev,
offset,
count,
with_scores,
} => {
let items = if by_score {
let (min_val, min_exclusive) = webc_cmd::parse::parse_score_boundary(&min);
let (max_val, max_exclusive) = webc_cmd::parse::parse_score_boundary(&max);
let spec = RangeScoreSpec {
min: min_val,
max: max_val,
minex: min_exclusive,
maxex: max_exclusive,
offset,
count,
};
if rev {
db.zrevrangebyscore(key.as_bytes(), &spec)?
} else {
db.zrangebyscore(key.as_bytes(), &spec)?
}
} else {
let start = min.parse::<i64>().unwrap_or(0);
let stop = max.parse::<i64>().unwrap_or(-1);
if rev {
db.zrevrange(key.as_bytes(), start, stop)?
} else {
db.zrange(key.as_bytes(), start, stop)?
}
};
let mut res = Vec::new();
for (m, s) in items {
res.push(RespValue::Blob(m));
if with_scores {
res.push(RespValue::Blob(format!("{s}").into_bytes()));
}
}
Ok(RespValue::Arr(res))
}
Cmd::ZRevRange(key, start, stop, with_scores) => {
let items = db.zrevrange(key.as_bytes(), start, stop)?;
let mut res = Vec::new();
for (m, s) in items {
res.push(RespValue::Blob(m));
if with_scores {
res.push(RespValue::Blob(format!("{s}").into_bytes()));
}
}
Ok(RespValue::Arr(res))
}
Cmd::ZCount(key, min, max) => {
let (min_val, min_exclusive) = webc_cmd::parse::parse_score_boundary(&min);
let (max_val, max_exclusive) = webc_cmd::parse::parse_score_boundary(&max);
let spec = RangeScoreSpec {
min: min_val,
max: max_val,
minex: min_exclusive,
maxex: max_exclusive,
offset: 0,
count: None,
};
let items = db.zrangebyscore(key.as_bytes(), &spec)?;
Ok(RespValue::Int(items.len() as i64))
}
Cmd::ZPopMin(key, count) => {
let c = count.unwrap_or(1);
let items = db.zpopmin(key.as_bytes(), c)?;
let mut res = Vec::new();
for (m, s) in items {
res.push(RespValue::Blob(m));
res.push(RespValue::Blob(format!("{s}").into_bytes()));
}
Ok(RespValue::Arr(res))
}
Cmd::ZPopMax(key, count) => {
let c = count.unwrap_or(1);
let items = db.zpopmax(key.as_bytes(), c)?;
let mut res = Vec::new();
for (m, s) in items {
res.push(RespValue::Blob(m));
res.push(RespValue::Blob(format!("{s}").into_bytes()));
}
Ok(RespValue::Arr(res))
}
Cmd::GeoAdd { key, items, .. } => {
let m: Vec<(f64, f64, &[u8])> = items
.iter()
.map(|(lon, lat, mem)| (*lon, *lat, mem.as_bytes()))
.collect();
let count = db.geoadd(key.as_bytes(), &m)?;
Ok(RespValue::Int(count as i64))
}
Cmd::GeoDist { key, m1, m2, unit } => {
let dist = db.geodist(
key.as_bytes(),
m1.as_bytes(),
m2.as_bytes(),
Some(unit.as_str()),
)?;
match dist {
Some(d) => Ok(RespValue::Blob(format!("{d}").into_bytes())),
None => Ok(RespValue::Null),
}
}
Cmd::GeoHash(key, members) => {
let m: Vec<&[u8]> = members.iter().map(|s| s.as_bytes()).collect();
let hashes = db.geohash(key.as_bytes(), &m)?;
Ok(RespValue::Arr(
hashes
.into_iter()
.map(|h| match h {
Some(s) => RespValue::Blob(s.into_bytes()),
None => RespValue::Null,
})
.collect(),
))
}
Cmd::GeoPos(key, members) => {
let m: Vec<&[u8]> = members.iter().map(|s| s.as_bytes()).collect();
let positions = db.geopos(key.as_bytes(), &m)?;
Ok(RespValue::Arr(
positions
.into_iter()
.map(|p| match p {
Some((lon, lat)) => RespValue::Arr(vec![
RespValue::Blob(format!("{lon}").into_bytes()),
RespValue::Blob(format!("{lat}").into_bytes()),
]),
None => RespValue::Null,
})
.collect(),
))
}
Cmd::BfAdd(key, item) => {
let added = db.bf_add(key.as_bytes(), &item)?;
Ok(RespValue::Int(if added { 1 } else { 0 }))
}
Cmd::BfExists(key, item) => {
let exists = db.bf_exists(key.as_bytes(), &item)?;
Ok(RespValue::Int(if exists { 1 } else { 0 }))
}
Cmd::CfAdd(key, item) => {
let added = db.cf_add(key.as_bytes(), &item)?;
Ok(RespValue::Int(if added { 1 } else { 0 }))
}
Cmd::CfExists(key, item) => {
let exists = db.cf_exists(key.as_bytes(), &item)?;
Ok(RespValue::Int(if exists { 1 } else { 0 }))
}
Cmd::PfAdd(key, elements) => {
let e: Vec<&[u8]> = elements.iter().map(|e| e.as_slice()).collect();
let changed = db.pfadd(key.as_bytes(), &e)?;
Ok(RespValue::Int(if changed { 1 } else { 0 }))
}
Cmd::PfCount(keys) => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
let count = db.pfcount(&k)?;
Ok(RespValue::Int(count as i64))
}
Cmd::PfMerge(dest, keys) => {
let k: Vec<&[u8]> = keys.iter().map(|k| k.as_bytes()).collect();
db.pfmerge(dest.as_bytes(), &k)?;
Ok(RespValue::ok())
}
Cmd::JsonGet {
key,
paths,
indent: _,
newline: _,
space: _,
} => {
let path = paths.first().map(|s| s.as_str());
match db.json_get(key.as_bytes(), path)? {
Some(v) => Ok(RespValue::Blob(v.into_bytes())),
None => Ok(RespValue::Null),
}
}
Cmd::JsonSet {
key, path, value, ..
} => {
db.json_set(key.as_bytes(), &path, &value)?;
Ok(RespValue::ok())
}
Cmd::TsAdd {
key,
timestamp,
value,
..
} => {
let ts = timestamp.unwrap_or_else(|| ts_::sec() * 1000);
let added_ts = db.ts_add(key.as_bytes(), ts, value)?;
Ok(RespValue::Int(added_ts as i64))
}
Cmd::TsGet { key, .. } => match db.ts_get(key.as_bytes())? {
Some((ts, val)) => Ok(RespValue::Arr(vec![
RespValue::Int(ts as i64),
RespValue::Blob(format!("{val}").into_bytes()),
])),
None => Ok(RespValue::Null),
},
Cmd::TDigestCreate { key, compression } => {
let comp = compression.unwrap_or(100.0);
db.tdigest_create(key.as_bytes(), comp)?;
Ok(RespValue::ok())
}
Cmd::TDigestAdd(key, values) => {
db.tdigest_add(key.as_bytes(), &values)?;
Ok(RespValue::ok())
}
Cmd::TDigestQuantile(key, quantiles) => {
let q_res = db.tdigest_quantile(key.as_bytes(), &quantiles)?;
let results = q_res
.into_iter()
.map(|v| match v {
Some(f) => RespValue::Blob(format!("{f}").into_bytes()),
None => RespValue::Blob(b"nan".to_vec()),
})
.collect();
Ok(RespValue::Arr(results))
}
Cmd::TDigestCdf(key, vals) => {
let cdf_res = db.tdigest_cdf(key.as_bytes(), &vals)?;
let results = cdf_res
.into_iter()
.map(|v| match v {
Some(f) => RespValue::Blob(format!("{f}").into_bytes()),
None => RespValue::Blob(b"nan".to_vec()),
})
.collect();
Ok(RespValue::Arr(results))
}
Cmd::TDigestMin(key) => {
let val = db.tdigest_min(key.as_bytes())?;
if val.is_nan() {
Ok(RespValue::Blob(b"nan".to_vec()))
} else {
Ok(RespValue::Blob(format!("{val}").into_bytes()))
}
}
Cmd::TDigestMax(key) => {
let val = db.tdigest_max(key.as_bytes())?;
if val.is_nan() {
Ok(RespValue::Blob(b"nan".to_vec()))
} else {
Ok(RespValue::Blob(format!("{val}").into_bytes()))
}
}
Cmd::TDigestRank(key, vals) => {
let ranks = db.tdigest_rank(key.as_bytes(), &vals)?;
let results = ranks.into_iter().map(RespValue::Int).collect();
Ok(RespValue::Arr(results))
}
Cmd::TDigestRevRank(key, vals) => {
let ranks = db.tdigest_revrank(key.as_bytes(), &vals)?;
let results = ranks.into_iter().map(RespValue::Int).collect();
Ok(RespValue::Arr(results))
}
Cmd::TDigestByRank(key, ranks) => {
let res = db.tdigest_byrank(key.as_bytes(), &ranks)?;
let results = res
.into_iter()
.map(|v| match v {
Some(f) => {
if f.is_infinite() {
if f > 0.0 {
RespValue::Blob(b"inf".to_vec())
} else {
RespValue::Blob(b"-inf".to_vec())
}
} else {
RespValue::Blob(format!("{f}").into_bytes())
}
}
None => RespValue::Blob(b"nan".to_vec()),
})
.collect();
Ok(RespValue::Arr(results))
}
Cmd::TDigestByRevRank(key, ranks) => {
let res = db.tdigest_byrevrank(key.as_bytes(), &ranks)?;
let results = res
.into_iter()
.map(|v| match v {
Some(f) => {
if f.is_infinite() {
if f > 0.0 {
RespValue::Blob(b"inf".to_vec())
} else {
RespValue::Blob(b"-inf".to_vec())
}
} else {
RespValue::Blob(format!("{f}").into_bytes())
}
}
None => RespValue::Blob(b"nan".to_vec()),
})
.collect();
Ok(RespValue::Arr(results))
}
Cmd::TDigestTrimmedMean(key, low_cut, high_cut) => {
let mean = db.tdigest_trimmed_mean(key.as_bytes(), low_cut, high_cut)?;
match mean {
Some(f) => Ok(RespValue::Blob(format!("{f}").into_bytes())),
None => Ok(RespValue::Blob(b"nan".to_vec())),
}
}
Cmd::TDigestReset(key) => {
db.tdigest_reset(key.as_bytes())?;
Ok(RespValue::ok())
}
Cmd::TDigestMerge {
dst,
sources,
compression,
override_flag,
} => {
let src_slices: Vec<&[u8]> = sources.iter().map(|s| s.as_bytes()).collect();
let opts = wedb_embed::tdigest::TDigestMerge {
compression: compression.map(|c| c as u32),
override_dest: override_flag,
};
db.tdigest_merge(dst.as_bytes(), &src_slices, opts)?;
Ok(RespValue::ok())
}
Cmd::TDigestInfo(key) => {
let info = db.tdigest_info(key.as_bytes())?;
let min_blob = match info.minimum {
Some(m) => RespValue::Blob(format!("{m}").into_bytes()),
None => RespValue::Blob(b"nan".to_vec()),
};
let max_blob = match info.maximum {
Some(m) => RespValue::Blob(format!("{m}").into_bytes()),
None => RespValue::Blob(b"nan".to_vec()),
};
let entries = vec![
RespValue::Simple("Compression".to_string()),
RespValue::Int(info.compression as i64),
RespValue::Simple("Capacity".to_string()),
RespValue::Int(info.capacity as i64),
RespValue::Simple("Merged nodes".to_string()),
RespValue::Int(info.merged_nodes as i64),
RespValue::Simple("Unmerged nodes".to_string()),
RespValue::Int(info.unmerged_nodes as i64),
RespValue::Simple("Merged weight".to_string()),
RespValue::Blob(format!("{}", info.merged_weight).into_bytes()),
RespValue::Simple("Unmerged weight".to_string()),
RespValue::Blob(format!("{}", info.unmerged_weight).into_bytes()),
RespValue::Simple("Total weight".to_string()),
RespValue::Blob(format!("{}", info.total_weight).into_bytes()),
RespValue::Simple("Observations".to_string()),
RespValue::Int(info.observations as i64),
RespValue::Simple("Total compressions".to_string()),
RespValue::Int(info.total_compressions as i64),
RespValue::Simple("Minimum".to_string()),
min_blob,
RespValue::Simple("Maximum".to_string()),
max_blob,
];
Ok(RespValue::Arr(entries))
}
_ => Ok(RespValue::ok()),
}
}