rullst-orm 12.1.1

Fast, un-opinionated, database-agnostic ORM and query builder for Rust
Documentation
use std::{
    cmp::Ordering,
    collections::{BTreeMap, BTreeSet},
};

use tokio::sync::RwLock;

use super::{
    RedisDataKey, RedisField, RedisMember, RedisScanLimit, RedisStructure, RedisValue,
    ScoredRedisMember, redis_error, validate_score,
};
use crate::polyglot::PolyglotError;

#[derive(Default)]
pub(super) struct MockRedisData {
    hashes: RwLock<BTreeMap<String, BTreeMap<String, String>>>,
    sets: RwLock<BTreeMap<String, BTreeSet<String>>>,
    sorted_sets: RwLock<BTreeMap<String, BTreeMap<String, f64>>>,
}

impl MockRedisData {
    pub(super) async fn hash_set(
        &self,
        key: &RedisDataKey,
        field: &RedisField,
        value: &RedisValue,
    ) -> Result<bool, PolyglotError> {
        let mut hashes = self.hashes.write().await;
        let inserted = hashes
            .entry(key.as_str().to_owned())
            .or_default()
            .insert(field.as_str().to_owned(), value.as_str().to_owned())
            .is_none();
        Ok(inserted)
    }

    pub(super) async fn hash_get(
        &self,
        key: &RedisDataKey,
        field: &RedisField,
    ) -> Result<Option<RedisValue>, PolyglotError> {
        self.hashes
            .read()
            .await
            .get(key.as_str())
            .and_then(|hash| hash.get(field.as_str()))
            .cloned()
            .map(RedisValue::new)
            .transpose()
    }

    pub(super) async fn hash_increment(
        &self,
        key: &RedisDataKey,
        field: &RedisField,
        amount: i64,
    ) -> Result<i64, PolyglotError> {
        let mut hashes = self.hashes.write().await;
        let hash = hashes.entry(key.as_str().to_owned()).or_default();
        let current = hash.get(field.as_str()).map_or(Ok(0), |value| {
            value.parse::<i64>().map_err(|_| redis_error())
        })?;
        let updated = current.checked_add(amount).ok_or_else(redis_error)?;
        hash.insert(field.as_str().to_owned(), updated.to_string());
        Ok(updated)
    }

    pub(super) async fn set_add(
        &self,
        key: &RedisDataKey,
        member: &RedisMember,
    ) -> Result<bool, PolyglotError> {
        Ok(self
            .sets
            .write()
            .await
            .entry(key.as_str().to_owned())
            .or_default()
            .insert(member.as_str().to_owned()))
    }

    pub(super) async fn set_contains(
        &self,
        key: &RedisDataKey,
        member: &RedisMember,
    ) -> Result<bool, PolyglotError> {
        Ok(self
            .sets
            .read()
            .await
            .get(key.as_str())
            .is_some_and(|set| set.contains(member.as_str())))
    }

    pub(super) async fn set_scan(
        &self,
        key: &RedisDataKey,
        limit: RedisScanLimit,
    ) -> Result<Vec<RedisMember>, PolyglotError> {
        self.sets
            .read()
            .await
            .get(key.as_str())
            .into_iter()
            .flat_map(|set| set.iter())
            .take(usize::from(limit.get()))
            .cloned()
            .map(RedisMember::new)
            .collect()
    }

    pub(super) async fn sorted_set_add(
        &self,
        key: &RedisDataKey,
        member: &RedisMember,
        score: f64,
    ) -> Result<bool, PolyglotError> {
        validate_score(score)?;
        Ok(self
            .sorted_sets
            .write()
            .await
            .entry(key.as_str().to_owned())
            .or_default()
            .insert(member.as_str().to_owned(), score)
            .is_none())
    }

    pub(super) async fn sorted_set_top(
        &self,
        key: &RedisDataKey,
        limit: RedisScanLimit,
    ) -> Result<Vec<ScoredRedisMember>, PolyglotError> {
        let sorted_sets = self.sorted_sets.read().await;
        let mut rows = sorted_sets
            .get(key.as_str())
            .into_iter()
            .flat_map(|set| set.iter())
            .map(|(member, score)| (member.clone(), *score))
            .collect::<Vec<_>>();
        rows.sort_by(|left, right| {
            right
                .1
                .partial_cmp(&left.1)
                .unwrap_or(Ordering::Equal)
                .then_with(|| left.0.cmp(&right.0))
        });
        rows.truncate(usize::from(limit.get()));
        rows.into_iter()
            .map(|(member, score)| ScoredRedisMember::new(RedisMember::new(member)?, score))
            .collect()
    }

    pub(super) async fn delete(
        &self,
        key: &RedisDataKey,
        structure: RedisStructure,
    ) -> Result<bool, PolyglotError> {
        let removed = match structure {
            RedisStructure::Hash => self.hashes.write().await.remove(key.as_str()).is_some(),
            RedisStructure::Set => self.sets.write().await.remove(key.as_str()).is_some(),
            RedisStructure::SortedSet => self
                .sorted_sets
                .write()
                .await
                .remove(key.as_str())
                .is_some(),
        };
        Ok(removed)
    }
}