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