Add set_options cmd (#879)

Add `set_options`, which supports the full set of options
available for Redis `SET` command.

redis/src/commands/mod.rs

@@ -3,7 +3,7 @@
 use crate::cmd::{cmd, Cmd, Iter};
 use crate::connection::{Connection, ConnectionLike, Msg};
 use crate::pipeline::Pipeline;
-use crate::types::{FromRedisValue, NumericBehavior, RedisResult, ToRedisArgs, RedisWrite, Expiry};
+use crate::types::{FromRedisValue, NumericBehavior, RedisResult, ToRedisArgs, RedisWrite, Expiry, SetExpiry, ExistenceCheck};
 
 #[macro_use]
 mod macros;
@@ -87,6 +87,11 @@ implement_commands! {
         cmd("SET").arg(key).arg(value)
     }
 
+    /// Set the string value of a key with options.
+    fn set_options<K: ToRedisArgs, V: ToRedisArgs>(key: K, value: V, options: SetOptions) {
+        cmd("SET").arg(key).arg(value).arg(options)
+    }
+
     /// Sets multiple keys to their values.
     #[allow(deprecated)]
     #[deprecated(since = "0.22.4", note = "Renamed to mset() to reflect Redis name")]
@@ -2064,3 +2069,91 @@ impl ToRedisArgs for Direction {
         out.write_arg(s);
     }
 }
+
+/// Options for the [SET](https://redis.io/commands/set) command
+///
+/// # Example
+/// ```rust,no_run
+/// use redis::{Commands, RedisResult, SetOptions, SetExpiry, ExistenceCheck};
+/// fn set_key_value(
+///     con: &mut redis::Connection,
+///     key: &str,
+///     value: &str,
+/// ) -> RedisResult<Vec<usize>> {
+///     let opts = SetOptions::default()
+///         .conditional_set(ExistenceCheck::NX)
+///         .get(true)
+///         .with_expiration(SetExpiry::EX(60));
+///     con.set_options(key, value, opts)
+/// }
+/// ```
+#[derive(Default)]
+pub struct SetOptions {
+    conditional_set: Option<ExistenceCheck>,
+    get: bool,
+    expiration: Option<SetExpiry>,
+}
+
+impl SetOptions {
+    /// Set the existence check for the SET command
+    pub fn conditional_set(mut self, existence_check: ExistenceCheck) -> Self {
+        self.conditional_set = Some(existence_check);
+        self
+    }
+
+    /// Set the GET option for the SET command
+    pub fn get(mut self, get: bool) -> Self {
+        self.get = get;
+        self
+    }
+
+    /// Set the expiration for the SET command
+    pub fn with_expiration(mut self, expiration: SetExpiry) -> Self {
+        self.expiration = Some(expiration);
+        self
+    }
+}
+
+impl ToRedisArgs for SetOptions {
+    fn write_redis_args<W>(&self, out: &mut W)
+    where
+        W: ?Sized + RedisWrite,
+    {
+        if let Some(ref conditional_set) = self.conditional_set {
+            match conditional_set {
+                ExistenceCheck::NX => {
+                    out.write_arg(b"NX");
+                }
+                ExistenceCheck::XX => {
+                    out.write_arg(b"XX");
+                }
+            }
+        }
+        if self.get {
+            out.write_arg(b"GET");
+        }
+        if let Some(ref expiration) = self.expiration {
+            match expiration {
+                SetExpiry::EX(secs) => {
+                    out.write_arg(b"EX");
+                    out.write_arg(format!("{}", secs).as_bytes());
+                }
+                SetExpiry::PX(millis) => {
+                    out.write_arg(b"PX");
+                    out.write_arg(format!("{}", millis).as_bytes());
+                }
+                SetExpiry::EXAT(unix_time) => {
+                    out.write_arg(b"EXAT");
+                    out.write_arg(format!("{}", unix_time).as_bytes());
+                }
+                SetExpiry::PXAT(unix_time) => {
+                    out.write_arg(b"PXAT");
+                    out.write_arg(format!("{}", unix_time).as_bytes());
+                }
+                SetExpiry::KEEPTTL => {
+                    out.write_arg(b"KEEPTTL");
+                }
+            }
+        }
+    }
+}

redis/src/lib.rs

