path: root/src/app/store/sqlite.rs

use std::str::FromStr;

use crate::db;

use super::records::{EquityChangeEvent, EquitySymbol, FiduciaryAccount};

use super::{ContextError, Store};
use bigdecimal::ToPrimitive;
use chrono::Local;
use miette::Result;
use sqlite::{Connection, Row, Value};

type DateTime = chrono::DateTime<chrono::Utc>;

pub struct SQLiteMigration {
    migration_id: usize,
    migration_filename: String,
    migration_timestamp: DateTime,
}

pub struct SQLiteStore {
    location: String,
    connection: Option<Connection>,
    counter: Option<u8>,
}

/// Define data conversion from sqlite::Row to SQLiteMigration
impl TryFrom<Result<Row, sqlite::Error>> for SQLiteMigration {
    type Error = ContextError;
    fn try_from(value: Result<Row, sqlite::Error>) -> Result<Self, Self::Error> {
        let result = value?;

        let id: i64 = result.try_read("migration_id")?;
        let filename: &str = result.try_read("migration_filename")?;
        let timestamp: &str = result.try_read("migration_timestamp")?;

        Ok(SQLiteMigration {
            migration_id: usize::try_from(id)?,
            migration_filename: filename.into(),
            migration_timestamp: DateTime::from_str(timestamp)?,
        })
    }
}

/// Define data conversion from sqlite::Row to FiduciaryAccount
impl TryFrom<Result<Row, sqlite::Error>> for FiduciaryAccount {
    type Error = ContextError;
    fn try_from(value: Result<Row, sqlite::Error>) -> Result<Self, Self::Error> {
        let result = value?;

        let id: i64 = result.try_read("fiduciary_account_id")?;
        let number: i64 = result.try_read("fiduciary_account_number")?;
        let name: &str = result.try_read("fiduciary_account_name")?;
        let description: &str = result.try_read("fiduciary_account_description")?;
        let timestamp: &str = result.try_read("fiduciary_account_created_timestamp")?;

        Ok(FiduciaryAccount {
            fiduciary_account_id: u8::try_from(id)?,
            fiduciary_account_number: usize::try_from(number)?,
            fiduciary_account_name: name.into(),
            fiduciary_account_description: description.into(),
            fiduciary_account_created_timestamp: DateTime::from_str(timestamp)?,
        })
    }
}

/// Define data conversion from sqlite::Row to EquitySymbol
impl TryFrom<Result<Row, sqlite::Error>> for EquitySymbol {
    type Error = ContextError;
    fn try_from(value: Result<Row, sqlite::Error>) -> Result<Self, Self::Error> {
        let result = value?;

        let id: i64 = result.try_read("equity_symbol_id")?;
        let name: &str = result.try_read("equity_symbol_name")?;
        let symbol_type: &str = result.try_read("equity_symbol_type")?;
        let company: &str = result.try_read("equity_symbol_managing_company")?;
        let timestamp: &str = result.try_read("equity_symbol_created_timestamp")?;

        Ok(EquitySymbol {
            equity_symbol_id: usize::try_from(id)?,
            equity_symbol_name: name.into(),
            equity_symbol_type: symbol_type.try_into()?,
            equity_symbol_managing_company: company.into(),
            equity_symbol_created_timestamp: DateTime::from_str(&timestamp)?,
        })
    }
}

/// Define data conversion from sqlite::Row to EquityChangeEvent
impl TryFrom<Result<Row, sqlite::Error>> for EquityChangeEvent {
    type Error = ContextError;
    fn try_from(value: Result<Row, sqlite::Error>) -> Result<Self, Self::Error> {
        let result = value?;

        let event_id: i64 = result.try_read("equity_change_event_id")?;
        let account_id: i64 = result.try_read("fiduciary_account_id")?;
        let symbol_id: i64 = result.try_read("equity_symbol_id")?;

        // FP is not to be trusted !!
        let quantity: &str = result.try_read("equity_change_quantity")?;
        let cost_basis: &str = result.try_read("equity_change_cost_basis_usd")?;

        let change_type: &str = result.try_read("equity_change_type")?;
        let timestamp: &str = result.try_read("equity_change_timestamp")?;

        Ok(EquityChangeEvent {
            equity_change_event_id: usize::try_from(event_id)?,
            fiduciary_account_id: u8::try_from(account_id)?,
            equity_symbol_id: usize::try_from(symbol_id)?,
            equity_change_quantity: bigdecimal::BigDecimal::from_str(quantity)?,
            equity_change_cost_basis: bigdecimal::BigDecimal::from_str(cost_basis)?,
            equity_change_type: change_type.try_into()?,
            equity_change_timestamp: DateTime::from_str(&timestamp)?,
        })
    }
}

