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};

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>,
}

/// 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()),
            }),
        }
    }

    /// Attempts to get the migration list from the database; if that fails, creates the migration table and returns an empty vec
    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::<&[(_, sqlite::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(())
    }

}

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

impl Store for SQLiteStore {
    /// Initialize SQLite database connection and execute any pending migrations
    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()?;

        Ok(())
    }

    /// Selects accounts from the fiduciary_account table
    fn get_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()
    }

    /// Inserts an account into the fiduciary_account table
    fn add_account(&mut self, account: &FiduciaryAccount) -> Result<(), ContextError> {
        let conn = self.get_connection()?;

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

        statement.bind::<&[(_, sqlite::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(())
    }

    /// Selects symbols from the equity_symbol table
    fn get_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()
    }

    /// Inserts a symbol into the equity_symbol table
    fn add_symbol(&mut self, symbol: &EquitySymbol) -> Result<(), ContextError> {
        let conn = self.get_connection()?;

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

        statement.bind::<&[(_, sqlite::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(())
    }

    /// Selects events from the equity_change_event table
    fn get_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()
    }

    /// Inserts an event into the equity_change_event table
    fn add_event(&mut self, event: &EquityChangeEvent) -> Result<(), ContextError> {
        let conn = self.get_connection()?;

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

        statement.bind::<&[(_, sqlite::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(())
    }
}

#[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());
    }
}