From 49d41791ab2866eb8e7896539ebd0510c0c6da50 Mon Sep 17 00:00:00 2001
From: Qu Chen <quchen@amazon.com>
Date: Fri, 3 Feb 2023 18:13:35 -0800
Subject: [PATCH] Add new API to remove expiration for RedisKeyWritable.

---
 Cargo.toml           |  4 ++++
 examples/expire.rs   | 30 ++++++++++++++++++++++++++++++
 src/key.rs           | 11 +++++++++++
 tests/integration.rs | 38 ++++++++++++++++++++++++++++++++++++++
 4 files changed, 83 insertions(+)
 create mode 100644 examples/expire.rs

diff --git a/Cargo.toml b/Cargo.toml
index a9b1c764..2c0e673a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -116,6 +116,10 @@ crate-type = ["cdylib"]
 name = "open_key_with_flags"
 crate-type = ["cdylib"]
 
+[[example]]
+name = "expire"
+crate-type = ["cdylib"]
+
 [dependencies]
 bitflags = "2"
 libc = "0.2"
diff --git a/examples/expire.rs b/examples/expire.rs
new file mode 100644
index 00000000..dfd0cacb
--- /dev/null
+++ b/examples/expire.rs
@@ -0,0 +1,30 @@
+use redis_module::{redis_module, Context, NextArg, RedisError, RedisResult, RedisString};
+use std::time::Duration;
+
+fn expire_cmd(ctx: &Context, args: Vec<RedisString>) -> RedisResult {
+    if args.len() < 3 {
+        return Err(RedisError::WrongArity);
+    }
+
+    let mut args = args.into_iter().skip(1);
+    let key_name = args.next_arg()?;
+    let ttl_sec = args.next_i64()?;
+    let key = ctx.open_key_writable(&key_name);
+    if ttl_sec >= 0 {
+        key.set_expire(Duration::new(ttl_sec as u64, 0))
+    } else {
+        key.remove_expire()
+    }
+}
+
+//////////////////////////////////////////////////////
+
+redis_module! {
+    name: "expire",
+    version: 1,
+    allocator: (redis_module::alloc::RedisAlloc, redis_module::alloc::RedisAlloc),
+    data_types: [],
+    commands: [
+        ["expire.cmd", expire_cmd, "write fast deny-oom", 1, 1, 1],
+    ],
+}
diff --git a/src/key.rs b/src/key.rs
index 9c4685ca..195b06e2 100644
--- a/src/key.rs
+++ b/src/key.rs
@@ -329,6 +329,17 @@ impl RedisKeyWritable {
         }
     }
 
+    /// Remove expiration from a key if it exists.
+    pub fn remove_expire(&self) -> RedisResult {
+        match raw::set_expire(self.key_inner, REDISMODULE_NO_EXPIRE.into()) {
+            raw::Status::Ok => REDIS_OK,
+
+            // Error may occur if the key wasn't open for writing or is an
+            // empty key.
+            raw::Status::Err => Err(RedisError::Str("Error while removing key expire")),
+        }
+    }
+
     pub fn write(&self, val: &str) -> RedisResult {
         let val_str = RedisString::create(NonNull::new(self.ctx), val);
         match raw::string_set(self.key_inner, val_str.inner) {
diff --git a/tests/integration.rs b/tests/integration.rs
index 7864e366..f40130ca 100644
--- a/tests/integration.rs
+++ b/tests/integration.rs
@@ -705,3 +705,41 @@ fn test_open_key_with_flags() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_expire() -> Result<()> {
+    let port: u16 = 6502;
+    let _guards = vec![start_redis_server_with_module("expire", port)
+        .with_context(|| "failed to start redis server")?];
+    let mut con =
+        get_redis_connection(port).with_context(|| "failed to connect to redis server")?;
+
+    // Create a key without TTL
+    redis::cmd("set")
+        .arg(&["key", "value"])
+        .query(&mut con)
+        .with_context(|| "failed to run set")?;
+
+    let ttl: i64 = redis::cmd("ttl").arg(&["key"]).query(&mut con)?;
+    assert_eq!(ttl, -1);
+
+    // Set TTL on the key
+    redis::cmd("expire.cmd")
+        .arg(&["key", "100"])
+        .query(&mut con)
+        .with_context(|| "failed to run expire.cmd")?;
+
+    let ttl: i64 = redis::cmd("ttl").arg(&["key"]).query(&mut con)?;
+    assert!(ttl > 0);
+
+    // Remove TTL on the key
+    redis::cmd("expire.cmd")
+        .arg(&["key", "-1"])
+        .query(&mut con)
+        .with_context(|| "failed to run expire.cmd")?;
+
+    let ttl: i64 = redis::cmd("ttl").arg(&["key"]).query(&mut con)?;
+    assert_eq!(ttl, -1);
+
+    Ok(())
+}