MOD-8720 command proc macro: add acl and command arguments parameters

Cargo.toml

@@ -134,6 +134,7 @@ nix = { version = "0.26", default-features = false }
 cfg-if = "1"
 redis-module-macros-internals = { path = "./redismodule-rs-macros-internals" }
 log = "0.4"
+common = { path = "./common" }
 
 [dev-dependencies]
 anyhow = "1"

common/Cargo.toml

@@ -0,0 +1,7 @@
+[package]
+name = "common"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+serde = { version = "1.0", features = ["derive"] }

common/src/lib.rs

@@ -0,0 +1,109 @@
+use serde::Deserialize;
+
+#[derive(Debug, Clone, Deserialize, PartialEq, Eq, Default)]
+pub enum AclCategory {
+    #[default]
+    None,
+    Keyspace,
+    Read,
+    Write,
+    Set,
+    SortedSet,
+    List,
+    Hash,
+    String,
+    Bitmap,
+    HyperLogLog,
+    Geo,
+    Stream,
+    PubSub,
+    Admin,
+    Fast,
+    Slow,
+    Blocking,
+    Dangerous,
+    Connection,
+    Transaction,
+    Scripting,
+    Single(String),
+    Multi(Vec<AclCategory>),
+}
+
+impl From<Vec<AclCategory>> for AclCategory {
+    fn from(value: Vec<AclCategory>) -> Self {
+        AclCategory::Multi(value)
+    }
+}
+
+impl From<&str> for AclCategory {
+    fn from(value: &str) -> Self {
+        match value {
+            "" => AclCategory::None,
+            "keyspace" => AclCategory::Keyspace,
+            "read" => AclCategory::Read,
+            "write" => AclCategory::Write,
+            "set" => AclCategory::Set,
+            "sortedset" => AclCategory::SortedSet,
+            "list" => AclCategory::List,
+            "hash" => AclCategory::Hash,
+            "string" => AclCategory::String,
+            "bitmap" => AclCategory::Bitmap,
+            "hyperloglog" => AclCategory::HyperLogLog,
+            "geo" => AclCategory::Geo,
+            "stream" => AclCategory::Stream,
+            "pubsub" => AclCategory::PubSub,
+            "admin" => AclCategory::Admin,
+            "fast" => AclCategory::Fast,
+            "slow" => AclCategory::Slow,
+            "blocking" => AclCategory::Blocking,
+            "dangerous" => AclCategory::Dangerous,
+            "connection" => AclCategory::Connection,
+            "transaction" => AclCategory::Transaction,
+            "scripting" => AclCategory::Scripting,
+            _ if !value.contains(" ") => AclCategory::Single(value.to_string()),
+            _ => AclCategory::Multi(value.split_whitespace().map(AclCategory::from).collect()),
+        }
+    }
+}
+
+impl From<AclCategory> for String {
+    fn from(value: AclCategory) -> Self {
+        match value {
+            AclCategory::None => "".to_string(),
+            AclCategory::Keyspace => "keyspace".to_string(),
+            AclCategory::Read => "read".to_string(),
+            AclCategory::Write => "write".to_string(),
+            AclCategory::Set => "set".to_string(),
+            AclCategory::SortedSet => "sortedset".to_string(),
+            AclCategory::List => "list".to_string(),
+            AclCategory::Hash => "hash".to_string(),
+            AclCategory::String => "string".to_string(),
+            AclCategory::Bitmap => "bitmap".to_string(),
+            AclCategory::HyperLogLog => "hyperloglog".to_string(),
+            AclCategory::Geo => "geo".to_string(),
+            AclCategory::Stream => "stream".to_string(),
+            AclCategory::PubSub => "pubsub".to_string(),
+            AclCategory::Admin => "admin".to_string(),
+            AclCategory::Fast => "fast".to_string(),
+            AclCategory::Slow => "slow".to_string(),
+            AclCategory::Blocking => "blocking".to_string(),
+            AclCategory::Dangerous => "dangerous".to_string(),
+            AclCategory::Connection => "connection".to_string(),
+            AclCategory::Transaction => "transaction".to_string(),
+            AclCategory::Scripting => "scripting".to_string(),
+            AclCategory::Single(s) => s,
+            AclCategory::Multi(v) => v
+                .into_iter()
+                .map(String::from)
+                .collect::<Vec<_>>()
+                .join(" "),
+        }
+    }
+}
+
+impl std::fmt::Display for AclCategory {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", String::from(self.clone()))
+    }
+}
+

redismodule-rs-macros/Cargo.toml

@@ -17,6 +17,7 @@ quote = "1.0"
 proc-macro2 = "1"
 serde = { version = "1", features = ["derive"] }
 serde_syn = "0.1.0"
+common = { path = "../common" }
 
 [lib]
 name = "redis_module_macros"

redismodule-rs-macros/src/command.rs

@@ -1,3 +1,4 @@
+use common::AclCategory;
 use proc_macro::TokenStream;
 use proc_macro2::Ident;
 use quote::quote;
