editoast: rename DbConnection to DieselConnection

Wadjetz · Wadjetz · commit 771fc0c4b8d8 · 2024-09-03T19:38:01.000+02:00
diff --git a/editoast/editoast_derive/src/model.rs b/editoast/editoast_derive/src/model.rs
@@ -84,7 +84,7 @@ fn create_functions(config: &Config) -> TokenStream {
     quote! {
         #[async_trait::async_trait]
         impl crate::models::Create for #model_name {
-            async fn create_conn(self, conn: &mut editoast_models::DbConnection) -> crate::error::Result<Self> {
+            async fn create_conn(self, conn: &mut editoast_models::DieselConnection) -> crate::error::Result<Self> {
                 use #table::table;
                 use diesel_async::RunQueryDsl;
 
@@ -126,7 +126,7 @@ fn retrieve_functions(config: &Config) -> TokenStream {
     quote! {
         #[async_trait::async_trait]
         impl crate::models::Retrieve for #model_name {
-            async fn retrieve_conn(conn: &mut editoast_models::DbConnection, obj_id: i64) -> crate::error::Result<Option<Self>> {
+            async fn retrieve_conn(conn: &mut editoast_models::DieselConnection, obj_id: i64) -> crate::error::Result<Option<Self>> {
                 use #table::table;
                 use #table::dsl;
                 use diesel_async::RunQueryDsl;
@@ -174,7 +174,7 @@ fn delete_functions(config: &Config) -> TokenStream {
     quote! {
         #[async_trait::async_trait]
         impl crate::models::Delete for #model_name {
-            async fn delete_conn(conn: &mut editoast_models::DbConnection, obj_id: i64) -> crate::error::Result<bool> {
+            async fn delete_conn(conn: &mut editoast_models::DieselConnection, obj_id: i64) -> crate::error::Result<bool> {
                 use #table::table;
                 use #table::dsl;
                 use diesel_async::RunQueryDsl;
@@ -207,7 +207,7 @@ fn update_functions(config: &Config) -> TokenStream {
     quote! {
         #[async_trait::async_trait]
         impl crate::models::Update for #model_name {
-            async fn update_conn(self, conn: &mut editoast_models::DbConnection, obj_id: i64) -> crate::error::Result<Option<Self>> {
+            async fn update_conn(self, conn: &mut editoast_models::DieselConnection, obj_id: i64) -> crate::error::Result<Option<Self>> {
                 use #table::table;
 
                 match diesel::update(table.find(obj_id)).set(&self).get_result(conn).await
diff --git a/editoast/editoast_models/src/db_connection_pool.rs b/editoast/editoast_models/src/db_connection_pool.rs
@@ -1,3 +1,5 @@
+use std::ops::Deref;
+use std::ops::DerefMut;
 use std::sync::Arc;
 
 use diesel::sql_query;
@@ -7,6 +9,8 @@ use diesel_async::pooled_connection::deadpool::Object;
 use diesel_async::pooled_connection::deadpool::Pool;
 use diesel_async::pooled_connection::AsyncDieselConnectionManager;
 use diesel_async::pooled_connection::ManagerConfig;
+use diesel_async::scoped_futures::ScopedBoxFuture;
+use diesel_async::AsyncConnection;
 use diesel_async::AsyncPgConnection;
 use diesel_async::RunQueryDsl;
 use futures::future::BoxFuture;
@@ -17,13 +21,12 @@ use openssl::ssl::SslMethod;
 use openssl::ssl::SslVerifyMode;
 use url::Url;
 
-#[cfg(feature = "testing")]
 use tokio::sync::OwnedRwLockWriteGuard;
-#[cfg(feature = "testing")]
 use tokio::sync::RwLock;
 
 use super::DbConnection;
 use super::DbConnectionPool;
+use super::DieselConnection;
 
 pub type DbConnectionConfig = AsyncDieselConnectionManager<AsyncPgConnection>;
 
@@ -33,6 +36,84 @@ pub type DbConnectionV2 = OwnedRwLockWriteGuard<Object<AsyncPgConnection>>;
 #[cfg(not(feature = "testing"))]
 pub type DbConnectionV2 = Object<AsyncPgConnection>;
 
+#[derive(Clone)]
+pub struct DbConnectionV3 {
+    inner: Arc<RwLock<Object<AsyncPgConnection>>>,
+}
+
+pub struct WriteHandle {
+    guard: OwnedRwLockWriteGuard<Object<AsyncPgConnection>>,
+}
+
+impl DbConnectionV3 {
+    pub fn new(inner: Arc<RwLock<Object<AsyncPgConnection>>>) -> Self {
+        Self { inner }
+    }
+
+    pub async fn write(&self) -> WriteHandle {
+        WriteHandle {
+            guard: self.inner.clone().write_owned().await,
+        }
+    }
+
+    pub async fn transaction<'a, R, E, F>(self, callback: F) -> std::result::Result<R, E>
+    where
+        F: FnOnce(Self) -> ScopedBoxFuture<'a, 'a, std::result::Result<R, E>> + Send + 'a,
+        E: From<diesel::result::Error> + Send + 'a,
+        R: Send + 'a,
+    {
+        use diesel_async::TransactionManager as _;
+
+        type TxManager = <AsyncPgConnection as AsyncConnection>::TransactionManager;
+
+        {
+            let mut handle = self.write().await;
+            TxManager::begin_transaction(
+                handle.deref_mut(),
+            )
+            .await?;
+        }
+
+        match callback(self.clone()).await {
+            Ok(result) => {
+                let mut handle = self.write().await;
+                TxManager::commit_transaction(
+                    handle.deref_mut(),
+                )
+                .await?;
+                Ok(result)
+            }
+            Err(callback_error) => {
+                let mut handle = self.write().await;
+                match TxManager::rollback_transaction(
+                    handle.deref_mut(),
+                )
+                .await
+                {
+                    Ok(()) | Err(diesel::result::Error::BrokenTransactionManager) => {
+                        Err(callback_error)
+                    }
+                    Err(rollback_error) => Err(rollback_error.into()),
+                }
+            }
+        }
+    }
+}
+
+impl Deref for WriteHandle {
+    type Target = AsyncPgConnection;
+
+    fn deref(&self) -> &Self::Target {
+        self.guard.deref()
+    }
+}
+
+impl DerefMut for WriteHandle {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.guard.deref_mut()
+    }
+}
+
 /// Wrapper for connection pooling with support for test connections on `cfg(test)`
 ///
 /// # Testing pool
@@ -47,6 +128,8 @@ pub struct DbConnectionPoolV2 {
     pool: Arc<Pool<AsyncPgConnection>>,
     #[cfg(feature = "testing")]
     test_connection: Option<Arc<RwLock<Object<AsyncPgConnection>>>>,
+    #[cfg(feature = "testing")]
+    test_connection_v3: Option<DbConnectionV3>,
 }
 
 #[cfg(feature = "testing")]
@@ -61,8 +144,20 @@ impl Default for DbConnectionPoolV2 {
 pub struct DatabasePoolBuildError(#[from] diesel_async::pooled_connection::deadpool::BuildError);
 
 #[derive(Debug, thiserror::Error)]
-#[error("an error occurred while getting a connection from the database pool: '{0}'")]
-pub struct DatabasePoolError(#[from] diesel_async::pooled_connection::deadpool::PoolError);
+pub enum DatabasePoolError {
+    #[error("an error occurred while getting a connection from the database pool: '{0}'")]
+    Pool(#[from] diesel_async::pooled_connection::deadpool::PoolError),
+    #[error("an error occured while querying the database: {0}")]
+    DieselError(#[from] diesel::result::Error),
+}
+
+#[derive(Debug, thiserror::Error)]
+pub enum DatabaseTransactionError {
+    #[error("an error occurred while getting a connection from the database pool: '{0}'")]
+    Pool(#[from] DatabasePoolError),
+    #[error("an error occured while querying the database: {0}")]
+    DieselError(#[from] diesel::result::Error),
+}
 
 impl DbConnectionPoolV2 {
     /// Get inner pool for retro compatibility
@@ -89,12 +184,27 @@ impl DbConnectionPoolV2 {
         Ok(connection)
     }
 
+    #[cfg(feature = "testing")]
+    async fn get_connection_v3(&self) -> Result<DbConnectionV3, DatabasePoolError> {
+        Ok(self
+            .test_connection_v3
+            .as_ref()
+            .expect("should already exist")
+            .clone())
+    }
+
     #[cfg(not(feature = "testing"))]
     async fn get_connection(&self) -> Result<DbConnectionV2, DatabasePoolError> {
         let connection = self.pool.get().await?;
         Ok(connection)
     }
 
+    #[cfg(not(feature = "testing"))]
+    async fn get_connection_v3(&self) -> Result<DbConnectionV3, DatabasePoolError> {
+        let connection = self.pool.get().await?;
+        Ok(DbConnectionV3::new(Arc::new(RwLock::new(connection))))
+    }
+
     /// Get a connection from the pool
     ///
     /// This function behaves differently in test mode.
@@ -228,6 +338,15 @@ impl DbConnectionPoolV2 {
         self.get_connection().await
     }
 
+    pub async fn get_v3(&self) -> Result<DbConnectionV3, DatabasePoolError> {
+        self.get_connection_v3().await
+    }
+
+    #[cfg(feature = "testing")]
+    pub fn get_ok_v3(&self) -> DbConnectionV3 {
+        futures::executor::block_on(self.get_v3()).expect("Failed to get test connection")
+    }
+
     /// Gets a test connection from the pool synchronously, failing if the connection is not available
     ///
     /// In unit tests, this is the preferred way to get a connection
@@ -300,9 +419,23 @@ impl DbConnectionPoolV2 {
         }
         let test_connection = Arc::new(RwLock::new(conn));
 
+        // Conn v3
+        let mut conn_v3 = pool
+            .get()
+            .await
+            .expect("cannot acquire a connection in the test pool");
+        if transaction {
+            conn_v3
+                .begin_test_transaction()
+                .await
+                .expect("cannot begin a test transaction");
+        }
+        let test_connection_v3 = Some(DbConnectionV3::new(Arc::new(RwLock::new(conn_v3))));
+
         Self {
             pool,
             test_connection: Some(test_connection),
+            test_connection_v3,
         }
     }
 
@@ -312,7 +445,7 @@ impl DbConnectionPoolV2 {
             .unwrap_or_else(|_| String::from("postgresql://osrd:password@localhost/osrd"));
         let url = Url::parse(&url).expect("Failed to parse postgresql url");
         let pool =
-            create_connection_pool(url, 1).expect("Failed to initialize test connection pool");
+            create_connection_pool(url, 2).expect("Failed to initialize test connection pool");
         futures::executor::block_on(Self::from_pool_test(Arc::new(pool), transaction))
     }
 
@@ -353,7 +486,7 @@ pub fn create_connection_pool(
     Ok(Pool::builder(manager).max_size(max_size).build()?)
 }
 
-fn establish_connection(config: &str) -> BoxFuture<ConnectionResult<DbConnection>> {
+fn establish_connection(config: &str) -> BoxFuture<ConnectionResult<DieselConnection>> {
     let fut = async {
         let mut connector_builder = SslConnector::builder(SslMethod::tls()).unwrap();
         connector_builder.set_verify(SslVerifyMode::NONE);
@@ -368,7 +501,7 @@ fn establish_connection(config: &str) -> BoxFuture<ConnectionResult<DbConnection
                 tracing::error!("connection error: {}", e);
             }
         });
-        DbConnection::try_from(client).await
+        DieselConnection::try_from(client).await
     };
     fut.boxed()
 }
diff --git a/editoast/editoast_models/src/lib.rs b/editoast/editoast_models/src/lib.rs
@@ -1,16 +1,24 @@
-use diesel_async::{pooled_connection::deadpool::Pool, AsyncPgConnection};
+use db_connection_pool::DatabasePoolError;
+use diesel_async::pooled_connection::deadpool::Pool;
+use diesel_async::AsyncPgConnection;
 
 pub mod db_connection_pool;
 pub mod tables;
 
 pub use db_connection_pool::DbConnectionPoolV2;
+pub use db_connection_pool::DbConnectionV3;
 
 pub type DbConnection = AsyncPgConnection;
-pub type DbConnectionPool = Pool<DbConnection>;
+pub type DieselConnection = AsyncPgConnection;
+pub type DbConnectionPool = Pool<DieselConnection>;
 
 /// Generic error type to forward errors from the database
 ///
 /// Useful for functions which only points of failure are the DB calls.
 #[derive(Debug, thiserror::Error)]
-#[error("an error occured while querying the database: {0}")]
-pub struct DatabaseError(#[from] diesel::result::Error);
+pub enum DatabaseError {
+    #[error("an error occured while querying the database: {0}")]
+    DieselError(#[from] diesel::result::Error),
+    #[error("an error occured while retrieving a connection from the pool: {0}")]
+    DatabasePoolError(#[from] DatabasePoolError),
+}
diff --git a/editoast/src/modelsv2/prelude/create.rs b/editoast/src/modelsv2/prelude/create.rs
@@ -1,8 +1,10 @@
 use std::fmt::Debug;
 
+use editoast_models::DbConnection;
+use editoast_models::DbConnectionV3;
+
 use crate::error::EditoastError;
 use crate::error::Result;
-use editoast_models::DbConnection;
 
 /// Describes how a [Model](super::Model) can be created in the database
 ///
diff --git a/editoast/src/modelsv2/prelude/delete.rs b/editoast/src/modelsv2/prelude/delete.rs
@@ -1,6 +1,6 @@
+use editoast_models::DbConnection;
 use crate::error::EditoastError;
 use crate::error::Result;
-use editoast_models::DbConnection;
 
 /// Describes how a [Model](super::Model) can be deleted from the database
 ///
