CODEWITHSUNDEEP
rustrust-dieselr2d2
balteo
balteo

Reputation: 24689

Imlementing connection pooling in a rust/diesel app with r2d2

I am trying to implement connection pooling in a rust/diesel/rocket application. I am not sure how to ensure the content of the establish_pooled_connection() method is invoked only once, so as to prepare the connection pool.

Here is my code.

from lib.rs:

pub fn establish_pooled_connection() -> PooledConnection<ConnectionManager<PgConnection>> {
    dotenv().ok();

    let database_url = env::var("DB_URL")
        .expect("DATABASE_URL must be set");

    let manager = ConnectionManager::<PgConnection>::new(&database_url);
    let pool = r2d2::Pool::builder().build(manager).expect("Failed to create pool.");
    let conn = pool.clone().get().unwrap();
    return conn;
}

Here I use the above method in main.rs:

#[get("/", format = "json")]
fn find_all() -> Json<Vec<Person>> {
    let connection = establish_pooled_connection();

    let all: QueryResult<Vec<Person>> = person::table().get_results(&connection);
    ...

The issue here is that each get method (e.g. above) calls the establish_pooled_connection() and everything is re-instantiated...

I come from the java world where dependency injection allows us to avoid re-instantiation.

What would be a proper way to implement connection pooling in an rust/diesel app?

Upvotes: 5

Views: 8061

Answers (2)

peaske
peaske

Reputation: 51

The issue here is that each get method (e.g. above) calls the establish_pooled_connection() and everything is re-instantiated...

I was having trouble setting up pooled connections as well, so I thought I'd leave behind what I did.

How I solved it

I managed to use pooled connections in a rocket/diesel app by adding the established pool to the server state and accessing the database via that server state from the individual functions.

// add the `r2d2` feature for diesel 
use diesel::{
    r2d2::{ConnectionManager, Pool, PooledConnection},
    MysqlConnection,
};

// set an alias, so we don't have to keep writing out this long type
pub type DbPool = Pool<ConnectionManager<MysqlConnection>>;

// a real-world app would have more fields in the server state like 
// CORS options, environment values needed for external APIs, etc.
pub struct ServerState {
  pub db_pool: DbPool
}

Rocket allows us to define the server state with the .manage() call when building the server with build().

#[rocket::main]
pub async fn main() {
  let db_pool: DbPool = establish_connection_pool();

  rocket::build()
    .mount("/", routes![test]) // mount routes
    .manage(ServerState { db_pool })
    .launch()
    .await
    .expect("Failed to launch rocket server.")
}

Because Rocket allows us to retrieve the server's state from functions as long as they have Rocket's macro (get, post, etc.), we can retrieve the database from the server state.

use rocket::{get, serde::json::Json, State};

#[get("/test?<user_id>")]
pub async fn test(
  state: &State<ServerState>,
  user_id: String
) -> Result<Json<User>, Error> {
  let pooled = &state.db_pool.get()?;

  pooled.transaction(|| {
    // some diesel query
  })
}

The versions I used when working on this are as follows.

diesel = { version = "1.4.8", features = ["r2d2", "mysql", "chrono"] }
rocket = { version = "0.5.0-rc.1", features = ["json"] }

What other things to consider

There are reasons why someone would want to set up pooled connections directly using r2d2 as I did above. I personally thought being able to pass in the pool as a field in the server state would allow for better control over writing integration tests. However, there are other ways of handling pooled connections you might want to look into.

Rocket provides its own solution of handling pooled database connections with rocket_sync_db_pools[1] and rocket_db_pools[2]. If you use ORMs like diesel[3] and sqlx[4] with databases covered in the library, then I highly recommend looking into them.

Just set up a Rocket.toml file in the root directory of the project like so.

[global.databases]
test_db = { url = "mysql://mysql:password@localhost:5432/test_db_name" }

And now, by harnessing the power of the database macro, you can easily establish connections.

#[macro_use]
extern crate rocket_sync_db_pools;

// as long as we set the `Rocket.toml` to include information about 
// what database we want to point to, the `database` macro does 
// the magic for us
#[database("test_db")]
pub struct TestDbConn(diesel::MysqlConnection);

// We can directly access the pooled connection as we did with the server state
#[get("/test?<user_id>")]
pub async fn test(
  conn: TestDbConn,
  user_id: String
) -> Result<Json<User>, Error> {
  conn.run(move || {
    // some diesel query
  })
}

[1] Use rocket_sync_db_pools with synchronous ORMs like diesel. https://api.rocket.rs/v0.5-rc/rocket_sync_db_pools/index.html#

[2] Use rocket_db_pools with asynchronous ORMs like deadpool-postgres and sqlx. https://api.rocket.rs/v0.5-rc/rocket_db_pools/index.html#

[3] https://api.rocket.rs/v0.5-rc/rocket_sync_db_pools/index.html#database-support

[4] https://api.rocket.rs/v0.5-rc/rocket_db_pools/index.html#supported-drivers

Upvotes: 4

musicmatze
musicmatze

Reputation: 4288

You wouldn't create a new connection in every handler, but rather use some state-sharing mechanisms of your framework.

See here for the guide on how to use state in rocket, that'll be a good starting point.

Upvotes: 2

Related Questions