redis_module_ext/
command.rs

1use std::borrow::{Borrow, Cow};
2use std::ffi::{CStr, CString};
3use std::fmt::Write;
4use std::marker::PhantomData;
5use std::ptr::NonNull;
6
7use redis_module::raw;
8
9pub trait RedisCommand: Sized {
10    const NAME: &'static CStr;
11
12    fn flags(ctx: &redis_module::Context) -> Vec<CommandFlag>;
13    fn command_info(ctx: &redis_module::Context) -> RedisModuleCommandInfo;
14    fn invoke(ctx: &redis_module::Context, args: CommandArgs) -> redis_module::RedisResult;
15
16    fn invoke_raw(ctx: *mut raw::RedisModuleCtx, argv: *mut *mut raw::RedisModuleString, argc: i32) -> i32 {
17        let wctx = redis_module::Context::new(ctx);
18        let resp = Self::invoke(&wctx, CommandArgs::new(ctx, argv, argc as usize));
19        wctx.reply(resp) as i32
20    }
21
22    fn register(ctx: &redis_module::Context) -> redis_module::RedisResult<()> {
23        unsafe extern "C" fn callback<C: RedisCommand>(
24            ctx: *mut raw::RedisModuleCtx,
25            argv: *mut *mut raw::RedisModuleString,
26            argc: i32,
27        ) -> i32 {
28            C::invoke_raw(ctx, argv, argc)
29        }
30
31        let name = Self::NAME.to_string_lossy();
32        let flags = fmtools::fmt(|f| {
33            let mut first = true;
34            for flag in Self::flags(ctx) {
35                if !first {
36                    f.write_char(' ')?;
37                }
38                first = false;
39
40                f.write_str(flag.as_str())?;
41            }
42            Ok(())
43        })
44        .to_string();
45
46        let flags = CString::new(flags).unwrap();
47
48        if unsafe {
49            redis_module::RedisModule_CreateCommand.unwrap()(
50                ctx.ctx,
51                Self::NAME.as_ptr(),
52                Some(callback::<Self>),
53                flags.as_ptr(),
54                0,
55                0,
56                0,
57            )
58        } == raw::Status::Err as i32
59        {
60            return Err(redis_module::RedisError::String(format!("Failed register command {name}.",)));
61        }
62
63        // Register the extra data of the command
64        let command = unsafe { raw::RedisModule_GetCommand.unwrap()(ctx.ctx, Self::NAME.as_ptr()) };
65
66        if command.is_null() {
67            return Err(redis_module::RedisError::String(format!(
68                "Failed finding command {name} after registration.",
69            )));
70        }
71
72        let command_info = Self::command_info(ctx);
73        let mut ffi = command_info.as_ffi();
74
75        if unsafe { raw::RedisModule_SetCommandInfo.unwrap()(command, ffi.ffi_mut()) } == raw::Status::Err as i32 {
76            return Err(redis_module::RedisError::String(format!(
77                "Failed setting info for command {name}.",
78            )));
79        }
80
81        Ok(())
82    }
83}
84
85pub struct CommandArgs {
86    ctx: Option<NonNull<raw::RedisModuleCtx>>,
87    argv: *mut *mut raw::RedisModuleString,
88    argc: usize,
89}
90
91impl CommandArgs {
92    fn new(ctx: *mut raw::RedisModuleCtx, argv: *mut *mut raw::RedisModuleString, argc: usize) -> Self {
93        CommandArgs {
94            ctx: NonNull::new(ctx),
95            argv,
96            argc,
97        }
98    }
99}
100
101impl Iterator for CommandArgs {
102    type Item = redis_module::RedisString;
103
104    fn next(&mut self) -> Option<Self::Item> {
105        if self.argc == 0 {
106            None
107        } else {
108            self.argc -= 1;
109            let arg = unsafe { *self.argv };
110            self.argv = unsafe { self.argv.add(1) };
111            Some(redis_module::RedisString::new(self.ctx, arg))
112        }
113    }
114}
115
116pub struct RedisModuleCommandInfo {
117    pub summary: Option<Cow<'static, CStr>>,
118    pub complexity: Option<Cow<'static, CStr>>,
119    pub since: Option<Cow<'static, CStr>>,
120    pub history: Vec<RedisModuleCommandHistoryEntry>,
121    pub tips: Option<Cow<'static, CStr>>,
122    pub arity: i32,
123    pub key_specs: Vec<RedisModuleCommandKeySpec>,
124    pub args: Vec<RedisModuleCommandArg>,
125}
126
127pub struct RedisModuleCommandHistoryEntry {
128    pub since: Option<Cow<'static, CStr>>,
129    pub changes: Option<Cow<'static, CStr>>,
130}
131
132pub struct RedisModuleCommandKeySpec {
133    pub notes: Option<Cow<'static, CStr>>,
134    pub flags: Vec<KeySpecFlag>,
135    pub begin_search: KeySpecBeginSearch,
136    pub find_keys: Option<KeySpecFindKeys>,
137}
138
139pub enum KeySpecBeginSearch {
140    Index(i32),
141    Keyword {
142        keyword: Cow<'static, CStr>,
143        start_from: i32,
144    },
145}
146
147pub enum KeySpecFindKeys {
148    Range {
149        last_key: i32,
150        key_step: i32,
151        limit: i32,
152    },
153    Keynum {
154        keynum_idx: i32,
155        first_key: i32,
156        key_step: i32,
157    },
158}
159
160#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash)]
161#[repr(u32)]
162pub enum RedisModuleCommandArgKind {
163    /// String argument.
164    String = raw::RedisModuleCommandArgType_REDISMODULE_ARG_TYPE_STRING,
165    /// Integer argument.
166    Integer = raw::RedisModuleCommandArgType_REDISMODULE_ARG_TYPE_INTEGER,
167    /// Double-precision float argument.
168    Double = raw::RedisModuleCommandArgType_REDISMODULE_ARG_TYPE_DOUBLE,
169    /// String argument representing a keyname.
170    Key = raw::RedisModuleCommandArgType_REDISMODULE_ARG_TYPE_KEY,
171    /// String, but regex pattern.
172    Pattern = raw::RedisModuleCommandArgType_REDISMODULE_ARG_TYPE_PATTERN,
173    /// Integer, but Unix timestamp.
174    UnixTime = raw::RedisModuleCommandArgType_REDISMODULE_ARG_TYPE_UNIX_TIME,
175    /// Argument doesn't have a placeholder. It's just a token without a value. Example: the KEEPTTL option of the SET command.
176    PureToken = raw::RedisModuleCommandArgType_REDISMODULE_ARG_TYPE_PURE_TOKEN,
177    /// Used when the user can choose only one of a few sub-arguments. Requires subargs. Example: the NX and XX options of SET.
178    OneOf = raw::RedisModuleCommandArgType_REDISMODULE_ARG_TYPE_ONEOF,
179    /// Used when one wants to group together several sub-arguments, usually to apply something on all of them,
180    /// like making the entire group "optional". Requires subargs. Example: the LIMIT offset count parameters in ZRANGE.
181    Block = raw::RedisModuleCommandArgType_REDISMODULE_ARG_TYPE_BLOCK,
182}
183
184#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash)]
185#[repr(u32)]
186pub enum RedisModuleCommandArgFlag {
187    /// The argument is optional (like GET in the SET command).
188    Optional = raw::REDISMODULE_CMD_ARG_OPTIONAL,
189    /// The argument may repeat itself (like key in DEL).
190    Multiple = raw::REDISMODULE_CMD_ARG_MULTIPLE,
191    /// The argument may repeat itself, and so does its token (like GET pattern in SORT).
192    MultipleToken = raw::REDISMODULE_CMD_ARG_MULTIPLE_TOKEN,
193}
194
195#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash)]
196pub enum CommandFlag {
197    /// The command may modify the data set (it may also read from it).
198    Write,
199
200    /// The command returns data from keys but never writes.
201    ReadOnly,
202
203    /// The command is an administrative command (may change replication or perform similar tasks).
204    Admin,
205
206    /// The command may use additional memory and should be denied during out of memory conditions.
207    DenyOOM,
208
209    /// Don't allow this command in Lua scripts.
210    DenyScript,
211
212    /// Allow this command while the server is loading data. Only commands not interacting with the data set
213    /// should be allowed to run in this mode. If not sure don't use this flag.
214    AllowLoading,
215
216    /// The command publishes things on Pub/Sub channels.
217    PubSub,
218
219    /// The command may have different outputs even starting from the same input arguments and key values.
220    /// Starting from Redis 7.0 this flag has been deprecated. Declaring a command as "random" can be done using
221    /// command tips, see <https://redis.io/topics/command-tips>.
222    #[deprecated = "Declaring a command as 'random' can be done using command tips, see https://redis.io/topics/command-tips."]
223    Random,
224
225    /// The command is allowed to run on slaves that don't serve stale data. Don't use if you don't know what
226    /// this means.
227    AllowStale,
228
229    /// Don't propagate the command on monitor. Use this if the command has sensitive data among the arguments.
230    NoMonitor,
231
232    /// Don't log this command in the slowlog. Use this if the command has sensitive data among the arguments.
233    NoSlowlog,
234
235    /// The command time complexity is not greater than O(log(N)) where N is the size of the collection or
236    /// anything else representing the normal scalability issue with the command.
237    Fast,
238
239    /// The command implements the interface to return the arguments that are keys. Used when start/stop/step
240    /// is not enough because of the command syntax.
241    GetkeysApi,
242
243    /// The command should not register in Redis Cluster since is not designed to work with it because, for
244    /// example, is unable to report the position of the keys, programmatically creates key names, or any
245    /// other reason.
246    NoCluster,
247
248    /// This command can be run by an un-authenticated client. Normally this is used by a command that is used
249    /// to authenticate a client.
250    NoAuth,
251
252    /// This command may generate replication traffic, even though it's not a write command.
253    MayReplicate,
254
255    /// All the keys this command may take are optional
256    NoMandatoryKeys,
257
258    /// The command has the potential to block the client.
259    Blocking,
260
261    /// Permit the command while the server is blocked either by a script or by a slow module command, see
262    /// RM_Yield.
263    AllowBusy,
264
265    /// The command implements the interface to return the arguments that are channels.
266    GetchannelsApi,
267
268    /// Internal command, one that should not be exposed to the user connections.
269    /// For example, module commands that are called by the modules, commands that do not perform ACL validations (relying on earlier checks)
270    Internal,
271}
272
273impl CommandFlag {
274    pub fn as_str(&self) -> &str {
275        match self {
276            Self::Write => "write",
277            Self::ReadOnly => "readonly",
278            Self::Admin => "admin",
279            Self::DenyOOM => "deny-oom",
280            Self::DenyScript => "deny-script",
281            Self::AllowLoading => "allow-loading",
282            Self::PubSub => "pubsub",
283            #[allow(deprecated)]
284            Self::Random => "random",
285            Self::AllowStale => "allow-stale",
286            Self::NoMonitor => "no-monitor",
287            Self::NoSlowlog => "no-slowlog",
288            Self::Fast => "fast",
289            Self::GetkeysApi => "getkeys-api",
290            Self::NoCluster => "no-cluster",
291            Self::NoAuth => "no-auth",
292            Self::MayReplicate => "may-replicate",
293            Self::NoMandatoryKeys => "no-mandatory-keys",
294            Self::Blocking => "blocking",
295            Self::AllowBusy => "allow-busy",
296            Self::GetchannelsApi => "getchannels-api",
297            Self::Internal => "internal",
298        }
299    }
300}
301
302impl std::fmt::Display for CommandFlag {
303    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
304        f.write_str(self.as_str())
305    }
306}
307
308#[derive(Debug, Clone, Copy, PartialEq, PartialOrd, Eq, Ord, Hash)]
309#[repr(u32)]
310pub enum KeySpecFlag {
311    /// Read-Only. Reads the value of the key, but doesn't necessarily return it.
312    ReadOnly = raw::REDISMODULE_CMD_KEY_RO,
313
314    /// Read-Write. Modifies the data stored in the value of the key or its metadata.
315    ReadWrite = raw::REDISMODULE_CMD_KEY_RW,
316
317    /// Overwrite. Overwrites the data stored in the value of the key.
318    Overwrite = raw::REDISMODULE_CMD_KEY_OW,
319
320    /// Deletes the key.
321    Remove = raw::REDISMODULE_CMD_KEY_RM,
322
323    /// Returns, copies or uses the user data from the value of the key.
324    Access = raw::REDISMODULE_CMD_KEY_ACCESS,
325
326    /// Updates data to the value, new value may depend on the old value.
327    Update = raw::REDISMODULE_CMD_KEY_UPDATE,
328
329    /// Adds data to the value with no chance of modification or deletion of existing data.
330    Insert = raw::REDISMODULE_CMD_KEY_INSERT,
331
332    /// Explicitly deletes some content from the value of the key.
333    Delete = raw::REDISMODULE_CMD_KEY_DELETE,
334
335    /// The key is not actually a key, but should be routed in cluster mode as if it was a key.
336    NotKey = raw::REDISMODULE_CMD_KEY_NOT_KEY,
337
338    /// The keyspec might not point out all the keys it should cover.
339    Incomplete = raw::REDISMODULE_CMD_KEY_INCOMPLETE,
340
341    /// Some keys might have different flags depending on arguments.
342    VariableFlags = raw::REDISMODULE_CMD_KEY_VARIABLE_FLAGS,
343}
344
345fn opt_cstr_ptr<R: AsRef<CStr>>(cstr: impl Borrow<Option<R>>) -> *const std::ffi::c_char {
346    cstr.borrow()
347        .as_ref()
348        .map(|cstr| cstr.as_ref().as_ptr())
349        .unwrap_or(std::ptr::null())
350}
351
352pub struct RedisModuleCommandArg {
353    pub name: Cow<'static, CStr>,
354    pub kind: RedisModuleCommandArgKind,
355    pub key_spec_idx: i32,
356    pub token: Option<Cow<'static, CStr>>,
357    pub summary: Option<Cow<'static, CStr>>,
358    pub since: Option<Cow<'static, CStr>>,
359    pub flags: i32,
360    pub deprecated_since: Option<Cow<'static, CStr>>,
361    pub sub_args: Vec<RedisModuleCommandArg>,
362    pub display_text: Option<Cow<'static, CStr>>,
363}
364
365struct RawRedisModuleCommandInfo<'a> {
366    ffi: raw::RedisModuleCommandInfo,
367    _history_storage: Box<[raw::RedisModuleCommandHistoryEntry]>,
368    _key_spec_storage: Box<[raw::RedisModuleCommandKeySpec]>,
369    _args_storage: Box<[Box<[raw::RedisModuleCommandArg]>]>,
370    phantom: PhantomData<&'a ()>,
371}
372
373impl RawRedisModuleCommandInfo<'_> {
374    fn ffi_mut(&mut self) -> &mut raw::RedisModuleCommandInfo {
375        &mut self.ffi
376    }
377}
378
379impl RedisModuleCommandInfo {
380    fn as_ffi(&self) -> RawRedisModuleCommandInfo<'_> {
381        static COMMNAD_INFO_VERSION: raw::RedisModuleCommandInfoVersion = raw::RedisModuleCommandInfoVersion {
382            version: 1,
383            sizeof_historyentry: std::mem::size_of::<raw::RedisModuleCommandHistoryEntry>(),
384            sizeof_keyspec: std::mem::size_of::<raw::RedisModuleCommandKeySpec>(),
385            sizeof_arg: std::mem::size_of::<raw::RedisModuleCommandArg>(),
386        };
387
388        let mut history_vec = self
389            .history
390            .iter()
391            .map(|h| raw::RedisModuleCommandHistoryEntry {
392                changes: opt_cstr_ptr(&h.changes),
393                since: opt_cstr_ptr(&h.since),
394            })
395            .collect::<Vec<_>>();
396        if !history_vec.is_empty() {
397            history_vec.push(unsafe { std::mem::zeroed() });
398        }
399
400        let mut history_storage = history_vec.into_boxed_slice();
401
402        let mut key_spec_vec = self
403            .key_specs
404            .iter()
405            .map(|spec| raw::RedisModuleCommandKeySpec {
406                notes: opt_cstr_ptr(&spec.notes),
407                flags: spec.flags.iter().fold(0, |mut flags, flag| {
408                    flags |= *flag as u64;
409                    flags
410                }),
411                begin_search_type: match spec.begin_search {
412                    KeySpecBeginSearch::Index(_) => raw::RedisModuleKeySpecBeginSearchType_REDISMODULE_KSPEC_BS_INDEX,
413                    KeySpecBeginSearch::Keyword { .. } => {
414                        raw::RedisModuleKeySpecBeginSearchType_REDISMODULE_KSPEC_BS_KEYWORD
415                    }
416                },
417                bs: match &spec.begin_search {
418                    KeySpecBeginSearch::Index(idx) => raw::RedisModuleCommandKeySpec__bindgen_ty_1 {
419                        index: raw::RedisModuleCommandKeySpec__bindgen_ty_1__bindgen_ty_1 { pos: *idx },
420                    },
421                    KeySpecBeginSearch::Keyword { keyword, start_from } => raw::RedisModuleCommandKeySpec__bindgen_ty_1 {
422                        keyword: raw::RedisModuleCommandKeySpec__bindgen_ty_1__bindgen_ty_2 {
423                            keyword: keyword.as_ptr(),
424                            startfrom: *start_from,
425                        },
426                    },
427                },
428                find_keys_type: match spec.find_keys {
429                    Some(KeySpecFindKeys::Keynum { .. }) => raw::RedisModuleKeySpecFindKeysType_REDISMODULE_KSPEC_FK_KEYNUM,
430                    Some(KeySpecFindKeys::Range { .. }) => raw::RedisModuleKeySpecFindKeysType_REDISMODULE_KSPEC_FK_RANGE,
431                    None => raw::RedisModuleKeySpecFindKeysType_REDISMODULE_KSPEC_FK_OMITTED,
432                },
433                fk: match spec.find_keys {
434                    Some(KeySpecFindKeys::Keynum {
435                        first_key,
436                        key_step,
437                        keynum_idx,
438                    }) => raw::RedisModuleCommandKeySpec__bindgen_ty_2 {
439                        keynum: raw::RedisModuleCommandKeySpec__bindgen_ty_2__bindgen_ty_2 {
440                            firstkey: first_key,
441                            keynumidx: keynum_idx,
442                            keystep: key_step,
443                        },
444                    },
445                    Some(KeySpecFindKeys::Range {
446                        key_step,
447                        last_key,
448                        limit,
449                    }) => raw::RedisModuleCommandKeySpec__bindgen_ty_2 {
450                        range: raw::RedisModuleCommandKeySpec__bindgen_ty_2__bindgen_ty_1 {
451                            keystep: key_step,
452                            lastkey: last_key,
453                            limit,
454                        },
455                    },
456                    None => unsafe { std::mem::zeroed() },
457                },
458            })
459            .collect::<Vec<_>>();
460        if !key_spec_vec.is_empty() {
461            key_spec_vec.push(unsafe { std::mem::zeroed() });
462        }
463
464        let mut key_spec_storage = key_spec_vec.into_boxed_slice();
465
466        fn convert(
467            arg_vec: &mut Vec<Box<[raw::RedisModuleCommandArg]>>,
468            arg: &RedisModuleCommandArg,
469        ) -> raw::RedisModuleCommandArg {
470            let mut sub_args = arg.sub_args.iter().map(|arg| convert(arg_vec, arg)).collect::<Vec<_>>();
471            let subargs = if !sub_args.is_empty() {
472                sub_args.push(unsafe { std::mem::zeroed() });
473                let mut sub_args = sub_args.into_boxed_slice();
474                let ptr = sub_args.as_mut_ptr();
475                arg_vec.push(sub_args);
476                ptr
477            } else {
478                std::ptr::null_mut()
479            };
480            raw::RedisModuleCommandArg {
481                name: arg.name.as_ptr(),
482                type_: arg.kind as u32,
483                key_spec_index: arg.key_spec_idx,
484                token: opt_cstr_ptr(&arg.token),
485                summary: opt_cstr_ptr(&arg.summary),
486                since: opt_cstr_ptr(&arg.since),
487                flags: arg.flags,
488                deprecated_since: opt_cstr_ptr(&arg.deprecated_since),
489                subargs,
490                display_text: opt_cstr_ptr(&arg.display_text),
491            }
492        }
493
494        let mut args_storage_vec = Vec::new();
495        let mut args_vec = self
496            .args
497            .iter()
498            .map(|arg| convert(&mut args_storage_vec, arg))
499            .collect::<Vec<_>>();
500        let args = if !args_vec.is_empty() {
501            args_vec.push(unsafe { std::mem::zeroed() });
502            let mut args_vec = args_vec.into_boxed_slice();
503            let ptr = args_vec.as_mut_ptr();
504            args_storage_vec.push(args_vec);
505            ptr
506        } else {
507            std::ptr::null_mut()
508        };
509
510        RawRedisModuleCommandInfo {
511            ffi: raw::RedisModuleCommandInfo {
512                version: &raw const COMMNAD_INFO_VERSION,
513                summary: opt_cstr_ptr(&self.summary),
514                complexity: opt_cstr_ptr(&self.complexity),
515                since: opt_cstr_ptr(&self.since),
516                history: if history_storage.is_empty() {
517                    std::ptr::null_mut()
518                } else {
519                    history_storage.as_mut_ptr()
520                },
521                tips: opt_cstr_ptr(&self.tips),
522                arity: self.arity,
523                key_specs: if key_spec_storage.is_empty() {
524                    std::ptr::null_mut()
525                } else {
526                    key_spec_storage.as_mut_ptr()
527                },
528                args,
529            },
530            _history_storage: history_storage,
531            _args_storage: args_storage_vec.into_boxed_slice(),
532            _key_spec_storage: key_spec_storage,
533            phantom: PhantomData,
534        }
535    }
536}