@@ -223,6 +224,68 @@ pub struct KeySpecArg {
     find_keys: FindKeys,
 }
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize)]
+pub enum CommandArgType {
+    String,
+    Integer,
+    Double,
+    Key,
+    Pattern,
+    UnixTime,
+    PureToken,
+    OneOf,
+    Block,
+}
+
+impl From<CommandArgType> for u32 {
+    fn from(arg_type: CommandArgType) -> Self {
+        match arg_type {
+            CommandArgType::String => 0,
+            CommandArgType::Integer => 1,
+            CommandArgType::Double => 2,
+            CommandArgType::Key => 3,
+            CommandArgType::Pattern => 4,
+            CommandArgType::UnixTime => 5,
+            CommandArgType::PureToken => 6,
+            CommandArgType::OneOf => 7,
+            CommandArgType::Block => 8,
+        }
+    }
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub enum CommandArgFlags {
+    None,
+    Optional,
+    Multiple,
+    MultipleToken,
+}
+
+impl From<&CommandArgFlags> for &'static str {
+    fn from(value: &CommandArgFlags) -> Self {
+        match value {
+            CommandArgFlags::None => "NONE",
+            CommandArgFlags::Optional => "OPTIONAL",
+            CommandArgFlags::Multiple => "MULTIPLE",
+            CommandArgFlags::MultipleToken => "MULTIPLE_TOKEN",
+        }
+    }
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct CommandArg {
+    pub name: String,
+    pub arg_type: CommandArgType,
+    pub key_spec_index: Option<u32>,
+    pub token: Option<String>,
+    pub summary: Option<String>,
+    pub since: Option<String>,
+    pub flags: Option<Vec<CommandArgFlags>>,
+    pub deprecated_since: Option<String>,
+    pub subargs: Option<Vec<CommandArg>>,
+    pub display_text: Option<String>,
+}
+
 #[derive(Debug, Deserialize)]
 struct Args {
     name: Option<String>,
@@ -234,6 +297,8 @@ struct Args {
     tips: Option<String>,
     arity: i64,
     key_spec: Vec<KeySpecArg>,
+    args: Option<Vec<CommandArg>>,
+    acl_categories: Option<Vec<AclCategory>>,
 }
 
 impl Parse for Args {
@@ -247,6 +312,52 @@ fn to_token_stream(s: Option<String>) -> proc_macro2::TokenStream {
         .unwrap_or(quote! {None})
 }
 
+fn generate_command_arg(arg: &CommandArg) -> proc_macro2::TokenStream {
+    let name = &arg.name;
+    let arg_type: u32 = arg.arg_type.into();
+    let key_spec_index = arg
+        .key_spec_index
+        .map(|v| quote! {Some(#v)})
+        .unwrap_or(quote! {None});
+    let token = to_token_stream(arg.token.clone());
+    let summary = to_token_stream(arg.summary.clone());
+    let since = to_token_stream(arg.since.clone());
+    let flags: Vec<&'static str> = arg
+        .flags
+        .as_ref()
+        .map(|v| v.iter().map(|v| v.into()).collect())
+        .unwrap_or_default();
+    let flags = quote! {
+        vec![#(redis_module::commands::CommandArgFlags::try_from(#flags)?, )*]
+    };
+    let deprecated_since = to_token_stream(arg.deprecated_since.clone());
+    let display_text = to_token_stream(arg.display_text.clone());
+
+    let subargs = if let Some(subargs_vec) = &arg.subargs {
+        let subargs_tokens: Vec<_> = subargs_vec.iter().map(generate_command_arg).collect();
+        quote! {
+            Some(vec![#(#subargs_tokens),*])
+        }
+    } else {
+        quote! { None }
+    };
+
+    quote! {
+        redis_module::commands::RedisModuleCommandArg::new(
+            #name.to_owned(),
+            #arg_type,
+            #key_spec_index,
+            #token,
+            #summary,
+            #since,
+            #flags.into(),
+            #deprecated_since,
+            #subargs,
+            #display_text,
+        )
+    }
+}
+
 pub(crate) fn redis_command(attr: TokenStream, item: TokenStream) -> TokenStream {
     let args = parse_macro_input!(attr as Args);
     let func: ItemFn = match syn::parse(item) {
@@ -358,6 +469,24 @@ pub(crate) fn redis_command(attr: TokenStream, item: TokenStream) -> TokenStream
         })
         .collect();
 
+    let command_args: Vec<_> = args
+        .args
+        .as_ref()
+        .map(|v| v.iter().map(generate_command_arg).collect())
+        .unwrap_or_default();
+
+    let acl_categories = args
+        .acl_categories
+        .map(|v| v.into_iter().map(String::from).collect::<Vec<_>>());
+
+    let acl_categories_tokens = if let Some(categories) = &acl_categories {
+        quote! {
+            Some(vec![#(#categories.to_owned()),*])
+        }
+    } else {
+        quote! { None }
+    };
+
     let gen = quote! {
         #func
 
@@ -385,6 +514,7 @@ pub(crate) fn redis_command(attr: TokenStream, item: TokenStream) -> TokenStream
                     ),
                 )*
             ];
+            let command_args = vec![#(#command_args),*];
             Ok(redis_module::commands::CommandInfo::new(
                 #name_literal.to_owned(),
                 Some(#flags_literal.to_owned()),
@@ -396,6 +526,8 @@ pub(crate) fn redis_command(attr: TokenStream, item: TokenStream) -> TokenStream
                 #arity_literal,
                 key_spec,
                 #c_function_name,
+                command_args,
+                #acl_categories_tokens,
             ))
         }
     };