feat: Add PostgreSQL and Redis client support
- Add PostgreSQL client functionality for database interactions. - Add Redis client functionality for cache and data store operations. - Extend Rhai scripting with PostgreSQL and Redis client modules. - Add documentation and test cases for both clients.
This commit is contained in:
356
src/postgresclient/postgresclient.rs
Normal file
356
src/postgresclient/postgresclient.rs
Normal file
@@ -0,0 +1,356 @@
|
||||
use lazy_static::lazy_static;
|
||||
use postgres::{Client, Error as PostgresError, NoTls, Row};
|
||||
use std::env;
|
||||
use std::sync::{Arc, Mutex, Once};
|
||||
|
||||
// Helper function to create a PostgreSQL error
|
||||
fn create_postgres_error(_message: &str) -> PostgresError {
|
||||
// Since we can't directly create a PostgresError, we'll create one by
|
||||
// attempting to connect to an invalid connection string and capturing the error
|
||||
let result = Client::connect("invalid-connection-string", NoTls);
|
||||
match result {
|
||||
Ok(_) => unreachable!(), // This should never happen
|
||||
Err(e) => {
|
||||
// We have a valid PostgresError now, but we want to customize the message
|
||||
// Unfortunately, PostgresError doesn't provide a way to modify the message
|
||||
// So we'll just return the error we got
|
||||
e
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Global PostgreSQL client instance using lazy_static
|
||||
lazy_static! {
|
||||
static ref POSTGRES_CLIENT: Mutex<Option<Arc<PostgresClientWrapper>>> = Mutex::new(None);
|
||||
static ref INIT: Once = Once::new();
|
||||
}
|
||||
|
||||
/// PostgreSQL connection configuration builder
|
||||
///
|
||||
/// This struct is used to build a PostgreSQL connection configuration.
|
||||
/// It follows the builder pattern to allow for flexible configuration.
|
||||
#[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>,
|
||||
}
|
||||
|
||||
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,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PostgresConfigBuilder {
|
||||
/// Create a new PostgreSQL connection configuration builder with default values
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
/// Set the host for the PostgreSQL connection
|
||||
pub fn host(mut self, host: &str) -> Self {
|
||||
self.host = host.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the port for the PostgreSQL connection
|
||||
pub fn port(mut self, port: u16) -> Self {
|
||||
self.port = port;
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the user for the PostgreSQL connection
|
||||
pub fn user(mut self, user: &str) -> Self {
|
||||
self.user = user.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the password for the PostgreSQL connection
|
||||
pub fn password(mut self, password: &str) -> Self {
|
||||
self.password = Some(password.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the database for the PostgreSQL connection
|
||||
pub fn database(mut self, database: &str) -> Self {
|
||||
self.database = database.to_string();
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the application name for the PostgreSQL connection
|
||||
pub fn application_name(mut self, application_name: &str) -> Self {
|
||||
self.application_name = Some(application_name.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the connection timeout in seconds
|
||||
pub fn connect_timeout(mut self, seconds: u64) -> Self {
|
||||
self.connect_timeout = Some(seconds);
|
||||
self
|
||||
}
|
||||
|
||||
/// Set the SSL mode for the PostgreSQL connection
|
||||
pub fn ssl_mode(mut self, ssl_mode: &str) -> Self {
|
||||
self.ssl_mode = Some(ssl_mode.to_string());
|
||||
self
|
||||
}
|
||||
|
||||
/// Build the connection string from the configuration
|
||||
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
|
||||
}
|
||||
|
||||
/// Build a PostgreSQL client from the configuration
|
||||
pub fn build(&self) -> Result<Client, PostgresError> {
|
||||
let conn_string = self.build_connection_string();
|
||||
Client::connect(&conn_string, NoTls)
|
||||
}
|
||||
}
|
||||
|
||||
/// Wrapper for PostgreSQL client to handle connection
|
||||
pub struct PostgresClientWrapper {
|
||||
connection_string: String,
|
||||
client: Mutex<Option<Client>>,
|
||||
}
|
||||
|
||||
impl PostgresClientWrapper {
|
||||
/// Create a new PostgreSQL client wrapper
|
||||
fn new(connection_string: String) -> Self {
|
||||
PostgresClientWrapper {
|
||||
connection_string,
|
||||
client: Mutex::new(None),
|
||||
}
|
||||
}
|
||||
|
||||
/// Get a reference to the PostgreSQL client, creating it if it doesn't exist
|
||||
fn get_client(&self) -> Result<&Mutex<Option<Client>>, PostgresError> {
|
||||
let mut client_guard = self.client.lock().unwrap();
|
||||
|
||||
// If we don't have a client or it's not working, create a new one
|
||||
if client_guard.is_none() {
|
||||
*client_guard = Some(Client::connect(&self.connection_string, NoTls)?);
|
||||
}
|
||||
|
||||
Ok(&self.client)
|
||||
}
|
||||
|
||||
/// Execute a query on the PostgreSQL connection
|
||||
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"))
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a query on the PostgreSQL connection and return the rows
|
||||
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"))
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a query on the PostgreSQL connection and return a single row
|
||||
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"))
|
||||
}
|
||||
}
|
||||
|
||||
/// Execute a query on the PostgreSQL connection and return an optional row
|
||||
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"))
|
||||
}
|
||||
}
|
||||
|
||||
/// Ping the PostgreSQL server to check if the connection is alive
|
||||
pub fn ping(&self) -> Result<bool, PostgresError> {
|
||||
let result = self.query("SELECT 1", &[]);
|
||||
match result {
|
||||
Ok(_) => Ok(true),
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the PostgreSQL client instance
|
||||
pub fn get_postgres_client() -> Result<Arc<PostgresClientWrapper>, PostgresError> {
|
||||
// Check if we already have a client
|
||||
{
|
||||
let guard = POSTGRES_CLIENT.lock().unwrap();
|
||||
if let Some(ref client) = &*guard {
|
||||
return Ok(Arc::clone(client));
|
||||
}
|
||||
}
|
||||
|
||||
// Create a new client
|
||||
let client = create_postgres_client()?;
|
||||
|
||||
// Store the client globally
|
||||
{
|
||||
let mut guard = POSTGRES_CLIENT.lock().unwrap();
|
||||
*guard = Some(Arc::clone(&client));
|
||||
}
|
||||
|
||||
Ok(client)
|
||||
}
|
||||
|
||||
/// Create a new PostgreSQL client
|
||||
fn create_postgres_client() -> Result<Arc<PostgresClientWrapper>, PostgresError> {
|
||||
// Try to get connection details from environment variables
|
||||
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"));
|
||||
|
||||
// Build the connection string
|
||||
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();
|
||||
|
||||
// Create the client wrapper
|
||||
let wrapper = Arc::new(PostgresClientWrapper::new(connection_string));
|
||||
|
||||
// Test the connection
|
||||
match wrapper.ping() {
|
||||
Ok(_) => Ok(wrapper),
|
||||
Err(e) => Err(e),
|
||||
}
|
||||
}
|
||||
|
||||
/// Reset the PostgreSQL client
|
||||
pub fn reset() -> Result<(), PostgresError> {
|
||||
// Clear the existing client
|
||||
{
|
||||
let mut client_guard = POSTGRES_CLIENT.lock().unwrap();
|
||||
*client_guard = None;
|
||||
}
|
||||
|
||||
// Create a new client, only return error if it fails
|
||||
get_postgres_client()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Execute a query on the PostgreSQL connection
|
||||
pub fn execute(
|
||||
query: &str,
|
||||
params: &[&(dyn postgres::types::ToSql + Sync)],
|
||||
) -> Result<u64, PostgresError> {
|
||||
let client = get_postgres_client()?;
|
||||
client.execute(query, params)
|
||||
}
|
||||
|
||||
/// Execute a query on the PostgreSQL connection and return the rows
|
||||
pub fn query(
|
||||
query: &str,
|
||||
params: &[&(dyn postgres::types::ToSql + Sync)],
|
||||
) -> Result<Vec<Row>, PostgresError> {
|
||||
let client = get_postgres_client()?;
|
||||
client.query(query, params)
|
||||
}
|
||||
|
||||
/// Execute a query on the PostgreSQL connection and return a single row
|
||||
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)
|
||||
}
|
||||
|
||||
/// Execute a query on the PostgreSQL connection and return an optional row
|
||||
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)
|
||||
}
|
||||
|
||||
/// Create a new PostgreSQL client with custom configuration
|
||||
pub fn with_config(config: PostgresConfigBuilder) -> Result<Client, PostgresError> {
|
||||
config.build()
|
||||
}
|
Reference in New Issue
Block a user