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