upend/src/database.rs

use crate::addressing::Address;
use crate::hash::{decode, hash, Hash, Hashable};
use crate::models;
use crate::util::LoggerSink;
use anyhow::{anyhow, Result};
use diesel::debug_query;
use diesel::prelude::*;
use diesel::r2d2::{self, ConnectionManager};
use diesel::sqlite::{Sqlite, SqliteConnection};
use log::{debug, trace};
use serde::export::Formatter;
use serde_json::{json, Value};
use std::convert::TryFrom;
use std::fs;
use std::io::{Cursor, Write};
use std::path::{Path, PathBuf};
use std::time::Duration;

#[derive(Debug, Clone)]
pub struct Entry {
    pub identity: Hash,
    pub target: Address,
    pub key: String,
    pub value: EntryValue,
}

#[derive(Debug, Clone)]
pub struct InnerEntry {
    pub target: Address,
    pub key: String,
    pub value: EntryValue,
}

#[derive(Debug, Clone, PartialEq)]
pub enum EntryValue {
    Value(serde_json::Value),
    Address(Address),
    Invalid,
}

impl Entry {
    pub fn as_json(&self) -> serde_json::Value {
        json!({
            "target": self.target.to_string(),
            "key": self.key,
            "value": match &self.value {
                EntryValue::Value(value) => ("VALUE", value.clone()),
                EntryValue::Address(address) => ("ADDR", json!(address.to_string())),
                EntryValue::Invalid => ("INVALID", json!("INVALID")),
            }
        })
    }
}

impl TryFrom<models::Entry> for Entry {
    type Error = anyhow::Error;

    fn try_from(e: models::Entry) -> Result<Self, Self::Error> {
        Ok(Entry {
            identity: Hash(e.identity),
            target: Address::decode(&e.target)?,
            key: e.key,
            value: e.value.parse().unwrap(),
        })
    }
}

impl std::fmt::Display for InnerEntry {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "{} | {} | {}", self.target, self.key, self.value)
    }
}

impl Hashable for InnerEntry {
    fn hash(self: &InnerEntry) -> Result<Hash> {
        let mut result = Cursor::new(vec![0u8; 0]);
        result.write_all(self.target.encode()?.as_slice())?;
        result.write_all(self.key.as_bytes())?;
        result.write_all(self.value.to_str()?.as_bytes())?;
        Ok(hash(result.get_ref()))
    }
}

impl EntryValue {
    pub fn to_str(&self) -> Result<String> {
        let (type_char, content) = match self {
            EntryValue::Value(value) => ('J', serde_json::to_string(value)?),
            EntryValue::Address(address) => ('O', address.to_string()),
            EntryValue::Invalid => return Err(anyhow!("Cannot serialize invalid Entity value.")),
        };

        Ok(format!("{}{}", type_char, content))
    }
}

// unsafe unwraps!
impl std::fmt::Display for EntryValue {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        let (entry_type, entry_value) = match self {
            EntryValue::Address(address) => ("ADDRESS", address.to_string()),
            EntryValue::Value(value) => ("VALUE", serde_json::to_string(value).unwrap()),
            EntryValue::Invalid => ("INVALID", "INVALID".to_string()),
        };
        write!(f, "{}: {}", entry_type, entry_value)
    }
}

impl std::str::FromStr for EntryValue {
    type Err = std::convert::Infallible;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.len() < 2 {
            Ok(EntryValue::Invalid)
        } else {
            let (type_char, content) = s.split_at(1);
            match (type_char, content) {
                ("J", content) => {
                    if let Ok(value) = serde_json::from_str(content) {
                        Ok(EntryValue::Value(value))
                    } else {
                        Ok(EntryValue::Invalid)
                    }
                }
                ("O", content) => {
                    if let Ok(addr) = decode(content).and_then(|v| Address::decode(&v)) {
                        Ok(EntryValue::Address(addr))
                    } else {
                        Ok(EntryValue::Invalid)
                    }
                }
                _ => Ok(EntryValue::Invalid),
            }
        }
    }
}

