greed 0.6.2

A rust tool to automate trades 📈
use crate::asset::AssetSymbol;
use crate::config::simple::SimpleConfig;
use crate::config::Config;
use crate::db;
use crate::error::GreedError;
use crate::platform;
use crate::platform::FinancialPlatform;
use crate::quote_accumulator::QuoteAccumulator;
use crate::strategy::factory::StrategyProviderFactory;
use crate::strategy::provider::StrategyRunnerProvider;
use crate::strategy::runner::{StrategyRunner, TacticStrategyRunner};
use duckdb::Connection;
use log::warn;
use std::ffi::OsStr;
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::time::sleep;

#[derive(Debug, Default, PartialEq)]
pub struct GreedRunnerArgs {
    pub config_path: PathBuf,
    pub db_path: Option<PathBuf>,
    pub is_simulated: bool,
    pub record_history: bool,
}

pub struct GreedRunner {
    run_interval: u64,
    config_strategy: Box<dyn StrategyRunner>,
    strategy_providers: Vec<Box<dyn StrategyRunnerProvider>>,
    config_assets: Vec<AssetSymbol>,
    platform: Arc<dyn FinancialPlatform>,
    connection: Arc<Mutex<Option<Connection>>>,
    quote_accumulator: QuoteAccumulator,
    db_path: PathBuf,
    record_history: bool,
}

impl GreedRunner {
    pub async fn new(
        config: Config,
        config_path: PathBuf,
        platform: Arc<dyn FinancialPlatform>,
        db_path: PathBuf,
        record_history: bool,
    ) -> Result<Self, GreedError> {
        let connection = Arc::new(Mutex::new(None));
        let quote_accumulator = QuoteAccumulator::new();
        let config_strategy = Box::new(TacticStrategyRunner::from_config(
            &config,
            &platform,
            &quote_accumulator,
        ));
        let factory =
            StrategyProviderFactory::new(&config, config_path, &platform, &quote_accumulator);
        let strategy_providers = factory.create_providers().await?;

        // Accumulate assets from all providers
        let mut config_assets = Vec::new();
        for provider in &strategy_providers {
            config_assets.extend(provider.config_assets());
        }
        // Add assets from config strategy
        config_assets.extend(config.tactics.iter().flat_map(|t| t.assets()));

        Ok(Self {
            run_interval: config.interval,
            config_strategy,
            strategy_providers,
            config_assets,
            platform,
            connection,
            quote_accumulator,
            db_path,
            record_history,
        })
    }

    pub async fn from_args(
        args: GreedRunnerArgs,
        db_path: PathBuf,
    ) -> Result<GreedRunner, GreedError> {
        let config_path = args.config_path.clone();
        let record_history = args.record_history;
        let config = GreedRunner::read_config(&args).await?;
        let platform = platform::for_type(&config.platform, args.into())?;
        Self::new(config, config_path, platform, db_path, record_history).await
    }

    async fn read_config(args: &GreedRunnerArgs) -> Result<Config, GreedError> {
        let ext = args.config_path.extension();
        if Some(OsStr::new("csv")) == ext {
            let simple_config = SimpleConfig::from_path(&args.config_path).await?;
            Ok(simple_config.into())
        } else {
            Config::from_path(&args.config_path).await
        }
    }

    pub async fn run_loop(&self) {
        let loop_interval = Duration::from_secs(self.run_interval);
        loop {
            self.run_tick(loop_interval).await;
            sleep(loop_interval).await;
        }
    }

    async fn run_tick(&self, loop_interval: Duration) {
        // Create strategy runners from providers
        let strategy_runners = match self.provide_strategy_runners().await {
            Ok(runners) => runners,
            Err(e) => {
                // If we fail we will sleep for a bit then try again
                warn!("error creating strategy runners: {}", e);
                return;
            }
        };
        // Open the DB for this tick's writes; held only for the writes below.
        self.open_tick_connection().await;
        // Run each strategy runner
        for runner in strategy_runners {
            let _ = runner.run(&self.config_assets).await;
            sleep(loop_interval).await;
        }
        // Run the general strategy from the top level config
        self.config_strategy.run(&self.config_assets).await;
        // Capture portfolio status history for this cycle.
        self.record_portfolio_status().await;
        // Flush this tick's deduped quotes in a single batch write.
        self.record_quotes();
        // Close the DB now that this tick's writes are done, releasing the lock.
        self.close_tick_connection();
    }

    async fn open_tick_connection(&self) {
        if !self.record_history {
            return;
        }
        match db::open_for_write_with_retry(&self.db_path).await {
            Ok(connection) => {
                *self.connection.lock().expect("connection slot lock") = Some(connection);
            }
            Err(e) => warn!("failed to open db for this tick's writes: {}", e),
        }
    }

    fn close_tick_connection(&self) {
        self.connection.lock().expect("connection slot lock").take();
    }

    async fn provide_strategy_runners(&self) -> Result<Vec<Box<dyn StrategyRunner>>, GreedError> {
        let mut runners = Vec::new();
        for provider in &self.strategy_providers {
            let runner = provider.provide_strategy_runner().await?;
            runners.push(runner);
        }
        Ok(runners)
    }

    // A DB hiccup must not stop trading automation, so failures are logged and swallowed here.
    // The insert is synchronous for now; switch to tokio::task::spawn_blocking if it proves slow.
    async fn record_portfolio_status(&self) {
        if !self.record_history {
            return;
        }
        match self.platform.account().await {
            Ok(account) => {
                let guard = self
                    .connection
                    .lock()
                    .expect("portfolio status connection lock");
                let Some(connection) = guard.as_ref() else {
                    return;
                };
                if let Err(e) = db::portfolio_status::insert(connection, &account) {
                    warn!("failed to write portfolio status: {}", e);
                }
            }
            Err(e) => warn!("failed to fetch account for portfolio status: {}", e),
        }
    }

