comunicapiv3/src/helpers/database.rs

use std::collections;
use std::error::Error;
use std::ops::Add;
use std::sync::{Arc, Mutex};

use mysql::{Binary, Pool, ResultSet, Value};
use mysql::prelude::Queryable;

use crate::data::config::DatabaseConfig;
use crate::data::error::{ExecError, ResultBoxError};
use std::collections::HashMap;

/// Database access helper
///
/// @author Pierre Hubert

pub type ProcessRowResult<E> = Result<E, Box<dyn Error>>;

// Pool shared across threads
static mut POOL: Option<Arc<Mutex<mysql::Pool>>> = None;

/// Connect to the database
pub fn connect(conf: &DatabaseConfig) -> Result<(), Box<dyn Error>> {
    let url = format!(
        "mysql://{}:{}@{}:3306/{}",
        conf.username, conf.password, conf.host, conf.name
    );

    let pool = Pool::new(url)?;
    let pool = Some(Arc::new(Mutex::new(pool)));

    unsafe {
        POOL = pool;
    }

    Ok(())
}

/// Get a connection to the database
pub fn get_connection() -> Result<mysql::PooledConn, Box<dyn Error>> {
    let pool: Pool;

    unsafe {
        let guard = POOL.as_ref().unwrap().lock();
        pool = guard.unwrap().clone()
    }

    Ok(pool.get_conn()?)
}


pub struct QueryInfo {
    /// Fetched table
    pub table: String,

    /// Query limits
    pub conditions: collections::HashMap<String, String>,

    /// Limit of the query (0 = no limit)
    pub limit: u64,

    /// Queried arguments
    ///
    /// If this attribute is empty, all the columns of the table set are fetched
    pub fields: Vec<String>,
}

impl QueryInfo {
    /// Initialize a new query on the database
    pub fn new(table: &str) -> QueryInfo {
        QueryInfo {
            table: table.to_string(),
            conditions: collections::HashMap::new(),
            limit: 0,
            fields: Vec::new(),
        }
    }

    pub fn cond(mut self, key: &str, val: &str) -> QueryInfo {
        self.conditions.insert(key.to_string(), val.to_string());
        self
    }

    pub fn cond_u32(mut self, key: &str, val: u32) -> QueryInfo {
        self.conditions.insert(key.to_string(), val.to_string());
        self
    }

    pub fn cond_u64(mut self, key: &str, val: u64) -> QueryInfo {
        self.conditions.insert(key.to_string(), val.to_string());
        self
    }

    pub fn cond_i64(mut self, key: &str, val: i64) -> QueryInfo {
        self.conditions.insert(key.to_string(), val.to_string());
        self
    }

    /// Append a field to the list of selected fields
    pub fn add_field(mut self, key: &str) -> QueryInfo {
        self.fields.push(key.to_string());
        self
    }

    /// Execute query
    pub fn exec<E, F: Fn(&RowResult) -> ProcessRowResult<E>>(self, process_function: F)
                                                             -> Result<Vec<E>, Box<dyn Error>> {
        query(self, process_function)
    }

    /// Execute count query
    pub fn exec_count(self) -> ResultBoxError<usize> {
        count(self)
    }
}

/// Struct used to read the result of a request
pub struct RowResult<'a> {
    row: &'a mysql::Row
}

impl<'a> RowResult<'a> {
    pub fn new(row: &mysql::Row) -> RowResult {
        RowResult {
            row
        }
    }

    /// Find a column in result set
    fn find_col(&self, name: &str) -> Result<usize, ExecError> {
        let name_bytes = name.as_bytes();
        let mut index = 0;
        for c in self.row.columns_ref() {
            if c.name_ref().eq(name_bytes) {
                return Ok(index);
            }

            index += 1;
        }

        Err(ExecError(format!("Column {} not found in database result set!", name)))
    }

    /// Find an integer included in the request
    pub fn get_int64(&self, name: &str) -> Result<i64, Box<dyn Error>> {
        let value = self.row.get_opt(self.find_col(name)?);

        match value {
            None => Err(ExecError::boxed_string(
                format!("Could not extract integer field {} !", name))),
            Some(s) => Ok(s?)
        }
    }

    /// Find an integer included in the request
    pub fn get_u64(&self, name: &str) -> Result<u64, Box<dyn Error>> {
        let value = self.row.get_opt(self.find_col(name)?);

        match value {
            None => Err(ExecError::boxed_string(
                format!("Could not extract integer field {} !", name))),
            Some(s) => Ok(s?)
        }
    }

    /// Find an integer included in the request
    pub fn get_usize(&self, name: &str) -> Result<usize, Box<dyn Error>> {
        let value = self.row.get_opt(self.find_col(name)?);

        match value {
            None => Err(ExecError::boxed_string(
                format!("Could not extract integer field {} !", name))),
            Some(s) => Ok(s?)
        }
    }

    /// Find a string included in the request
    pub fn get_str(&self, name: &str) -> Result<String, Box<dyn Error>> {
        let value = self.row.get_opt(self.find_col(name)?);

        match value {
            None => Err(ExecError::boxed_string(
                format!("Could not extract string field {} !", name))),
            Some(s) => Ok(s?)
        }
    }

    /// Check out whether a given value is null or not
    pub fn is_null(&self, name: &str) -> ResultBoxError<bool> {
        let value = self.row.get_opt(self.find_col(name)?);

        match value {
            None => Ok(true),
            Some(Ok(Value::NULL)) => Ok(true),
            _ => Ok(false)
        }
    }

    /// Get an optional string => Set to None if string is null / empty
    pub fn get_optional_str(&self, name: &str) -> ResultBoxError<Option<String>> {
        match self.is_null(name)? {
            true => Ok(None),
            false => Ok(Some(self.get_str(name)?).map_or(None, |d| {
                match d.is_empty() {
                    true => None,
                    false => Some(d)
                }
            }))
        }
    }

