1 #![cfg(feature = "script")]
2 use sha1::Sha1;
3 
4 use crate::cmd::cmd;
5 use crate::connection::ConnectionLike;
6 use crate::types::{ErrorKind, FromRedisValue, RedisResult, ToRedisArgs};
7 
8 /// Represents a lua script.
9 #[derive(Debug, Clone)]
10 pub struct Script {
11     code: String,
12     hash: String,
13 }
14 
15 /// The script object represents a lua script that can be executed on the
16 /// redis server.  The object itself takes care of automatic uploading and
17 /// execution.  The script object itself can be shared and is immutable.
18 ///
19 /// Example:
20 ///
21 /// ```rust,no_run
22 /// # let client = redis::Client::open("redis://127.0.0.1/").unwrap();
23 /// # let mut con = client.get_connection().unwrap();
24 /// let script = redis::Script::new(r"
25 ///     return tonumber(ARGV[1]) + tonumber(ARGV[2]);
26 /// ");
27 /// let result = script.arg(1).arg(2).invoke(&mut con);
28 /// assert_eq!(result, Ok(3));
29 /// ```
30 impl Script {
31     /// Creates a new script object.
new(code: &str) -> Script32     pub fn new(code: &str) -> Script {
33         let mut hash = Sha1::new();
34         hash.update(code.as_bytes());
35         Script {
36             code: code.to_string(),
37             hash: hash.digest().to_string(),
38         }
39     }
40 
41     /// Returns the script's SHA1 hash in hexadecimal format.
get_hash(&self) -> &str42     pub fn get_hash(&self) -> &str {
43         &self.hash
44     }
45 
46     /// Creates a script invocation object with a key filled in.
47     #[inline]
key<T: ToRedisArgs>(&self, key: T) -> ScriptInvocation<'_>48     pub fn key<T: ToRedisArgs>(&self, key: T) -> ScriptInvocation<'_> {
49         ScriptInvocation {
50             script: self,
51             args: vec![],
52             keys: key.to_redis_args(),
53         }
54     }
55 
56     /// Creates a script invocation object with an argument filled in.
57     #[inline]
arg<T: ToRedisArgs>(&self, arg: T) -> ScriptInvocation<'_>58     pub fn arg<T: ToRedisArgs>(&self, arg: T) -> ScriptInvocation<'_> {
59         ScriptInvocation {
60             script: self,
61             args: arg.to_redis_args(),
62             keys: vec![],
63         }
64     }
65 
66     /// Returns an empty script invocation object.  This is primarily useful
67     /// for programmatically adding arguments and keys because the type will
68     /// not change.  Normally you can use `arg` and `key` directly.
69     #[inline]
prepare_invoke(&self) -> ScriptInvocation<'_>70     pub fn prepare_invoke(&self) -> ScriptInvocation<'_> {
71         ScriptInvocation {
72             script: self,
73             args: vec![],
74             keys: vec![],
75         }
76     }
77 
78     /// Invokes the script directly without arguments.
79     #[inline]
invoke<T: FromRedisValue>(&self, con: &mut dyn ConnectionLike) -> RedisResult<T>80     pub fn invoke<T: FromRedisValue>(&self, con: &mut dyn ConnectionLike) -> RedisResult<T> {
81         ScriptInvocation {
82             script: self,
83             args: vec![],
84             keys: vec![],
85         }
86         .invoke(con)
87     }
88 }
89 
90 /// Represents a prepared script call.
91 pub struct ScriptInvocation<'a> {
92     script: &'a Script,
93     args: Vec<Vec<u8>>,
94     keys: Vec<Vec<u8>>,
95 }
96 
97 /// This type collects keys and other arguments for the script so that it
98 /// can be then invoked.  While the `Script` type itself holds the script,
99 /// the `ScriptInvocation` holds the arguments that should be invoked until
100 /// it's sent to the server.
101 impl<'a> ScriptInvocation<'a> {
102     /// Adds a regular argument to the invocation.  This ends up as `ARGV[i]`
103     /// in the script.
104     #[inline]
arg<'b, T: ToRedisArgs>(&'b mut self, arg: T) -> &'b mut ScriptInvocation<'a> where 'a: 'b,105     pub fn arg<'b, T: ToRedisArgs>(&'b mut self, arg: T) -> &'b mut ScriptInvocation<'a>
106     where
107         'a: 'b,
108     {
109         arg.write_redis_args(&mut self.args);
110         self
111     }
112 
113     /// Adds a key argument to the invocation.  This ends up as `KEYS[i]`
114     /// in the script.
115     #[inline]
key<'b, T: ToRedisArgs>(&'b mut self, key: T) -> &'b mut ScriptInvocation<'a> where 'a: 'b,116     pub fn key<'b, T: ToRedisArgs>(&'b mut self, key: T) -> &'b mut ScriptInvocation<'a>
117     where
118         'a: 'b,
119     {
120         key.write_redis_args(&mut self.keys);
121         self
122     }
123 
124     /// Invokes the script and returns the result.
125     #[inline]
invoke<T: FromRedisValue>(&self, con: &mut dyn ConnectionLike) -> RedisResult<T>126     pub fn invoke<T: FromRedisValue>(&self, con: &mut dyn ConnectionLike) -> RedisResult<T> {
127         loop {
128             match cmd("EVALSHA")
129                 .arg(self.script.hash.as_bytes())
130                 .arg(self.keys.len())
131                 .arg(&*self.keys)
132                 .arg(&*self.args)
133                 .query(con)
134             {
135                 Ok(val) => {
136                     return Ok(val);
137                 }
138                 Err(err) => {
139                     if err.kind() == ErrorKind::NoScriptError {
140                         cmd("SCRIPT")
141                             .arg("LOAD")
142                             .arg(self.script.code.as_bytes())
143                             .query(con)?;
144                     } else {
145                         fail!(err);
146                     }
147                 }
148             }
149         }
150     }
151 
152     /// Asynchronously invokes the script and returns the result.
153     #[inline]
154     #[cfg(feature = "aio")]
invoke_async<C, T>(&self, con: &mut C) -> RedisResult<T> where C: crate::aio::ConnectionLike, T: FromRedisValue,155     pub async fn invoke_async<C, T>(&self, con: &mut C) -> RedisResult<T>
156     where
157         C: crate::aio::ConnectionLike,
158         T: FromRedisValue,
159     {
160         let mut eval_cmd = cmd("EVALSHA");
161         eval_cmd
162             .arg(self.script.hash.as_bytes())
163             .arg(self.keys.len())
164             .arg(&*self.keys)
165             .arg(&*self.args);
166 
167         let mut load_cmd = cmd("SCRIPT");
168         load_cmd.arg("LOAD").arg(self.script.code.as_bytes());
169         match eval_cmd.query_async(con).await {
170             Ok(val) => {
171                 // Return the value from the script evaluation
172                 Ok(val)
173             }
174             Err(err) => {
175                 // Load the script into Redis if the script hash wasn't there already
176                 if err.kind() == ErrorKind::NoScriptError {
177                     load_cmd.query_async(con).await?;
178                     eval_cmd.query_async(con).await
179                 } else {
180                     Err(err)
181                 }
182             }
183         }
184     }
185 }
186