db_redis/
pool.rs

1use deadpool::managed::Object;
2use redis::{
3    Client, FromRedisValue, RedisError, RedisResult, ToRedisArgs,
4    aio::{ConnectionManager, ConnectionManagerConfig},
5    cluster::ClusterClientBuilder,
6    cmd,
7};
8use std::fmt::Display;
9use crate::connection::RedisConnection;
10
11pub struct Redis {
12    pub connection: RedisConnection,
13    pub db_name: String,
14}
15
16impl Redis {
17    pub async fn init(uri: &str, db_name: &str) -> RedisResult<Redis> {
18        let nodes = uri.split(",").collect::<Vec<&str>>();
19        if nodes.len() == 1 {
20            let client = Client::open(uri)?;
21            let config = ConnectionManagerConfig::default().set_number_of_retries(usize::MAX);
22            let connection = ConnectionManager::new_with_config(client.clone(), config).await?;
23            let connection = RedisConnection::Single(connection);
24            Ok(Redis {
25                connection,
26                db_name: db_name.to_string(),
27            })
28        } else {
29            let client = ClusterClientBuilder::new(nodes)
30                .use_protocol(redis::ProtocolVersion::RESP3)
31                .retries(u32::MAX)
32                .build()?;
33            let connection = client.get_async_connection().await?;
34            let connection = RedisConnection::Cluster(connection);
35            Ok(Redis {
36                connection,
37                db_name: db_name.to_string(),
38            })
39        }
40    }
41}
42
43#[derive(Debug)]
44pub struct Manager {
45    uri: String,
46    db_name: String,
47}
48
49impl deadpool::managed::Manager for Manager {
50    type Type = Redis;
51    type Error = RedisError;
52
53    async fn create(&self) -> Result<Redis, RedisError> {
54        Redis::init(&self.uri, &self.db_name).await
55    }
56
57    async fn recycle(
58        &self,
59        _: &mut Redis,
60        _: &deadpool::managed::Metrics,
61    ) -> deadpool::managed::RecycleResult<RedisError> {
62        Ok(())
63    }
64}
65
66pub struct RedisPool(pub deadpool::managed::Pool<Manager>);
67
68impl RedisPool {
69    pub fn init(uri: &str, db_name: &str) -> anyhow::Result<Self> {
70        let manager = Manager {
71            uri: uri.to_string(),
72            db_name: db_name.to_string(),
73        };
74        let pool = deadpool::managed::Pool::<Manager>::builder(manager).build()?;
75        Ok(Self(pool))
76    }
77
78    pub async fn instance(&self) -> anyhow::Result<Object<Manager>> {
79        let redis = self.0.get().await?;
80        Ok(redis)
81    }
82
83    /// SET key value
84    pub async fn set<K, V>(&self, key: K, value: V) -> anyhow::Result<()>
85    where
86        K: ToRedisArgs + Display,
87        V: ToRedisArgs,
88    {
89        let mut redis = self.instance().await?;
90        let key = format!("{}:{}", redis.db_name, key);
91        redis
92            .connection
93            .exec(cmd("SET").arg(key).arg(value))
94            .await?;
95        Ok(())
96    }
97
98    /// SET key value EX seconds
99    pub async fn set_ex<K, V>(&self, key: K, value: V, ex: u32) -> anyhow::Result<()>
100    where
101        K: ToRedisArgs + Display,
102        V: ToRedisArgs,
103    {
104        let mut redis = self.instance().await?;
105        let key = format!("{}:{}", redis.db_name, key);
106        redis
107            .connection
108            .exec(cmd("SET").arg(key).arg(value).arg("EX").arg(ex))
109            .await?;
110        Ok(())
111    }
112
113    /// GET key
114    pub async fn get<K, V>(&self, key: K) -> anyhow::Result<Option<V>>
115    where
116        K: ToRedisArgs + Display,
117        V: FromRedisValue,
118    {
119        let mut redis = self.instance().await?;
120        let key = format!("{}:{}", redis.db_name, key);
121        let value = redis
122            .connection
123            .query::<Option<V>>(cmd("GET").arg(key))
124            .await?;
125        Ok(value)
126    }
127
128    /// HSET table key value
129    pub async fn hset<K, V>(&self, table: &str, key: K, value: V) -> anyhow::Result<()>
130    where
131        K: ToRedisArgs,
132        V: ToRedisArgs,
133    {
134        let mut redis = self.instance().await?;
135        let table = format!("{}:{}", redis.db_name, table);
136        redis
137            .connection
138            .exec(cmd("HSET").arg(table).arg(key).arg(value))
139            .await?;
140        Ok(())
141    }
142
143    /// HSET table key2 value2 key2 value2
144    pub async fn hmset<K, V>(&self, table: &str, pairs: &[(K, V)]) -> anyhow::Result<()>
145    where
146        K: ToRedisArgs,
147        V: ToRedisArgs,
148    {
149        let mut redis = self.instance().await?;
150        let table = format!("{}:{}", redis.db_name, table);
151        let mut command = cmd("HSET");
152        command.arg(table);
153        for (key, value) in pairs {
154            command.arg(key);
155            command.arg(value);
156        }
157        redis.connection.exec(&mut command).await?;
158        Ok(())
159    }
160
161    /// HSET table key value EX seconds
162    pub async fn hset_ex<K, V>(&self, table: &str, key: K, value: V, ex: u32) -> anyhow::Result<()>
163    where
164        K: ToRedisArgs,
165        V: ToRedisArgs,
166    {
167        let mut redis = self.instance().await?;
168        let table = format!("{}:{}", redis.db_name, table);
169        redis
170            .connection
171            .exec(cmd("HSET").arg(table).arg(key).arg(value).arg("EX").arg(ex))
172            .await?;
173        Ok(())
174    }
175
176    /// HGET table key
177    pub async fn hget<K, V>(&self, table: &str, key: K) -> anyhow::Result<Option<V>>
178    where
179        K: ToRedisArgs,
180        V: FromRedisValue,
181    {
182        let mut redis = self.instance().await?;
183        let table = format!("{}:{}", redis.db_name, table);
184        let value = redis
185            .connection
186            .query::<Option<V>>(cmd("HGET").arg(table).arg(key))
187            .await?;
188        Ok(value)
189    }
190
191    /// HMGET table key1 key2 key3
192    pub async fn hmget<K, V>(&self, table: &str, keys: &[K]) -> anyhow::Result<Vec<V>>
193    where
194        K: ToRedisArgs,
195        V: FromRedisValue,
196    {
197        let mut redis = self.instance().await?;
198        let table = format!("{}:{}", redis.db_name, table);
199        let items: Vec<V> = redis
200            .connection
201            .query(cmd("HMGET").arg(table).arg(keys))
202            .await?;
203        Ok(items)
204    }
205
206    /// HGETALL table
207    pub async fn hget_all<K, V>(&self, table: &str) -> anyhow::Result<Vec<(K, V)>>
208    where
209        K: FromRedisValue,
210        V: FromRedisValue,
211    {
212        let mut redis = self.instance().await?;
213        let table = format!("{}:{}", redis.db_name, table);
214        let value = redis
215            .connection
216            .query::<Vec<(K, V)>>(cmd("HGETALL").arg(table))
217            .await?;
218        Ok(value)
219    }
220
221    /// HDEL table key
222    pub async fn hdel<K>(&self, table: &str, key: K) -> anyhow::Result<()>
223    where
224        K: ToRedisArgs,
225    {
226        let mut redis = self.instance().await?;
227        let table = format!("{}:{}", redis.db_name, table);
228        redis
229            .connection
230            .exec(cmd("HDEL").arg(table).arg(key))
231            .await?;
232        Ok(())
233    }
234
235    /// HDEL table key1 key2 key3
236    pub async fn hmdel<K>(&self, table: &str, keys: &[K]) -> anyhow::Result<()>
237    where
238        K: ToRedisArgs,
239    {
240        let mut redis = self.instance().await?;
241        let table = format!("{}:{}", redis.db_name, table);
242        let mut command = cmd("HDEL");
243        command.arg(table);
244        for key in keys {
245            command.arg(key);
246        }
247        redis.connection.exec(&mut command).await?;
248        Ok(())
249    }
250
251    /// DEL table
252    pub async fn hdel_all(&self, table: &str) -> anyhow::Result<()> {
253        let mut redis = self.instance().await?;
254        let table = format!("{}:{}", redis.db_name, table);
255        redis.connection.exec(cmd("DEL").arg(table)).await?;
256        Ok(())
257    }
258
259    pub async fn publish<C, M>(&self, channel: C, message: M) -> anyhow::Result<()>
260    where
261        C: ToRedisArgs + Display,
262        M: ToRedisArgs + Send + Sync,
263    {
264        let mut redis = self.instance().await?;
265        let channel = format!("{}:{}", redis.db_name, channel);
266        redis
267            .connection
268            .publish(channel.as_bytes(), message)
269            .await?;
270        Ok(())
271    }
272}
db_redis/pool.rs

db_redis/
pool.rs
