Skip to main content

pyra_redis/drift/
typed.rs

1//! Typed fetch helpers for Drift-specific Redis data.
2
3use std::collections::HashMap;
4
5use pyra_tokens::AssetId;
6use solana_pubkey::Pubkey;
7
8use pyra_types::{Cache, DriftUser, SpotMarket};
9
10use crate::{RedisClient, RedisKey, RedisResult};
11
12impl RedisClient {
13    /// Fetch a Drift user account by the vault (authority) address.
14    pub async fn fetch_drift_user(&self, authority: &Pubkey) -> RedisResult<Option<DriftUser>> {
15        let key = RedisKey::drift_user(authority);
16        let cached: Option<Cache<DriftUser>> = self.get_json(key.as_str()).await?;
17        Ok(cached.map(|c| c.account))
18    }
19
20    /// Fetch a single Drift spot market by asset ID.
21    pub async fn fetch_spot_market(&self, asset_id: AssetId) -> RedisResult<Option<SpotMarket>> {
22        let key = RedisKey::drift_spot_market(asset_id);
23        let cached: Option<Cache<SpotMarket>> = self.get_json(key.as_str()).await?;
24        Ok(cached.map(|c| c.account))
25    }
26
27    /// Fetch multiple spot markets by their asset IDs in a single MGET round-trip.
28    ///
29    /// Returns only the markets that were found and successfully deserialized.
30    pub async fn fetch_spot_markets(
31        &self,
32        asset_ids: &[AssetId],
33    ) -> RedisResult<HashMap<AssetId, SpotMarket>> {
34        if asset_ids.is_empty() {
35            return Ok(HashMap::new());
36        }
37
38        let keys: Vec<String> = asset_ids
39            .iter()
40            .map(|&id| RedisKey::drift_spot_market(id).to_string())
41            .collect();
42
43        let values = self.mget(&keys).await?;
44        let mut result = HashMap::with_capacity(asset_ids.len());
45
46        for (id, val) in asset_ids.iter().zip(values.into_iter()) {
47            if let Some(raw) = val {
48                if let Ok(cached) = serde_json::from_str::<Cache<SpotMarket>>(&raw) {
49                    result.insert(*id, cached.account);
50                }
51            }
52        }
53
54        Ok(result)
55    }
56
57    /// Fetch all spot markets by scanning `account:drift:spot_market:*` keys.
58    ///
59    /// Uses SCAN + MGET for efficiency. Returns asset_id -> SpotMarket.
60    /// Only includes tokens known to pyra-tokens.
61    pub async fn fetch_all_spot_markets(&self) -> RedisResult<HashMap<AssetId, SpotMarket>> {
62        let pattern = RedisKey::pattern(RedisKey::DRIFT_SPOT_MARKET_PREFIX);
63        let keys = self.scan_keys(&pattern).await?;
64
65        if keys.is_empty() {
66            return Ok(HashMap::new());
67        }
68
69        let prefix_with_colon = format!("{}:", RedisKey::DRIFT_SPOT_MARKET_PREFIX);
70        let values = self.mget(&keys).await?;
71        let mut result = HashMap::with_capacity(keys.len());
72
73        for (key, val) in keys.iter().zip(values.into_iter()) {
74            let Some(raw) = val else { continue };
75            let Some(suffix) = key.strip_prefix(prefix_with_colon.as_str()) else {
76                continue;
77            };
78            let Ok(id_raw) = suffix.parse::<u16>() else {
79                continue;
80            };
81            let Ok(asset_id) = AssetId::new(id_raw) else {
82                continue;
83            };
84            if let Ok(cached) = serde_json::from_str::<Cache<SpotMarket>>(&raw) {
85                result.insert(asset_id, cached.account);
86            }
87        }
88
89        Ok(result)
90    }
91}