    // A DB hiccup must not stop trading automation, so failures are logged and swallowed here.
    // The accumulator is drained unconditionally so it never grows unbounded across ticks,
    // even when `record_history` is disabled.
    fn record_quotes(&self) {
        let quotes = self.quote_accumulator.drain();
        if !self.record_history || quotes.is_empty() {
            return;
        }
        let mut guard = self.connection.lock().expect("stock quote connection lock");
        let Some(connection) = guard.as_mut() else {
            return;
        };
        if let Err(e) = db::stock_quotes::insert_many(connection, &quotes) {
            warn!("failed to write stock quotes: {}", e);
        }
    }
}

#[cfg(test)]
mod test {
    use crate::asset::AssetSymbol;
    use crate::config::Config;
    use crate::db::bootstrap_at;
    use crate::fixture;
    use crate::platform::account::Account;
    use crate::platform::quote::Quote;
    use crate::platform::MockPlatform;
    use crate::run::GreedRunner;
    use duckdb::Connection;
    use std::path::PathBuf;
    use std::time::Duration;
    use uuid::Uuid;

    #[tokio::test]
    async fn run_tick_writes_portfolio_status_row_per_tick() {
        let db_path = temp_db_path();
        bootstrap_at(&db_path).expect("bootstrap should succeed");
        let platform = MockPlatform::new().with_account(Account::fixture()).arc();
        let config = Config {
            interval: 1,
            ..Default::default()
        };
        let runner = GreedRunner::new(
            config,
            PathBuf::from("unused.toml"),
            platform,
            db_path.clone(),
            true,
        )
        .await
        .expect("runner should build");

        for _ in 0..10 {
            runner.run_tick(Duration::from_secs(1)).await;
        }

        assert_eq!(table_row_count(&db_path, "portfolio_status"), 10);
        std::fs::remove_file(&db_path).ok();
    }

    #[tokio::test]
    async fn run_tick_writes_stock_quote_row_per_tick() {
        let db_path = temp_db_path();
        bootstrap_at(&db_path).expect("bootstrap should succeed");
        let mut config = fixture::config("config_single_tactic.toml").await;
        config.interval = 0;
        let platform = MockPlatform::new()
            .with_account(Account::fixture())
            .with_quotes(vec![Quote::fixture(AssetSymbol::new("VTI"))])
            .arc();
        let runner = GreedRunner::new(
            config,
            PathBuf::from("unused.toml"),
            platform,
            db_path.clone(),
            true,
        )
        .await
        .expect("runner should build");

        for _ in 0..10 {
            runner.run_tick(Duration::from_secs(0)).await;
        }

        assert_eq!(table_row_count(&db_path, "portfolio_status"), 10);
        assert_eq!(table_row_count(&db_path, "stock_quotes"), 10);
        std::fs::remove_file(&db_path).ok();
    }

    #[tokio::test]
    async fn run_tick_dedupes_quotes_from_multiple_tactics_watching_same_symbol() {
        let db_path = temp_db_path();
        bootstrap_at(&db_path).expect("bootstrap should succeed");
        let mut config = fixture::config("config_multi_tactic.toml").await;
        config.interval = 0;
        // MockPlatform::latest_quotes ignores the requested symbols, so both of
        // this config's tactics fetch (and would, pre-fix, each write) this same quote.
        let platform = MockPlatform::new()
            .with_account(Account::fixture())
            .with_quotes(vec![Quote::fixture(AssetSymbol::new("VTI"))])
            .arc();
        let runner = GreedRunner::new(
            config,
            PathBuf::from("unused.toml"),
            platform,
            db_path.clone(),
            true,
        )
        .await
        .expect("runner should build");

        for _ in 0..10 {
            runner.run_tick(Duration::from_secs(0)).await;
        }

        assert_eq!(table_row_count(&db_path, "stock_quotes"), 10);
        std::fs::remove_file(&db_path).ok();
    }

    #[tokio::test]
    async fn run_tick_skips_writes_when_record_history_disabled() {
        let db_path = temp_db_path();
        bootstrap_at(&db_path).expect("bootstrap should succeed");
        let mut config = fixture::config("config_strategic.toml").await;
        config.interval = 0;
        let platform = MockPlatform::new()
            .with_account(Account::fixture())
            .with_quotes(vec![Quote::fixture(AssetSymbol::new("VTI"))])
            .arc();
        let runner = GreedRunner::new(
            config,
            fixture::path("config_strategic.toml"),
            platform,
            db_path.clone(),
            false,
        )
        .await
        .expect("runner should build");

        for _ in 0..10 {
            runner.run_tick(Duration::from_secs(0)).await;
        }

        assert_eq!(table_row_count(&db_path, "portfolio_status"), 0);
        assert_eq!(table_row_count(&db_path, "stock_quotes"), 0);
        std::fs::remove_file(&db_path).ok();
    }

    fn table_row_count(db_path: &PathBuf, table: &str) -> i64 {
        let connection = Connection::open(db_path).expect("db should be queryable");
        connection
            .query_row(&format!("SELECT COUNT(*) FROM {table}"), [], |row| {
                row.get(0)
            })
            .expect("table should be queryable")
    }

    fn temp_db_path() -> PathBuf {
        std::env::temp_dir().join(format!("greed_test_{}.duckdb", Uuid::new_v4()))
    }
}