    /// Get legacy boolean value : 1 = true / 0 = false
    pub fn get_legacy_bool(&self, name: &str) -> ResultBoxError<bool> {
        Ok(self.get_int64(name)? == 1)
    }
}


/// Query a single row of the database
pub fn query_row<E, F: Fn(&RowResult) -> ProcessRowResult<E>>(mut info: QueryInfo,
                                                              process_function: F) -> Result<E, Box<dyn Error>> {
    info.limit = 1;
    let mut list = query(info, process_function)?;

    match list.len() {
        0 => Err(ExecError::boxed_new("Request did not return a result!")),
        _ => Ok(list.remove(0))
    }
}

/// Make a simple query
pub fn query<E, F: Fn(&RowResult) -> ProcessRowResult<E>>(info: QueryInfo, process_function: F)
                                                          -> Result<Vec<E>, Box<dyn Error>> {
    let mut params = vec![];

    let select_elements = match info.fields.len() {
        0 => "*".to_string(),
        _ => info.fields.join(", ")
    };

    // Build query
    let mut query = format!("SELECT {} FROM {} ", select_elements, info.table);


    // WHERE clause
    if !info.conditions.is_empty() {
        let mut where_args = vec![];

        for (k, v) in info.conditions {
            where_args.push(format!("{} = ?", k));
            params.push(v)
        }

        let where_args = format!(" WHERE {} ", where_args.join(" AND "));
        query = query.add(&where_args);
    }

    // LIMIT clause
    if info.limit > 0 {
        query = query.add(&format!(" LIMIT {}", info.limit));
    }

    // Execute query
    let mut con = get_connection()?;
    let stmt = con.prep(&query)?;

    let mut res = con.exec_iter(stmt, params)?;

    // This system is made to support only one dataset
    let result_set = res.next_set();
    if let None = result_set {
        return Err(Box::new(ExecError::new("No result set in a query!")));
    }
    let result_set: ResultSet<Binary> = result_set.unwrap()?;


    // Parse each result of the dataset
    let mut res: Vec<E> = vec![];
    for row in result_set {
        let row = row?;
        let result = RowResult::new(&row);

        let parsed = process_function(&result)?;
        res.push(parsed);
    }


    Ok(res)
}

/// Count the number of results a query would have produced
pub fn count(mut q: QueryInfo) -> ResultBoxError<usize> {
    &q.fields.clear();
    &q.fields.push("COUNT(*) as count".to_string());

    query(q, |res| res.get_usize("count"))?.pop()
        .ok_or(ExecError::boxed_new("database::count did not return a result!"))
}

/// Structure used to execute a insert query
pub struct InsertQuery {
    pub table: String,
    pub values: HashMap<String, mysql::Value>,
}

impl InsertQuery {
    /// Construct a new InsertQuery instance
    pub fn new(table: &str) -> InsertQuery {
        InsertQuery {
            table: table.to_string(),
            values: HashMap::new(),
        }
    }

    /// Add a string
    pub fn add_str(mut self, key: &str, value: &str) -> InsertQuery {
        self.values.insert(key.to_string(), Value::from(value));
        self
    }

    /// Add an integer
    pub fn add_i64(mut self, key: &str, value: i64) -> InsertQuery {
        self.values.insert(key.to_string(), Value::from(value));
        self
    }

    pub fn add_u32(mut self, key: &str, value: u32) -> InsertQuery {
        self.values.insert(key.to_string(), Value::from(value));
        self
    }
}

/// Insert a new entry into the database
pub fn insert(query: InsertQuery) -> ResultBoxError<()> {

    // Collect keys
    let keys = query.values
        .keys()
        .map(|f| f.to_string())
        .collect::<Vec<String>>();

    let query_sql = format!(
        "INSERT INTO {} ({}) VALUES({})",
        query.table,
        keys.join(", "),
        keys.iter().map(|_| "?").collect::<Vec<&str>>().join(", ")
    );

    get_connection()?.exec_drop(
        query_sql,
        query.values.values().collect::<Vec<&mysql::Value>>(),
    )?;

    Ok(())
}


/// Structure used to build delete queries
pub struct DeleteQuery {
    table: String,
    conditions: HashMap<String, Value>,
}

impl DeleteQuery {
    /// Construct a new delete table query
    pub fn new(table: &str) -> DeleteQuery {
        DeleteQuery {
            table: table.to_string(),
            conditions: HashMap::new(),
        }
    }

    /// Add a string condition
    pub fn cond_str(mut self, key: &str, value: &str) -> DeleteQuery {
        self.conditions.insert(key.to_string(), Value::from(value));
        self
    }

    /// Add an integer condition
    pub fn cond_u32(mut self, key: &str, value: u32) -> DeleteQuery {
        self.conditions.insert(key.to_string(), Value::from(value));
        self
    }

    pub fn cond_i64(mut self, key: &str, value: i64) -> DeleteQuery {
        self.conditions.insert(key.to_string(), Value::from(value));
        self
    }
}

/// Delete an entry from the database
pub fn delete(query: DeleteQuery) -> ResultBoxError<()> {
    if query.conditions.is_empty() {
        return Err(ExecError::boxed_new("DELETE without WHERE condition blocked for security reasons!"));
    }

    let query_sql = format!(
        "DELETE FROM {} WHERE {}",
        query.table,
        query.conditions.keys()
            .map(|f| format!("{} = ?", f))
            .collect::<Vec<String>>()
            .join(" AND ")
    );

    get_connection()?.exec_drop(
        query_sql,
        query.conditions.values().collect::<Vec<&Value>>()
    )?;

    Ok(())
}