shared/models/
mod.rs

1use crate::database::DatabaseError;
2use compact_str::CompactStringExt;
3use futures_util::{StreamExt, TryStreamExt};
4use garde::Validate;
5use parking_lot::RwLock;
6use serde::{Deserialize, Serialize, de::DeserializeOwned};
7use sha2::Digest;
8use sqlx::{
9    Arguments, Postgres, QueryBuilder, Row,
10    encode::IsNull,
11    error::BoxDynError,
12    postgres::{PgArgumentBuffer, PgArguments, PgRow, PgTypeInfo},
13};
14use std::{
15    collections::{BTreeMap, HashSet},
16    marker::PhantomData,
17    pin::Pin,
18    sync::{Arc, LazyLock},
19};
20use utoipa::ToSchema;
21
22pub mod admin_activity;
23pub mod announcement;
24pub mod backup_configuration;
25pub mod database_agent_host;
26pub mod database_agent_template;
27pub mod database_host;
28pub mod egg_configuration;
29pub mod egg_repository;
30pub mod egg_repository_egg;
31pub mod location;
32pub mod location_database_agent_host;
33pub mod location_database_host;
34pub mod mount;
35pub mod nest;
36pub mod nest_egg;
37pub mod nest_egg_mount;
38pub mod nest_egg_variable;
39pub mod node;
40pub mod node_allocation;
41pub mod node_database_agent_host;
42pub mod node_database_host;
43pub mod node_mount;
44pub mod node_tunnel;
45pub mod oauth_provider;
46pub mod oauth_provider_mapping;
47pub mod role;
48pub mod server;
49pub mod server_activity;
50pub mod server_allocation;
51pub mod server_backup;
52pub mod server_backup_group;
53pub mod server_database;
54pub mod server_database_instance;
55pub mod server_mount;
56pub mod server_schedule;
57pub mod server_schedule_step;
58pub mod server_subuser;
59pub mod server_tunnel;
60pub mod server_variable;
61pub mod system_backup_policy;
62pub mod system_backup_policy_database_agent_host;
63pub mod system_backup_policy_location;
64pub mod system_backup_policy_node;
65pub mod system_backup_policy_server;
66pub mod user;
67pub mod user_activity;
68pub mod user_api_key;
69pub mod user_command_snippet;
70pub mod user_email_verification;
71pub mod user_oauth_link;
72pub mod user_password_reset;
73pub mod user_recovery_code;
74pub mod user_security_key;
75pub mod user_server_group;
76pub mod user_session;
77pub mod user_ssh_key;
78pub mod user_two_factor_code;
79
80#[derive(ToSchema, Validate, Deserialize, Serialize)]
81pub struct PaginationParams {
82    #[garde(range(min = 1))]
83    #[schema(minimum = 1)]
84    #[serde(default = "Pagination::default_page")]
85    pub page: i64,
86    #[garde(range(min = 1, max = 100))]
87    #[schema(minimum = 1, maximum = 100)]
88    #[serde(default = "Pagination::default_per_page")]
89    pub per_page: i64,
90}
91
92#[derive(ToSchema, Validate, Deserialize, Serialize)]
93pub struct PaginationParamsWithSearch {
94    #[garde(range(min = 1))]
95    #[schema(minimum = 1)]
96    #[serde(default = "Pagination::default_page")]
97    pub page: i64,
98    #[garde(range(min = 1, max = 100))]
99    #[schema(minimum = 1, maximum = 100)]
100    #[serde(default = "Pagination::default_per_page")]
101    pub per_page: i64,
102    #[garde(length(chars, min = 1, max = 128))]
103    #[schema(min_length = 1, max_length = 128)]
104    #[serde(
105        default,
106        deserialize_with = "crate::deserialize::deserialize_string_option"
107    )]
108    pub search: Option<compact_str::CompactString>,
109}
110
111#[derive(ToSchema, Deserialize, Serialize)]
112pub struct Pagination<T: Serialize = serde_json::Value> {
113    pub total: i64,
114    pub per_page: i64,
115    pub page: i64,
116
117    pub data: Vec<T>,
118}
119
120impl Pagination {
121    #[inline]
122    pub const fn default_page() -> i64 {
123        1
124    }
125
126    #[inline]
127    pub const fn default_per_page() -> i64 {
128        25
129    }
130}
131
132impl<T: Serialize> Pagination<T> {
133    pub async fn async_map<R: serde::Serialize, Fut: Future<Output = R>>(
134        self,
135        mapper: impl Fn(T) -> Fut,
136    ) -> Pagination<R> {
137        let mut results = Vec::new();
138        results.reserve_exact(self.data.len());
139        let mut result_stream =
140            futures_util::stream::iter(self.data.into_iter().map(mapper)).buffered(25);
141
142        while let Some(result) = result_stream.next().await {
143            results.push(result);
144        }
145
146        Pagination {
147            total: self.total,
148            per_page: self.per_page,
149            page: self.page,
150            data: results,
151        }
152    }
153
154    pub async fn try_async_map<R: serde::Serialize, E, Fut: Future<Output = Result<R, E>>>(
155        self,
156        mapper: impl Fn(T) -> Fut,
157    ) -> Result<Pagination<R>, E> {
158        let mut results = Vec::new();
159        results.reserve_exact(self.data.len());
160        let mut result_stream =
161            futures_util::stream::iter(self.data.into_iter().map(mapper)).buffered(25);
162
163        while let Some(result) = result_stream.try_next().await? {
164            results.push(result);
165        }
166
167        Ok(Pagination {
168            total: self.total,
169            per_page: self.per_page,
170            page: self.page,
171            data: results,
172        })
173    }
174}
175
176pub type ModelExtensionList = parking_lot::RwLock<Vec<Box<dyn ModelExtension + Send + Sync>>>;
177pub type ModelExtensionData = Vec<(compact_str::CompactString, Vec<u8>)>;
178pub type ModelExtensionMapType = Box<dyn erased_serde::Serialize>;
179
180pub trait ModelExtension {
181    fn extension_name(&self) -> &'static str;
182
183    fn extended_columns(&self, prefix: &str) -> BTreeMap<&'static str, compact_str::CompactString>;
184
185    fn map_extended(
186        &self,
187        prefix: &str,
188        row: &PgRow,
189    ) -> Result<ModelExtensionMapType, crate::database::DatabaseError>;
190}
191
192pub trait SafeModelExtension: ModelExtension {
193    type Value: Serialize + DeserializeOwned;
194
195    fn name() -> &'static str;
196}
197
198pub trait BaseModel: Serialize + DeserializeOwned {
199    const NAME: &'static str;
200
201    fn get_extension_list() -> &'static ModelExtensionList;
202    fn get_extension_data(&self) -> &ModelExtensionData;
203
204    /// Registers a model extension. If an extension with the same name is already registered, this function will do nothing.
205    fn register_model_extension(extension: impl ModelExtension + Send + Sync + 'static) {
206        let mut extensions = Self::get_extension_list().write();
207
208        if extensions
209            .iter()
210            .any(|e| e.extension_name() == extension.extension_name())
211        {
212            return;
213        }
214
215        extensions.push(Box::new(extension));
216    }
217
218    /// Parses a model extension from the model's extension data. If the extension is not found, or if the data cannot be deserialized, an error is returned.
219    ///
220    /// This can be costly depending on what is stored, so use sparingly.
221    fn parse_model_extension<Extension: SafeModelExtension>(
222        &self,
223    ) -> Result<Extension::Value, crate::database::DatabaseError>
224    where
225        Extension::Value: Serialize + DeserializeOwned,
226    {
227        let data = self.get_extension_data();
228
229        for (name, value) in data.iter() {
230            if name.as_str() == Extension::name() {
231                let deserialized =
232                    rmp_serde::from_slice::<Extension::Value>(value).map_err(anyhow::Error::new)?;
233
234                return Ok(deserialized);
235            }
236        }
237
238        Err(crate::database::DatabaseError::Any(anyhow::anyhow!(
239            "model extension not found"
240        )))
241    }
242
243    fn base_columns(prefix: Option<&str>) -> BTreeMap<&'static str, compact_str::CompactString>;
244    fn columns(prefix: Option<&str>) -> BTreeMap<&'static str, compact_str::CompactString> {
245        let extensions = Self::get_extension_list().read();
246
247        let mut columns = Self::base_columns(prefix);
248
249        for extension in extensions.iter() {
250            columns.extend(extension.extended_columns(prefix.unwrap_or_default()));
251        }
252
253        columns
254    }
255
256    #[inline]
257    fn columns_sql(prefix: Option<&str>) -> compact_str::CompactString {
258        Self::columns(prefix)
259            .iter()
260            .map(|(key, value)| compact_str::format_compact!("{key} as {value}"))
261            .join_compact(", ")
262    }
263
264    fn map_extensions(
265        prefix: &str,
266        row: &PgRow,
267    ) -> Result<ModelExtensionData, crate::database::DatabaseError> {
268        let mut data = Vec::new();
269
270        let extensions = Self::get_extension_list().read();
271        for extension in extensions.iter() {
272            let value = extension.map_extended(prefix, row)?;
273            let serialized = rmp_serde::to_vec(&value).map_err(anyhow::Error::new)?;
274
275            data.push((
276                compact_str::CompactString::const_new(extension.extension_name()),
277                serialized,
278            ));
279        }
280
281        Ok(data)
282    }
283
284    fn map(prefix: Option<&str>, row: &PgRow) -> Result<Self, crate::database::DatabaseError>;
285
286    fn cache_invalidation_keys(&self) -> Vec<compact_str::CompactString> {
287        Vec::new()
288    }
289}
290
291pub async fn invalidate_cache_keys(
292    database: &crate::database::Database,
293    keys: &[compact_str::CompactString],
294) {
295    for key in keys {
296        if let Err(err) = database.cache.invalidate(key).await {
297            tracing::warn!(key = %key, "failed to invalidate cache entry: {err:#?}");
298        }
299    }
300}
301
302pub trait EventEmittingModel: BaseModel {
303    type Event: Send + Sync + 'static;
304
305    fn get_event_emitter() -> &'static crate::events::EventEmitter<Self::Event>;
306
307    fn register_event_handler<
308        F: Fn(crate::State, Arc<Self::Event>) -> Fut + Send + Sync + 'static,
309        Fut: Future<Output = Result<(), anyhow::Error>> + Send + 'static,
310    >(
311        listener: F,
312    ) -> crate::events::EventHandlerHandle {
313        Self::get_event_emitter().register_event_handler(listener)
314    }
315}
316
317type CreateHandlerResult<'a> =
318    Pin<Box<dyn Future<Output = Result<(), crate::database::DatabaseError>> + Send + 'a>>;
319type CreateHandler<M> = dyn for<'a> Fn(
320        &'a mut <M as CreatableModel>::CreateOptions<'_>,
321        &'a mut InsertQueryBuilder,
322        &'a crate::State,
323        &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
324    ) -> CreateHandlerResult<'a>
325    + Send
326    + Sync;
327type CreateAfterHandler<M> = dyn for<'a> Fn(
328        &'a mut <M as CreatableModel>::CreateResult,
329        &'a <M as CreatableModel>::CreateOptions<'_>,
330        &'a crate::State,
331        &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
332    ) -> CreateHandlerResult<'a>
333    + Send
334    + Sync;
335pub type CreateListenerList<M> =
336    Arc<ModelHandlerList<Arc<CreateHandler<M>>, Arc<CreateAfterHandler<M>>>>;
337
338#[async_trait::async_trait]
339pub trait CreatableModel: BaseModel + Send + Sync + 'static {
340    type CreateOptions<'a>: Send + Sync + Validate;
341    type CreateResult: Send;
342
343    fn get_create_handlers() -> &'static LazyLock<CreateListenerList<Self>>;
344
345    fn register_create_handler<
346        F: for<'a> Fn(
347                &'a mut Self::CreateOptions<'_>,
348                &'a mut InsertQueryBuilder,
349                &'a crate::State,
350                &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
351            ) -> Pin<
352                Box<dyn Future<Output = Result<(), crate::database::DatabaseError>> + Send + 'a>,
353            > + Send
354            + Sync
355            + 'static,
356    >(
357        priority: ListenerPriority,
358        callback: F,
359    ) {
360        let erased = Arc::new(callback) as Arc<CreateHandler<Self>>;
361
362        Self::get_create_handlers().register_handler(priority, erased);
363    }
364
365    fn register_after_create_handler<
366        F: for<'a> Fn(
367                &'a mut Self::CreateResult,
368                &'a Self::CreateOptions<'_>,
369                &'a crate::State,
370                &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
371            ) -> Pin<
372                Box<dyn Future<Output = Result<(), crate::database::DatabaseError>> + Send + 'a>,
373            > + Send
374            + Sync
375            + 'static,
376    >(
377        priority: ListenerPriority,
378        callback: F,
379    ) {
380        let erased = Arc::new(callback) as Arc<CreateAfterHandler<Self>>;
381
382        Self::get_create_handlers().register_after_handler(priority, erased);
383    }
384
385    async fn run_create_handlers(
386        options: &mut Self::CreateOptions<'_>,
387        query_builder: &mut InsertQueryBuilder,
388        state: &crate::State,
389        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
390    ) -> Result<(), crate::database::DatabaseError> {
391        let callbacks = Self::get_create_handlers()
392            .before_handlers
393            .read()
394            .iter()
395            .map(|l| l.callback.clone())
396            .collect::<Vec<_>>();
397
398        for callback in callbacks.iter() {
399            (*callback)(options, query_builder, state, transaction).await?;
400        }
401
402        Ok(())
403    }
404
405    async fn run_after_create_handlers(
406        result: &mut Self::CreateResult,
407        options: &Self::CreateOptions<'_>,
408        state: &crate::State,
409        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
410    ) -> Result<(), crate::database::DatabaseError> {
411        let callbacks = Self::get_create_handlers()
412            .after_handlers
413            .read()
414            .iter()
415            .map(|l| l.callback.clone())
416            .collect::<Vec<_>>();
417
418        for callback in callbacks.iter() {
419            (*callback)(result, options, state, transaction).await?;
420        }
421
422        Ok(())
423    }
424
425    async fn create_with_transaction(
426        state: &crate::State,
427        options: Self::CreateOptions<'_>,
428        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
429    ) -> Result<Self::CreateResult, crate::database::DatabaseError>;
430
431    async fn create(
432        state: &crate::State,
433        options: Self::CreateOptions<'_>,
434    ) -> Result<Self::CreateResult, crate::database::DatabaseError> {
435        let mut transaction = state.database.write().begin().await?;
436
437        let result = match Self::create_with_transaction(state, options, &mut transaction).await {
438            Ok(result) => result,
439            Err(err) => {
440                transaction.rollback().await?;
441                return Err(err);
442            }
443        };
444
445        transaction.commit().await?;
446
447        Ok(result)
448    }
449}
450
451type UpdateHandlerResult<'a> =
452    Pin<Box<dyn Future<Output = Result<(), crate::database::DatabaseError>> + Send + 'a>>;
453type UpdateHandler<M> = dyn for<'a> Fn(
454        &'a mut M,
455        &'a mut <M as UpdatableModel>::UpdateOptions,
456        &'a mut UpdateQueryBuilder,
457        &'a crate::State,
458        &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
459    ) -> UpdateHandlerResult<'a>
460    + Send
461    + Sync;
462type UpdateAfterHandler<M> = dyn for<'a> Fn(
463        &'a mut M,
464        &'a crate::State,
465        &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
466    ) -> UpdateHandlerResult<'a>
467    + Send
468    + Sync;
469pub type UpdateHandlerList<M> =
470    Arc<ModelHandlerList<Arc<UpdateHandler<M>>, Arc<UpdateAfterHandler<M>>>>;
471
472#[async_trait::async_trait]
473pub trait UpdatableModel: BaseModel + Send + Sync + 'static {
474    type UpdateOptions: Send + Sync + Default + ToSchema + DeserializeOwned + Serialize + Validate;
475
476    fn get_update_handlers() -> &'static LazyLock<UpdateHandlerList<Self>>;
477
478    fn register_update_handler<
479        F: for<'a> Fn(
480                &'a mut Self,
481                &'a mut Self::UpdateOptions,
482                &'a mut UpdateQueryBuilder,
483                &'a crate::State,
484                &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
485            ) -> Pin<
486                Box<dyn Future<Output = Result<(), crate::database::DatabaseError>> + Send + 'a>,
487            > + Send
488            + Sync
489            + 'static,
490    >(
491        priority: ListenerPriority,
492        callback: F,
493    ) {
494        let erased = Arc::new(callback) as Arc<UpdateHandler<Self>>;
495
496        Self::get_update_handlers().register_handler(priority, erased);
497    }
498
499    fn register_after_update_handler<
500        F: for<'a> Fn(
501                &'a mut Self,
502                &'a crate::State,
503                &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
504            ) -> Pin<
505                Box<dyn Future<Output = Result<(), crate::database::DatabaseError>> + Send + 'a>,
506            > + Send
507            + Sync
508            + 'static,
509    >(
510        priority: ListenerPriority,
511        callback: F,
512    ) {
513        let erased = Arc::new(callback) as Arc<UpdateAfterHandler<Self>>;
514
515        Self::get_update_handlers().register_after_handler(priority, erased);
516    }
517
518    async fn run_update_handlers(
519        &mut self,
520        options: &mut Self::UpdateOptions,
521        query_builder: &mut UpdateQueryBuilder,
522        state: &crate::State,
523        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
524    ) -> Result<(), crate::database::DatabaseError> {
525        let callbacks = Self::get_update_handlers()
526            .before_handlers
527            .read()
528            .iter()
529            .map(|l| l.callback.clone())
530            .collect::<Vec<_>>();
531
532        for callback in callbacks.iter() {
533            (*callback)(self, options, query_builder, state, transaction).await?;
534        }
535
536        Ok(())
537    }
538
539    async fn run_after_update_handlers(
540        &mut self,
541        state: &crate::State,
542        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
543    ) -> Result<(), crate::database::DatabaseError> {
544        let callbacks = Self::get_update_handlers()
545            .after_handlers
546            .read()
547            .iter()
548            .map(|l| l.callback.clone())
549            .collect::<Vec<_>>();
550
551        for callback in callbacks.iter() {
552            (*callback)(self, state, transaction).await?;
553        }
554
555        Ok(())
556    }
557
558    async fn update_with_transaction(
559        &mut self,
560        state: &crate::State,
561        options: Self::UpdateOptions,
562        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
563    ) -> Result<(), crate::database::DatabaseError>;
564
565    async fn update(
566        &mut self,
567        state: &crate::State,
568        options: Self::UpdateOptions,
569    ) -> Result<(), crate::database::DatabaseError> {
570        let mut transaction = state.database.write().begin().await?;
571
572        if let Err(err) = self
573            .update_with_transaction(state, options, &mut transaction)
574            .await
575        {
576            transaction.rollback().await?;
577            return Err(err);
578        }
579
580        transaction.commit().await?;
581
582        invalidate_cache_keys(&state.database, &self.cache_invalidation_keys()).await;
583
584        Ok(())
585    }
586}
587
588type DeleteHandlerResult<'a> = Pin<Box<dyn Future<Output = Result<(), anyhow::Error>> + Send + 'a>>;
589type DeleteHandler<M> = dyn for<'a> Fn(
590        &'a M,
591        &'a <M as DeletableModel>::DeleteOptions,
592        &'a crate::State,
593        &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
594    ) -> DeleteHandlerResult<'a>
595    + Send
596    + Sync;
597type DeleteAfterHandler<M> = dyn for<'a> Fn(
598        &'a M,
599        &'a <M as DeletableModel>::DeleteOptions,
600        &'a crate::State,
601        &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
602    ) -> DeleteHandlerResult<'a>
603    + Send
604    + Sync;
605pub type DeleteHandlerList<M> =
606    Arc<ModelHandlerList<Arc<DeleteHandler<M>>, Arc<DeleteAfterHandler<M>>>>;
607
608#[async_trait::async_trait]
609pub trait DeletableModel: BaseModel + Send + Sync + 'static {
610    type DeleteOptions: Send + Sync + Default + Clone;
611
612    fn get_delete_handlers() -> &'static LazyLock<DeleteHandlerList<Self>>;
613
614    fn register_delete_handler<
615        F: for<'a> Fn(
616                &'a Self,
617                &'a Self::DeleteOptions,
618                &'a crate::State,
619                &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
620            )
621                -> Pin<Box<dyn Future<Output = Result<(), anyhow::Error>> + Send + 'a>>
622            + Send
623            + Sync
624            + 'static,
625    >(
626        priority: ListenerPriority,
627        callback: F,
628    ) {
629        let erased = Arc::new(callback) as Arc<DeleteHandler<Self>>;
630
631        Self::get_delete_handlers().register_handler(priority, erased);
632    }
633
634    fn register_after_delete_handler<
635        F: for<'a> Fn(
636                &'a Self,
637                &'a Self::DeleteOptions,
638                &'a crate::State,
639                &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
640            )
641                -> Pin<Box<dyn Future<Output = Result<(), anyhow::Error>> + Send + 'a>>
642            + Send
643            + Sync
644            + 'static,
645    >(
646        priority: ListenerPriority,
647        callback: F,
648    ) {
649        let erased = Arc::new(callback) as Arc<DeleteAfterHandler<Self>>;
650
651        Self::get_delete_handlers().register_after_handler(priority, erased);
652    }
653
654    async fn run_delete_handlers(
655        &self,
656        options: &Self::DeleteOptions,
657        state: &crate::State,
658        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
659    ) -> Result<(), anyhow::Error> {
660        let callbacks = Self::get_delete_handlers()
661            .before_handlers
662            .read()
663            .iter()
664            .map(|l| l.callback.clone())
665            .collect::<Vec<_>>();
666
667        for callback in callbacks.iter() {
668            (*callback)(self, options, state, transaction).await?;
669        }
670
671        Ok(())
672    }
673
674    async fn run_after_delete_handlers(
675        &self,
676        options: &Self::DeleteOptions,
677        state: &crate::State,
678        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
679    ) -> Result<(), anyhow::Error> {
680        let callbacks = Self::get_delete_handlers()
681            .after_handlers
682            .read()
683            .iter()
684            .map(|l| l.callback.clone())
685            .collect::<Vec<_>>();
686
687        for callback in callbacks.iter() {
688            (*callback)(self, options, state, transaction).await?;
689        }
690
691        Ok(())
692    }
693
694    async fn delete_with_transaction(
695        &self,
696        state: &crate::State,
697        options: Self::DeleteOptions,
698        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
699    ) -> Result<(), anyhow::Error>;
700
701    async fn delete(
702        &self,
703        state: &crate::State,
704        options: Self::DeleteOptions,
705    ) -> Result<(), anyhow::Error> {
706        let mut transaction = state.database.write().begin().await?;
707
708        if let Err(err) = self
709            .delete_with_transaction(state, options, &mut transaction)
710            .await
711        {
712            transaction.rollback().await?;
713            return Err(err);
714        }
715
716        transaction.commit().await?;
717
718        invalidate_cache_keys(&state.database, &self.cache_invalidation_keys()).await;
719
720        Ok(())
721    }
722}
723
724type DuplicateHandlerResult<'a> =
725    Pin<Box<dyn Future<Output = Result<(), crate::database::DatabaseError>> + Send + 'a>>;
726type DuplicateHandler<M> = dyn for<'a> Fn(
727        &'a M,
728        &'a <M as DuplicableModel>::DuplicateOptions<'_>,
729        &'a crate::State,
730        &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
731    ) -> DuplicateHandlerResult<'a>
732    + Send
733    + Sync;
734type DuplicateAfterHandler<M> = dyn for<'a> Fn(
735        &'a M,
736        &'a mut M,
737        &'a <M as DuplicableModel>::DuplicateOptions<'_>,
738        &'a crate::State,
739        &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
740    ) -> DuplicateHandlerResult<'a>
741    + Send
742    + Sync;
743pub type DuplicateHandlerList<M> =
744    Arc<ModelHandlerList<Arc<DuplicateHandler<M>>, Arc<DuplicateAfterHandler<M>>>>;
745
746#[async_trait::async_trait]
747pub trait DuplicableModel: BaseModel + Send + Sync + 'static {
748    type DuplicateOptions<'a>: Send + Sync + Validate;
749
750    fn get_duplicate_handlers() -> &'static LazyLock<DuplicateHandlerList<Self>>;
751
752    fn register_duplicate_handler<
753        F: for<'a> Fn(
754                &'a Self,
755                &'a Self::DuplicateOptions<'_>,
756                &'a crate::State,
757                &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
758            ) -> DuplicateHandlerResult<'a>
759            + Send
760            + Sync
761            + 'static,
762    >(
763        priority: ListenerPriority,
764        callback: F,
765    ) {
766        let erased = Arc::new(callback) as Arc<DuplicateHandler<Self>>;
767
768        Self::get_duplicate_handlers().register_handler(priority, erased);
769    }
770
771    fn register_after_duplicate_handler<
772        F: for<'a> Fn(
773                &'a Self,
774                &'a mut Self,
775                &'a Self::DuplicateOptions<'_>,
776                &'a crate::State,
777                &'a mut sqlx::Transaction<'_, sqlx::Postgres>,
778            ) -> DuplicateHandlerResult<'a>
779            + Send
780            + Sync
781            + 'static,
782    >(
783        priority: ListenerPriority,
784        callback: F,
785    ) {
786        let erased = Arc::new(callback) as Arc<DuplicateAfterHandler<Self>>;
787
788        Self::get_duplicate_handlers().register_after_handler(priority, erased);
789    }
790
791    async fn run_duplicate_handlers(
792        &self,
793        options: &Self::DuplicateOptions<'_>,
794        state: &crate::State,
795        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
796    ) -> Result<(), crate::database::DatabaseError> {
797        let callbacks = Self::get_duplicate_handlers()
798            .before_handlers
799            .read()
800            .iter()
801            .map(|l| l.callback.clone())
802            .collect::<Vec<_>>();
803
804        for callback in callbacks.iter() {
805            (*callback)(self, options, state, transaction).await?;
806        }
807
808        Ok(())
809    }
810
811    async fn run_after_duplicate_handlers(
812        &self,
813        duplicated: &mut Self,
814        options: &Self::DuplicateOptions<'_>,
815        state: &crate::State,
816        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
817    ) -> Result<(), crate::database::DatabaseError> {
818        let callbacks = Self::get_duplicate_handlers()
819            .after_handlers
820            .read()
821            .iter()
822            .map(|l| l.callback.clone())
823            .collect::<Vec<_>>();
824
825        for callback in callbacks.iter() {
826            (*callback)(self, duplicated, options, state, transaction).await?;
827        }
828
829        Ok(())
830    }
831
832    async fn duplicate_with_transaction(
833        &self,
834        state: &crate::State,
835        options: Self::DuplicateOptions<'_>,
836        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
837    ) -> Result<Self, crate::database::DatabaseError>;
838
839    async fn duplicate(
840        &self,
841        state: &crate::State,
842        options: Self::DuplicateOptions<'_>,
843    ) -> Result<Self, crate::database::DatabaseError> {
844        let mut transaction = state.database.write().begin().await?;
845
846        let duplicated = match self
847            .duplicate_with_transaction(state, options, &mut transaction)
848            .await
849        {
850            Ok(duplicated) => duplicated,
851            Err(err) => {
852                transaction.rollback().await?;
853                return Err(err);
854            }
855        };
856
857        transaction.commit().await?;
858
859        Ok(duplicated)
860    }
861}
862
863pub const BY_UUID_CACHE_TTL: u64 = 10;
864
865#[async_trait::async_trait]
866pub trait ByUuid: BaseModel {
867    async fn by_uuid(
868        database: &crate::database::Database,
869        uuid: uuid::Uuid,
870    ) -> Result<Self, DatabaseError>;
871
872    async fn by_uuid_with_transaction(
873        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
874        uuid: uuid::Uuid,
875    ) -> Result<Self, DatabaseError>;
876
877    async fn by_uuid_cached(
878        database: &crate::database::Database,
879        uuid: uuid::Uuid,
880    ) -> Result<Self, anyhow::Error> {
881        database
882            .cache
883            .cached(
884                &format!("{}::{uuid}", Self::NAME),
885                BY_UUID_CACHE_TTL,
886                || Self::by_uuid(database, uuid),
887            )
888            .await
889    }
890
891    async fn by_uuid_optional(
892        database: &crate::database::Database,
893        uuid: uuid::Uuid,
894    ) -> Result<Option<Self>, DatabaseError> {
895        match Self::by_uuid(database, uuid).await {
896            Ok(res) => Ok(Some(res)),
897            Err(DatabaseError::Sqlx(sqlx::Error::RowNotFound)) => Ok(None),
898            Err(err) => Err(err),
899        }
900    }
901
902    async fn by_uuid_optional_with_transaction(
903        transaction: &mut sqlx::Transaction<'_, sqlx::Postgres>,
904        uuid: uuid::Uuid,
905    ) -> Result<Option<Self>, DatabaseError> {
906        match Self::by_uuid_with_transaction(transaction, uuid).await {
907            Ok(res) => Ok(Some(res)),
908            Err(DatabaseError::Sqlx(sqlx::Error::RowNotFound)) => Ok(None),
909            Err(err) => Err(err),
910        }
911    }
912
913    async fn by_uuid_optional_cached(
914        database: &crate::database::Database,
915        uuid: uuid::Uuid,
916    ) -> Result<Option<Self>, anyhow::Error> {
917        match Self::by_uuid_cached(database, uuid).await {
918            Ok(res) => Ok(Some(res)),
919            Err(err)
920                if err.chain().any(|err| {
921                    matches!(
922                        err.downcast_ref::<DatabaseError>(),
923                        Some(DatabaseError::Sqlx(sqlx::Error::RowNotFound))
924                    )
925                }) =>
926            {
927                Ok(None)
928            }
929            Err(err) => Err(err),
930        }
931    }
932
933    #[inline]
934    fn get_fetchable(uuid: uuid::Uuid) -> Fetchable<Self> {
935        Fetchable {
936            uuid,
937            _model: PhantomData,
938        }
939    }
940
941    #[inline]
942    fn get_fetchable_from_row(row: &PgRow, column: impl AsRef<str>) -> Option<Fetchable<Self>> {
943        match row.try_get(column.as_ref()) {
944            Ok(uuid) => Some(Fetchable {
945                uuid,
946                _model: PhantomData,
947            }),
948            Err(_) => None,
949        }
950    }
951
952    async fn invalidate_cached(database: &crate::database::Database, uuid: uuid::Uuid) {
953        let key = compact_str::format_compact!("{}::{uuid}", Self::NAME);
954
955        invalidate_cache_keys(database, &[key]).await;
956    }
957}
958
959#[derive(Debug)]
960struct UnresolvedIdentifier;
961
962impl std::fmt::Display for UnresolvedIdentifier {
963    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
964        f.write_str("identifier does not resolve to a row")
965    }
966}
967
968impl std::error::Error for UnresolvedIdentifier {}
969
970#[async_trait::async_trait]
971pub trait ResolvableModel: ByUuid + Serialize + DeserializeOwned + Send + Sync {
972    type Fingerprint: Serialize + Send;
973
974    fn uuid(&self) -> uuid::Uuid;
975
976    fn fingerprint(&self) -> Self::Fingerprint;
977
978    async fn resolve(
979        database: &crate::database::Database,
980        identifier: &str,
981    ) -> Result<Option<Self>, anyhow::Error>;
982
983    async fn resolve_cached(
984        database: &crate::database::Database,
985        identifier: &str,
986    ) -> Result<Option<Self>, anyhow::Error> {
987        let key = format!(
988            "{}::resolve::{}",
989            Self::NAME,
990            hex::encode(sha2::Sha256::digest(identifier.as_bytes()))
991        );
992
993        if let Some(resolution) = database.cache.resolution(&key).await {
994            if let Some(row) = Self::by_uuid_optional_cached(database, resolution.uuid).await?
995                && rmp_serde::to_vec(&row.fingerprint())? == *resolution.fingerprint
996            {
997                return Ok(Some(row));
998            }
999
1000            database.cache.remove_resolution(&key).await;
1001        }
1002
1003        let resolution = database
1004            .cache
1005            .resolve(&key, || async {
1006                let Some(row) = Self::resolve(database, identifier).await? else {
1007                    return Err(UnresolvedIdentifier.into());
1008                };
1009
1010                database
1011                    .cache
1012                    .set(
1013                        &format!("{}::{}", Self::NAME, row.uuid()),
1014                        BY_UUID_CACHE_TTL,
1015                        &row,
1016                    )
1017                    .await?;
1018
1019                Ok(crate::cache::Resolution {
1020                    uuid: row.uuid(),
1021                    fingerprint: Arc::new(rmp_serde::to_vec(&row.fingerprint())?),
1022                })
1023            })
1024            .await;
1025
1026        match resolution {
1027            Ok(resolution) => Self::by_uuid_optional_cached(database, resolution.uuid).await,
1028            Err(err) if err.chain().any(|err| err.is::<UnresolvedIdentifier>()) => Ok(None),
1029            Err(err) => Err(err),
1030        }
1031    }
1032}
1033
1034#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
1035pub enum ListenerPriority {
1036    Highest,
1037    High,
1038    #[default]
1039    Normal,
1040    Low,
1041    Lowest,
1042}
1043
1044impl ListenerPriority {
1045    #[inline]
1046    fn rank(self) -> u8 {
1047        match self {
1048            Self::Highest => 5,
1049            Self::High => 4,
1050            Self::Normal => 3,
1051            Self::Low => 2,
1052            Self::Lowest => 1,
1053        }
1054    }
1055}
1056
1057impl PartialOrd for ListenerPriority {
1058    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
1059        Some(self.cmp(other))
1060    }
1061}
1062
1063impl Ord for ListenerPriority {
1064    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
1065        let self_rank = self.rank();
1066        let other_rank = other.rank();
1067
1068        other_rank.cmp(&self_rank)
1069    }
1070}
1071
1072impl<F: Send + Sync, AfterF: Send + Sync> crate::events::DisconnectEventHandler
1073    for ModelHandlerList<F, AfterF>
1074{
1075    #[inline]
1076    fn disconnect(&self, id: uuid::Uuid) {
1077        self.before_handlers.write().retain(|l| l.uuid != id);
1078        self.after_handlers.write().retain(|l| l.uuid != id);
1079    }
1080}
1081
1082pub struct ModelHandlerList<F: Send + Sync + 'static, AfterF: Send + Sync + 'static> {
1083    before_handlers: RwLock<Vec<ModelHandler<F>>>,
1084    after_handlers: RwLock<Vec<ModelHandler<AfterF>>>,
1085}
1086
1087impl<F: Send + Sync + 'static, AfterF: Send + Sync + 'static> Default
1088    for ModelHandlerList<F, AfterF>
1089{
1090    fn default() -> Self {
1091        Self {
1092            before_handlers: RwLock::new(Vec::new()),
1093            after_handlers: RwLock::new(Vec::new()),
1094        }
1095    }
1096}
1097
1098impl<F: Send + Sync + 'static, AfterF: Send + Sync + 'static> ModelHandlerList<F, AfterF> {
1099    pub fn register_handler(
1100        self: &Arc<Self>,
1101        priority: ListenerPriority,
1102        callback: F,
1103    ) -> ModelHandlerHandle {
1104        let (listener, aborter) = ModelHandler::new(callback, priority, self.clone());
1105
1106        let mut self_listeners = self.before_handlers.write();
1107        self_listeners.push(listener);
1108        self_listeners.sort_by_key(|a| a.priority);
1109
1110        aborter
1111    }
1112
1113    pub fn register_after_handler(
1114        self: &Arc<Self>,
1115        priority: ListenerPriority,
1116        callback: AfterF,
1117    ) -> ModelHandlerHandle {
1118        let (listener, aborter) = ModelHandler::new(callback, priority, self.clone());
1119
1120        let mut self_listeners = self.after_handlers.write();
1121        self_listeners.push(listener);
1122        self_listeners.sort_by_key(|a| a.priority);
1123
1124        aborter
1125    }
1126}
1127
1128pub struct ModelHandler<F: Send + Sync + 'static> {
1129    uuid: uuid::Uuid,
1130    priority: ListenerPriority,
1131
1132    pub callback: F,
1133}
1134
1135impl<F: Send + Sync + 'static> ModelHandler<F> {
1136    pub(crate) fn new(
1137        callback: F,
1138        priority: ListenerPriority,
1139        list: Arc<dyn crate::events::DisconnectEventHandler + Send + Sync>,
1140    ) -> (Self, ModelHandlerHandle) {
1141        let handler = Self {
1142            uuid: uuid::Uuid::new_v4(),
1143            priority,
1144            callback,
1145        };
1146        let handle = ModelHandlerHandle {
1147            list_ref: list,
1148            id: handler.uuid,
1149        };
1150        (handler, handle)
1151    }
1152}
1153
1154pub struct ModelHandlerHandle {
1155    list_ref: Arc<dyn crate::events::DisconnectEventHandler + Send + Sync>,
1156    id: uuid::Uuid,
1157}
1158
1159impl ModelHandlerHandle {
1160    pub fn disconnect(&self) {
1161        self.list_ref.disconnect(self.id);
1162    }
1163}
1164
1165#[derive(Serialize, Deserialize, Clone, Copy)]
1166pub struct Fetchable<M: ByUuid> {
1167    pub uuid: uuid::Uuid,
1168    #[serde(skip)]
1169    _model: PhantomData<M>,
1170}
1171
1172impl<M: ByUuid + Send> Fetchable<M> {
1173    #[inline]
1174    pub async fn fetch(&self, database: &crate::database::Database) -> Result<M, DatabaseError> {
1175        M::by_uuid(database, self.uuid).await
1176    }
1177
1178    #[inline]
1179    pub async fn fetch_cached(
1180        &self,
1181        database: &crate::database::Database,
1182    ) -> Result<M, anyhow::Error> {
1183        M::by_uuid_cached(database, self.uuid).await
1184    }
1185
1186    #[inline]
1187    pub async fn fetch_optional(
1188        &self,
1189        database: &crate::database::Database,
1190    ) -> Result<Option<M>, DatabaseError> {
1191        M::by_uuid_optional(database, self.uuid).await
1192    }
1193
1194    #[inline]
1195    pub async fn fetch_optional_cached(
1196        &self,
1197        database: &crate::database::Database,
1198    ) -> Result<Option<M>, anyhow::Error> {
1199        M::by_uuid_optional_cached(database, self.uuid).await
1200    }
1201}
1202
1203pub struct InsertQueryBuilder<'a> {
1204    table: &'a str,
1205    columns: Vec<&'a str>,
1206    expressions: Vec<String>,
1207    arguments: PgArguments,
1208    returning_clause: Option<&'a str>,
1209}
1210
1211impl<'a> InsertQueryBuilder<'a> {
1212    pub fn new(table: &'a str) -> Self {
1213        Self {
1214            table,
1215            columns: Vec::new(),
1216            expressions: Vec::new(),
1217            arguments: PgArguments::default(),
1218            returning_clause: None,
1219        }
1220    }
1221
1222    pub fn set<T: 'a + sqlx::Encode<'a, Postgres> + sqlx::Type<Postgres> + Send>(
1223        &mut self,
1224        column: &'a str,
1225        value: T,
1226    ) -> &mut Self {
1227        if self.columns.contains(&column) {
1228            return self;
1229        }
1230
1231        if self.arguments.add(value).is_ok() {
1232            self.columns.push(column);
1233            let idx = self.arguments.len();
1234            self.expressions.push(format!("${}", idx));
1235        }
1236
1237        self
1238    }
1239
1240    pub fn set_expr<T: 'a + sqlx::Encode<'a, Postgres> + sqlx::Type<Postgres> + Send>(
1241        &mut self,
1242        column: &'a str,
1243        expression: &str,
1244        values: Vec<T>,
1245    ) -> &mut Self {
1246        if self.columns.contains(&column) {
1247            return self;
1248        }
1249
1250        let start_len = self.arguments.len();
1251
1252        for value in values {
1253            if self.arguments.add(value).is_err() {
1254                return self;
1255            }
1256        }
1257
1258        let mut expr = expression.to_string();
1259        let added_count = self.arguments.len() - start_len;
1260
1261        for i in (1..=added_count).rev() {
1262            let global_idx = start_len + i;
1263            expr = expr.replace(&format!("${}", i), &format!("${}", global_idx));
1264        }
1265
1266        self.columns.push(column);
1267        self.expressions.push(expr);
1268
1269        self
1270    }
1271
1272    pub fn returning(mut self, clause: &'a str) -> Self {
1273        self.returning_clause = Some(clause);
1274        self
1275    }
1276
1277    fn build_sql(&self) -> String {
1278        let columns_sql = self.columns.join(", ");
1279        let values_sql = self.expressions.join(", ");
1280
1281        let mut sql = format!(
1282            "INSERT INTO {} ({}) VALUES ({})",
1283            self.table, columns_sql, values_sql
1284        );
1285
1286        if let Some(clause) = self.returning_clause {
1287            sql.push_str(" RETURNING ");
1288            sql.push_str(clause);
1289        }
1290
1291        sql
1292    }
1293
1294    pub async fn execute(
1295        self,
1296        executor: impl sqlx::Executor<'a, Database = Postgres>,
1297    ) -> Result<sqlx::postgres::PgQueryResult, sqlx::Error> {
1298        let sql = self.build_sql();
1299        sqlx::query_with(sqlx::AssertSqlSafe(sql), self.arguments)
1300            .execute(executor)
1301            .await
1302    }
1303
1304    pub async fn fetch_one(
1305        self,
1306        executor: impl sqlx::Executor<'a, Database = Postgres>,
1307    ) -> Result<sqlx::postgres::PgRow, sqlx::Error> {
1308        let sql = self.build_sql();
1309        sqlx::query_with(sqlx::AssertSqlSafe(sql), self.arguments)
1310            .fetch_one(executor)
1311            .await
1312    }
1313}
1314
1315pub struct UpdateQueryBuilder<'a> {
1316    builder: QueryBuilder<Postgres>,
1317    updated_fields: HashSet<&'a str>,
1318    has_set_fields: bool,
1319    has_where: bool,
1320}
1321
1322impl<'a> UpdateQueryBuilder<'a> {
1323    pub fn new(table: &'a str) -> Self {
1324        let mut builder = QueryBuilder::new("UPDATE ");
1325        builder.push(table);
1326        builder.push(" SET ");
1327
1328        Self {
1329            builder,
1330            updated_fields: HashSet::new(),
1331            has_set_fields: false,
1332            has_where: false,
1333        }
1334    }
1335
1336    /// Adds a field to be updated, if `None`, will not add the field
1337    /// To set a field to null (`None`), you need a `Some(None)`
1338    pub fn set<T: 'a + sqlx::Encode<'a, Postgres> + sqlx::Type<Postgres> + Send>(
1339        &mut self,
1340        column: &'a str,
1341        value: Option<T>,
1342    ) -> &mut Self {
1343        let Some(value) = value else {
1344            return self;
1345        };
1346
1347        if !self.updated_fields.insert(column) {
1348            return self;
1349        }
1350
1351        if self.has_set_fields {
1352            self.builder.push(", ");
1353        }
1354
1355        self.builder.push(column);
1356        self.builder.push(" = ");
1357        self.builder.push_bind(value);
1358
1359        self.has_set_fields = true;
1360        self
1361    }
1362
1363    pub fn where_eq<T: 'a + sqlx::Encode<'a, Postgres> + sqlx::Type<Postgres> + Send>(
1364        &mut self,
1365        column: &'a str,
1366        value: T,
1367    ) -> &mut Self {
1368        if self.has_where {
1369            self.builder.push(" AND ");
1370        } else {
1371            self.builder.push(" WHERE ");
1372            self.has_where = true;
1373        }
1374
1375        self.builder.push(column);
1376        self.builder.push(" = ");
1377        self.builder.push_bind(value);
1378        self
1379    }
1380
1381    pub async fn execute(
1382        mut self,
1383        executor: impl sqlx::Executor<'a, Database = Postgres>,
1384    ) -> Result<sqlx::any::AnyQueryResult, sqlx::Error> {
1385        if !self.has_set_fields {
1386            return Ok(sqlx::any::AnyQueryResult::default());
1387        }
1388
1389        let query = self.builder.build();
1390        query.execute(executor).await.map(|r| r.into())
1391    }
1392}
1393
1394/// SQLx helper type to preserve order of keys when encoding JSON. By default, SQLx encodes JSON using `serde_json::Value`, which does not preserve order of keys. This type allows you to encode any serializable type as JSON while preserving the order of keys.
1395pub struct OrderedJson<T>(pub T);
1396
1397impl<T: Serialize> sqlx::Encode<'_, sqlx::Postgres> for OrderedJson<T> {
1398    fn encode_by_ref(&self, buf: &mut PgArgumentBuffer) -> Result<IsNull, BoxDynError> {
1399        serde_json::to_writer(&mut **buf, &self.0)?;
1400        Ok(IsNull::No)
1401    }
1402}
1403
1404impl<T> sqlx::Type<sqlx::Postgres> for OrderedJson<T> {
1405    fn type_info() -> PgTypeInfo {
1406        // JSON, not JSONB, to preserve order of keys
1407        PgTypeInfo::with_oid(sqlx::postgres::types::Oid(114))
1408    }
1409}
1410
1411#[async_trait::async_trait]
1412pub trait IntoApiObject {
1413    type ApiObject: Send;
1414    type ExtraArgs<'a>: Send;
1415
1416    async fn into_api_object<'a>(
1417        self,
1418        state: &crate::State,
1419        args: Self::ExtraArgs<'a>,
1420    ) -> Result<Self::ApiObject, DatabaseError>;
1421}
1422
1423#[async_trait::async_trait]
1424pub trait IntoAdminApiObject {
1425    type AdminApiObject: Send;
1426    type ExtraArgs<'a>: Send;
1427
1428    async fn into_admin_api_object<'a>(
1429        self,
1430        state: &crate::State,
1431        args: Self::ExtraArgs<'a>,
1432    ) -> Result<Self::AdminApiObject, DatabaseError>;
1433}
shared/models/mod.rs

shared/models/
mod.rs
