db_redis/
redis.rs

1use std::fmt::Display;
2use deadpool::managed::Object;
3use redis::{
4    aio::{ConnectionManager, ConnectionManagerConfig}, cluster::ClusterClientBuilder, cluster_async::ClusterConnection, cmd, Client, Cmd, FromRedisValue, PushInfo, RedisError, RedisResult, ToRedisArgs
5};
6use tokio::sync::OnceCell;
7
8static REDIS: OnceCell<Pool> = OnceCell::const_new();
9static DB_NAME: OnceCell<String> = OnceCell::const_new();
10
11pub struct Redis {
12    pub connection: RedisConnection,
13}
14
15impl Redis {
16    async fn init(uri: &str) -> RedisResult<Redis> {
17        let nodes = uri.split(",").collect::<Vec<&str>>();
18        if nodes.len() == 1 {
19            let client = Client::open(uri)?;
20            let config = ConnectionManagerConfig::default().set_number_of_retries(usize::MAX);
21            let connection = ConnectionManager::new_with_config(client.clone(), config).await?;
22            let connection = RedisConnection::Single(connection);
23            Ok(Redis { connection })
24        } else {
25            let client = ClusterClientBuilder::new(nodes)
26                .use_protocol(redis::ProtocolVersion::RESP3)
27                .retries(u32::MAX)
28                .build()?;
29            let connection = client.get_async_connection().await?;
30            let connection = RedisConnection::Cluster(connection);
31            Ok(Redis { connection })
32        }
33    }
34
35    pub async fn init_pubsub(uri: &str) -> RedisResult<RedisPubsub> {
36        let nodes = uri.split(",").collect::<Vec<&str>>();
37        let (tx, stream) = tokio::sync::mpsc::unbounded_channel();
38        if nodes.len() == 1 {
39            let client = Client::open(uri)?;
40            let config = ConnectionManagerConfig::default()
41                .set_number_of_retries(usize::MAX)
42                .set_factor(3000)
43                .set_exponent_base(1)
44                .set_push_sender(tx)
45                .set_automatic_resubscription();
46            let connection = ConnectionManager::new_with_config(client.clone(), config).await?;
47            let connection = RedisConnection::Single(connection);
48            Ok(RedisPubsub { connection, stream })
49        } else {
50            let client = ClusterClientBuilder::new(nodes)
51                .use_protocol(redis::ProtocolVersion::RESP3)
52                .push_sender(tx)
53                .retries(u32::MAX)
54                .retry_wait_formula(3000, 1)
55                .build()?;
56            let connection = client.get_async_connection().await?;
57            let connection = RedisConnection::Cluster(connection);
58            Ok(RedisPubsub { connection, stream })
59        }
60    }
61
62    pub async fn instance() -> anyhow::Result<Object<Manager>> {
63        let pool = REDIS.get().unwrap();
64        let redis = pool.get().await?;
65        Ok(redis)
66    }
67
68    /// SET key value
69    pub async fn set<K, V>(key: K, value: V) -> anyhow::Result<()>
70    where
71        K: ToRedisArgs + Display,
72        V: ToRedisArgs,
73    {
74        let key = format!("{}:{}", DB_NAME.get().unwrap(), key);
75        let mut redis = Redis::instance().await?;
76        redis
77            .connection
78            .exec(cmd("SET").arg(key).arg(value))
79            .await?;
80        Ok(())
81    }
82
83    /// SET key value EX seconds
84    pub async fn set_ex<K, V>(key: K, value: V, ex: u32) -> anyhow::Result<()>
85    where
86        K: ToRedisArgs + Display,
87        V: ToRedisArgs,
88    {
89        let key = format!("{}:{}", DB_NAME.get().unwrap(), key);
90        let mut redis = Redis::instance().await?;
91        redis
92            .connection
93            .exec(cmd("SET").arg(key).arg(value).arg("EX").arg(ex))
94            .await?;
95        Ok(())
96    }
97
98    /// GET key
99    pub async fn get<K, V>(key: K) -> anyhow::Result<Option<V>>
100    where
101        K: ToRedisArgs + Display,
102        V: FromRedisValue,
103    {
104        let key = format!("{}:{}", DB_NAME.get().unwrap(), key);
105        let mut redis = Redis::instance().await?;
106        let value = redis
107            .connection
108            .query::<Option<V>>(cmd("GET").arg(key))
109            .await?;
110        Ok(value)
111    }
112
113    /// HSET table key value
114    pub async fn hset<K, V>(table: &str, key: K, value: V) -> anyhow::Result<()>
115    where
116        K: ToRedisArgs,
117        V: ToRedisArgs,
118    {
119        let table = format!("{}:{}", DB_NAME.get().unwrap(), table);
120        let mut redis = Redis::instance().await?;
121        redis
122            .connection
123            .exec(cmd("HSET").arg(table).arg(key).arg(value))
124            .await?;
125        Ok(())
126    }
127
128    /// HSET table key2 value2 key2 value2
129    pub async fn hmset<K, V>(table: &str, pairs: &[(K, V)]) -> anyhow::Result<()>
130    where
131        K: ToRedisArgs,
132        V: ToRedisArgs,
133    {
134        let table = format!("{}:{}", DB_NAME.get().unwrap(), table);
135        let mut redis = Redis::instance().await?;
136        let mut command = cmd("HSET");
137        command.arg(table);
138        for (key, value) in pairs {
139            command.arg(key);
140            command.arg(value);
141        }
142        redis
143            .connection
144            .exec(&mut command)
145            .await?;
146        Ok(())
147    }
148
149    /// HSET table key value EX seconds
150    pub async fn hset_ex<K, V>(table: &str, key: K, value: V, ex: u32) -> anyhow::Result<()>
151    where
152        K: ToRedisArgs,
153        V: ToRedisArgs,
154    {
155        let table = format!("{}:{}", DB_NAME.get().unwrap(), table);
156        let mut redis = Redis::instance().await?;
157        redis
158            .connection
159            .exec(cmd("HSET").arg(table).arg(key).arg(value).arg("EX").arg(ex))
160            .await?;
161        Ok(())
162    }
163
164    /// HGET table key
165    pub async fn hget<K, V>(table: &str, key: K) -> anyhow::Result<Option<V>>
166    where
167        K: ToRedisArgs,
168        V: FromRedisValue,
169    {
170        let table = format!("{}:{}", DB_NAME.get().unwrap(), table);
171        let mut redis = Redis::instance().await?;
172        let value = redis
173            .connection
174            .query::<Option<V>>(cmd("HGET").arg(table).arg(key))
175            .await?;
176        Ok(value)
177    }
178
179    /// HMGET table key1 key2 key3
180    pub async fn hmget<K, V>(table: &str, keys: &[K]) -> anyhow::Result<Vec<V>>
181    where
182        K: ToRedisArgs,
183        V: FromRedisValue,
184    {
185        let table = format!("{}:{}", DB_NAME.get().unwrap(), table);
186        let mut redis = Redis::instance().await?;
187        let items: Vec<V> = redis
188            .connection
189            .query(cmd("HMGET").arg(table).arg(keys))
190            .await?;
191        Ok(items)
192    }
193
194    /// HGETALL table
195    pub async fn hget_all<K, V>(table: &str) -> anyhow::Result<Vec<(K, V)>>
196    where
197        K: FromRedisValue,
198        V: FromRedisValue,
199    {
200        let table = format!("{}:{}", DB_NAME.get().unwrap(), table);
201        let mut redis = Redis::instance().await?;
202        let value = redis
203            .connection
204            .query::<Vec<(K, V)>>(cmd("HGETALL").arg(table))
205            .await?;
206        Ok(value)
207    }
208
209    /// HDEL table key
210    pub async fn hdel<K>(table: &str, key: K) -> anyhow::Result<()>
211    where
212        K: ToRedisArgs,
213    {
214        let table = format!("{}:{}", DB_NAME.get().unwrap(), table);
215        let mut redis = Redis::instance().await?;
216        redis
217            .connection
218            .exec(cmd("HDEL").arg(table).arg(key))
219            .await?;
220        Ok(())
221    }
222
223    /// HDEL table key1 key2 key3
224    pub async fn hmdel<K>(table: &str, keys: &[K]) -> anyhow::Result<()>
225    where
226        K: ToRedisArgs,
227    {
228        let table = format!("{}:{}", DB_NAME.get().unwrap(), table);
229        let mut redis = Redis::instance().await?;
230        let mut command = cmd("HDEL");
231        command.arg(table);
232        for key in keys {
233             command.arg(key);
234        }
235        redis
236            .connection
237            .exec(&mut command)
238            .await?;
239        Ok(())
240    }
241
242    /// DEL table
243    pub async fn hdel_all(table: &str) -> anyhow::Result<()> {
244        let table = format!("{}:{}", DB_NAME.get().unwrap(), table);
245        let mut redis = Redis::instance().await?;
246        redis
247            .connection
248            .exec(cmd("DEL").arg(table))
249            .await?;
250        Ok(())
251    }
252}
253
254pub struct RedisPubsub {
255    pub connection: RedisConnection,
256    pub stream: tokio::sync::mpsc::UnboundedReceiver<PushInfo>,
257}
258
259impl RedisPubsub {
260    pub async fn subscribe(&mut self, channel_name: impl ToRedisArgs + Display) -> RedisResult<()> {
261        let channel_name = format!("{}:{}", DB_NAME.get().unwrap(), channel_name);
262        match &mut self.connection {
263            RedisConnection::Single(c) => c.subscribe(channel_name).await,
264            RedisConnection::Cluster(c) => c.subscribe(channel_name).await,
265        }
266    }
267}
268
269#[derive(Debug)]
270pub struct Manager {
271    uri: String,
272}
273
274impl deadpool::managed::Manager for Manager {
275    type Type = Redis;
276    type Error = RedisError;
277
278    async fn create(&self) -> Result<Redis, RedisError> {
279        Redis::init(&self.uri).await
280    }
281
282    async fn recycle(
283        &self,
284        _: &mut Redis,
285        _: &deadpool::managed::Metrics,
286    ) -> deadpool::managed::RecycleResult<RedisError> {
287        Ok(())
288    }
289}
290
291type Pool = deadpool::managed::Pool<Manager>;
292
293pub enum RedisConnection {
294    Single(ConnectionManager),
295    Cluster(ClusterConnection),
296}
297
298impl RedisConnection {
299    pub async fn query<T: FromRedisValue>(&mut self, cmd: &mut Cmd) -> RedisResult<T> {
300        match self {
301            RedisConnection::Single(client) => cmd.query_async(client).await,
302            RedisConnection::Cluster(client) => cmd.query_async(client).await,
303        }
304    }
305    pub async fn exec(&mut self, cmd: &mut Cmd) -> RedisResult<()> {
306        match self {
307            RedisConnection::Single(client) => cmd.exec_async(client).await,
308            RedisConnection::Cluster(client) => cmd.exec_async(client).await,
309        }
310    }
311}
312
313/// Initialize Redis connection
314pub async fn init(uri: &str, db_name: &str) -> anyhow::Result<()> {
315    tracing::info!("Initializing Redis...");
316    let mgr = Manager {
317        uri: uri.to_string(),
318    };
319    let pool = Pool::builder(mgr).build()?;
320    REDIS.set(pool).map_err(|e| anyhow::anyhow!("Failed to set Redis pool: {}", e))?;
321    DB_NAME.set(db_name.to_string()).map_err(|e| anyhow::anyhow!("Failed to set DB_NAME: {}", e))?;
322    tracing::info!("Redis initialized with uri: {}", uri);
323    Ok(())
324}
db_redis/redis.rs

db_redis/
redis.rs
