scouter_sql/sql/
cache.rs

1use crate::sql::error::SqlError;
2use crate::sql::traits::EntitySqlLogic;
3use crate::PostgresClient;
4use mini_moka::sync::Cache;
5use scouter_types::DriftType;
6use sqlx::{Pool, Postgres};
7use std::sync::OnceLock;
8
9#[derive(Clone, Debug)]
10pub struct EntityCache {
11    cache: Cache<String, i32>,
12}
13
14impl EntityCache {
15    pub fn new(max_capacity: u64) -> Self {
16        let cache = Cache::new(max_capacity);
17        Self { cache }
18    }
19
20    /// Get entity ID from UID with caching
21    /// # Arguments
22    /// * `pool` - Database pool to use for queries
23    /// * `uid` - The UID of the entity
24    /// # Returns
25    /// * `Result<i32, SqlError>` - Result of the query returning the entity
26    pub async fn get_entity_id_from_uid(
27        &self,
28        pool: &Pool<Postgres>,
29        uid: &String,
30    ) -> Result<i32, SqlError> {
31        match self.cache.get(uid) {
32            Some(cached_id) => Ok(cached_id),
33            None => {
34                let entity_id = PostgresClient::get_entity_id_from_uid(pool, uid).await?;
35                self.cache.insert(uid.to_string(), entity_id);
36                Ok(entity_id)
37            }
38        }
39    }
40
41    pub async fn get_optional_entity_id_from_uid(
42        &self,
43        pool: &Pool<Postgres>,
44        uid: &String,
45    ) -> Result<Option<i32>, SqlError> {
46        match self.cache.get(uid) {
47            Some(cached_id) => Ok(Some(cached_id)),
48            None => {
49                let entity_id = PostgresClient::get_optional_entity_id_from_uid(pool, uid).await?;
50                if let Some(id) = entity_id {
51                    self.cache.insert(uid.to_string(), id);
52                }
53                Ok(entity_id)
54            }
55        }
56    }
57
58    pub async fn get_entity_id_from_space_name_version_drift_type(
59        &self,
60        pool: &Pool<Postgres>,
61        space: &str,
62        name: &str,
63        version: &str,
64        drift_type: &DriftType,
65    ) -> Result<i32, SqlError> {
66        let id = PostgresClient::get_entity_id_from_space_name_version_drift_type(
67            pool,
68            space,
69            name,
70            version,
71            drift_type.to_string(),
72        )
73        .await?;
74
75        Ok(id)
76    }
77}
78
79static INSTANCE: OnceLock<EntityCache> = OnceLock::new();
80
81pub fn init_entity_cache(max_capacity: u64) {
82    INSTANCE.get_or_init(|| {
83        tracing::info!("Initializing EntityCache");
84        EntityCache::new(max_capacity)
85    });
86}
87
88pub fn entity_cache() -> &'static EntityCache {
89    INSTANCE
90        .get()
91        .expect("EntityCache is not initialized - call init_entity_cache first")
92}
scouter_sql/sql/cache.rs

scouter_sql/sql/
cache.rs