pub fn insert_file<C: Connection<Backend = Sqlite>>(
    connection: &C,
    file: models::NewFile,
) -> Result<usize> {
    use crate::schema::files;

    debug!(
        "Inserting {} ({})...",
        &file.path,
        Address::Hash(Hash((&file.hash).clone()))
    );

    Ok(diesel::insert_into(files::table)
        .values(file)
        .execute(connection)?)
}

pub fn retrieve_all_files<C: Connection<Backend = Sqlite>>(
    connection: &C,
) -> Result<Vec<models::File>> {
    use crate::schema::files::dsl::*;
    let matches = files.load::<models::File>(connection)?;
    Ok(matches)
}

pub fn retrieve_file<C: Connection<Backend = Sqlite>>(
    connection: &C,
    obj_hash: Hash,
) -> Result<Option<String>> {
    use crate::schema::files::dsl::*;

    let matches = files
        .filter(valid.eq(true))
        .filter(hash.eq(obj_hash.0))
        .load::<models::File>(connection)?;

    Ok(matches.get(0).map(|f| f.path.clone()))
}

pub fn file_set_valid<C: Connection<Backend = Sqlite>>(
    connection: &C,
    file_id: i32,
    is_valid: bool,
) -> Result<usize> {
    use crate::schema::files::dsl::*;

    debug!("Setting file ID {} to valid = {}", file_id, is_valid);

    Ok(diesel::update(files.filter(id.eq(file_id)))
        .set(valid.eq(is_valid))
        .execute(connection)?)
}

pub fn retrieve_object<C: Connection<Backend = Sqlite>>(
    connection: &C,
    object_address: Address,
) -> Result<Vec<Entry>> {
    use crate::schema::data::dsl::*;

    let matches = data
        .filter(target.eq(object_address.encode()?))
        .or_filter(value.eq(EntryValue::Address(object_address).to_str()?))
        .load::<models::Entry>(connection)?;
    let entries = matches
        .into_iter()
        .map(Entry::try_from)
        .filter_map(Result::ok)
        .collect();

    Ok(entries)
}

pub fn bulk_retrieve_objects<C: Connection<Backend = Sqlite>>(
    connection: &C,
    object_addresses: Vec<Address>,
) -> Result<Vec<Entry>> {
    use crate::schema::data::dsl::*;

    let matches = data
        .filter(
            target.eq_any(
                object_addresses
                    .iter()
                    .filter_map(|addr| addr.encode().ok()),
            ),
        )
        // .or_filter(value.eq(EntryValue::Address(object_address).to_str()?))
        .load::<models::Entry>(connection)?;
    let entries = matches
        .into_iter()
        .map(Entry::try_from)
        .filter_map(Result::ok)
        .collect();

    Ok(entries)
}

pub fn remove_object<C: Connection<Backend = Sqlite>>(
    connection: &C,
    object_address: Address,
) -> Result<usize> {
    use crate::schema::data::dsl::*;

    debug!("Deleting {}!", object_address);

    let matches = data
        .filter(target.eq(object_address.encode()?))
        .or_filter(value.eq(EntryValue::Address(object_address).to_str()?));

    Ok(diesel::delete(matches).execute(connection)?)
}

pub enum QueryComponent<T> {
    Exact(T),
    ILike(T),
    In(Vec<T>),
    Any,
}

pub struct EntryQuery {
    pub target: QueryComponent<Address>,
    pub key: QueryComponent<String>,
    pub value: QueryComponent<EntryValue>,
}

