Skip to main content

rustvello_mongo/
client_data_store.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use mongodb::bson::doc;
5
6use rustvello_core::client_data_store::ClientDataStore;
7use rustvello_core::error::{RustvelloError, RustvelloResult};
8
9use crate::connection::{mongo_err, MongoPool};
10
11const COLLECTION: &str = "client_data_store";
12
13/// MongoDB-backed client data store.
14///
15/// Uses a simple collection with `{ _id: key, value: string }` documents.
16#[non_exhaustive]
17pub struct MongoClientDataStore {
18    pool: Arc<MongoPool>,
19}
20
21impl MongoClientDataStore {
22    pub fn new(pool: Arc<MongoPool>) -> Self {
23        Self { pool }
24    }
25}
26
27#[async_trait]
28impl ClientDataStore for MongoClientDataStore {
29    async fn store(&self, key: &str, value: &str) -> RustvelloResult<()> {
30        let db = self.pool.db().await?;
31        let col = db.collection::<mongodb::bson::Document>(COLLECTION);
32        let filter = doc! { "_id": key };
33        let update = doc! { "$set": { "value": value } };
34        col.update_one(filter, update)
35            .upsert(true)
36            .await
37            .map_err(mongo_err)?;
38        Ok(())
39    }
40
41    async fn retrieve(&self, key: &str) -> RustvelloResult<String> {
42        let db = self.pool.db().await?;
43        let col = db.collection::<mongodb::bson::Document>(COLLECTION);
44        let filter = doc! { "_id": key };
45        let result = col.find_one(filter).await.map_err(mongo_err)?;
46        match result {
47            Some(d) => {
48                let val = d
49                    .get_str("value")
50                    .map_err(|e| RustvelloError::state_backend(e.to_string()))?;
51                Ok(val.to_string())
52            }
53            None => Err(RustvelloError::state_backend(format!(
54                "CDS key not found: {}",
55                key
56            ))),
57        }
58    }
59
60    async fn purge(&self) -> RustvelloResult<()> {
61        let db = self.pool.db().await?;
62        let col = db.collection::<mongodb::bson::Document>(COLLECTION);
63        col.delete_many(doc! {}).await.map_err(mongo_err)?;
64        Ok(())
65    }
66}