hive_console_sdk/
persisted_documents.rs

1use std::time::Duration;
2
3use moka::future::Cache;
4use reqwest::header::HeaderMap;
5use reqwest::header::HeaderValue;
6use reqwest_middleware::ClientBuilder;
7use reqwest_middleware::ClientWithMiddleware;
8use reqwest_retry::{policies::ExponentialBackoff, RetryTransientMiddleware};
9use tracing::{debug, info, warn};
10
11#[derive(Debug)]
12pub struct PersistedDocumentsManager {
13    agent: ClientWithMiddleware,
14    cache: Cache<String, String>,
15    endpoint: String,
16}
17
18#[derive(Debug, thiserror::Error)]
19pub enum PersistedDocumentsError {
20    #[error("Failed to read body: {0}")]
21    FailedToReadBody(String),
22    #[error("Failed to parse body: {0}")]
23    FailedToParseBody(serde_json::Error),
24    #[error("Persisted document not found.")]
25    DocumentNotFound,
26    #[error("Failed to locate the persisted document key in request.")]
27    KeyNotFound,
28    #[error("Failed to validate persisted document")]
29    FailedToFetchFromCDN(reqwest_middleware::Error),
30    #[error("Failed to read CDN response body")]
31    FailedToReadCDNResponse(reqwest::Error),
32    #[error("No persisted document provided, or document id cannot be resolved.")]
33    PersistedDocumentRequired,
34}
35
36impl PersistedDocumentsError {
37    pub fn message(&self) -> String {
38        self.to_string()
39    }
40
41    pub fn code(&self) -> String {
42        match self {
43            PersistedDocumentsError::FailedToReadBody(_) => "FAILED_TO_READ_BODY".into(),
44            PersistedDocumentsError::FailedToParseBody(_) => "FAILED_TO_PARSE_BODY".into(),
45            PersistedDocumentsError::DocumentNotFound => "PERSISTED_DOCUMENT_NOT_FOUND".into(),
46            PersistedDocumentsError::KeyNotFound => "PERSISTED_DOCUMENT_KEY_NOT_FOUND".into(),
47            PersistedDocumentsError::FailedToFetchFromCDN(_) => "FAILED_TO_FETCH_FROM_CDN".into(),
48            PersistedDocumentsError::FailedToReadCDNResponse(_) => {
49                "FAILED_TO_READ_CDN_RESPONSE".into()
50            }
51            PersistedDocumentsError::PersistedDocumentRequired => {
52                "PERSISTED_DOCUMENT_REQUIRED".into()
53            }
54        }
55    }
56}
57
58impl PersistedDocumentsManager {
59    #[allow(clippy::too_many_arguments)]
60    pub fn new(
61        key: String,
62        endpoint: String,
63        accept_invalid_certs: bool,
64        connect_timeout: Duration,
65        request_timeout: Duration,
66        retry_count: u32,
67        cache_size: u64,
68        user_agent: String,
69    ) -> Self {
70        let retry_policy = ExponentialBackoff::builder().build_with_max_retries(retry_count);
71
72        let mut default_headers = HeaderMap::new();
73        default_headers.insert("X-Hive-CDN-Key", HeaderValue::from_str(&key).unwrap());
74        let reqwest_agent = reqwest::Client::builder()
75            .danger_accept_invalid_certs(accept_invalid_certs)
76            .connect_timeout(connect_timeout)
77            .timeout(request_timeout)
78            .user_agent(user_agent)
79            .default_headers(default_headers)
80            .build()
81            .expect("Failed to create reqwest client");
82        let agent = ClientBuilder::new(reqwest_agent)
83            .with(RetryTransientMiddleware::new_with_policy(retry_policy))
84            .build();
85
86        let cache = Cache::<String, String>::new(cache_size);
87
88        Self {
89            agent,
90            cache,
91            endpoint,
92        }
93    }
94
95    /// Resolves the document from the cache, or from the CDN
96    pub async fn resolve_document(
97        &self,
98        document_id: &str,
99    ) -> Result<String, PersistedDocumentsError> {
100        let cached_record = self.cache.get(document_id).await;
101
102        match cached_record {
103            Some(document) => {
104                debug!("Document {} found in cache: {}", document_id, document);
105
106                Ok(document)
107            }
108            None => {
109                debug!(
110                    "Document {} not found in cache. Fetching from CDN",
111                    document_id
112                );
113                let cdn_document_id = str::replace(document_id, "~", "/");
114                let cdn_artifact_url = format!("{}/apps/{}", &self.endpoint, cdn_document_id);
115                info!(
116                    "Fetching document {} from CDN: {}",
117                    document_id, cdn_artifact_url
118                );
119                let cdn_response = self.agent.get(cdn_artifact_url).send().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