@@ -326,7 +326,7 @@ In addition to the synchronous interface that's been explained above there also
 asynchronous interface based on [`futures`][] and [`tokio`][].
 
 This interface exists under the `aio` (async io) module (which requires that the `aio` feature
-is enabled) and largely mirrors the synchronous with a few concessions to make it fit the 
+is enabled) and largely mirrors the synchronous with a few concessions to make it fit the
 constraints of `futures`.
 
 ```rust,no_run
@@ -363,7 +363,9 @@ assert_eq!(result, Ok(("foo".to_string(), b"bar".to_vec())));
 // public api
 pub use crate::client::Client;
 pub use crate::cmd::{cmd, pack_command, pipe, Arg, Cmd, Iter};
-pub use crate::commands::{Commands, ControlFlow, Direction, LposOptions, PubSubCommands};
+pub use crate::commands::{
+    Commands, ControlFlow, Direction, LposOptions, PubSubCommands, SetOptions,
+};
 pub use crate::connection::{
     parse_redis_url, transaction, Connection, ConnectionAddr, ConnectionInfo, ConnectionLike,
     IntoConnectionInfo, Msg, PubSub, RedisConnectionInfo,
@@ -391,6 +393,8 @@ pub use crate::types::{
     InfoDict,
     NumericBehavior,
     Expiry,
+    SetExpiry,
+    ExistenceCheck,
 
     // error and result types
     RedisError,

redis/src/types.rs

@@ -45,6 +45,28 @@ pub enum Expiry {
     PERSIST,
 }
 
+/// Helper enum that is used to define expiry time for SET command
+pub enum SetExpiry {
+    /// EX seconds -- Set the specified expire time, in seconds.
+    EX(usize),
+    /// PX milliseconds -- Set the specified expire time, in milliseconds.
+    PX(usize),
+    /// EXAT timestamp-seconds -- Set the specified Unix time at which the key will expire, in seconds.
+    EXAT(usize),
+    /// PXAT timestamp-milliseconds -- Set the specified Unix time at which the key will expire, in milliseconds.
+    PXAT(usize),
+    /// KEEPTTL -- Retain the time to live associated with the key.
+    KEEPTTL,
+}
+
+/// Helper enum that is used to define existence checks
+pub enum ExistenceCheck {
+    /// NX -- Only set the key if it does not already exist.
+    NX,
+    /// XX -- Only set the key if it already exists.
+    XX,
+}
+
 /// Helper enum that is used in some situations to describe
 /// the behavior of arguments in a numeric context.
 #[derive(PartialEq, Eq, Clone, Debug, Copy)]

redis/tests/support/mod.rs

@@ -41,6 +41,8 @@ mod cluster;
 #[cfg(any(feature = "cluster", feature = "cluster-async"))]
 mod mock_cluster;
 
+mod util;
+
 #[cfg(any(feature = "cluster", feature = "cluster-async"))]
 pub use self::cluster::*;
 

redis/tests/support/util.rs

@@ -0,0 +1,10 @@
+#[macro_export]
+macro_rules! assert_args {
+    ($value:expr, $($args:expr),+) => {
+        let args = $value.to_redis_args();
+        let strings: Vec<_> = args.iter()
+                                .map(|a| std::str::from_utf8(a.as_ref()).unwrap())
+                                .collect();
+        assert_eq!(strings, vec![$($args),+]);
+    }
+}

redis/tests/test_basic.rs

@@ -1,8 +1,8 @@
 #![allow(clippy::let_unit_value)]
 
 use redis::{
-    Commands, ConnectionInfo, ConnectionLike, ControlFlow, ErrorKind, Expiry, PubSubCommands,
-    RedisResult,
+    Commands, ConnectionInfo, ConnectionLike, ControlFlow, ErrorKind, ExistenceCheck, Expiry,
+    PubSubCommands, RedisResult, SetExpiry, SetOptions, ToRedisArgs,
 };
 
 use std::collections::{BTreeMap, BTreeSet};
@@ -1184,3 +1184,53 @@ fn test_multi_generics() {
     let _: () = con.rename(999_i64, b"set2").unwrap();
     assert_eq!(con.sunionstore("res", &[b"set1", b"set2"]), Ok(3));
 }
+
+#[test]
+fn test_set_options_with_get() {
+    let ctx = TestContext::new();
+    let mut con = ctx.connection();
+
+    let opts = SetOptions::default().get(true);
+    let data: Option<String> = con.set_options(1, "1", opts).unwrap();
+    assert_eq!(data, None);
+
+    let opts = SetOptions::default().get(true);
+    let data: Option<String> = con.set_options(1, "1", opts).unwrap();
+    assert_eq!(data, Some("1".to_string()));
+}
+
+#[test]
+fn test_set_options_options() {
+    let empty = SetOptions::default();
+    assert_eq!(ToRedisArgs::to_redis_args(&empty).len(), 0);
+
+    let opts = SetOptions::default()
+        .conditional_set(ExistenceCheck::NX)
+        .get(true)
+        .with_expiration(SetExpiry::PX(1000));
+
+    assert_args!(&opts, "NX", "GET", "PX", "1000");
+
+    let opts = SetOptions::default()
+        .conditional_set(ExistenceCheck::XX)
+        .get(true)
+        .with_expiration(SetExpiry::PX(1000));
+
+    assert_args!(&opts, "XX", "GET", "PX", "1000");
+
+    let opts = SetOptions::default()
+        .conditional_set(ExistenceCheck::XX)
+        .with_expiration(SetExpiry::KEEPTTL);
+
+    assert_args!(&opts, "XX", "KEEPTTL");
+
+    let opts = SetOptions::default()
+        .conditional_set(ExistenceCheck::XX)
+        .with_expiration(SetExpiry::EXAT(100));
+
+    assert_args!(&opts, "XX", "EXAT", "100");
+
+    let opts = SetOptions::default().with_expiration(SetExpiry::EX(1000));
+
+    assert_args!(&opts, "EX", "1000");
+}

redis/tests/test_streams.rs

@@ -11,16 +11,6 @@ use std::str;
 use std::thread::sleep;
 use std::time::Duration;
 
-macro_rules! assert_args {
-    ($value:expr, $($args:expr),+) => {
-        let args = $value.to_redis_args();
-        let strings: Vec<_> = args.iter()
-                                .map(|a| str::from_utf8(a.as_ref()).unwrap())
-                                .collect();
-        assert_eq!(strings, vec![$($args),+]);
-    }
-}
-
 fn xadd(con: &mut Connection) {
     let _: RedisResult<String> =
         con.xadd("k1", "1000-0", &[("hello", "world"), ("redis", "streams")]);