impl SQLiteStore {
    const FIDUCIARY_ACCOUNT_INSERT_QUERY: &str = "
        INSERT INTO fiduciary_account VALUES (
            :fiduciary_account_id,
            :fiduciary_account_number,
            :fiduciary_account_name,
            :fiduciary_account_description,
            :fiduciary_account_created_timestamp
        )";
    const FIDUCIARY_ACCOUNT_SELECT_QUERY: &str = "
        SELECT
            fiduciary_account_id,
            fiduciary_account_number,
            fiduciary_account_name,
            fiduciary_account_description,
            fiduciary_account_created_timestamp
        FROM fiduciary_account";
    const EQUITY_SYMBOL_INSERT_QUERY: &str = "
        INSERT INTO equity_symbol VALUES (
            :equity_symbol_id,
            :equity_symbol_name,
            :equity_symbol_type,
            :equity_symbol_managing_company,
            :equity_symbol_created_timestamp
        )";
    const EQUITY_SYMBOL_SELECT_QUERY: &str = "
        SELECT
            equity_symbol_id,
            equity_symbol_name,
            equity_symbol_type,
            equity_symbol_managing_company,
            equity_symbol_created_timestamp
        FROM equity_symbol";
    const EQUITY_CHANGE_EVENT_INSERT_QUERY: &str = "
        INSERT INTO equity_change_event VALUES (
            :equity_change_event_id,
            :fiduciary_account_id,
            :equity_symbol_id,
            :equity_change_quantity,
            :equity_change_cost_basis_usd,
            :equity_change_type,
            :equity_change_timestamp
        )";
    const EQUITY_CHANGE_EVENT_SELECT_QUERY: &str = "
        SELECT
            equity_change_event_id,
            fiduciary_account_id,
            equity_symbol_id,
            equity_change_quantity,
            equity_change_cost_basis_usd,
            equity_change_type,
            equity_change_timestamp
        FROM equity_change_event";

    const MIGRATION_CONTROL_CREATE_QUERY: &str = "
        CREATE TABLE migration_control (
            migration_id        UNSIGNED INT PRIMARY KEY,
            migration_filename  NVARCHAR(512) NOT NULL,
            migration_timestamp TIMESTAMP NOT NULL
        )";
    const MIGRATION_CONTROL_SELECT_QUERY: &str = "
        SELECT
            migration_id,
            migration_filename,
            migration_timestamp
        FROM migration_control";
    const MIGRATION_CONTROL_INSERT_QUERY: &str = "
        INSERT INTO migration_control VALUES (
            :migration_id,
            :migration_filename,
            :migration_timestamp
        )";

    /// Wrap connection in a Result to make it easier to unwrap later
    fn get_connection(&self) -> Result<&Connection, sqlite::Error> {
        match &self.connection {
            Some(conn) => Ok(conn),
            None => Err(sqlite::Error {
                code: None,
                message: Some("Uninitialized".into()),
            }),
        }
    }

    fn exec_migration_list(&mut self) -> Result<Vec<SQLiteMigration>, ContextError> {
        let conn = self.get_connection()?;

        let statement = conn.prepare(Self::MIGRATION_CONTROL_SELECT_QUERY);

        if statement.is_err() {
            conn.execute(Self::MIGRATION_CONTROL_CREATE_QUERY)?;
            Ok(vec![])
        } else {
            statement?
                .into_iter()
                .map(SQLiteMigration::try_from)
                .collect()
        }
    }

    /// Gets list of available migrations, asserts whether they need to be run, and executes them against the sqlite db.
    fn run_migrations(&mut self) -> Result<(), ContextError> {
        // - get list of sql files
        let migrations = db::get_migration_scripts()?;

        let run_migrations = self.exec_migration_list()?;
        let mut seed_id = run_migrations.len();
        tracing::event!(
            tracing::Level::DEBUG,
            "Found {} already run migrations.",
            seed_id
        );

        let conn = self.get_connection()?;

        // - against list of migration records, run any scripts absent
        for migration in &migrations {
            let prev_migration = run_migrations
                .iter()
                .find(|m| m.migration_filename == migration.filename);
            if prev_migration.is_none() {
                let timestamp = Local::now().to_string();
                tracing::event!(
                    tracing::Level::INFO,
                    migration_id = seed_id,
                    directory = migration.directory,
                    migration_filename = migration.filename,
                    migration_timestamp = timestamp,
                    "Performing migration."
                );

                conn.execute(migration.script_contents.clone())?;
                let mut statement = conn.prepare(Self::MIGRATION_CONTROL_INSERT_QUERY)?;
                statement.bind::<&[(_, Value)]>(&[
                    (":migration_id", i64::try_from(seed_id)?.into()),
                    (":migration_filename", (&*migration.filename).into()),
                    (":migration_timestamp", (&*timestamp).into()),
                ])?;

                statement.next()?;
                seed_id += 1;
            } else {
                let actual_migration = prev_migration.unwrap();
                tracing::event!(
                    tracing::Level::DEBUG,
                    migration_id = actual_migration.migration_id,
                    migration_filename = actual_migration.migration_filename,
                    migration_timestamp = actual_migration.migration_timestamp.to_string(),
                    "Migration already occurred."
                )
            }
        }

        Ok(())
    }

