Federated blogging application, thanks to ActivityPub https://joinplu.me
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 

90 lines
2.9 KiB

  1. use crate::{instance::Instance, Connection, CONFIG};
  2. use diesel::r2d2::{
  3. ConnectionManager, CustomizeConnection, Error as ConnError, Pool, PooledConnection,
  4. };
  5. #[cfg(feature = "sqlite")]
  6. use diesel::{dsl::sql_query, ConnectionError, RunQueryDsl};
  7. use rocket::{
  8. http::Status,
  9. request::{self, FromRequest},
  10. Outcome, Request, State,
  11. };
  12. use std::ops::Deref;
  13. pub type DbPool = Pool<ConnectionManager<Connection>>;
  14. /// Initializes a database pool.
  15. pub fn init_pool() -> Option<DbPool> {
  16. let manager = ConnectionManager::<Connection>::new(CONFIG.database_url.as_str());
  17. let mut builder = DbPool::builder()
  18. .connection_customizer(Box::new(PragmaForeignKey))
  19. .min_idle(CONFIG.db_min_idle);
  20. if let Some(max_size) = CONFIG.db_max_size {
  21. builder = builder.max_size(max_size);
  22. };
  23. let pool = builder.build(manager).ok()?;
  24. Instance::cache_local(&pool.get().unwrap());
  25. Some(pool)
  26. }
  27. // From rocket documentation
  28. // Connection request guard type: a wrapper around an r2d2 pooled connection.
  29. pub struct DbConn(pub PooledConnection<ConnectionManager<Connection>>);
  30. /// Attempts to retrieve a single connection from the managed database pool. If
  31. /// no pool is currently managed, fails with an `InternalServerError` status. If
  32. /// no connections are available, fails with a `ServiceUnavailable` status.
  33. impl<'a, 'r> FromRequest<'a, 'r> for DbConn {
  34. type Error = ();
  35. fn from_request(request: &'a Request<'r>) -> request::Outcome<Self, Self::Error> {
  36. let pool = request.guard::<State<'_, DbPool>>()?;
  37. match pool.get() {
  38. Ok(conn) => Outcome::Success(DbConn(conn)),
  39. Err(_) => Outcome::Failure((Status::ServiceUnavailable, ())),
  40. }
  41. }
  42. }
  43. // For the convenience of using an &DbConn as an &Connection.
  44. impl Deref for DbConn {
  45. type Target = Connection;
  46. fn deref(&self) -> &Self::Target {
  47. &self.0
  48. }
  49. }
  50. // Execute a pragma for every new sqlite connection
  51. #[derive(Debug)]
  52. pub struct PragmaForeignKey;
  53. impl CustomizeConnection<Connection, ConnError> for PragmaForeignKey {
  54. #[cfg(feature = "sqlite")] // will default to an empty function for postgres
  55. fn on_acquire(&self, conn: &mut Connection) -> Result<(), ConnError> {
  56. sql_query("PRAGMA foreign_keys = on;")
  57. .execute(conn)
  58. .map(|_| ())
  59. .map_err(|_| {
  60. ConnError::ConnectionError(ConnectionError::BadConnection(String::from(
  61. "PRAGMA foreign_keys = on failed",
  62. )))
  63. })
  64. }
  65. }
  66. #[cfg(test)]
  67. pub(crate) mod tests {
  68. use super::*;
  69. use diesel::Connection as _;
  70. #[derive(Debug)]
  71. pub struct TestConnectionCustomizer;
  72. impl CustomizeConnection<Connection, ConnError> for TestConnectionCustomizer {
  73. fn on_acquire(&self, conn: &mut Connection) -> Result<(), ConnError> {
  74. PragmaForeignKey.on_acquire(conn)?;
  75. Ok(conn.begin_test_transaction().unwrap())
  76. }
  77. }
  78. }