redis_lua/
script.rs

1use crate::types::ScriptArg;
2use futures::prelude::*;
3
4/// Script information which is generated by proc-macro.
5#[derive(Clone, Debug)]
6pub struct Info {
7    /// The entire script including arguments initialization.
8    script: &'static str,
9    /// The script excluding arguments initialization.
10    body: &'static str,
11    /// The list of arguments.
12    args: &'static [&'static str],
13}
14
15impl Info {
16    /// Create the new script information.
17    pub fn new(script: &'static str, body: &'static str, args: &'static [&'static str]) -> Self {
18        Self { script, body, args }
19    }
20}
21
22/// To make sure `Script` be object safe.
23fn _object_safe(_: &dyn Script) {}
24
25/// Represents a complete invocable script which has a complete set of arguments.
26pub trait Script {
27    /// Retrieve all the script information.
28    fn info(&self, _: &mut Vec<Info>, _: &mut Vec<ScriptArg>);
29
30    /// Join another script making self as inner.
31    fn join<T: Script>(self, other: T) -> ScriptJoin<Self, T>
32    where
33        Self: Sized,
34    {
35        ScriptJoin(self, other)
36    }
37
38    /// Invoke the script.
39    fn invoke<T>(self, con: &mut dyn redis::ConnectionLike) -> redis::RedisResult<T>
40    where
41        T: redis::FromRedisValue,
42        Self: Sized,
43    {
44        let mut info = vec![];
45        let mut args = vec![];
46        self.info(&mut info, &mut args);
47        let script = gen_script(&info, &args);
48        let mut invoke = script.prepare_invoke();
49        for wr in args {
50            invoke.arg(wr);
51        }
52        invoke.invoke(con)
53    }
54
55    /// Invoke the script asynchronously.
56    fn invoke_async<'a, C, T>(self, con: &'a mut C) -> redis::RedisFuture<'a, T>
57    where
58        C: redis::aio::ConnectionLike + Send,
59        T: redis::FromRedisValue + Send,
60        Self: Sized + Send + 'a,
61    {
62        async move {
63            let mut info = vec![];
64            let mut args = vec![];
65            self.info(&mut info, &mut args);
66            let script = gen_script(&info, &args);
67            let mut invoke = script.prepare_invoke();
68            for wr in args {
69                invoke.arg(wr);
70            }
71            invoke.invoke_async(con).await
72        }
73        .boxed()
74    }
75}
76
77impl<S: Script + ?Sized> Script for Box<S> {
78    fn info(&self, infos: &mut Vec<Info>, args: &mut Vec<ScriptArg>) {
79        (**self).info(infos, args);
80    }
81}
82
83impl Script for () {
84    fn info(&self, _: &mut Vec<Info>, _: &mut Vec<ScriptArg>) {}
85}
86
87/// Represents the set of two scripts which are joined.
88pub struct ScriptJoin<S, T>(S, T);
89
90impl<S, T> Script for ScriptJoin<S, T>
91where
92    S: Script,
93    T: Script,
94{
95    fn info(&self, info: &mut Vec<Info>, args: &mut Vec<ScriptArg>) {
96        self.0.info(info, args);
97        self.1.info(info, args);
98    }
99}
100
101/// Take another script as the inner of the script.
102pub trait TakeScript<I> {
103    type Item;
104
105    /// Take the inner script.
106    fn take(self, inner: I) -> Self::Item;
107}
108
109/// Generate a script from a list of script information.
110pub fn gen_script(info: &[Info], args: &[ScriptArg]) -> redis::Script {
111    assert!(info.len() > 0, "No script information");
112
113    // Generate the joined script.
114    let mut arg_index = 0;
115    let mut script = String::new();
116    let last = info.len() - 1;
117    for (index, info) in info.iter().enumerate() {
118        let prefix = if index == last { "return " } else { "" };
119        let mut init = String::new();
120
121        for arg in info.args {
122            let pack = args[arg_index].pack();
123
124            arg_index += 1;
125
126            if pack {
127                init += &format!("local {} = cmsgpack.unpack(ARGV[{}]) ", arg, arg_index);
128            } else {
129                init += &format!("local {} = ARGV[{}] ", arg, arg_index);
130            }
131        }
132
133        script += &format!("{}(function() {} {} end)();\n", prefix, init, info.body);
134    }
135    redis::Script::new(&script)
136}