pyra-redis 0.4.1

//! Typed fetch helpers for Drift-specific Redis data.

use std::collections::HashMap;

use pyra_tokens::AssetId;
use solana_pubkey::Pubkey;

use pyra_types::{Cache, DriftUser, SpotMarket};

use crate::{RedisClient, RedisKey, RedisResult};

impl RedisClient {
    /// Fetch a Drift user account by the vault (authority) address.
    pub async fn fetch_drift_user(&self, authority: &Pubkey) -> RedisResult<Option<DriftUser>> {
        let key = RedisKey::drift_user(authority);
        let cached: Option<Cache<DriftUser>> = self.get_json(key.as_str()).await?;
        Ok(cached.map(|c| c.account))
    }

    /// Fetch a single Drift spot market by asset ID.
    ///
    /// Returns `Ok(None)` if the asset has no Drift market index.
    pub async fn fetch_spot_market(&self, asset_id: AssetId) -> RedisResult<Option<SpotMarket>> {
        let Some(key) = RedisKey::drift_spot_market(asset_id) else {
            return Ok(None);
        };
        let cached: Option<Cache<SpotMarket>> = self.get_json(key.as_str()).await?;
        Ok(cached.map(|c| c.account))
    }

    /// Fetch multiple spot markets by their asset IDs in a single MGET round-trip.
    ///
    /// Returns only the markets that were found and successfully deserialized.
    /// Assets without a Drift market index are silently skipped.
    pub async fn fetch_spot_markets(
        &self,
        asset_ids: &[AssetId],
    ) -> RedisResult<HashMap<AssetId, SpotMarket>> {
        if asset_ids.is_empty() {
            return Ok(HashMap::new());
        }

        let keyed: Vec<(AssetId, String)> = asset_ids
            .iter()
            .filter_map(|&id| RedisKey::drift_spot_market(id).map(|k| (id, k.to_string())))
            .collect();

        if keyed.is_empty() {
            return Ok(HashMap::new());
        }

        let keys: Vec<String> = keyed.iter().map(|(_, k)| k.clone()).collect();
        let values = self.mget(&keys).await?;
        let mut result = HashMap::with_capacity(keyed.len());

        for ((id, _), val) in keyed.iter().zip(values.into_iter()) {
            if let Some(raw) = val {
                if let Ok(cached) = serde_json::from_str::<Cache<SpotMarket>>(&raw) {
                    result.insert(*id, cached.account);
                }
            }
        }

        Ok(result)
    }

    /// Fetch all spot markets by scanning `account:drift:spot_market:*` keys.
    ///
    /// Uses SCAN + MGET for efficiency. Returns asset_id -> SpotMarket.
    /// Only includes tokens known to pyra-tokens.
    pub async fn fetch_all_spot_markets(&self) -> RedisResult<HashMap<AssetId, SpotMarket>> {
        let pattern = RedisKey::pattern(RedisKey::DRIFT_SPOT_MARKET_PREFIX);
        let keys = self.scan_keys(&pattern).await?;

        if keys.is_empty() {
            return Ok(HashMap::new());
        }

        let values = self.mget(&keys).await?;
        let mut result = HashMap::with_capacity(keys.len());

        for val in values.into_iter().flatten() {
            if let Ok(cached) = serde_json::from_str::<Cache<SpotMarket>>(&val) {
                if let Some(token) =
                    pyra_tokens::Token::find_by_drift_market_index(cached.account.market_index)
                {
                    result.insert(token.asset_id, cached.account);
                }
            }
        }

        Ok(result)
    }
}