upend/src/database.rs

use crate::addressing::Address;
use crate::hash::{decode, hash, Hash, Hashable};
use crate::models;
use actix::prelude::*;
use actix_derive::Message;
use anyhow::{anyhow, Result};
use diesel::prelude::*;
use diesel::r2d2::{self, ConnectionManager};
use diesel::sqlite::SqliteConnection;
use log::debug;
use serde::export::Formatter;
use serde_json::json;
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 type DbPool = r2d2::Pool<ConnectionManager<SqliteConnection>>;

pub struct DbExecutor(pub DbPool);

impl Actor for DbExecutor {
    type Context = SyncContext<Self>;
}

#[derive(Message)]
#[rtype(result = "Result<usize>")]
pub struct InsertFile {
    pub file: models::NewFile,
}

impl Handler<InsertFile> for DbExecutor {
    type Result = Result<usize>;

    fn handle(&mut self, msg: InsertFile, _: &mut Self::Context) -> Self::Result {
        use crate::schema::files;

        let connection = &self.0.get()?;

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

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

#[derive(Message)]
#[rtype(result = "Result<Option<String>>")]
pub struct RetrieveByHash {
    pub hash: Hash,
}

impl Handler<RetrieveByHash> for DbExecutor {
    type Result = Result<Option<String>>;

    fn handle(&mut self, msg: RetrieveByHash, _: &mut Self::Context) -> Self::Result {
        use crate::schema::files::dsl::*;

        let connection = &self.0.get()?;

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

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

#[derive(Message)]
#[rtype(result = "Result<Vec<models::File>>")]
pub struct LookupByFilename {
    pub query: String,
}

impl Handler<LookupByFilename> for DbExecutor {
    type Result = Result<Vec<models::File>>;

    fn handle(&mut self, msg: LookupByFilename, _: &mut Self::Context) -> Self::Result {
        use crate::schema::files::dsl::*;

        let connection = &self.0.get()?;

        let matches = files
            .filter(path.like(format!("%{}%", msg.query)))
            .filter(valid.eq(true))
            .load::<models::File>(connection)?;

        Ok(matches)
    }
}

#[derive(Message)]
#[rtype(result = "Result<Vec<Entry>>")]
pub struct RetrieveObject {
    pub target: Address,
}

impl Handler<RetrieveObject> for DbExecutor {
    type Result = Result<Vec<Entry>>;

    fn handle(&mut self, msg: RetrieveObject, _: &mut Self::Context) -> Self::Result {
        use crate::schema::data::dsl::*;

        let connection = &self.0.get()?;

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

        Ok(entries)
    }
}

#[derive(Message)]
#[rtype(result = "Result<Vec<Entry>>")]
pub struct QueryEntries {
    pub target: Option<Address>,
    pub key: Option<String>,
    pub value: Option<EntryValue>,
}

impl Handler<QueryEntries> for DbExecutor {
    type Result = Result<Vec<Entry>>;

    fn handle(&mut self, msg: QueryEntries, _: &mut Self::Context) -> Self::Result {
        use crate::schema::data::dsl::*;

        let connection = &self.0.get()?;

        let mut query = data.into_boxed();

        if let Some(q_target) = msg.target {
            query = query.filter(target.eq(q_target.encode()?));
        }

        if let Some(q_key) = msg.key {
            query = query.filter(key.eq(q_key));
        }

        if let Some(q_value) = msg.value {
            query = query.filter(value.eq(q_value.to_str()?));
        }

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

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

        Ok(entries)
    }
}

#[derive(Message)]
#[rtype(result = "Result<usize>")]
pub struct InsertEntry {
    pub entry: InnerEntry,
}

impl Handler<InsertEntry> for DbExecutor {
    type Result = Result<usize>;

    fn handle(&mut self, msg: InsertEntry, _: &mut Self::Context) -> Self::Result {
        use crate::schema::data;

        let connection = &self.0.get()?;

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

        let insert_entry = models::Entry {
            identity: msg.entry.hash()?.0,
            target: msg.entry.target.encode()?,
            key: msg.entry.key,
            value: msg.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)
    }
}

#[derive(Default)]
struct LoggerSink {
    buffer: Vec<u8>,
}

impl std::io::Write for LoggerSink {
    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
        self.buffer.extend(buf.iter());

        if self.buffer.ends_with(b"\n") {
            self.flush()?;
        }

        Ok(buf.len())
    }

    fn flush(&mut self) -> std::io::Result<()> {
        use std::str;

        debug!(
            "{}",
            str::from_utf8(self.buffer.as_mut())
                .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?
                .trim()
        );

        Ok(())
    }
}

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

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(3)),
        }))
        .build(manager)
        .expect("Failed to create pool.");

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

    Ok(OpenResult { pool, new })
}