    fn select_accounts(&self) -> Result<Vec<FiduciaryAccount>, ContextError> {
        let conn = self.get_connection()?;

        let statement = conn.prepare(Self::FIDUCIARY_ACCOUNT_SELECT_QUERY)?;

        statement
            .into_iter()
            .map(FiduciaryAccount::try_from)
            .collect()
    }

    fn select_account(&self, account_id: u8) -> Result<FiduciaryAccount, ContextError> {
        let conn = self.get_connection()?;

        let mut statement = conn.prepare(
            Self::FIDUCIARY_ACCOUNT_SELECT_QUERY.to_owned() + " WHERE fiduciary_account_id = ?",
        )?;
        statement.bind((1, i64::from(account_id)))?;

        statement
            .into_iter()
            .map(FiduciaryAccount::try_from)
            .last()
            .unwrap_or(Err(ContextError::ItemNotFoundError {}))
    }

    fn insert_account(&mut self, account: &FiduciaryAccount) -> Result<(), ContextError> {
        let conn = self.get_connection()?;

        let mut statement = conn.prepare(Self::FIDUCIARY_ACCOUNT_INSERT_QUERY)?;

        statement.bind::<&[(_, Value)]>(&[
            (":fiduciary_account_id", i64::from(account.fiduciary_account_id).into()),
            (":fiduciary_account_name", account.fiduciary_account_name.clone().into()),
            (":fiduciary_account_number", i64::try_from(account.fiduciary_account_number)?.into()),
            (":fiduciary_account_description", account.fiduciary_account_description.clone().into()),
            (":fiduciary_account_created_timestamp", account.fiduciary_account_created_timestamp.to_string().into()),
        ])?;
        Ok(())
    }

    fn select_symbols(&self) -> Result<Vec<EquitySymbol>, ContextError> {
        let conn = self.get_connection()?;

        let statement = conn.prepare(Self::EQUITY_SYMBOL_SELECT_QUERY)?;

        statement.into_iter().map(EquitySymbol::try_from).collect()
    }

    fn select_symbol(&self, symbol_id: usize) -> Result<EquitySymbol, ContextError> {
        let conn = self.get_connection()?;

        let mut statement = conn
            .prepare(Self::EQUITY_SYMBOL_SELECT_QUERY.to_owned() + " WHERE equity_symbol_id = ?")?;
        statement.bind((1, i64::try_from(symbol_id)?))?;

        statement
            .into_iter()
            .map(EquitySymbol::try_from)
            .last()
            .unwrap_or(Err(ContextError::ItemNotFoundError()))
    }

    fn insert_symbol(&mut self, symbol: &EquitySymbol) -> Result<(), ContextError> {
        let conn = self.get_connection()?;

        let mut statement = conn.prepare(Self::EQUITY_SYMBOL_INSERT_QUERY)?;

        statement.bind::<&[(_, Value)]>(&[
            (":equity_symbol_id", i64::try_from(symbol.equity_symbol_id)?.into()),
            (":equity_symbol_name", symbol.equity_symbol_name.clone().into()),
            (":equity_symbol_type", symbol.equity_symbol_type.to_string().into()),
            (":equity_symbol_managing_company", symbol.equity_symbol_managing_company.clone().into()),
            (":equity_symbol_created_timestamp", symbol.equity_symbol_created_timestamp.to_string().into()),
        ])?;
        Ok(())
    }

    fn select_events(&self) -> Result<Vec<EquityChangeEvent>, ContextError> {
        let conn = self.get_connection()?;

        let statement = conn.prepare(Self::EQUITY_CHANGE_EVENT_SELECT_QUERY)?;

        statement
            .into_iter()
            .map(EquityChangeEvent::try_from)
            .collect()
    }

