mini_redis/cmd/
set.rs

1use crate::cmd::{Parse, ParseError};
2use crate::{Connection, Db, Frame};
3
4use bytes::Bytes;
5use std::time::Duration;
6use tracing::{debug, instrument};
7
8/// Set `key` to hold the string `value`.
9///
10/// If `key` already holds a value, it is overwritten, regardless of its type.
11/// Any previous time to live associated with the key is discarded on successful
12/// SET operation.
13///
14/// # Options
15///
16/// Currently, the following options are supported:
17///
18/// * EX `seconds` -- Set the specified expire time, in seconds.
19/// * PX `milliseconds` -- Set the specified expire time, in milliseconds.
20#[derive(Debug)]
21pub struct Set {
22    /// the lookup key
23    key: String,
24
25    /// the value to be stored
26    value: Bytes,
27
28    /// When to expire the key
29    expire: Option<Duration>,
30}
31
32impl Set {
33    /// Create a new `Set` command which sets `key` to `value`.
34    ///
35    /// If `expire` is `Some`, the value should expire after the specified
36    /// duration.
37    pub fn new(key: impl ToString, value: Bytes, expire: Option<Duration>) -> Set {
38        Set {
39            key: key.to_string(),
40            value,
41            expire,
42        }
43    }
44
45    /// Get the key
46    pub fn key(&self) -> &str {
47        &self.key
48    }
49
50    /// Get the value
51    pub fn value(&self) -> &Bytes {
52        &self.value
53    }
54
55    /// Get the expire
56    pub fn expire(&self) -> Option<Duration> {
57        self.expire
58    }
59
60    /// Parse a `Set` instance from a received frame.
61    ///
62    /// The `Parse` argument provides a cursor-like API to read fields from the
63    /// `Frame`. At this point, the entire frame has already been received from
64    /// the socket.
65    ///
66    /// The `SET` string has already been consumed.
67    ///
68    /// # Returns
69    ///
70    /// Returns the `Set` value on success. If the frame is malformed, `Err` is
71    /// returned.
72    ///
73    /// # Format
74    ///
75    /// Expects an array frame containing at least 3 entries.
76    ///
77    /// ```text
78    /// SET key value [EX seconds|PX milliseconds]
79    /// ```
80    pub(crate) fn parse_frames(parse: &mut Parse) -> crate::Result<Set> {
81        use ParseError::EndOfStream;
82
83        // Read the key to set. This is a required field
84        let key = parse.next_string()?;
85
86        // Read the value to set. This is a required field.
87        let value = parse.next_bytes()?;
88
89        // The expiration is optional. If nothing else follows, then it is
90        // `None`.
91        let mut expire = None;
92
93        // Attempt to parse another string.
94        match parse.next_string() {
95            Ok(s) if s.to_uppercase() == "EX" => {
96                // An expiration is specified in seconds. The next value is an
97                // integer.
98                let secs = parse.next_int()?;
99                expire = Some(Duration::from_secs(secs));
100            }
101            Ok(s) if s.to_uppercase() == "PX" => {
102                // An expiration is specified in milliseconds. The next value is
103                // an integer.
104                let ms = parse.next_int()?;
105                expire = Some(Duration::from_millis(ms));
106            }
107            // Currently, mini-redis does not support any of the other SET
108            // options. An error here results in the connection being
109            // terminated. Other connections will continue to operate normally.
110            Ok(_) => return Err("currently `SET` only supports the expiration option".into()),
111            // The `EndOfStream` error indicates there is no further data to
112            // parse. In this case, it is a normal run time situation and
113            // indicates there are no specified `SET` options.
114            Err(EndOfStream) => {}
115            // All other errors are bubbled up, resulting in the connection
116            // being terminated.
117            Err(err) => return Err(err.into()),
118        }
119
120        Ok(Set { key, value, expire })
121    }
122
123    /// Apply the `Set` command to the specified `Db` instance.
124    ///
125    /// The response is written to `dst`. This is called by the server in order
126    /// to execute a received command.
127    #[instrument(skip(self, db, dst))]
128    pub(crate) async fn apply(self, db: &Db, dst: &mut Connection) -> crate::Result<()> {
129        // Set the value in the shared database state.
130        db.set(self.key, self.value, self.expire);
131
132        // Create a success response and write it to `dst`.
133        let response = Frame::Simple("OK".to_string());
134        debug!(?response);
135        dst.write_frame(&response).await?;
136
137        Ok(())
138    }
139
140    /// Converts the command into an equivalent `Frame`.
141    ///
142    /// This is called by the client when encoding a `Set` command to send to
143    /// the server.
144    pub(crate) fn into_frame(self) -> Frame {
145        let mut frame = Frame::array();
146        frame.push_bulk(Bytes::from("set".as_bytes()));
147        frame.push_bulk(Bytes::from(self.key.into_bytes()));
148        frame.push_bulk(self.value);
149        if let Some(ms) = self.expire {
150            // Expirations in Redis procotol can be specified in two ways
151            // 1. SET key value EX seconds
152            // 2. SET key value PX milliseconds
153            // We the second option because it allows greater precision and
154            // src/bin/cli.rs parses the expiration argument as milliseconds
155            // in duration_from_ms_str()
156            frame.push_bulk(Bytes::from("px".as_bytes()));
157            frame.push_int(ms.as_millis() as u64);
158        }
159        frame
160    }
161}