hive_console_sdk/
persisted_documents.rs

1use std::time::Duration;
2
3use moka::future::Cache;
4use reqwest_middleware::ClientBuilder;
5use reqwest_middleware::ClientWithMiddleware;
6use reqwest_retry::{policies::ExponentialBackoff, RetryTransientMiddleware};
7use tracing::{debug, info, warn};
8
9#[derive(Debug)]
10pub struct PersistedDocumentsManager {
11    agent: ClientWithMiddleware,
12    cache: Cache<String, String>,
13    endpoint: String,
14    key: String,
15}
16
17#[derive(Debug, thiserror::Error)]
18pub enum PersistedDocumentsError {
19    #[error("Failed to read body: {0}")]
20    FailedToReadBody(String),
21    #[error("Failed to parse body: {0}")]
22    FailedToParseBody(serde_json::Error),
23    #[error("Persisted document not found.")]
24    DocumentNotFound,
25    #[error("Failed to locate the persisted document key in request.")]
26    KeyNotFound,
27    #[error("Failed to validate persisted document")]
28    FailedToFetchFromCDN(reqwest_middleware::Error),
29    #[error("Failed to read CDN response body")]
30    FailedToReadCDNResponse(reqwest::Error),
31    #[error("No persisted document provided, or document id cannot be resolved.")]
32    PersistedDocumentRequired,
33}
34
35impl PersistedDocumentsError {
36    pub fn message(&self) -> String {
37        self.to_string()
38    }
39
40    pub fn code(&self) -> String {
41        match self {
42            PersistedDocumentsError::FailedToReadBody(_) => "FAILED_TO_READ_BODY".into(),
43            PersistedDocumentsError::FailedToParseBody(_) => "FAILED_TO_PARSE_BODY".into(),
44            PersistedDocumentsError::DocumentNotFound => "PERSISTED_DOCUMENT_NOT_FOUND".into(),
45            PersistedDocumentsError::KeyNotFound => "PERSISTED_DOCUMENT_KEY_NOT_FOUND".into(),
46            PersistedDocumentsError::FailedToFetchFromCDN(_) => "FAILED_TO_FETCH_FROM_CDN".into(),
47            PersistedDocumentsError::FailedToReadCDNResponse(_) => {
48                "FAILED_TO_READ_CDN_RESPONSE".into()
49            }
50            PersistedDocumentsError::PersistedDocumentRequired => {
51                "PERSISTED_DOCUMENT_REQUIRED".into()
52            }
53        }
54    }
55}
56
57impl PersistedDocumentsManager {
58    #[allow(clippy::too_many_arguments)]
59    pub fn new(
60        key: String,
61        endpoint: String,
62        accept_invalid_certs: bool,
63        connect_timeout: Duration,
64        request_timeout: Duration,
65        retry_count: u32,
66        cache_size: u64,
67    ) -> Self {
68        let retry_policy = ExponentialBackoff::builder().build_with_max_retries(retry_count);
69
70        let reqwest_agent = reqwest::Client::builder()
71            .danger_accept_invalid_certs(accept_invalid_certs)
72            .connect_timeout(connect_timeout)
73            .timeout(request_timeout)
74            .build()
75            .expect("Failed to create reqwest client");
76        let agent = ClientBuilder::new(reqwest_agent)
77            .with(RetryTransientMiddleware::new_with_policy(retry_policy))
78            .build();
79
80        let cache = Cache::<String, String>::new(cache_size);
81
82        Self {
83            agent,
84            cache,
85            endpoint,
86            key,
87        }
88    }
89
90    /// Resolves the document from the cache, or from the CDN
91    pub async fn resolve_document(
92        &self,
93        document_id: &str,
94    ) -> Result<String, PersistedDocumentsError> {
95        let cached_record = self.cache.get(document_id).await;
96
97        match cached_record {
98            Some(document) => {
99                debug!("Document {} found in cache: {}", document_id, document);
100
101                Ok(document)
102            }
103            None => {
104                debug!(
105                    "Document {} not found in cache. Fetching from CDN",
106                    document_id
107                );
108                let cdn_document_id = str::replace(document_id, "~", "/");
109                let cdn_artifact_url = format!("{}/apps/{}", &self.endpoint, cdn_document_id);
110                info!(
111                    "Fetching document {} from CDN: {}",
112                    document_id, cdn_artifact_url
113                );
114                let cdn_response = self
115                    .agent
116                    .get(cdn_artifact_url)
117                    .header("X-Hive-CDN-Key", self.key.to_string())
118                    .send()
119                    .await;
120
121                match cdn_response {
122                    Ok(response) => {
123                        if response.status().is_success() {
124                            let document = response
125                                .text()
126                                .await
127                                .map_err(PersistedDocumentsError::FailedToReadCDNResponse)?;
128                            debug!(
129                                "Document fetched from CDN: {}, storing in local cache",
130                                document
131                            );
132                            self.cache
133                                .insert(document_id.into(), document.clone())
134                                .await;
135
136                            return Ok(document);
137                        }
138
139                        warn!(
140                            "Document fetch from CDN failed: HTTP {}, Body: {:?}",
141                            response.status(),
142                            response
143                                .text()
144                                .await
145                                .unwrap_or_else(|_| "Unavailable".to_string())
146                        );
147
148                        Err(PersistedDocumentsError::DocumentNotFound)
149                    }
150                    Err(e) => {
151                        warn!("Failed to fetch document from CDN: {:?}", e);
152
153                        Err(PersistedDocumentsError::FailedToFetchFromCDN(e))
154                    }
155                }
156            }
157        }
158    }
159}
hive_console_sdk/persisted_documents.rs

hive_console_sdk/
persisted_documents.rs
