featherbit/plugins/native/
store_incr.rs1use async_trait::async_trait;
9use std::collections::HashMap;
10use std::sync::Arc;
11
12use crate::context::Context;
13use crate::plugins::resources::PluginResources;
14use crate::plugins::util::store_kv::{self, StoreHandle};
15use crate::plugins::{Plugin, PluginResult};
16use crate::vars::template::Template;
17
18#[cfg(feature = "redis-store")]
19use crate::plugins::PluginOutput;
20
21#[cfg(feature = "redis-store")]
29const INCR_SCRIPT: &str = r#"
30local n = redis.call('INCRBY', KEYS[1], ARGV[1])
31if tonumber(ARGV[2]) > 0 and (tonumber(ARGV[3]) == 1 or redis.call('TTL', KEYS[1]) < 0) then
32 redis.call('EXPIRE', KEYS[1], ARGV[2])
33end
34return n
35"#;
36
37pub struct StoreIncrPlugin {
38 store: StoreHandle,
39 key: Template,
40 by: i64,
41 ttl_seconds: Option<u64>,
42 refresh_ttl: bool,
43 name: String,
44 #[cfg(feature = "redis-store")]
45 script: redis::Script,
46}
47
48impl std::fmt::Debug for StoreIncrPlugin {
57 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58 f.debug_struct("StoreIncrPlugin")
59 .field("store", &self.store)
60 .field("key", &self.key)
61 .field("by", &self.by)
62 .field("ttl_seconds", &self.ttl_seconds)
63 .field("refresh_ttl", &self.refresh_ttl)
64 .field("name", &self.name)
65 .finish()
66 }
67}
68
69fn parse_refresh_ttl(config: &HashMap<String, serde_json::Value>) -> Result<bool, String> {
77 match config.get("refresh_ttl") {
78 None | Some(serde_json::Value::Null) => Ok(false),
79 Some(v) => v
80 .as_bool()
81 .ok_or_else(|| "store-incr: 'refresh_ttl' must be a boolean".to_string()),
82 }
83}
84
85fn parse_by(config: &HashMap<String, serde_json::Value>) -> Result<i64, String> {
87 match config.get("by") {
88 None | Some(serde_json::Value::Null) => Ok(1),
89 Some(v) => v
90 .as_i64()
91 .ok_or_else(|| "store-incr: 'by' must be an integer".to_string()),
92 }
93}
94
95impl StoreIncrPlugin {
96 pub fn from_config(
104 config: &HashMap<String, serde_json::Value>,
105 resources: &Arc<PluginResources>,
106 ) -> Result<Self, String> {
107 let name = config
108 .get("name")
109 .and_then(|v| v.as_str())
110 .filter(|s| !s.is_empty())
111 .ok_or_else(|| {
112 "store-incr: 'name' is required (the context.message key receiving the new value)"
113 .to_string()
114 })?
115 .to_string();
116 let key = store_kv::required_template(config, "key", "store-incr")?;
117 let by = parse_by(config)?;
118 let ttl_seconds = store_kv::optional_ttl(config, "store-incr")?;
119 let refresh_ttl = parse_refresh_ttl(config)?;
120 if refresh_ttl && ttl_seconds.is_none() {
121 return Err(
122 "store-incr: 'refresh_ttl' requires 'ttl_seconds'; there is no expiry to refresh"
123 .to_string(),
124 );
125 }
126 Ok(Self {
127 store: store_kv::resolve(config, resources, "store-incr")?,
128 key,
129 by,
130 ttl_seconds,
131 refresh_ttl,
132 name,
133 #[cfg(feature = "redis-store")]
134 script: redis::Script::new(INCR_SCRIPT),
135 })
136 }
137}
138
139#[async_trait]
140impl Plugin for StoreIncrPlugin {
141 fn plugin_type(&self) -> &str {
142 "store-incr"
143 }
144
145 fn reads_response_body(&self) -> bool {
146 self.key.references_response_body()
147 }
148
149 #[cfg(feature = "redis-store")]
150 async fn execute(&self, mut ctx: Context) -> PluginResult {
151 let rendered = self.key.render(&ctx).to_string();
152 if rendered.is_empty() {
153 return Err(store_kv::key_invalid(
154 ctx,
155 "store-incr",
156 "INCRBY",
157 &self.store.name,
158 ));
159 }
160 let key = self.store.key_for(&rendered);
161
162 let mut conn = match self.store.conn().await {
163 Ok(c) => c,
164 Err(e) => {
165 return Err(store_kv::store_error(
166 ctx,
167 "store-incr",
168 "INCRBY",
169 &self.store.name,
170 e,
171 ))
172 }
173 };
174
175 let n: i64 = match self
176 .script
177 .key(key.as_str())
178 .arg(self.by)
179 .arg(self.ttl_seconds.unwrap_or(0))
180 .arg(i64::from(self.refresh_ttl))
181 .invoke_async(&mut conn)
182 .await
183 {
184 Ok(n) => n,
185 Err(e) => {
186 return Err(if store_kv::is_value_type_error(&e) {
192 store_kv::value_invalid(
193 ctx,
194 "store-incr",
195 "INCRBY",
196 &self.store.name,
197 format!("value at '{}' is not usable as an integer", rendered),
198 )
199 } else {
200 store_kv::store_error(
201 ctx,
202 "store-incr",
203 "INCRBY",
204 &self.store.name,
205 e.to_string(),
206 )
207 });
208 }
209 };
210
211 ctx.message
212 .insert(self.name.clone(), serde_json::Value::from(n));
213 Ok(PluginOutput::success(ctx))
214 }
215
216 #[cfg(not(feature = "redis-store"))]
217 async fn execute(&self, ctx: Context) -> PluginResult {
218 Err(store_kv::store_error(
219 ctx,
220 "store-incr",
221 "INCRBY",
222 &self.store.name,
223 "built without the redis-store feature".to_string(),
224 ))
225 }
226}
227
228#[cfg(test)]
229mod tests {
230 use super::*;
231 use crate::plugins::resources::PluginResources;
232 use std::collections::HashMap;
233
234 fn cfg(json: serde_json::Value) -> HashMap<String, serde_json::Value> {
235 serde_json::from_value(json).unwrap()
236 }
237
238 #[test]
239 fn test_requires_a_name() {
240 let r = PluginResources::empty();
241 let err =
242 StoreIncrPlugin::from_config(&cfg(serde_json::json!({ "store": "s", "key": "k" })), &r)
243 .unwrap_err();
244 assert!(err.contains("name"), "{err}");
245 }
246
247 #[test]
248 fn test_by_defaults_to_one_and_accepts_negatives() {
249 assert_eq!(parse_by(&cfg(serde_json::json!({}))).unwrap(), 1);
250 assert_eq!(parse_by(&cfg(serde_json::json!({ "by": -2 }))).unwrap(), -2);
251 assert!(parse_by(&cfg(serde_json::json!({ "by": "x" }))).is_err());
252 }
253
254 #[test]
258 fn test_refresh_ttl_without_a_ttl_is_a_config_error() {
259 let r = PluginResources::empty();
260 let err = StoreIncrPlugin::from_config(
261 &cfg(serde_json::json!({
262 "store": "s", "key": "k", "name": "n", "refresh_ttl": true
263 })),
264 &r,
265 )
266 .unwrap_err();
267 assert!(err.contains("refresh_ttl"), "{err}");
268 assert!(err.contains("ttl_seconds"), "{err}");
269 }
270
271 #[test]
275 fn test_refresh_ttl_defaults_to_false_and_parses() {
276 assert!(!parse_refresh_ttl(&cfg(serde_json::json!({}))).unwrap());
277 assert!(parse_refresh_ttl(&cfg(serde_json::json!({ "refresh_ttl": true }))).unwrap());
278 assert!(parse_refresh_ttl(&cfg(serde_json::json!({ "refresh_ttl": "yes" }))).is_err());
279 }
280}