Skip to content

Commit 3792bb0

Browse files
arthur-zhangarthur-zhang
committed
refract cmd type
1 parent 866eeaa commit 3792bb0

File tree

18 files changed

+165
-117
lines changed

18 files changed

+165
-117
lines changed

‎redis-command-gen/Cargo.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@ serde_json = "1.0.114"
1414
smol_str = "0.2.1"
1515
lazy_static = "1.4.0"
1616
redis-command = { path = "../redis-command" }
17+
strum = { version = "0.26.2",features = ["derive"] }
18+
strum_macros = "0.26.2"
1719

1820
[build-dependencies]
1921
serde = { version = "1.0.197", features = ["derive"] }

‎redis-command-gen/build_redis_cmd.rs‎

Lines changed: 36 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -91,14 +91,29 @@ fn main() {
9191
}
9292
});
9393

94+
95+
let mut cmd_type_str = r#"
96+
#[allow(non_camel_case_types)]
97+
#[derive(EnumString, AsRefStr, Debug, Eq, PartialEq, Hash, Clone, Copy)]
98+
pub enum CmdType {
99+
ALL,
100+
UNKNOWN,
101+
"#.to_string();
102+
for (k, _) in &all_cmd_map {
103+
cmd_type_str.push_str(&format!(" #[strum(serialize = \"{}\")]\n", k));
104+
cmd_type_str.push_str(&format!(" {},\n", k.replace(" ", "_").replace("-", "_").to_ascii_uppercase()));
105+
}
106+
107+
cmd_type_str.push_str("}\n");
108+
94109
let mut multi_cmd_map = String::new();
95110
for (container, _) in all_multi_cmd_map {
96-
multi_cmd_map.push_str(&format!(" (SmolStr::from(\"{}\"), true),\n", container.to_ascii_lowercase()));
111+
multi_cmd_map.push_str(&format!(" (CmdType::from_str(\"{}\").unwrap(), true),\n", container.to_ascii_uppercase()));
97112
}
98113

99114
let mut cmd_map = String::new();
100115
for (k, v) in all_cmd_map {
101-
let head = format!(" (SmolStr::from(\"{}\"), Command {{\n", k.to_ascii_lowercase());
116+
let head = format!(" (CmdType::from_str(\"{}\").unwrap(), Command {{\n", k.to_ascii_uppercase());
102117
let name = format!(" name: String::from(\"{}\"),\n", v.name);
103118
let container = if let Some(c) = v.container {
104119
format!(" container: Some(String::from(\"{}\")),\n", c)
@@ -173,25 +188,34 @@ fn main() {
173188
}
174189

175190
let mut cmd_holder_rs_content = String::new();
176-
cmd_holder_rs_content.push_str("/**\n");
177-
cmd_holder_rs_content.push_str(" * warning: This file was generated by build_redis_cmd.rs\n");
178-
cmd_holder_rs_content.push_str(" * do not modify it manually!\n");
179-
cmd_holder_rs_content.push_str(" */\n");
180-
cmd_holder_rs_content.push_str("use std::collections::HashMap;\n\n");
181-
cmd_holder_rs_content.push_str("use lazy_static::lazy_static;\n");
182-
cmd_holder_rs_content.push_str("use smol_str::SmolStr;\n\n");
183-
cmd_holder_rs_content.push_str("use redis_command::{BeginSearch, Command, CommandFlags, FindKeys, Group, Index, KeyNum, KeySpecs, Keyword, Range};\n\n");
191+
192+
193+
cmd_holder_rs_content.push_str(r#"/**
194+
* warning: This file was generated by build_redis_cmd.rs
195+
* do not modify it manually!
196+
*/
197+
use std::collections::HashMap;
198+
use std::str::FromStr;
199+
use lazy_static::lazy_static;
200+
use strum::EnumString;
201+
use std::convert::AsRef;
202+
use strum_macros::AsRefStr;
203+
use redis_command::{BeginSearch, Command, CommandFlags, FindKeys, Group, Index, KeyNum, KeySpecs, Keyword, Range};
204+
"#);
205+
206+
cmd_holder_rs_content.push_str(&cmd_type_str);
207+
184208
cmd_holder_rs_content.push_str("lazy_static! {\n");
185209
cmd_holder_rs_content.push_str(" /**\n");
186210
cmd_holder_rs_content.push_str(" * Redis command's name(or container) and whether it is a multipart command\n");
187211
cmd_holder_rs_content.push_str(" */\n");
188-
cmd_holder_rs_content.push_str(" pub static ref MULTIPART_COMMANDS: HashMap<SmolStr, bool> = HashMap::from([\n");
212+
cmd_holder_rs_content.push_str(" pub static ref MULTIPART_COMMANDS: HashMap<CmdType, bool> = HashMap::from([\n");
189213
cmd_holder_rs_content.push_str(&multi_cmd_map);
190214
cmd_holder_rs_content.push_str(" ]);\n\n");
191215
cmd_holder_rs_content.push_str(" /**\n");
192216
cmd_holder_rs_content.push_str(" * Redis command's full name and its description\n");
193217
cmd_holder_rs_content.push_str(" */\n");
194-
cmd_holder_rs_content.push_str(" pub static ref COMMAND_ATTRIBUTES: HashMap<SmolStr, Command> = HashMap::from([\n");
218+
cmd_holder_rs_content.push_str(" pub static ref COMMAND_ATTRIBUTES: HashMap<CmdType, Command> = HashMap::from([\n");
195219
cmd_holder_rs_content.push_str(&cmd_map);
196220
cmd_holder_rs_content.push_str(" ]);\n");
197221
cmd_holder_rs_content.push_str("}\n");

‎redis-command/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -115,4 +115,4 @@ pub struct KeyNum {
115115
pub keynumidx: usize,
116116
pub firstkey: usize,
117117
pub step: usize
118-
}
118+
}

‎redis-proxy-common/src/command/utils.rs‎

Lines changed: 48 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,15 @@ use bytes::Bytes;
55
use lazy_static::lazy_static;
66
use smol_str::SmolStr;
77
use redis_command::{Command, CommandFlags, Group, Index, KeyNum, Keyword, Range};
8-
use redis_command_gen::COMMAND_ATTRIBUTES;
8+
use redis_command_gen::{CmdType, COMMAND_ATTRIBUTES};
99

10-
lazy_static!{
10+
lazy_static! {
1111
// some special command types
12-
pub static ref CMD_TYPE_ALL: SmolStr = SmolStr::from("*");
13-
pub static ref CMD_TYPE_UNKNOWN: SmolStr = SmolStr::from("unknown");
14-
15-
pub static ref CMD_TYPE_AUTH: SmolStr = SmolStr::from("auth");
16-
pub static ref CMD_TYPE_SELECT: SmolStr = SmolStr::from("select");
12+
// pub static ref CMD_TYPE_ALL: SmolStr = SmolStr::from("*");
13+
// pub static ref CMD_TYPE_UNKNOWN: SmolStr = SmolStr::from("unknown");
14+
15+
// pub static ref CMD_TYPE_AUTH: CmdType = CmdType::from_str("AUTH").unwrap();
16+
// pub static ref CMD_TYPE_SELECT: CmdType = CmdType::from_str("SELECT");
1717
}
1818

1919
#[inline]
@@ -25,42 +25,41 @@ pub fn to_lower_effective(origin: &[u8]) -> Vec<u8> {
2525
target
2626
}
2727

28-
pub fn is_write_cmd(cmd: &SmolStr) -> bool {
28+
pub fn is_write_cmd(cmd: &CmdType) -> bool {
2929
if let Some(cmd) = COMMAND_ATTRIBUTES.get(cmd) {
3030
return cmd.command_flags.contains(CommandFlags::Write);
3131
}
32-
return false
32+
return false;
3333
}
3434

35-
pub fn is_readonly_cmd(cmd: &SmolStr) -> bool {
35+
pub fn is_readonly_cmd(cmd: &CmdType) -> bool {
3636
if let Some(cmd) = COMMAND_ATTRIBUTES.get(cmd) {
3737
return cmd.command_flags.contains(CommandFlags::Readonly);
3838
}
39-
return false
39+
return false;
4040
}
4141

42-
pub fn has_key(cmd: &SmolStr) -> bool {
42+
pub fn has_key(cmd: &CmdType) -> bool {
4343
if let Some(cmd) = COMMAND_ATTRIBUTES.get(cmd) {
44-
return cmd.key_specs.is_some()
44+
return cmd.key_specs.is_some();
4545
}
46-
return false
46+
return false;
4747
}
4848

49-
pub fn is_connection_cmd(cmd: &SmolStr) -> bool {
49+
pub fn is_connection_cmd(cmd: &CmdType) -> bool {
5050
if let Some(cmd) = COMMAND_ATTRIBUTES.get(cmd) {
51-
return cmd.group == Group::Connection
51+
return cmd.group == Group::Connection;
5252
}
53-
return false
53+
return false;
5454
}
5555

56-
pub fn get_cmd_key_bulks<'a>(cmd: &SmolStr, bulk_args: &'a Vec<Bytes>) -> Option<Vec<&'a [u8]>> {
56+
pub fn get_cmd_key_bulks<'a>(cmd: &CmdType, bulk_args: &'a Vec<Bytes>) -> Option<Vec<&'a [u8]>> {
5757
if bulk_args.len() < 2 {
58-
return None
58+
return None;
5959
}
60-
// let cmd = COMMAND_ATTRIBUTES.get(cmd);
61-
let cmd :Option<Command> = None;
60+
let cmd = COMMAND_ATTRIBUTES.get(cmd);
6261
if cmd.is_none() {
63-
return None
62+
return None;
6463
}
6564
let cmd = cmd.unwrap();
6665

@@ -77,7 +76,7 @@ pub fn get_cmd_key_bulks<'a>(cmd: &SmolStr, bulk_args: &'a Vec<Bytes>) -> Option
7776
i += range.step;
7877
}
7978
}
80-
},
79+
}
8180
(Some(index), None, None, Some(keynum)) => {
8281
if let Some((first, key_nums)) = first_and_keynums_from_index(index, keynum, bulk_args) {
8382
let mut nums = 0;
@@ -88,7 +87,7 @@ pub fn get_cmd_key_bulks<'a>(cmd: &SmolStr, bulk_args: &'a Vec<Bytes>) -> Option
8887
i += keynum.step;
8988
}
9089
}
91-
},
90+
}
9291
(None, Some(keyword), Some(find_keys_range), None) => {
9392
if let Some((keyword_index, find_last_index)) = range_from_keyword(keyword, find_keys_range, bulk_args) {
9493
let mut i = keyword_index + 1;
@@ -97,16 +96,16 @@ pub fn get_cmd_key_bulks<'a>(cmd: &SmolStr, bulk_args: &'a Vec<Bytes>) -> Option
9796
i += find_keys_range.step;
9897
}
9998
}
100-
},
99+
}
101100
_ => { //other case do not exist in redis command
102101
continue;
103102
}
104103
}
105104
}
106105
if key_bulks.is_empty() {
107-
return None
106+
return None;
108107
}
109-
return Some(key_bulks)
108+
return Some(key_bulks);
110109
}
111110
None
112111
}
@@ -137,7 +136,7 @@ fn first_and_keynums_from_index(index: &Index, keynum: &KeyNum, bulk_args: &Vec<
137136
}
138137

