Switch to pooled postgres client

Signed-off-by: trivernis <trivernis@protonmail.com>

Cargo.lock

@@ -382,6 +382,8 @@ dependencies = [
  "mime 0.3.16",
  "msgrpc",
  "postgres",
+ "r2d2",
+ "r2d2_postgres",
  "rand 0.7.3",
  "rmp",
  "rmp-serde",
@@ -1036,6 +1038,27 @@ dependencies = [
  "proc-macro2",
 ]
 
+[[package]]
+name = "r2d2"
+version = "0.8.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "545c5bc2b880973c9c10e4067418407a0ccaa3091781d1671d46eb35107cb26f"
+dependencies = [
+ "log 0.4.11",
+ "parking_lot",
+ "scheduled-thread-pool",
+]
+
+[[package]]
+name = "r2d2_postgres"
+version = "0.16.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "707d27f66f43bac1081141f6d9611fffcce7da2841ae97c7ac53619d098efe8f"
+dependencies = [
+ "postgres",
+ "r2d2",
+]
+
 [[package]]
 name = "rand"
 version = "0.4.6"
@@ -1332,6 +1355,15 @@ version = "0.3.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ef703b7cb59335eae2eb93ceb664c0eb7ea6bf567079d843e09420219668e072"
 
+[[package]]
+name = "scheduled-thread-pool"
+version = "0.2.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "dc6f74fd1204073fa02d5d5d68bec8021be4c38690b61264b2fdb48083d0e7d7"
+dependencies = [
+ "parking_lot",
+]
+
 [[package]]
 name = "scopeguard"
 version = "1.1.0"

Cargo.toml

@@ -26,4 +26,6 @@ mime = "0.3.16"
 serde_json = "1.0.57"
 rouille = "3.0.0"
 base64 = "0.12.3"
-chrono = "0.4.15"
+chrono = "0.4.15"
+r2d2 = "0.8.9"
+r2d2_postgres = "0.16.0"

src/database/mod.rs

@@ -3,10 +3,11 @@ use crate::database::role_permissions::RolePermissions;
 use crate::database::roles::Roles;
 use crate::database::user_roles::UserRoles;
 use crate::database::users::Users;
-use crate::utils::error::{DBError, DatabaseClient, DatabaseResult, PostgresError};
+use crate::utils::error::DatabaseResult;
 use dotenv;
-use postgres::{Client, NoTls};
+use postgres::NoTls;
-use std::sync::{Arc, Mutex};
+use r2d2::Pool;
+use r2d2_postgres::PostgresConnectionManager;
 
 pub mod database_error;
 pub mod models;
@@ -21,13 +22,13 @@ const DB_CONNECTION_URL: &str = "POSTGRES_CONNECTION_URL";
 const DEFAULT_CONNECTION: &str = "postgres://postgres:postgres@localhost/postgres";
 
 pub trait Table {
-    fn new(database_connection: Arc<Mutex<DatabaseClient>>) -> Self;
+    fn new(pool: PostgresPool) -> Self;
     fn init(&self) -> DatabaseResult<()>;
 }
 
 #[derive(Clone)]
 pub struct Database {
-    database_connection: Arc<Mutex<Client>>,
+    pool: PostgresPool,
     pub users: Users,
     pub roles: Roles,
     pub permissions: Permissions,
@@ -37,16 +38,14 @@ pub struct Database {
 
 impl Database {
     pub fn new() -> DatabaseResult<Self> {
-        let database_connection = Arc::new(Mutex::new(
+        let pool = get_database_connection()?;
-            get_database_connection().map_err(|e| DBError::Postgres(e))?,
-        ));
         Ok(Self {
-            users: Users::new(Arc::clone(&database_connection)),
+            users: Users::new(PostgresPool::clone(&pool)),
-            roles: Roles::new(Arc::clone(&database_connection)),
+            roles: Roles::new(PostgresPool::clone(&pool)),
-            permissions: Permissions::new(Arc::clone(&database_connection)),
+            permissions: Permissions::new(PostgresPool::clone(&pool)),
-            user_roles: UserRoles::new(Arc::clone(&database_connection)),
+            user_roles: UserRoles::new(PostgresPool::clone(&pool)),
-            role_permission: RolePermissions::new(Arc::clone(&database_connection)),
+            role_permission: RolePermissions::new(PostgresPool::clone(&pool)),
-            database_connection,
+            pool,
         })
     }
 
@@ -67,8 +66,15 @@ impl Database {
         Ok(())
     }
 }
+
+pub type PostgresPool = Pool<PostgresConnectionManager<NoTls>>;
+
 /// Returns a database connection
-fn get_database_connection() -> Result<DatabaseClient, PostgresError> {
+fn get_database_connection() -> Result<PostgresPool, r2d2::Error> {
     let conn_url = dotenv::var(DB_CONNECTION_URL).unwrap_or(DEFAULT_CONNECTION.to_string());
-    Client::connect(conn_url.as_str(), NoTls)
+
+    Pool::new(PostgresConnectionManager::new(
+        conn_url.parse().unwrap(),
+        NoTls,
+    ))
 }

src/database/permissions.rs

@@ -1,25 +1,20 @@
 use crate::database::models::{CreatePermissionsEntry, Permission};
-use crate::database::{DatabaseClient, DatabaseResult, Table};
+use crate::database::{DatabaseResult, PostgresPool, Table};
 use crate::utils::error::DBError;
-use postgres::Client;
-use std::sync::{Arc, Mutex};
 
 #[derive(Clone)]
 pub struct Permissions {
-    database_connection: Arc<Mutex<DatabaseClient>>,
+    pool: PostgresPool,
 }
 
 impl Table for Permissions {
-    fn new(database_connection: Arc<Mutex<Client>>) -> Self {
+    fn new(pool: PostgresPool) -> Self {
-        Self {
+        Self { pool }
-            database_connection,
-        }
     }
 
     fn init(&self) -> DatabaseResult<()> {
-        self.database_connection
+        self.pool
-            .lock()
+            .get()?
-            .unwrap()
             .batch_execute(
                 "CREATE TABLE IF NOT EXISTS permissions (
                         id              SERIAL PRIMARY KEY,
@@ -36,7 +31,7 @@ impl Permissions {
         &self,
         permissions: Vec<CreatePermissionsEntry>,
     ) -> DatabaseResult<Vec<Permission>> {
-        let mut connection = self.database_connection.lock().unwrap();
+        let mut connection = self.pool.get()?;
         let mut transaction = connection.transaction()?;
         let mut created_permissions = Vec::new();
         let _: Vec<DatabaseResult<()>> = permissions

src/database/role_permissions.rs

@@ -1,24 +1,20 @@
 use crate::database::models::Permission;
-use crate::database::{DatabaseClient, DatabaseResult, Table};
+use crate::database::{DatabaseResult, PostgresPool, Table};
 use crate::utils::error::DBError;
-use std::sync::{Arc, Mutex};
 
 #[derive(Clone)]
 pub struct RolePermissions {
-    database_connection: Arc<Mutex<DatabaseClient>>,
+    pool: PostgresPool,
 }
 
 impl Table for RolePermissions {
-    fn new(database_connection: Arc<Mutex<DatabaseClient>>) -> Self {
+    fn new(pool: PostgresPool) -> Self {
-        Self {
+        Self { pool }
-            database_connection,
-        }
     }
 
     fn init(&self) -> DatabaseResult<()> {
-        self.database_connection
+        self.pool
-            .lock()
+            .get()?
-            .unwrap()
             .batch_execute(
                 "
             CREATE TABLE IF NOT EXISTS role_permissions (
@@ -33,7 +29,7 @@ impl Table for RolePermissions {
 
 impl RolePermissions {
     pub fn by_role(&self, role_id: i32) -> DatabaseResult<Vec<Permission>> {
-        let mut connection = self.database_connection.lock().unwrap();
+        let mut connection = self.pool.get()?;
         let rows = connection.query(
             "SELECT * FROM role_permissions, permissions WHERE role_id = $1 AND role_permissions.permission_id = permissions.id", 
             &[&role_id])?;

src/database/roles.rs

@@ -1,28 +1,25 @@
 use crate::database::models::Role;
 use crate::database::role_permissions::RolePermissions;
-use crate::database::{DatabaseResult, Table};
+use crate::database::{DatabaseResult, PostgresPool, Table};
 use crate::utils::error::DBError;
-use postgres::Client;
-use std::sync::{Arc, Mutex};
 
 #[derive(Clone)]
 pub struct Roles {
-    database_connection: Arc<Mutex<Client>>,
+    pool: PostgresPool,
     role_permission: RolePermissions,
 }
 
 impl Table for Roles {
-    fn new(database_connection: Arc<Mutex<Client>>) -> Self {
+    fn new(pool: PostgresPool) -> Self {
         Self {
-            role_permission: RolePermissions::new(Arc::clone(&database_connection)),
+            role_permission: RolePermissions::new(PostgresPool::clone(&pool)),
-            database_connection,
+            pool,
         }
     }
 
     fn init(&self) -> DatabaseResult<()> {
-        self.database_connection
+        self.pool
-            .lock()
+            .get()?
-            .unwrap()
             .batch_execute(
                 "
             CREATE TABLE IF NOT EXISTS roles (
@@ -42,7 +39,7 @@ impl Roles {
         description: Option<String>,
         permissions: Vec<i32>,
     ) -> DatabaseResult<Role> {
-        let mut connection = self.database_connection.lock().unwrap();
+        let mut connection = self.pool.get()?;
         let exists = connection.query_opt("SELECT id FROM roles WHERE name = $1", &[&name])?;
 
         if exists.is_some() {

src/database/user_roles.rs

@@ -1,25 +1,20 @@
 use crate::database::models::Role;
-use crate::database::{DatabaseResult, Table};
+use crate::database::{DatabaseResult, PostgresPool, Table};
 use crate::utils::error::DBError;
-use postgres::Client;
-use std::sync::{Arc, Mutex};
 
 #[derive(Clone)]
 pub struct UserRoles {
-    database_connection: Arc<Mutex<Client>>,
+    pool: PostgresPool,
 }
 
 impl Table for UserRoles {
-    fn new(database_connection: Arc<Mutex<Client>>) -> Self {
+    fn new(pool: PostgresPool) -> Self {
-        Self {
+        Self { pool }
-            database_connection,
-        }
     }
 
     fn init(&self) -> DatabaseResult<()> {
-        self.database_connection
+        self.pool
-            .lock()
+            .get()?
-            .unwrap()
             .batch_execute(
                 "
         CREATE TABLE IF NOT EXISTS user_roles (
@@ -34,7 +29,7 @@ impl Table for UserRoles {
 
 impl UserRoles {
     pub fn by_user(&self, user_id: i32) -> DatabaseResult<Vec<Role>> {
-        let mut connection = self.database_connection.lock().unwrap();
+        let mut connection = self.pool.get()?;
         let rows = connection.query(
             "SELECT * FROM user_roles, roles WHERE user_id = $1 AND roles.id = user_roles.role_id",
             &[&user_id],

src/database/users.rs

@@ -1,11 +1,10 @@
 use crate::database::models::UserRecord;
 use crate::database::tokens::{SessionTokens, TokenStore};
 use crate::database::user_roles::UserRoles;
-use crate::database::{DatabaseResult, Table};
+use crate::database::{DatabaseResult, PostgresPool, Table};
 use crate::utils::error::DBError;
 use crate::utils::{create_salt, hash_password};
 
-use postgres::Client;
 use std::sync::{Arc, Mutex};
 use zeroize::{Zeroize, Zeroizing};
 
@@ -16,22 +15,22 @@ const ENV_ADMIN_EMAIL: &str = "ADMIN_EMAIL";
 
 #[derive(Clone)]
 pub struct Users {
-    database_connection: Arc<Mutex<Client>>,
+    pool: PostgresPool,
     user_roles: UserRoles,
     token_store: Arc<Mutex<TokenStore>>,
 }
 
 impl Table for Users {
-    fn new(database_connection: Arc<Mutex<Client>>) -> Self {
+    fn new(pool: PostgresPool) -> Self {
         Self {
-            user_roles: UserRoles::new(Arc::clone(&database_connection)),
+            user_roles: UserRoles::new(PostgresPool::clone(&pool)),
-            database_connection,
+            pool,
             token_store: Arc::new(Mutex::new(TokenStore::new())),
         }
     }
 
     fn init(&self) -> DatabaseResult<()> {
-        self.database_connection.lock().unwrap().batch_execute(
+        self.pool.get()?.batch_execute(
             "CREATE TABLE IF NOT EXISTS users (
             id              SERIAL PRIMARY KEY,
             name            VARCHAR(255) NOT NULL,
@@ -62,7 +61,7 @@ impl Users {
         email: String,
         password: String,
     ) -> DatabaseResult<UserRecord> {
-        let mut connection = self.database_connection.lock().unwrap();
+        let mut connection = self.pool.get()?;
         let mut password = Zeroizing::new(password);
         log::trace!("Creating user {} with email  {}", name, email);
 
@@ -90,7 +89,7 @@ impl Users {
         password: &String,
     ) -> DatabaseResult<SessionTokens> {
         if self.validate_login(&email, password)? {
-            let mut connection = self.database_connection.lock().unwrap();
+            let mut connection = self.pool.get()?;
             let row = connection.query_one("SELECT id FROM users WHERE email = $1", &[&email])?;
             let id: i32 = row.get(0);
 
@@ -139,7 +138,7 @@ impl Users {
     }
 
     fn validate_login(&self, email: &String, password: &String) -> DatabaseResult<bool> {
-        let mut connection = self.database_connection.lock().unwrap();
+        let mut connection = self.pool.get()?;
         let row = connection
             .query_opt(
                 "SELECT password_hash, salt FROM users WHERE email = $1",

src/server/messages.rs

@@ -61,7 +61,7 @@ impl InfoEntry {
 
 #[derive(Deserialize)]
 pub struct GetPermissionsRequest {
-    pub role_ids: Vec<i32>,
+    pub roles: Vec<i32>,
 }
 
 #[derive(Deserialize)]

src/server/user_rpc.rs

@@ -97,7 +97,7 @@ impl UserRpcServer {
                     "get permissions",
                     GET_ROLE_PERMISSIONS,
                     "Returns all permissions the given roles are assigned to",
-                    "{role_ids: [i32]}",
+                    "{roles: [i32]}",
                 ),
                 InfoEntry::new(
                     "create role",
@@ -121,7 +121,7 @@ impl UserRpcServer {
             GetPermissionsRequest::deserialize(&mut Deserializer::new(&mut data.as_slice()))
                 .map_err(|e| ErrorMessage::new(e.to_string()))?;
         let mut response_data = HashMap::new();
-        for role_id in message.role_ids {
+        for role_id in message.roles {
             let permissions = self.database.role_permission.by_role(role_id)?;
             response_data.insert(role_id.to_string(), permissions);
         }

src/utils/error.rs

@@ -1,3 +1,4 @@
+use r2d2::Error;
 use serde_postgres::DeError;
 use std::error;
 use std::fmt::{self, Display, Formatter};
@@ -5,6 +6,7 @@ use std::fmt::{self, Display, Formatter};
 #[derive(Debug)]
 pub enum DBError {
     Postgres(PostgresError),
+    Pool(r2d2::Error),
     RecordExists,
     BCryptError,
     DeserializeError(serde_postgres::DeError),
@@ -27,6 +29,7 @@ impl DBError {
             DBError::Postgres(p) => p.to_string(),
             DBError::DeserializeError(de) => de.to_string(),
             DBError::BCryptError => "BCrypt Hash creation error".to_string(),
+            DBError::Pool(p) => p.to_string(),
         }
     }
 }
@@ -39,6 +42,12 @@ impl From<PostgresError> for DBError {
     }
 }
 
+impl From<r2d2::Error> for DBError {
+    fn from(other: Error) -> Self {
+        Self::Pool(other)
+    }
+}
+
 impl From<serde_postgres::DeError> for DBError {
     fn from(other: DeError) -> Self {
         Self::DeserializeError(other)