From 94ae69ec6036b43b9778ec25119c2c3352903656 Mon Sep 17 00:00:00 2001 From: trivernis Date: Sat, 12 Sep 2020 13:09:41 +0200 Subject: [PATCH] Switch to pooled postgres client Signed-off-by: trivernis --- Cargo.lock | 32 +++++++++++++++++++++++++++ Cargo.toml | 4 +++- src/database/mod.rs | 38 ++++++++++++++++++-------------- src/database/permissions.rs | 19 ++++++---------- src/database/role_permissions.rs | 18 ++++++--------- src/database/roles.rs | 19 +++++++--------- src/database/user_roles.rs | 19 ++++++---------- src/database/users.rs | 19 ++++++++-------- src/server/messages.rs | 2 +- src/server/user_rpc.rs | 4 ++-- src/utils/error.rs | 9 ++++++++ 11 files changed, 107 insertions(+), 76 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 38d3f7b..c2e7fe2 100644 --- a/Cargo.lock +++ b/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" diff --git a/Cargo.toml b/Cargo.toml index 0a70280..f650c82 100644 --- a/Cargo.toml +++ b/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" \ No newline at end of file +chrono = "0.4.15" +r2d2 = "0.8.9" +r2d2_postgres = "0.16.0" \ No newline at end of file diff --git a/src/database/mod.rs b/src/database/mod.rs index 24155c7..fedf31f 100644 --- a/src/database/mod.rs +++ b/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 std::sync::{Arc, Mutex}; +use postgres::NoTls; +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>) -> Self; + fn new(pool: PostgresPool) -> Self; fn init(&self) -> DatabaseResult<()>; } #[derive(Clone)] pub struct Database { - database_connection: Arc>, + pool: PostgresPool, pub users: Users, pub roles: Roles, pub permissions: Permissions, @@ -37,16 +38,14 @@ pub struct Database { impl Database { pub fn new() -> DatabaseResult { - let database_connection = Arc::new(Mutex::new( - get_database_connection().map_err(|e| DBError::Postgres(e))?, - )); + let pool = get_database_connection()?; Ok(Self { - users: Users::new(Arc::clone(&database_connection)), - roles: Roles::new(Arc::clone(&database_connection)), - permissions: Permissions::new(Arc::clone(&database_connection)), - user_roles: UserRoles::new(Arc::clone(&database_connection)), - role_permission: RolePermissions::new(Arc::clone(&database_connection)), - database_connection, + users: Users::new(PostgresPool::clone(&pool)), + roles: Roles::new(PostgresPool::clone(&pool)), + permissions: Permissions::new(PostgresPool::clone(&pool)), + user_roles: UserRoles::new(PostgresPool::clone(&pool)), + role_permission: RolePermissions::new(PostgresPool::clone(&pool)), + pool, }) } @@ -67,8 +66,15 @@ impl Database { Ok(()) } } + +pub type PostgresPool = Pool>; + /// Returns a database connection -fn get_database_connection() -> Result { +fn get_database_connection() -> Result { 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, + )) } diff --git a/src/database/permissions.rs b/src/database/permissions.rs index 62c9ba3..4ef98b9 100644 --- a/src/database/permissions.rs +++ b/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>, + pool: PostgresPool, } impl Table for Permissions { - fn new(database_connection: Arc>) -> Self { - Self { - database_connection, - } + fn new(pool: PostgresPool) -> Self { + Self { pool } } fn init(&self) -> DatabaseResult<()> { - self.database_connection - .lock() - .unwrap() + self.pool + .get()? .batch_execute( "CREATE TABLE IF NOT EXISTS permissions ( id SERIAL PRIMARY KEY, @@ -36,7 +31,7 @@ impl Permissions { &self, permissions: Vec, ) -> DatabaseResult> { - 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> = permissions diff --git a/src/database/role_permissions.rs b/src/database/role_permissions.rs index 60f3ec5..4e98cd4 100644 --- a/src/database/role_permissions.rs +++ b/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>, + pool: PostgresPool, } impl Table for RolePermissions { - fn new(database_connection: Arc>) -> Self { - Self { - database_connection, - } + fn new(pool: PostgresPool) -> Self { + Self { pool } } fn init(&self) -> DatabaseResult<()> { - self.database_connection - .lock() - .unwrap() + self.pool + .get()? .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> { - 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])?; diff --git a/src/database/roles.rs b/src/database/roles.rs index b1c303c..9f06357 100644 --- a/src/database/roles.rs +++ b/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>, + pool: PostgresPool, role_permission: RolePermissions, } impl Table for Roles { - fn new(database_connection: Arc>) -> Self { + fn new(pool: PostgresPool) -> Self { Self { - role_permission: RolePermissions::new(Arc::clone(&database_connection)), - database_connection, + role_permission: RolePermissions::new(PostgresPool::clone(&pool)), + pool, } } fn init(&self) -> DatabaseResult<()> { - self.database_connection - .lock() - .unwrap() + self.pool + .get()? .batch_execute( " CREATE TABLE IF NOT EXISTS roles ( @@ -42,7 +39,7 @@ impl Roles { description: Option, permissions: Vec, ) -> DatabaseResult { - 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() { diff --git a/src/database/user_roles.rs b/src/database/user_roles.rs index 0018f78..f4b07c6 100644 --- a/src/database/user_roles.rs +++ b/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>, + pool: PostgresPool, } impl Table for UserRoles { - fn new(database_connection: Arc>) -> Self { - Self { - database_connection, - } + fn new(pool: PostgresPool) -> Self { + Self { pool } } fn init(&self) -> DatabaseResult<()> { - self.database_connection - .lock() - .unwrap() + self.pool + .get()? .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> { - 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], diff --git a/src/database/users.rs b/src/database/users.rs index 1f21a86..dc7ef06 100644 --- a/src/database/users.rs +++ b/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>, + pool: PostgresPool, user_roles: UserRoles, token_store: Arc>, } impl Table for Users { - fn new(database_connection: Arc>) -> Self { + fn new(pool: PostgresPool) -> Self { Self { - user_roles: UserRoles::new(Arc::clone(&database_connection)), - database_connection, + user_roles: UserRoles::new(PostgresPool::clone(&pool)), + 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 { - 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 { 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 { - 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", diff --git a/src/server/messages.rs b/src/server/messages.rs index cb17d09..91adea8 100644 --- a/src/server/messages.rs +++ b/src/server/messages.rs @@ -61,7 +61,7 @@ impl InfoEntry { #[derive(Deserialize)] pub struct GetPermissionsRequest { - pub role_ids: Vec, + pub roles: Vec, } #[derive(Deserialize)] diff --git a/src/server/user_rpc.rs b/src/server/user_rpc.rs index a4671df..2ff41fc 100644 --- a/src/server/user_rpc.rs +++ b/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); } diff --git a/src/utils/error.rs b/src/utils/error.rs index d17b4c0..fc84032 100644 --- a/src/utils/error.rs +++ b/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 for DBError { } } +impl From for DBError { + fn from(other: Error) -> Self { + Self::Pool(other) + } +} + impl From for DBError { fn from(other: DeError) -> Self { Self::DeserializeError(other)