pub fn query_entries<C: Connection<Backend = Sqlite>>(
    connection: &C,
    entry_query: EntryQuery,
) -> Result<Vec<Entry>> {
    use crate::schema::data::dsl::*;

    let mut query = data.into_boxed();

    query = match entry_query.target {
        QueryComponent::Exact(q_target) => query.filter(target.eq(q_target.encode()?)),
        QueryComponent::In(q_targets) => {
            let targets: Result<Vec<_>, _> = q_targets.into_iter().map(|t| t.encode()).collect();
            query.filter(target.eq_any(targets?))
        }
        QueryComponent::ILike(_) => return Err(anyhow!("Cannot query Address alike.")),
        QueryComponent::Any => query,
    };

    query = match entry_query.key {
        QueryComponent::Exact(q_key) => query.filter(key.eq(q_key)),
        QueryComponent::In(q_keys) => query.filter(key.eq_any(q_keys)),
        QueryComponent::ILike(q_key) => query.filter(key.like(format!("%{}%", q_key))),
        QueryComponent::Any => query,
    };

    query = match entry_query.value {
        QueryComponent::Exact(q_value) => query.filter(value.eq(q_value.to_str()?)),
        QueryComponent::In(q_values) => {
            let values: Result<Vec<_>, _> = q_values.into_iter().map(|v| v.to_str()).collect();
            query.filter(value.eq_any(values?))
        }
        QueryComponent::ILike(EntryValue::Value(Value::String(q_value_string))) => {
            query.filter(value.like(format!("%{}%", q_value_string)))
        }
        QueryComponent::ILike(_) => {
            return Err(anyhow!("Only string Values can be queried alike."))
        }
        QueryComponent::Any => query,
    };

    trace!("Querying: {}", debug_query(&query));

    let matches = query.load::<models::Entry>(connection)?;

    let entries = matches
        .into_iter()
        .map(Entry::try_from)
        .filter_map(Result::ok)
        .collect();

    Ok(entries)
}

pub fn insert_entry<C: Connection<Backend = Sqlite>>(
    connection: &C,
    entry: InnerEntry,
) -> Result<usize> {
    use crate::schema::data;

    debug!("Inserting: {}", entry);

    let insert_entry = models::Entry {
        identity: entry.hash()?.0,
        target: entry.target.encode()?,
        key: entry.key,
        value: entry.value.to_str()?,
    };

    Ok(diesel::insert_into(data::table)
        .values(insert_entry)
        .execute(connection)?)
}

#[derive(Debug)]
pub struct ConnectionOptions {
    pub enable_foreign_keys: bool,
    pub busy_timeout: Option<Duration>,
}

impl ConnectionOptions {
    pub fn apply(&self, conn: &SqliteConnection) -> QueryResult<()> {
        if self.enable_foreign_keys {
            conn.execute("PRAGMA foreign_keys = ON;")?;
        }
        if let Some(duration) = self.busy_timeout {
            conn.execute(&format!("PRAGMA busy_timeout = {};", duration.as_millis()))?;
        }
        Ok(())
    }
}

impl diesel::r2d2::CustomizeConnection<SqliteConnection, diesel::r2d2::Error>
    for ConnectionOptions
{
    fn on_acquire(&self, conn: &mut SqliteConnection) -> Result<(), diesel::r2d2::Error> {
        self.apply(conn).map_err(diesel::r2d2::Error::QueryError)
    }
}

pub type DbPool = r2d2::Pool<ConnectionManager<SqliteConnection>>;

pub struct OpenResult {
    pub pool: DbPool,
    pub new: bool,
}

pub const DATABASE_FILENAME: &str = "upend.sqlite3";

pub fn open_upend<P: AsRef<Path>>(dirpath: P, reinitialize: bool) -> Result<OpenResult> {
    embed_migrations!("./migrations/upend/");

    let database_path: PathBuf = dirpath.as_ref().join(DATABASE_FILENAME);
    if reinitialize {
        let _ = fs::remove_file(&database_path);
    }
    let new = !database_path.exists();

    let manager = ConnectionManager::<SqliteConnection>::new(database_path.to_str().unwrap());
    let pool = r2d2::Pool::builder()
        .connection_customizer(Box::new(ConnectionOptions {
            enable_foreign_keys: true,
            busy_timeout: Some(Duration::from_secs(30)),
        }))
        .build(manager)
        .expect("Failed to create pool.");

    embedded_migrations::run_with_output(
        &pool.get().unwrap(),
        &mut LoggerSink {
            ..Default::default()
        },
    )?;

    Ok(OpenResult { pool, new })
}