    fn select_event(&self, event_id: usize) -> Result<EquityChangeEvent, ContextError> {
        let conn = self.get_connection()?;

        let mut statement = conn.prepare(
            Self::EQUITY_CHANGE_EVENT_SELECT_QUERY.to_owned() + " WHERE equity_change_event_id = ?",
        )?;
        statement.bind((1, i64::try_from(event_id)?))?;

        statement
            .into_iter()
            .map(EquityChangeEvent::try_from)
            .last()
            .unwrap_or(Err(ContextError::ItemNotFoundError()))
    }

    fn insert_event(&mut self, event: &EquityChangeEvent) -> Result<(), ContextError> {
        let conn = self.get_connection()?;

        let mut statement = conn.prepare(Self::EQUITY_SYMBOL_INSERT_QUERY)?;

        statement.bind::<&[(_, Value)]>(&[
            (":equity_change_event_id", i64::try_from(event.equity_change_event_id)?.into()),
            (":fiduciary_account_id", i64::from(event.fiduciary_account_id).into()),
            (":equity_symbol_id", i64::try_from(event.equity_symbol_id)?.into()),
            (":equity_change_quantity", event.equity_change_quantity.to_f64().into()),
            (":equity_change_cost_basis", event.equity_change_cost_basis.to_f64().into()),
            (":equity_change_type", event.equity_change_type.to_string().into()),
            (":equity_change_timestamp", event.equity_change_timestamp.to_string().into()),
        ])?;
        Ok(())
    }
}

impl Default for SQLiteStore {
    fn default() -> Self {
        SQLiteStore {
            location: ":memory:".into(),
            connection: Option::None,
            counter: None,
        }
    }
}

impl Store for SQLiteStore {
    /// Init SQLite for counting. Creates `RATAT` table with single row, with
    /// counter set to `0`
    fn init(&mut self) -> Result<(), ContextError> {
        let conn_path = self.location.clone();
        let conn = sqlite::open(conn_path)?;
        self.connection = Option::Some(conn);

        self.run_migrations()?;
        self.counter = Some(0_u8);

        Ok(())
    }

    /// Selects counter from `RATAT` table. Since sqlite doesn't scrutinize the
    /// value and we have a more restricted datatype (i64 -> u8), we use
    /// `try_from` which may produce an error result.
    fn get_counter(&self) -> Result<u8, ContextError> {
        self.counter.ok_or(ContextError::DataError(sqlite::Error {
            code: None,
            message: None,
        }))
    }

    /// Updates row with counter value on `RATAT` table
    fn set_counter(&mut self, new_count: u8) -> Result<u8, ContextError> {
        if self.counter.is_none() {
            Err(sqlite::Error {
                code: None,
                message: None,
            })?;
        }
        self.counter = Some(new_count);

        tracing::event!(tracing::Level::DEBUG, new_count, "Updated counter.");

        Ok(new_count)
    }
}

#[cfg(test)]
mod tests {

    use super::*;

    // Practically, the `Store` trait is what we interface with externally, and
    // provides the more general `Result` interface we want to work with so
    // that's what we test here

    #[test]
    fn default() {
        let context = SQLiteStore::default();

        assert_eq!(context.location, ":memory:");
        assert!(context.connection.is_none());
    }

    #[test]
    fn init() {
        let mut context = SQLiteStore::default();

        let res = context.init();

        assert!(res.is_ok());
        assert!(context.connection.is_some());
    }
}

#[cfg(test)]
mod counter_tests {
    use super::*;

    #[test]
    fn get_counter() {
        let mut context = SQLiteStore::default();
        assert!(context.init().is_ok());

        let res = context.get_counter();

        assert!(res.is_ok());
        assert_eq!(res.unwrap(), 0);
    }

    #[test]
    fn get_counter_error() {
        let context = SQLiteStore::default();
        // don't initialize

        let res = context.get_counter();

        assert!(res.is_err());
    }

    #[test]
    fn set_counter() {
        let mut context = SQLiteStore::default();
        assert!(context.init().is_ok());

        let expected = 240;
        let res = context.set_counter(expected);

        assert!(res.is_ok());
        assert_eq!(res.unwrap(), expected);
    }

    #[test]
    fn set_counter_error() {
        let mut context = SQLiteStore::default();
        // don't initialize

        let res = context.set_counter(1);

        assert!(res.is_err());
    }

    #[test]
    fn set_and_get() {
        let mut context = SQLiteStore::default();
        assert!(context.init().is_ok());

        let count_0 = context.get_counter().unwrap();

        assert_eq!(count_0, 0);

        let new_num = 244;
        let res = context.set_counter(new_num);

        assert!(res.is_ok());
        assert_eq!(res.unwrap(), new_num);

        let count_1 = context.get_counter().unwrap();

        assert_eq!(count_1, new_num);
        assert_ne!(count_0, count_1);
    }
}