use sa_token_adapter::storage::{SaStorage, StorageError, StorageResult};
use std::time::Duration;
use summer::async_trait;
use summer_redis::redis::AsyncCommands;
use summer_redis::Redis;
pub struct SummerRedisStorage {
client: Redis,
prefix: Option<String>,
rewrite_prefix: bool,
}
impl SummerRedisStorage {
pub fn new(client: Redis, prefix: Option<String>, rewrite_prefix: bool) -> Self {
let prefix = normalize_storage_prefix(prefix);
Self {
client,
prefix,
rewrite_prefix,
}
}
}
fn normalize_storage_prefix(prefix: Option<String>) -> Option<String> {
match prefix {
Some(prefix) if prefix.is_empty() => None,
Some(prefix) if prefix.ends_with(':') => Some(prefix),
Some(prefix) => Some(format!("{prefix}:")),
None => None,
}
}
fn apply_storage_prefix(prefix: Option<&str>, rewrite_prefix: bool, key: &str) -> String {
match prefix {
Some(prefix) if !prefix.is_empty() && rewrite_prefix => {
if let Some(stripped) = key.strip_prefix("sa:") {
format!("{prefix}{stripped}")
} else {
format!("{prefix}{key}")
}
}
Some(prefix) if !prefix.is_empty() => format!("{prefix}{key}"),
_ => key.to_string(),
}
}
fn apply_storage_prefix_to_pattern(
prefix: Option<&str>,
rewrite_prefix: bool,
pattern: &str,
) -> String {
apply_storage_prefix(prefix, rewrite_prefix, pattern)
}
fn strip_storage_prefix(prefix: Option<&str>, rewrite_prefix: bool, key: &str) -> String {
match prefix {
Some(prefix) if !prefix.is_empty() && rewrite_prefix => {
if let Some(stripped) = key.strip_prefix(prefix) {
if stripped.starts_with("sa:") {
stripped.to_string()
} else {
format!("sa:{stripped}")
}
} else {
key.to_string()
}
}
Some(prefix) if !prefix.is_empty() => key.strip_prefix(prefix).unwrap_or(key).to_string(),
_ => key.to_string(),
}
}
#[async_trait]
impl SaStorage for SummerRedisStorage {
async fn get(&self, key: &str) -> StorageResult<Option<String>> {
let mut conn = self.client.clone();
let key = apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key);
tracing::debug!("SummerRedisStorage GET key: {}", key);
conn.get(key)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
}
async fn set(&self, key: &str, value: &str, ttl: Option<Duration>) -> StorageResult<()> {
let mut conn = self.client.clone();
let key = apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key);
tracing::debug!("SummerRedisStorage SET key: {}", key);
if let Some(ttl) = ttl {
conn.set_ex(key, value, ttl.as_secs())
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
} else {
conn.set(key, value)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
}
}
async fn delete(&self, key: &str) -> StorageResult<()> {
let mut conn = self.client.clone();
let key = apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key);
conn.del(key)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
}
async fn exists(&self, key: &str) -> StorageResult<bool> {
let mut conn = self.client.clone();
let key = apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key);
conn.exists(key)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
}
async fn expire(&self, key: &str, ttl: Duration) -> StorageResult<()> {
let mut conn = self.client.clone();
let key = apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key);
conn.expire(key, ttl.as_secs() as i64)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
}
async fn ttl(&self, key: &str) -> StorageResult<Option<Duration>> {
let mut conn = self.client.clone();
let key = apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key);
let ttl_secs: i64 = conn
.ttl(key)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))?;
match ttl_secs {
-2 => Ok(None), -1 => Ok(None), secs if secs > 0 => Ok(Some(Duration::from_secs(secs as u64))),
_ => Ok(Some(Duration::from_secs(0))),
}
}
async fn mget(&self, keys: &[&str]) -> StorageResult<Vec<Option<String>>> {
let mut conn = self.client.clone();
let keys: Vec<String> = keys
.iter()
.map(|key| apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key))
.collect();
summer_redis::redis::cmd("MGET")
.arg(keys)
.query_async(&mut conn)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
}
async fn mset(&self, items: &[(&str, &str)], ttl: Option<Duration>) -> StorageResult<()> {
let mut conn = self.client.clone();
let mut pipe = summer_redis::redis::pipe();
for (key, value) in items {
let key = apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key);
if let Some(ttl) = ttl {
pipe.set_ex(key, *value, ttl.as_secs());
} else {
pipe.set(key, *value);
}
}
pipe.query_async(&mut conn)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
}
async fn mdel(&self, keys: &[&str]) -> StorageResult<()> {
let mut conn = self.client.clone();
let keys: Vec<String> = keys
.iter()
.map(|key| apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key))
.collect();
conn.del(keys)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
}
async fn incr(&self, key: &str) -> StorageResult<i64> {
let mut conn = self.client.clone();
let key = apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key);
conn.incr(key, 1)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
}
async fn decr(&self, key: &str) -> StorageResult<i64> {
let mut conn = self.client.clone();
let key = apply_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, key);
conn.decr(key, 1)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))
}
async fn clear(&self) -> StorageResult<()> {
let mut conn = self.client.clone();
let pattern =
apply_storage_prefix_to_pattern(self.prefix.as_deref(), self.rewrite_prefix, "sa:*");
let keys: Vec<String> = conn
.keys(pattern)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))?;
if !keys.is_empty() {
conn.del::<_, ()>(&keys)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))?;
}
Ok(())
}
async fn keys(&self, pattern: &str) -> StorageResult<Vec<String>> {
let mut conn = self.client.clone();
let pattern =
apply_storage_prefix_to_pattern(self.prefix.as_deref(), self.rewrite_prefix, pattern);
let keys: Vec<String> = conn
.keys(pattern)
.await
.map_err(|e| StorageError::OperationFailed(e.to_string()))?;
Ok(keys
.into_iter()
.map(|key| strip_storage_prefix(self.prefix.as_deref(), self.rewrite_prefix, &key))
.collect())
}
}
#[cfg(test)]
mod tests {
use super::{
apply_storage_prefix, apply_storage_prefix_to_pattern, normalize_storage_prefix,
strip_storage_prefix,
};
#[test]
fn normalizes_storage_prefix() {
assert_eq!(
normalize_storage_prefix(Some("demo".to_string())),
Some("demo:".to_string())
);
assert_eq!(
normalize_storage_prefix(Some("demo:".to_string())),
Some("demo:".to_string())
);
assert_eq!(normalize_storage_prefix(Some(String::new())), None);
assert_eq!(normalize_storage_prefix(None), None);
}
#[test]
fn applies_prefix_to_key() {
assert_eq!(
apply_storage_prefix(Some("demo:"), false, "sa:token:abc"),
"demo:sa:token:abc"
);
assert_eq!(
apply_storage_prefix(Some("demo:"), true, "sa:token:abc"),
"demo:token:abc"
);
assert_eq!(
apply_storage_prefix(None, false, "sa:token:abc"),
"sa:token:abc"
);
}
#[test]
fn applies_prefix_to_pattern() {
assert_eq!(
apply_storage_prefix_to_pattern(Some("demo:"), false, "sa:*"),
"demo:sa:*"
);
assert_eq!(
apply_storage_prefix_to_pattern(Some("demo:"), true, "sa:*"),
"demo:*"
);
assert_eq!(apply_storage_prefix_to_pattern(None, false, "sa:*"), "sa:*");
}
#[test]
fn strips_prefix_from_key() {
assert_eq!(
strip_storage_prefix(Some("demo:"), false, "demo:sa:token:abc"),
"sa:token:abc"
);
assert_eq!(
strip_storage_prefix(Some("demo:"), true, "demo:token:abc"),
"sa:token:abc"
);
assert_eq!(
strip_storage_prefix(Some("demo:"), true, "sa:token:abc"),
"sa:token:abc"
);
assert_eq!(
strip_storage_prefix(None, false, "sa:token:abc"),
"sa:token:abc"
);
}
}