139138
let key_num_bytes = &bulk_args[index.pos + keynum.keynumidx];
140-
let key_nums = unsafe {std::str::from_utf8_unchecked(key_num_bytes.as_ref())};
139+
let key_nums = unsafe { std::str::from_utf8_unchecked(key_num_bytes.as_ref()) };
141140
let key_nums = usize::from_str(key_nums);
142141
if key_nums.is_err() {
143142
return None;
@@ -158,11 +157,11 @@ fn range_from_keyword(keyword: &Keyword, range: &Range, bulk_args: &Vec<Bytes>)
158157
};
159158

160159
let mut keyword_bytes = &bulk_args[keyword_index];
161-
let mut keyword_str = unsafe {std::str::from_utf8_unchecked(keyword_bytes.as_ref())};
160+
let mut keyword_str = unsafe { std::str::from_utf8_unchecked(keyword_bytes.as_ref()) };
162161
while !keyword_str.eq_ignore_ascii_case(&keyword.keyword) && keyword_index < (bulks_length - 1) {
163162
keyword_index = if reverse { keyword_index - 1 } else { keyword_index + 1 };
164163
keyword_bytes = &bulk_args[keyword_index];
165-
keyword_str = unsafe {std::str::from_utf8_unchecked(keyword_bytes.as_ref())};
164+
keyword_str = unsafe { std::str::from_utf8_unchecked(keyword_bytes.as_ref()) };
166165
}
167166
if keyword_index >= (bulks_length - 1) {
168167
return None;
@@ -181,8 +180,12 @@ fn range_from_keyword(keyword: &Keyword, range: &Range, bulk_args: &Vec<Bytes>)
181180

182181
#[cfg(test)]
183182
mod test {
183+
use std::str::FromStr;
184+
184185
use bytes::Bytes;
185186

187+
use redis_command_gen::CmdType;
188+
186189
#[test]
187190
fn test_to_lower_effective() {
188191
let origin = b"HELLO_WORLD-";
@@ -192,80 +195,79 @@ mod test {
192195

193196
#[test]
194197
fn test_is_write_cmd() {
195-
let cmd = smol_str::SmolStr::from("set");
198+
let cmd = CmdType::from_str("set").unwrap();
196199
assert_eq!(super::is_write_cmd(&cmd), true);
197-
let cmd = smol_str::SmolStr::from("get");
200+
let cmd = CmdType::from_str("get").unwrap();
198201
assert_eq!(super::is_write_cmd(&cmd), false);
199202
}
200203

201204
#[test]
202205
fn test_is_connection_cmd() {
203-
let cmd = smol_str::SmolStr::from("auth");
206+
let cmd = CmdType::from_str("auth").unwrap();
204207
assert_eq!(super::is_connection_cmd(&cmd), true);
205-
let cmd = smol_str::SmolStr::from("get");
208+
let cmd = CmdType::from_str("get").unwrap();
206209
assert_eq!(super::is_connection_cmd(&cmd), false);
207210
}
208211

209212
#[test]
210213
fn test_get_cmd_key_bulks() {
211-
212214
/* index and range
213215
*/
214-
let cmd = smol_str::SmolStr::from("set");
216+
let cmd = CmdType::from_str("set").unwrap();
215217
let bulk_args = vec![Bytes::from("set"), Bytes::from("key"), Bytes::from("value")];
216218
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
217219
assert_eq!(key_bulks, Some(vec![b"key".as_ref()]));
218220

219-
let cmd = smol_str::SmolStr::from("lcs");
221+
let cmd = CmdType::from_str("lcs").unwrap();
220222
let bulk_args = vec![Bytes::from("lcs"), Bytes::from("key1"), Bytes::from("key2"), Bytes::from("IDX"), Bytes::from("MINMATCHLEN"), Bytes::from("4"), Bytes::from("WITHMATCHLEN")];
221223
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
222224
assert_eq!(key_bulks, Some(vec![b"key1".as_ref(), b"key2".as_ref()]));
223225

224-
let cmd = smol_str::SmolStr::from("mget");
226+
let cmd = CmdType::from_str("mget").unwrap();
225227
let bulk_args = vec![Bytes::from("mget"), Bytes::from("key1"), Bytes::from("key2")];
226228
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
227229
assert_eq!(key_bulks, Some(vec![b"key1".as_ref(), b"key2".as_ref()]));
228230

229-
let cmd = smol_str::SmolStr::from("mset");
231+
let cmd = CmdType::from_str("mset").unwrap();
230232
let bulk_args = vec![Bytes::from("mset"), Bytes::from("key1"), Bytes::from("value1"), Bytes::from("key2"), Bytes::from("value2")];
231233
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
232234
assert_eq!(key_bulks, Some(vec![b"key1".as_ref(), b"key2".as_ref()]));
233235

234-
let cmd = smol_str::SmolStr::from("blpop");
236+
let cmd = CmdType::from_str("blpop").unwrap();
235237
let bulk_args = vec![Bytes::from("blpop"), Bytes::from("key1"), Bytes::from("key2"), Bytes::from("key3"), Bytes::from("0")];
236238
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
237239
assert_eq!(key_bulks, Some(vec![b"key1".as_ref(), b"key2".as_ref(), b"key3".as_ref()]));
238240

239-
let cmd = smol_str::SmolStr::from("pfmerge");
241+
let cmd = CmdType::from_str("pfmerge").unwrap();
240242
let bulk_args = vec![Bytes::from("pfmerge"), Bytes::from("key1"), Bytes::from("key2"), Bytes::from("key3")];
241243
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
242244
assert_eq!(key_bulks, Some(vec![b"key1".as_ref(), b"key2".as_ref(), b"key3".as_ref()]));
243245

244246
/* index and keynum
245247
*/
246-
let cmd = smol_str::SmolStr::from("lmpop");
248+
let cmd = CmdType::from_str("lmpop").unwrap();
247249
let bulk_args = vec![Bytes::from("lmpop"), Bytes::from("3"), Bytes::from("key1"), Bytes::from("key2"), Bytes::from("key3"), Bytes::from("LEFT"), Bytes::from("COUNT"), Bytes::from("3")];
248250
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
249251
assert_eq!(key_bulks, Some(vec![b"key1".as_ref(), b"key2".as_ref(), b"key3".as_ref()]));
250252

251-
let cmd = smol_str::SmolStr::from("blmpop");
253+
let cmd = CmdType::from_str("blmpop").unwrap();
252254
let bulk_args = vec![Bytes::from("blmpop"), Bytes::from("1000"), Bytes::from("3"), Bytes::from("key1"), Bytes::from("key2"), Bytes::from("key3"), Bytes::from("LEFT"), Bytes::from("COUNT"), Bytes::from("3")];
253255
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
254256
assert_eq!(key_bulks, Some(vec![b"key1".as_ref(), b"key2".as_ref(), b"key3".as_ref()]));
255257

256258
/* keyword and range
257259
*/
258-
let cmd = smol_str::SmolStr::from("georadius");
260+
let cmd = CmdType::from_str("georadius").unwrap();
259261
let bulk_args = vec![Bytes::from("georadius"), Bytes::from("key1"), Bytes::from("15"), Bytes::from("37"), Bytes::from("300"), Bytes::from("km"), Bytes::from("WITHCOORD"), Bytes::from("WITHDIST"), Bytes::from("WITHHASH"), Bytes::from("COUNT"), Bytes::from("3"), Bytes::from("ASC"), Bytes::from("store"), Bytes::from("key2"), Bytes::from("STOREDIST"), Bytes::from("key3")];
260262
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
261263
assert_eq!(key_bulks, Some(vec![b"key1".as_ref(), b"key2".as_ref(), b"key3".as_ref()]));
262264

263-
let cmd = smol_str::SmolStr::from("xread");
265+
let cmd = CmdType::from_str("xread").unwrap();
264266
let bulk_args = vec![Bytes::from("xread"), Bytes::from("count"), Bytes::from("1"), Bytes::from("block"), Bytes::from("1000"), Bytes::from("streams"), Bytes::from("key1"), Bytes::from("key2"), Bytes::from("key3"), Bytes::from("id1"), Bytes::from("id2"), Bytes::from("id3")];
265267
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
266268
assert_eq!(key_bulks, Some(vec![b"key1".as_ref(), b"key2".as_ref(), b"key3".as_ref()]));
267269

268-
let cmd = smol_str::SmolStr::from("migrate");
270+
let cmd = CmdType::from_str("migrate").unwrap();
269271
let bulk_args = vec![Bytes::from("migrate"), Bytes::from("127.0.0.1"), Bytes::from("9001"), Bytes::from("key1"), Bytes::from("0"), Bytes::from("1000"), Bytes::from("COPY"), Bytes::from("REPLACE"), Bytes::from("AUTH"), Bytes::from("123456"), Bytes::from("AUTH2"), Bytes::from("admin"), Bytes::from("root"), Bytes::from("keys"), Bytes::from("key2"), Bytes::from("key3")];
270272
let key_bulks = super::get_cmd_key_bulks(&cmd, &bulk_args);
271273
assert_eq!(key_bulks, Some(vec![b"key1".as_ref(), b"key2".as_ref(), b"key3".as_ref()]));

‎redis-proxy-common/src/lib.rs‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
1+
use std::str::FromStr;
2+
13
use bytes::Bytes;
24
use smol_str::SmolStr;
3-
use redis_command_gen::MULTIPART_COMMANDS;
5+
6+
use redis_command_gen::{CmdType, MULTIPART_COMMANDS};
47

58
// use crate::command::holder::MULTIPART_COMMANDS;
69
use crate::command::utils::get_cmd_key_bulks;
@@ -10,19 +13,22 @@ pub mod command;
1013

1114
#[derive(Debug)]
1215
pub struct ReqPkt {
13-
pub cmd_type: SmolStr,
16+
pub cmd_type: CmdType,
1417
pub bulk_args: Vec<Bytes>,
1518
pub bytes_total: usize,
1619
}
1720

1821
impl ReqPkt {
1922
pub fn new(bulk_args: Vec<Bytes>, bytes_total: usize) -> Self {
20-
let mut cmd_type = bulk_args[0].iter().map(|it| it.to_ascii_lowercase() as char).collect::<SmolStr>();
23+
let mut cmd_type_str = bulk_args[0].iter().map(|it| it.to_ascii_uppercase() as char).collect::<SmolStr>();
24+
// todo, add unknown
25+
let mut cmd_type = CmdType::from_str(&cmd_type_str).unwrap();
2126
if MULTIPART_COMMANDS.contains_key(&cmd_type) && bulk_args.len() > 1 {
22-
cmd_type = cmd_type.chars()
27+
let cmd_with_sub_cmd: SmolStr = cmd_type_str.chars()
2328
.chain([' ' as char])
2429
.chain(bulk_args[1].iter().map(|it| it.to_ascii_lowercase() as char))
2530
.collect();
31+
cmd_type = CmdType::from_str(&cmd_with_sub_cmd).unwrap();
2632
}
2733
return ReqPkt {
2834
cmd_type,
@@ -31,7 +37,7 @@ impl ReqPkt {
3137
};
3238
}
3339

34-
pub fn keys(&self)-> Option<Vec<&[u8]>> {
40+
pub fn keys(&self) -> Option<Vec<&[u8]>> {
3541
get_cmd_key_bulks(&self.cmd_type, &self.bulk_args)
3642
}
3743
}

0 commit comments

Comments
 (0)