shared/
database.rs

1use base64::Engine;
2use sha2::Digest;
3use sqlx::postgres::PgPoolOptions;
4use std::{collections::HashMap, fmt::Display, pin::Pin, sync::Arc};
5use tokio::sync::Mutex;
6
7pub static BASE64_ENGINE: base64::engine::GeneralPurpose = base64::engine::GeneralPurpose::new(
8    &base64::alphabet::STANDARD,
9    base64::engine::GeneralPurposeConfig::new()
10        .with_decode_allow_trailing_bits(true)
11        .with_decode_padding_mode(base64::engine::DecodePaddingMode::Indifferent),
12);
13
14type BatchFuture = Pin<Box<dyn Future<Output = Result<(), anyhow::Error>> + Send>>;
15
16pub struct Database {
17    pub cache: Arc<crate::cache::Cache>,
18
19    write: sqlx::PgPool,
20    read: Option<sqlx::PgPool>,
21
22    encryption_key: Arc<str>,
23    decryption_cache: Option<moka::future::Cache<[u8; 32], Arc<Vec<u8>>>>,
24    batch_actions: Arc<Mutex<HashMap<(&'static str, uuid::Uuid), BatchFuture>>>,
25}
26
27async fn connect_pool(label: &str, url: &str, max_connections: u32) -> sqlx::PgPool {
28    crate::retry::startup_connect(label, || {
29        PgPoolOptions::new()
30            .min_connections(10)
31            .max_connections(max_connections)
32            .test_before_acquire(false)
33            .connect(url)
34    })
35    .await
36}
37
38impl Database {
39    pub async fn new(env: &crate::env::Env, cache: Arc<crate::cache::Cache>) -> Self {
40        let start = std::time::Instant::now();
41
42        let instance = Self {
43            cache,
44
45            write: match &env.database_url_primary {
46                Some(url) => connect_pool("primary database", url, 20).await,
47                None => connect_pool("database", &env.database_url, 50).await,
48            },
49            read: if env.database_url_primary.is_some() {
50                Some(connect_pool("read database", &env.database_url, 50).await)
51            } else {
52                None
53            },
54
55            encryption_key: env.app_encryption_key.clone().into(),
56            decryption_cache: env.app_use_decryption_cache.then(|| {
57                moka::future::Cache::builder()
58                    .max_capacity(16384)
59                    .time_to_live(std::time::Duration::from_secs(30))
60                    .build()
61            }),
62            batch_actions: Arc::new(Mutex::new(HashMap::new())),
63        };
64
65        let version = instance
66            .version()
67            .await
68            .unwrap_or_else(|_| "unknown".into());
69
70        tracing::info!(
71            "database connected (postgres@{}, {}ms)",
72            version,
73            start.elapsed().as_millis()
74        );
75
76        tokio::spawn({
77            let batch_actions = instance.batch_actions.clone();
78
79            async move {
80                let mut interval = tokio::time::interval(std::time::Duration::from_secs(5));
81                interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);
82
83                loop {
84                    interval.tick().await;
85
86                    let actions = batch_actions.lock().await.drain().collect::<Vec<_>>();
87
88                    for (key, action) in actions {
89                        tracing::debug!("executing batch action for {}:{}", key.0, key.1);
90                        if let Err(err) = action.await {
91                            tracing::error!(
92                                "error executing batch action for {}:{} - {:?}",
93                                key.0,
94                                key.1,
95                                err
96                            );
97                            sentry_anyhow::capture_anyhow(&err);
98                        }
99                    }
100                }
101            }
102        });
103
104        instance
105    }
106
107    pub async fn flush_batch_actions(&self) {
108        let actions = self.batch_actions.lock().await.drain().collect::<Vec<_>>();
109
110        for (key, action) in actions {
111            tracing::debug!("executing batch action for {}:{}", key.0, key.1);
112            if let Err(err) = action.await {
113                tracing::error!(
114                    "error executing batch action for {}:{} - {:?}",
115                    key.0,
116                    key.1,
117                    err
118                );
119                sentry_anyhow::capture_anyhow(&err);
120            }
121        }
122    }
123
124    pub async fn version(&self) -> Result<compact_str::CompactString, sqlx::Error> {
125        let version: (compact_str::CompactString,) =
126            sqlx::query_as("SELECT split_part(version(), ' ', 2)")
127                .fetch_one(self.read())
128                .await?;
129
130        Ok(version.0)
131    }
132
133    pub async fn size(&self) -> Result<u64, sqlx::Error> {
134        let size: (i64,) = sqlx::query_as("SELECT pg_database_size(current_database())")
135            .fetch_one(self.read())
136            .await?;
137
138        Ok(size.0 as u64)
139    }
140
141    #[inline]
142    pub fn write(&self) -> &sqlx::PgPool {
143        &self.write
144    }
145
146    #[inline]
147    pub fn read(&self) -> &sqlx::PgPool {
148        self.read.as_ref().unwrap_or(&self.write)
149    }
150
151    pub async fn encrypt(
152        &self,
153        data: impl AsRef<[u8]> + Send + 'static,
154    ) -> Result<Vec<u8>, anyhow::Error> {
155        let encryption_key = self.encryption_key.clone();
156
157        tokio::task::spawn_blocking(move || {
158            simple_crypt::encrypt(data.as_ref(), encryption_key.as_bytes())
159        })
160        .await?
161    }
162
163    pub async fn encrypt_with_input<P: AsRef<[u8]> + Send + 'static>(
164        &self,
165        data: P,
166    ) -> Result<(P, Vec<u8>), anyhow::Error> {
167        let encryption_key = self.encryption_key.clone();
168
169        tokio::task::spawn_blocking(move || {
170            simple_crypt::encrypt(data.as_ref(), encryption_key.as_bytes())
171                .map(|bytes| (data, bytes))
172        })
173        .await?
174    }
175
176    pub async fn encrypt_base64(
177        &self,
178        data: impl AsRef<[u8]> + Send + 'static,
179    ) -> Result<compact_str::CompactString, anyhow::Error> {
180        let encrypted = self.encrypt(data).await?;
181        Ok(BASE64_ENGINE.encode(&encrypted).into())
182    }
183
184    #[inline]
185    pub fn blocking_encrypt(&self, data: impl AsRef<[u8]>) -> Result<Vec<u8>, anyhow::Error> {
186        simple_crypt::encrypt(data.as_ref(), self.encryption_key.as_bytes())
187    }
188
189    #[inline]
190    pub fn blocking_encrypt_base64(
191        &self,
192        data: impl AsRef<[u8]>,
193    ) -> Result<compact_str::CompactString, anyhow::Error> {
194        let encrypted = self.blocking_encrypt(data)?;
195        Ok(BASE64_ENGINE.encode(&encrypted).into())
196    }
197
198    pub async fn decrypt(
199        &self,
200        data: impl AsRef<[u8]> + Send + 'static,
201    ) -> Result<compact_str::CompactString, anyhow::Error> {
202        let decrypted = self.decrypt_raw(data).await?;
203
204        Ok(compact_str::CompactString::from_utf8_lossy(&decrypted))
205    }
206
207    pub async fn decrypt_raw(
208        &self,
209        data: impl AsRef<[u8]> + Send + 'static,
210    ) -> Result<Vec<u8>, anyhow::Error> {
211        let Some(decryption_cache) = &self.decryption_cache else {
212            return self.decrypt_uncached(data.as_ref().to_vec()).await;
213        };
214
215        let digest: [u8; 32] = sha2::Sha256::digest(data.as_ref()).into();
216        let decrypted = decryption_cache
217            .try_get_with(digest, async {
218                self.decrypt_uncached(data.as_ref().to_vec())
219                    .await
220                    .map(Arc::new)
221            })
222            .await
223            .map_err(crate::cache::SharedComputeError)?;
224
225        Ok(decrypted.as_ref().clone())
226    }
227
228    async fn decrypt_uncached(&self, data: Vec<u8>) -> Result<Vec<u8>, anyhow::Error> {
229        let encryption_key = self.encryption_key.clone();
230
231        tokio::task::spawn_blocking(move || simple_crypt::decrypt(&data, encryption_key.as_bytes()))
232            .await?
233    }
234
235    pub async fn decrypt_base64(
236        &self,
237        data: impl AsRef<str>,
238    ) -> Result<compact_str::CompactString, anyhow::Error> {
239        let decoded = BASE64_ENGINE.decode(data.as_ref())?;
240        self.decrypt(decoded).await
241    }
242
243    pub async fn decrypt_base64_raw(
244        &self,
245        data: impl AsRef<str>,
246    ) -> Result<Vec<u8>, anyhow::Error> {
247        let decoded = BASE64_ENGINE.decode(data.as_ref())?;
248        self.decrypt_raw(decoded).await
249    }
250
251    pub async fn decrypt_base64_optional(
252        &self,
253        data: impl AsRef<str>,
254    ) -> Result<Option<compact_str::CompactString>, anyhow::Error> {
255        match BASE64_ENGINE.decode(data.as_ref()) {
256            Ok(decoded) => Ok(Some(self.decrypt(decoded).await?)),
257            Err(_) => Ok(None),
258        }
259    }
260
261    pub async fn decrypt_base64_raw_optional(
262        &self,
263        data: impl AsRef<str>,
264    ) -> Result<Option<Vec<u8>>, anyhow::Error> {
265        match BASE64_ENGINE.decode(data.as_ref()) {
266            Ok(decoded) => Ok(Some(self.decrypt_raw(decoded).await?)),
267            Err(_) => Ok(None),
268        }
269    }
270
271    #[inline]
272    pub fn blocking_decrypt(
273        &self,
274        data: impl AsRef<[u8]>,
275    ) -> Result<compact_str::CompactString, anyhow::Error> {
276        simple_crypt::decrypt(data.as_ref(), self.encryption_key.as_bytes())
277            .map(|s| compact_str::CompactString::from_utf8_lossy(&s))
278    }
279
280    #[inline]
281    pub fn blocking_decrypt_raw(&self, data: impl AsRef<[u8]>) -> Result<Vec<u8>, anyhow::Error> {
282        simple_crypt::decrypt(data.as_ref(), self.encryption_key.as_bytes())
283    }
284
285    #[inline]
286    pub fn blocking_decrypt_base64(
287        &self,
288        data: impl AsRef<str>,
289    ) -> Result<compact_str::CompactString, anyhow::Error> {
290        let decoded = BASE64_ENGINE.decode(data.as_ref())?;
291        self.blocking_decrypt(decoded)
292    }
293
294    #[inline]
295    pub fn blocking_decrypt_base64_raw(
296        &self,
297        data: impl AsRef<str>,
298    ) -> Result<Vec<u8>, anyhow::Error> {
299        let decoded = BASE64_ENGINE.decode(data.as_ref())?;
300        self.blocking_decrypt_raw(decoded)
301    }
302
303    #[inline]
304    pub fn blocking_decrypt_base64_optional(
305        &self,
306        data: impl AsRef<str>,
307    ) -> Result<Option<compact_str::CompactString>, anyhow::Error> {
308        match BASE64_ENGINE.decode(data.as_ref()) {
309            Ok(decoded) => Ok(Some(self.blocking_decrypt(decoded)?)),
310            Err(_) => Ok(None),
311        }
312    }
313
314    #[inline]
315    pub fn blocking_decrypt_base64_raw_optional(
316        &self,
317        data: impl AsRef<str>,
318    ) -> Result<Option<Vec<u8>>, anyhow::Error> {
319        match BASE64_ENGINE.decode(data.as_ref()) {
320            Ok(decoded) => Ok(Some(self.blocking_decrypt_raw(decoded)?)),
321            Err(_) => Ok(None),
322        }
323    }
324
325    #[inline]
326    pub async fn batch_action(
327        &self,
328        key: &'static str,
329        uuid: uuid::Uuid,
330        action: impl Future<Output = Result<(), anyhow::Error>> + Send + 'static,
331    ) {
332        let mut actions = self.batch_actions.lock().await;
333        actions.insert((key, uuid), Box::pin(action));
334    }
335}
336
337#[derive(Debug)]
338pub enum DatabaseError {
339    Sqlx(sqlx::Error),
340    Mongodb(mongodb::error::Error),
341    Serde(serde_json::Error),
342    Any(anyhow::Error),
343    Validation(garde::Report),
344    InvalidRelation(InvalidRelationError),
345}
346
347impl Display for DatabaseError {
348    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
349        match self {
350            Self::Sqlx(sqlx_value) => sqlx_value.fmt(f),
351            Self::Mongodb(mongodb_value) => mongodb_value.fmt(f),
352            Self::Serde(serde_value) => serde_value.fmt(f),
353            Self::Any(any_value) => any_value.fmt(f),
354            Self::Validation(validation_value) => validation_value.fmt(f),
355            Self::InvalidRelation(relation_value) => relation_value.fmt(f),
356        }
357    }
358}
359
360impl From<wings_api::client::ApiHttpError> for DatabaseError {
361    #[inline]
362    fn from(value: wings_api::client::ApiHttpError) -> Self {
363        Self::Any(value.into())
364    }
365}
366
367impl From<anyhow::Error> for DatabaseError {
368    #[inline]
369    fn from(value: anyhow::Error) -> Self {
370        Self::Any(value)
371    }
372}
373
374impl From<serde_json::Error> for DatabaseError {
375    #[inline]
376    fn from(value: serde_json::Error) -> Self {
377        Self::Serde(value)
378    }
379}
380
381impl From<sqlx::Error> for DatabaseError {
382    #[inline]
383    fn from(value: sqlx::Error) -> Self {
384        Self::Sqlx(value)
385    }
386}
387
388impl From<mongodb::error::Error> for DatabaseError {
389    #[inline]
390    fn from(value: mongodb::error::Error) -> Self {
391        Self::Mongodb(value)
392    }
393}
394
395impl From<garde::Report> for DatabaseError {
396    #[inline]
397    fn from(value: garde::Report) -> Self {
398        Self::Validation(value)
399    }
400}
401
402impl From<InvalidRelationError> for DatabaseError {
403    fn from(value: InvalidRelationError) -> Self {
404        Self::InvalidRelation(value)
405    }
406}
407
408impl DatabaseError {
409    #[inline]
410    pub fn is_unique_constraint_violation(&self, constraint: &str) -> bool {
411        match self {
412            Self::Sqlx(sqlx_value) => sqlx_value
413                .as_database_error()
414                .is_some_and(|e| e.is_unique_violation() && e.constraint() == Some(constraint)),
415            _ => false,
416        }
417    }
418
419    #[inline]
420    pub fn is_unique_violation(&self) -> bool {
421        match self {
422            Self::Sqlx(sqlx_value) => sqlx_value
423                .as_database_error()
424                .is_some_and(|e| e.is_unique_violation()),
425            _ => false,
426        }
427    }
428
429    #[inline]
430    pub fn is_foreign_key_violation(&self) -> bool {
431        match self {
432            Self::Sqlx(sqlx_value) => sqlx_value
433                .as_database_error()
434                .is_some_and(|e| e.is_foreign_key_violation()),
435            _ => false,
436        }
437    }
438
439    #[inline]
440    pub fn is_check_violation(&self) -> bool {
441        match self {
442            Self::Sqlx(sqlx_value) => sqlx_value
443                .as_database_error()
444                .is_some_and(|e| e.is_check_violation()),
445            _ => false,
446        }
447    }
448
449    #[inline]
450    pub const fn is_validation_error(&self) -> bool {
451        matches!(self, Self::Validation(_))
452    }
453
454    #[inline]
455    pub const fn is_invalid_relation(&self) -> bool {
456        matches!(self, Self::InvalidRelation(_))
457    }
458}
459
460impl std::error::Error for DatabaseError {}
461
462#[derive(Debug)]
463pub struct InvalidRelationError(pub &'static str);
464
465impl Display for InvalidRelationError {
466    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
467        write!(f, "invalid relation `{}` provided", self.0)
468    }
469}
shared/database.rs

shared/
database.rs
