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