Add optional info function to redis_module macro (#211)

* Add optional info function to redis_module macro

* Add InfoContext to abstract raw RedisModuleInfoCtx

* Add InfoContext::add_info_field_long_long

Author: oshadmi

examples/hello.rs

@@ -1,8 +1,10 @@
 #[macro_use]
 extern crate redis_module;
 
-use redis_module::{Context, RedisError, RedisResult, RedisString};
+use redis_module::InfoContext;
+use redis_module::Status;
 
+use redis_module::{Context, RedisError, RedisResult, RedisString};
 fn hello_mul(_: &Context, args: Vec<RedisString>) -> RedisResult {
     if args.len() < 2 {
         return Err(RedisError::WrongArity);
@@ -22,12 +24,19 @@ fn hello_mul(_: &Context, args: Vec<RedisString>) -> RedisResult {
     return Ok(response.into());
 }
 
+fn add_info(ctx: &InfoContext, _for_crash_report: bool) {
+    if ctx.add_info_section(Some("hello")) == Status::Ok {
+        ctx.add_info_field_str("field", "hello_value");
+    }
+}
+
 //////////////////////////////////////////////////////
 
 redis_module! {
     name: "hello",
     version: 1,
     data_types: [],
+    info: add_info,
     commands: [
         ["hello.mul", hello_mul, "", 0, 0, 0],
     ],

src/context/mod.rs

@@ -1,11 +1,11 @@
 use std::ffi::CString;
-use std::os::raw::{c_char, c_int, c_long};
+use std::os::raw::{c_char, c_int, c_long, c_longlong};
 use std::ptr;
 
 use crate::key::{RedisKey, RedisKeyWritable};
-use crate::raw;
 use crate::raw::ModuleOptions;
-use crate::LogLevel;
+use crate::{add_info_field_long_long, add_info_field_str, raw, Status};
+use crate::{add_info_section, LogLevel};
 use crate::{RedisError, RedisResult, RedisString, RedisValue};
 
 #[cfg(feature = "experimental-api")]
@@ -286,3 +286,36 @@ impl Context {
         unsafe { raw::RedisModule_SetModuleOptions.unwrap()(self.ctx, options.bits()) };
     }
 }
+
+pub struct InfoContext {
+    pub ctx: *mut raw::RedisModuleInfoCtx,
+}
+
+impl InfoContext {
+    pub fn new(ctx: *mut raw::RedisModuleInfoCtx) -> Self {
+        Self { ctx }
+    }
+
+    pub fn add_info_section(
+        &self,
+        name: Option<&str>, // assume NULL terminated
+    ) -> Status {
+        add_info_section(self.ctx, name)
+    }
+
+    pub fn add_info_field_str(
+        &self,
+        name: &str, // assume NULL terminated
+        content: &str,
+    ) -> Status {
+        add_info_field_str(self.ctx, name, content)
+    }
+
+    pub fn add_info_field_long_long(
+        &self,
+        name: &str, // assume NULL terminated
+        value: c_longlong,
+    ) -> Status {
+        add_info_field_long_long(self.ctx, name, value)
+    }
+}

src/lib.rs

@@ -1,9 +1,9 @@
 //#![allow(dead_code)]
 
+pub use crate::context::InfoContext;
 use std::os::raw::c_char;
 use std::str::Utf8Error;
 use strum_macros::AsRefStr;
-
 extern crate num_traits;
 
 use libc::size_t;
@@ -61,14 +61,21 @@ fn from_byte_string(byte_str: *const c_char, length: size_t) -> Result<String, U
     String::from_utf8(vec_str).map_err(|e| e.utf8_error())
 }
 
-pub fn base_info_func(ctx: *mut RedisModuleInfoCtx, for_crash_report: bool) {
+pub fn base_info_func(
+    ctx: &InfoContext,
+    for_crash_report: bool,
+    extended_info_func: Option<fn(&InfoContext, bool)>,
+) {
     if !for_crash_report {
-        return;
+        if let Some(func) = extended_info_func {
+            func(ctx, for_crash_report);
+            return;
+        }
     }
     // add rust trace into the crash report
-    if add_info_section(ctx, Some("trace")) == Status::Ok {
+    if ctx.add_info_section(Some("trace")) == Status::Ok {
         let current_backtrace = Backtrace::new();
         let trace = format!("{:?}", current_backtrace);
-        raw::add_info_field_str(ctx, "trace", &trace);
+        ctx.add_info_field_str("trace", &trace);
     }
 }

src/macros.rs

@@ -92,6 +92,7 @@ macro_rules! redis_module {
         ],
         $(init: $init_func:ident,)* $(,)*
         $(deinit: $deinit_func:ident,)* $(,)*
+        $(info: $info_func:ident,)?
         commands: [
             $([
                 $name:expr,
@@ -113,7 +114,10 @@ macro_rules! redis_module {
             ctx: *mut $crate::raw::RedisModuleInfoCtx,
             for_crash_report: i32,
         ) {
-            $crate::base_info_func(ctx, for_crash_report == 1);
+            use $crate::InfoContext;
+            let mut __info_func__cb : Option<fn(&InfoContext, bool)> = None;
+            $( __info_func__cb = Some($info_func); )?
+            $crate::base_info_func(&$crate::InfoContext::new(ctx), for_crash_report == 1, __info_func__cb);
         }
 
         #[no_mangle]

src/raw.rs

@@ -604,6 +604,18 @@ pub fn add_info_field_str(
     }
 }
 
+#[allow(clippy::not_unsafe_ptr_arg_deref)]
+pub fn add_info_field_long_long(
+    ctx: *mut RedisModuleInfoCtx,
+    name: &str, // assume NULL terminated
+    value: c_longlong,
+) -> Status {
+    let name = CString::new(name).unwrap();
+    unsafe {
+        RedisModule_InfoAddFieldLongLong.unwrap()(ctx, name.as_ptr() as *mut c_char, value).into()
+    }
+}
+
 #[cfg(feature = "experimental-api")]
 pub fn export_shared_api(
     ctx: *mut RedisModuleCtx,

tests/integration.rs

@@ -37,14 +37,30 @@ fn test_keys_pos() -> Result<()> {
     let res: Vec<String> = redis::cmd("keys_pos")
         .arg(&["a", "1", "b", "2"])
         .query(&mut con)
-        .with_context(|| "failed to run hello.mul")?;
+        .with_context(|| "failed to run keys_pos")?;
     assert_eq!(res, vec!["a", "b"]);
 
     let res: Result<Vec<String>, RedisError> =
         redis::cmd("keys_pos").arg(&["a", "1", "b"]).query(&mut con);
     if res.is_ok() {
-        return Err(anyhow::Error::msg("Shold return an error"));
+        return Err(anyhow::Error::msg("Shuold return an error"));
     }
 
     Ok(())
 }
+
+#[test]
+fn test_hello_info() -> Result<()> {
+    let _guards = vec![start_redis_server_with_module("hello", 6481)
+        .with_context(|| "failed to start redis server")?];
+    let mut con =
+        get_redis_connection(6481).with_context(|| "failed to connect to redis server")?;
+
+    let res: String = redis::cmd("INFO")
+        .arg("HELLO")
+        .query(&mut con)
+        .with_context(|| "failed to run hello.mul")?;
+    assert!(res.contains("hello_field:hello_value"));
+
+    Ok(())
+}