Skip to main content

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