use lazy_static::lazy_static;
use postgres::types::ToSql;
use postgres::{Client, Error as PostgresError, NoTls, Row};
use r2d2::Pool;
use r2d2_postgres::PostgresConnectionManager;
use std::env;
use std::sync::{Arc, Mutex, Once};
use std::time::Duration;
fn create_postgres_error(_message: &str) -> PostgresError {
let result = Client::connect("invalid-connection-string", NoTls);
match result {
Ok(_) => unreachable!(), Err(e) => {
e
}
}
}
lazy_static! {
static ref POSTGRES_CLIENT: Mutex<Option<Arc<PostgresClientWrapper>>> = Mutex::new(None);
static ref POSTGRES_POOL: Mutex<Option<Arc<Pool<PostgresConnectionManager<NoTls>>>>> =
Mutex::new(None);
static ref INIT: Once = Once::new();
}
#[derive(Debug)]
pub struct PostgresConfigBuilder {
pub host: String,
pub port: u16,
pub user: String,
pub password: Option<String>,
pub database: String,
pub application_name: Option<String>,
pub connect_timeout: Option<u64>,
pub ssl_mode: Option<String>,
pub pool_max_size: Option<u32>,
pub pool_min_idle: Option<u32>,
pub pool_idle_timeout: Option<Duration>,
pub pool_connection_timeout: Option<Duration>,
pub pool_max_lifetime: Option<Duration>,
pub use_pool: bool,
}
impl Default for PostgresConfigBuilder {
fn default() -> Self {
Self {
host: "localhost".to_string(),
port: 5432,
user: "postgres".to_string(),
password: None,
database: "postgres".to_string(),
application_name: None,
connect_timeout: None,
ssl_mode: None,
pool_max_size: Some(10),
pool_min_idle: Some(1),
pool_idle_timeout: Some(Duration::from_secs(300)),
pool_connection_timeout: Some(Duration::from_secs(30)),
pool_max_lifetime: Some(Duration::from_secs(1800)),
use_pool: false,
}
}
}
impl PostgresConfigBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn host(mut self, host: &str) -> Self {
self.host = host.to_string();
self
}
pub fn port(mut self, port: u16) -> Self {
self.port = port;
self
}
pub fn user(mut self, user: &str) -> Self {
self.user = user.to_string();
self
}
pub fn password(mut self, password: &str) -> Self {
self.password = Some(password.to_string());
self
}
pub fn database(mut self, database: &str) -> Self {
self.database = database.to_string();
self
}
pub fn application_name(mut self, application_name: &str) -> Self {
self.application_name = Some(application_name.to_string());
self
}
pub fn connect_timeout(mut self, seconds: u64) -> Self {
self.connect_timeout = Some(seconds);
self
}
pub fn ssl_mode(mut self, ssl_mode: &str) -> Self {
self.ssl_mode = Some(ssl_mode.to_string());
self
}
pub fn use_pool(mut self, use_pool: bool) -> Self {
self.use_pool = use_pool;
self
}
pub fn pool_max_size(mut self, size: u32) -> Self {
self.pool_max_size = Some(size);
self
}
pub fn pool_min_idle(mut self, size: u32) -> Self {
self.pool_min_idle = Some(size);
self
}
pub fn pool_idle_timeout(mut self, timeout: Duration) -> Self {
self.pool_idle_timeout = Some(timeout);
self
}
pub fn pool_connection_timeout(mut self, timeout: Duration) -> Self {
self.pool_connection_timeout = Some(timeout);
self
}
pub fn pool_max_lifetime(mut self, lifetime: Duration) -> Self {
self.pool_max_lifetime = Some(lifetime);
self
}
pub fn build_connection_string(&self) -> String {
let mut conn_string = format!(
"host={} port={} user={} dbname={}",
self.host, self.port, self.user, self.database
);
if let Some(password) = &self.password {
conn_string.push_str(&format!(" password={}", password));
}
if let Some(app_name) = &self.application_name {
conn_string.push_str(&format!(" application_name={}", app_name));
}
if let Some(timeout) = self.connect_timeout {
conn_string.push_str(&format!(" connect_timeout={}", timeout));
}
if let Some(ssl_mode) = &self.ssl_mode {
conn_string.push_str(&format!(" sslmode={}", ssl_mode));
}
conn_string
}
pub fn build(&self) -> Result<Client, PostgresError> {
let conn_string = self.build_connection_string();
Client::connect(&conn_string, NoTls)
}
pub fn build_pool(&self) -> Result<Pool<PostgresConnectionManager<NoTls>>, r2d2::Error> {
let conn_string = self.build_connection_string();
let manager = PostgresConnectionManager::new(conn_string.parse().unwrap(), NoTls);
let mut pool_builder = r2d2::Pool::builder();
if let Some(max_size) = self.pool_max_size {
pool_builder = pool_builder.max_size(max_size);
}
if let Some(min_idle) = self.pool_min_idle {
pool_builder = pool_builder.min_idle(Some(min_idle));
}
if let Some(idle_timeout) = self.pool_idle_timeout {
pool_builder = pool_builder.idle_timeout(Some(idle_timeout));
}
if let Some(connection_timeout) = self.pool_connection_timeout {
pool_builder = pool_builder.connection_timeout(connection_timeout);
}
if let Some(max_lifetime) = self.pool_max_lifetime {
pool_builder = pool_builder.max_lifetime(Some(max_lifetime));
}
pool_builder.build(manager)
}
}
pub struct PostgresClientWrapper {
connection_string: String,
client: Mutex<Option<Client>>,
}
pub fn transaction<F, T>(operations: F) -> Result<T, PostgresError>
where
F: FnOnce(&mut Client) -> Result<T, PostgresError>,
{
let client = get_postgres_client()?;
let client_mutex = client.get_client()?;
let mut client_guard = client_mutex.lock().unwrap();
if let Some(client) = client_guard.as_mut() {
client.execute("BEGIN", &[])?;
match operations(client) {
Ok(result) => {
client.execute("COMMIT", &[])?;
Ok(result)
}
Err(e) => {
let _ = client.execute("ROLLBACK", &[]);
Err(e)
}
}
} else {
Err(create_postgres_error("Failed to get PostgreSQL client"))
}
}
pub fn transaction_with_pool<F, T>(operations: F) -> Result<T, PostgresError>
where
F: FnOnce(&mut Client) -> Result<T, PostgresError>,
{
let pool = get_postgres_pool()?;
let mut client = pool.get().map_err(|e| {
create_postgres_error(&format!("Failed to get connection from pool: {}", e))
})?;
client.execute("BEGIN", &[])?;
match operations(&mut client) {
Ok(result) => {
client.execute("COMMIT", &[])?;
Ok(result)
}
Err(e) => {
let _ = client.execute("ROLLBACK", &[]);
Err(e)
}
}
}
impl PostgresClientWrapper {
fn new(connection_string: String) -> Self {
PostgresClientWrapper {
connection_string,
client: Mutex::new(None),
}
}
fn get_client(&self) -> Result<&Mutex<Option<Client>>, PostgresError> {
let mut client_guard = self.client.lock().unwrap();
if client_guard.is_none() {
*client_guard = Some(Client::connect(&self.connection_string, NoTls)?);
}
Ok(&self.client)
}
pub fn execute(
&self,
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<u64, PostgresError> {
let client_mutex = self.get_client()?;
let mut client_guard = client_mutex.lock().unwrap();
if let Some(client) = client_guard.as_mut() {
client.execute(query, params)
} else {
Err(create_postgres_error("Failed to get PostgreSQL client"))
}
}
pub fn query(
&self,
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<Vec<Row>, PostgresError> {
let client_mutex = self.get_client()?;
let mut client_guard = client_mutex.lock().unwrap();
if let Some(client) = client_guard.as_mut() {
client.query(query, params)
} else {
Err(create_postgres_error("Failed to get PostgreSQL client"))
}
}
pub fn query_one(
&self,
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<Row, PostgresError> {
let client_mutex = self.get_client()?;
let mut client_guard = client_mutex.lock().unwrap();
if let Some(client) = client_guard.as_mut() {
client.query_one(query, params)
} else {
Err(create_postgres_error("Failed to get PostgreSQL client"))
}
}
pub fn query_opt(
&self,
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<Option<Row>, PostgresError> {
let client_mutex = self.get_client()?;
let mut client_guard = client_mutex.lock().unwrap();
if let Some(client) = client_guard.as_mut() {
client.query_opt(query, params)
} else {
Err(create_postgres_error("Failed to get PostgreSQL client"))
}
}
pub fn ping(&self) -> Result<bool, PostgresError> {
let result = self.query("SELECT 1", &[]);
match result {
Ok(_) => Ok(true),
Err(e) => Err(e),
}
}
}
pub fn get_postgres_client() -> Result<Arc<PostgresClientWrapper>, PostgresError> {
{
let guard = POSTGRES_CLIENT.lock().unwrap();
if let Some(ref client) = &*guard {
return Ok(Arc::clone(client));
}
}
let client = create_postgres_client()?;
{
let mut guard = POSTGRES_CLIENT.lock().unwrap();
*guard = Some(Arc::clone(&client));
}
Ok(client)
}
fn create_postgres_client() -> Result<Arc<PostgresClientWrapper>, PostgresError> {
let host = env::var("POSTGRES_HOST").unwrap_or_else(|_| String::from("localhost"));
let port = env::var("POSTGRES_PORT")
.ok()
.and_then(|p| p.parse::<u16>().ok())
.unwrap_or(5432);
let user = env::var("POSTGRES_USER").unwrap_or_else(|_| String::from("postgres"));
let password = env::var("POSTGRES_PASSWORD").ok();
let database = env::var("POSTGRES_DB").unwrap_or_else(|_| String::from("postgres"));
let mut builder = PostgresConfigBuilder::new()
.host(&host)
.port(port)
.user(&user)
.database(&database);
if let Some(pass) = password {
builder = builder.password(&pass);
}
let connection_string = builder.build_connection_string();
let wrapper = Arc::new(PostgresClientWrapper::new(connection_string));
match wrapper.ping() {
Ok(_) => Ok(wrapper),
Err(e) => Err(e),
}
}
pub fn reset() -> Result<(), PostgresError> {
{
let mut client_guard = POSTGRES_CLIENT.lock().unwrap();
*client_guard = None;
}
get_postgres_client()?;
Ok(())
}
pub fn execute(
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<u64, PostgresError> {
let client = get_postgres_client()?;
client.execute(query, params)
}
pub fn query(
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<Vec<Row>, PostgresError> {
let client = get_postgres_client()?;
client.query(query, params)
}
pub fn query_one(
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<Row, PostgresError> {
let client = get_postgres_client()?;
client.query_one(query, params)
}
pub fn query_opt(
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<Option<Row>, PostgresError> {
let client = get_postgres_client()?;
client.query_opt(query, params)
}
pub fn with_config(config: PostgresConfigBuilder) -> Result<Client, PostgresError> {
config.build()
}
pub fn with_pool_config(
config: PostgresConfigBuilder,
) -> Result<Pool<PostgresConnectionManager<NoTls>>, r2d2::Error> {
config.build_pool()
}
pub fn get_postgres_pool() -> Result<Arc<Pool<PostgresConnectionManager<NoTls>>>, PostgresError> {
{
let guard = POSTGRES_POOL.lock().unwrap();
if let Some(ref pool) = &*guard {
return Ok(Arc::clone(pool));
}
}
let pool = create_postgres_pool()?;
{
let mut guard = POSTGRES_POOL.lock().unwrap();
*guard = Some(Arc::clone(&pool));
}
Ok(pool)
}
fn create_postgres_pool() -> Result<Arc<Pool<PostgresConnectionManager<NoTls>>>, PostgresError> {
let host = env::var("POSTGRES_HOST").unwrap_or_else(|_| String::from("localhost"));
let port = env::var("POSTGRES_PORT")
.ok()
.and_then(|p| p.parse::<u16>().ok())
.unwrap_or(5432);
let user = env::var("POSTGRES_USER").unwrap_or_else(|_| String::from("postgres"));
let password = env::var("POSTGRES_PASSWORD").ok();
let database = env::var("POSTGRES_DB").unwrap_or_else(|_| String::from("postgres"));
let mut builder = PostgresConfigBuilder::new()
.host(&host)
.port(port)
.user(&user)
.database(&database)
.use_pool(true);
if let Some(pass) = password {
builder = builder.password(&pass);
}
match builder.build_pool() {
Ok(pool) => {
match pool.get() {
Ok(_) => Ok(Arc::new(pool)),
Err(e) => Err(create_postgres_error(&format!(
"Failed to connect to PostgreSQL: {}",
e
))),
}
}
Err(e) => Err(create_postgres_error(&format!(
"Failed to create PostgreSQL connection pool: {}",
e
))),
}
}
pub fn reset_pool() -> Result<(), PostgresError> {
{
let mut pool_guard = POSTGRES_POOL.lock().unwrap();
*pool_guard = None;
}
get_postgres_pool()?;
Ok(())
}
pub fn execute_with_pool(
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<u64, PostgresError> {
let pool = get_postgres_pool()?;
let mut client = pool.get().map_err(|e| {
create_postgres_error(&format!("Failed to get connection from pool: {}", e))
})?;
client.execute(query, params)
}
pub fn query_with_pool(
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<Vec<Row>, PostgresError> {
let pool = get_postgres_pool()?;
let mut client = pool.get().map_err(|e| {
create_postgres_error(&format!("Failed to get connection from pool: {}", e))
})?;
client.query(query, params)
}
pub fn query_one_with_pool(
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<Row, PostgresError> {
let pool = get_postgres_pool()?;
let mut client = pool.get().map_err(|e| {
create_postgres_error(&format!("Failed to get connection from pool: {}", e))
})?;
client.query_one(query, params)
}
pub fn query_opt_with_pool(
query: &str,
params: &[&(dyn postgres::types::ToSql + Sync)],
) -> Result<Option<Row>, PostgresError> {
let pool = get_postgres_pool()?;
let mut client = pool.get().map_err(|e| {
create_postgres_error(&format!("Failed to get connection from pool: {}", e))
})?;
client.query_opt(query, params)
}
#[derive(Default)]
pub struct QueryParams {
params: Vec<Box<dyn ToSql + Sync>>,
}
impl QueryParams {
pub fn new() -> Self {
Self { params: Vec::new() }
}
pub fn add<T: 'static + ToSql + Sync>(&mut self, value: T) -> &mut Self {
self.params.push(Box::new(value));
self
}
pub fn add_str(&mut self, value: &str) -> &mut Self {
self.add(value.to_string())
}
pub fn add_int(&mut self, value: i32) -> &mut Self {
self.add(value)
}
pub fn add_float(&mut self, value: f64) -> &mut Self {
self.add(value)
}
pub fn add_bool(&mut self, value: bool) -> &mut Self {
self.add(value)
}
pub fn add_opt<T: 'static + ToSql + Sync>(&mut self, value: Option<T>) -> &mut Self {
if let Some(v) = value {
self.add(v);
} else {
self.params.push(Box::new(None::<String>));
}
self
}
pub fn as_slice(&self) -> Vec<&(dyn ToSql + Sync)> {
self.params
.iter()
.map(|p| p.as_ref() as &(dyn ToSql + Sync))
.collect()
}
}
pub fn execute_with_params(query_str: &str, params: &QueryParams) -> Result<u64, PostgresError> {
let client = get_postgres_client()?;
client.execute(query_str, ¶ms.as_slice())
}
pub fn query_with_params(query_str: &str, params: &QueryParams) -> Result<Vec<Row>, PostgresError> {
let client = get_postgres_client()?;
client.query(query_str, ¶ms.as_slice())
}
pub fn query_one_with_params(query_str: &str, params: &QueryParams) -> Result<Row, PostgresError> {
let client = get_postgres_client()?;
client.query_one(query_str, ¶ms.as_slice())
}
pub fn query_opt_with_params(
query_str: &str,
params: &QueryParams,
) -> Result<Option<Row>, PostgresError> {
let client = get_postgres_client()?;
client.query_opt(query_str, ¶ms.as_slice())
}
pub fn execute_with_pool_params(
query_str: &str,
params: &QueryParams,
) -> Result<u64, PostgresError> {
execute_with_pool(query_str, ¶ms.as_slice())
}
pub fn query_with_pool_params(
query_str: &str,
params: &QueryParams,
) -> Result<Vec<Row>, PostgresError> {
query_with_pool(query_str, ¶ms.as_slice())
}
pub fn query_one_with_pool_params(
query_str: &str,
params: &QueryParams,
) -> Result<Row, PostgresError> {
query_one_with_pool(query_str, ¶ms.as_slice())
}
pub fn query_opt_with_pool_params(
query_str: &str,
params: &QueryParams,
) -> Result<Option<Row>, PostgresError> {
query_opt_with_pool(query_str, ¶ms.as_slice())
}
pub fn notify(channel: &str, payload: &str) -> Result<(), PostgresError> {
let client = get_postgres_client()?;
client.execute(&format!("NOTIFY {}, '{}'", channel, payload), &[])?;
Ok(())
}
pub fn notify_with_pool(channel: &str, payload: &str) -> Result<(), PostgresError> {
let pool = get_postgres_pool()?;
let mut client = pool.get().map_err(|e| {
create_postgres_error(&format!("Failed to get connection from pool: {}", e))
})?;
client.execute(&format!("NOTIFY {}, '{}'", channel, payload), &[])?;
Ok(())
}