add initial support for diesel-async in aws-rds

shuttle-hq · May 4, 2024 · f82cecc · f82cecc
1 parent 28f4011
commit f82cecc
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 5 deletions.
diff --git a/resources/aws-rds/Cargo.toml b/resources/aws-rds/Cargo.toml
@@ -9,6 +9,7 @@ keywords = ["shuttle-service", "rds"]
 
 [dependencies]
 async-trait = "0.1.56"
+diesel-async = { version = "0.4.1", optional = true }
 paste = "1.0.7"
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
@@ -19,10 +20,15 @@ sqlx = { version = "0.7.1", optional = true }
 default = []
 
 # Database
-postgres = ["sqlx?/postgres"]
+postgres = ["diesel-async?/postgres", "sqlx?/postgres"]
 mysql = ["sqlx?/mysql"]
 mariadb = ["sqlx?/mysql"]
 
+# Support for diesel-async
+diesel-async = ["dep:diesel-async"]
+diesel-async-bb8 = ["diesel-async", "diesel-async/bb8"]
+diesel-async-deadpool = ["diesel-async", "diesel-async/deadpool"]
+
 # Add an sqlx Pool as a resource output type
 sqlx = ["dep:sqlx", "sqlx/runtime-tokio", "sqlx/tls-rustls"]
 sqlx-native-tls = ["dep:sqlx", "sqlx/runtime-tokio", "sqlx/tls-native-tls"]
diff --git a/resources/aws-rds/src/lib.rs b/resources/aws-rds/src/lib.rs
@@ -8,6 +8,20 @@ use shuttle_service::{
     DatabaseResource, DbInput, Error, IntoResource, ResourceFactory, ResourceInputBuilder,
 };
 
+#[cfg(any(feature = "diesel-async-bb8", feature = "diesel-async-deadpool"))]
+use diesel_async::pooled_connection::AsyncDieselConnectionManager;
+
+#[cfg(feature = "diesel-async-bb8")]
+use diesel_async::pooled_connection::bb8 as diesel_bb8;
+
+#[cfg(feature = "diesel-async-deadpool")]
+use diesel_async::pooled_connection::deadpool as diesel_deadpool;
+
+#[allow(dead_code)]
+const MIN_CONNECTIONS: u32 = 1;
+#[allow(dead_code)]
+const MAX_CONNECTIONS: u32 = 5;
+
 macro_rules! aws_engine {
     ($feature:expr, $struct_ident:ident) => {
         paste::paste! {
@@ -87,8 +101,8 @@ mod _sqlx {
             let connection_string: String = self.into_resource().await.unwrap();
 
             Ok(sqlx::postgres::PgPoolOptions::new()
-                .min_connections(1)
-                .max_connections(5)
+                .min_connections(MIN_CONNECTIONS)
+                .max_connections(MAX_CONNECTIONS)
                 .connect(&connection_string)
                 .await
                 .map_err(shuttle_service::error::CustomError::new)?)
@@ -102,11 +116,65 @@ mod _sqlx {
             let connection_string: String = self.into_resource().await.unwrap();
 
             Ok(sqlx::mysql::MySqlPoolOptions::new()
-                .min_connections(1)
-                .max_connections(5)
+                .min_connections(MIN_CONNECTIONS)
+                .max_connections(MAX_CONNECTIONS)
                 .connect(&connection_string)
                 .await
                 .map_err(shuttle_service::error::CustomError::new)?)
         }
     }
 }
+
+#[cfg(feature = "postgres")]
+mod _diesel {
+    use super::*;
+
+    #[cfg(feature = "diesel-async")]
+    #[async_trait]
+    impl IntoResource<diesel_async::AsyncPgConnection> for OutputWrapper {
+        async fn into_resource(self) -> Result<diesel_async::AsyncPgConnection, Error> {
+            use diesel_async::{AsyncConnection, AsyncPgConnection};
+
+            let connection_string: String = self.into_resource().await.unwrap();
+            Ok(AsyncPgConnection::establish(&connection_string)
+                .await
+                .map_err(shuttle_service::error::CustomError::new)?)
+        }
+    }
+
+    #[cfg(feature = "diesel-async-bb8")]
+    #[async_trait]
+    impl IntoResource<diesel_bb8::Pool<diesel_async::AsyncPgConnection>> for OutputWrapper {
+        async fn into_resource(
+            self,
+        ) -> Result<diesel_bb8::Pool<diesel_async::AsyncPgConnection>, Error> {
+            let connection_string: String = self.into_resource().await.unwrap();
+
+            Ok(diesel_bb8::Pool::builder()
+                .min_idle(Some(MIN_CONNECTIONS))
+                .max_size(MAX_CONNECTIONS)
+                .build(AsyncDieselConnectionManager::new(connection_string))
+                .await
+                .map_err(shuttle_service::error::CustomError::new)?)
+        }
+    }
+
+    #[cfg(feature = "diesel-async-deadpool")]
+    #[async_trait]
+    impl IntoResource<diesel_deadpool::Pool<diesel_async::AsyncPgConnection>> for OutputWrapper {
+        async fn into_resource(
+            self,
+        ) -> Result<diesel_deadpool::Pool<diesel_async::AsyncPgConnection>, Error> {
+            let connection_string: String = self.into_resource().await.unwrap();
+
+            Ok(
+                diesel_deadpool::Pool::builder(AsyncDieselConnectionManager::new(
+                    connection_string,
+                ))
+                .max_size(MAX_CONNECTIONS as usize)
+                .build()
+                .map_err(shuttle_service::error::CustomError::new)?,
+            )
+        }
+    }
+}