mod api; mod cli; mod db; mod ingest; mod types; use std::path::Path; use std::sync::Arc; use std::time::Instant; use axum::serve; use clap::Parser; use tokio::net::TcpListener; use tracing_subscriber::EnvFilter; struct LoggerGuard { _guard: Option, } fn init_logging(log_file_path: Option<&str>) -> LoggerGuard { let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); let is_json = std::env::var("HAGFISH_LOG_FORMAT").as_deref() == Ok("json"); match log_file_path { Some(path) => { let file = std::fs::File::create(path).expect("Failed to create log file"); let (non_blocking, guard) = tracing_appender::non_blocking(file); let subscriber = tracing_subscriber::fmt() .with_env_filter(filter.clone()) .with_writer(non_blocking); if is_json { subscriber.json().init(); } else { subscriber.init(); } LoggerGuard { _guard: Some(guard), } } None => { let subscriber = tracing_subscriber::fmt().with_env_filter(filter.clone()); if is_json { subscriber.json().init(); } else { subscriber.init(); } LoggerGuard { _guard: None } } } } fn cleanup_stale_parquet_files() { let tmp_dir = std::env::temp_dir(); match std::fs::read_dir(&tmp_dir) { Ok(entries) => { for entry in entries.filter_map(|e| e.ok()) { if let Some(name) = entry.file_name().to_str() { if name.starts_with("hagfish-") && name.ends_with(".parquet") { if let Ok(metadata) = entry.metadata() { if let Ok(modified) = metadata.modified() { let age = std::time::SystemTime::now().duration_since(modified); if age.map(|d| d.as_secs() > 3600).unwrap_or(false) { let _ = std::fs::remove_file(entry.path()); } } } } } } } Err(e) => tracing::warn!(error = %e, "Failed to list temp dir"), } } #[tokio::main] async fn main() -> anyhow::Result<()> { let cli = cli::Cli::parse(); let config = match &cli.config { Some(path) => types::Config::load(path)?, None => { let default_path = std::path::Path::new("config.json"); if default_path.exists() { types::Config::load(default_path)? } else { tracing::warn!("No config file found, using defaults"); types::Config::default() } } }; // Initialize logging AFTER config is loaded so we have log_file_path let _logger_guard = init_logging(config.log_file_path.as_deref()); match cli.command { cli::Command::Serve => run_serve(&config).await, cli::Command::Ingest { full } => run_ingest(&config, full).await, cli::Command::Export { out } => run_export(&config, &out), } } async fn run_serve(config: &types::Config) -> anyhow::Result<()> { cleanup_stale_parquet_files(); let conn = db::init(&config.duckdb_path)?; let state = api::AppState { conn: Arc::new(tokio::sync::Mutex::new(conn)), config: Arc::new(config.clone()), }; let app = api::build_router(state); let listener = TcpListener::bind(&config.bind_address).await?; tracing::info!(address = %config.bind_address, "hagfish server started"); serve(listener, app).await?; Ok(()) } async fn run_ingest(config: &types::Config, full: bool) -> anyhow::Result<()> { let run_start = Instant::now(); tracing::info!( mode = if full { "full" } else { "incremental" }, "starting ingestion run" ); let conn = db::init(&config.duckdb_path)?; let client = reqwest::Client::new(); let (lookup_maps, metadata) = ingest::fetch_metadata(&client, &config.data_source_url) .await .inspect_err(|e| { tracing::error!(error = %e, stage = "fetch_metadata", "ingestion failed"); })?; db::update_lookups(&conn, &lookup_maps)?; let all_months = ingest::extract_available_months(&metadata); let months_to_fetch: Vec = if full { all_months } else { let last = db::get_last_month(&conn)?; match last { Some(last_month) => all_months .into_iter() .filter(|m| m.as_str() > last_month.as_str()) .collect(), None => all_months, } }; if months_to_fetch.is_empty() { let elapsed = run_start.elapsed(); tracing::info!( mode = if full { "full" } else { "incremental" }, months_requested = 0, total_rows = 0, batches = 0, errors = 0, elapsed_ms = elapsed.as_millis() as u64, "ingestion complete — no new months" ); return Ok(()); } const BATCH_SIZE: usize = 12; let batches = ingest::chunk_months(&months_to_fetch, BATCH_SIZE); tracing::info!( mode = if full { "full" } else { "incremental" }, months_requested = months_to_fetch.len(), batches = batches.len(), "beginning batched ingestion" ); let mut total_rows = 0; let mut errors = 0; for (i, batch) in batches.iter().enumerate() { let batch_start = Instant::now(); tracing::info!( batch = i + 1, total_batches = batches.len(), months_in_batch = batch.len(), months = %batch.join(", "), "fetching batch" ); let query = ingest::build_query(batch)?; let data = match ingest::fetch_data(&client, &config.data_source_url, &query).await { Ok(d) => d, Err(e) => { tracing::error!( batch = i + 1, error = %e, stage = "fetch_data", "batch failed, continuing to next batch" ); errors += 1; continue; } }; let mut landings = Vec::with_capacity(data.dataset.value.len()); let mut parse_errors = 0; for j in 0..data.dataset.value.len() { match ingest::parse_row(j, &data.dataset, &lookup_maps) { Ok(row) => { landings.push(ingest::data_row_to_landing(&row)); } Err(e) => { tracing::warn!( batch = i + 1, row_index = j, error = %e, "failed to parse row" ); parse_errors += 1; errors += 1; } } } let count = match db::upsert_landings(&conn, &landings) { Ok(c) => c, Err(e) => { tracing::error!( batch = i + 1, error = %e, stage = "upsert_landings", "batch upsert failed" ); errors += 1; continue; } }; total_rows += count; let batch_elapsed = batch_start.elapsed(); tracing::info!( batch = i + 1, total_batches = batches.len(), rows_inserted = count, parse_errors = parse_errors, elapsed_ms = batch_elapsed.as_millis() as u64, "batch complete" ); // Backoff between batches to avoid rate limiting if i < batches.len() - 1 { tokio::time::sleep(std::time::Duration::from_secs(1)).await; } } let elapsed = run_start.elapsed(); tracing::info!( mode = if full { "full" } else { "incremental" }, months_requested = months_to_fetch.len(), batches = batches.len(), total_rows = total_rows, errors = errors, elapsed_ms = elapsed.as_millis() as u64, "ingestion run complete" ); Ok(()) } fn run_export(config: &types::Config, out: &Path) -> anyhow::Result<()> { let run_start = Instant::now(); let conn = db::init(&config.duckdb_path)?; let path_str = out .to_str() .ok_or_else(|| anyhow::anyhow!("Output path contains invalid UTF-8"))?; db::export_parquet(&conn, path_str)?; let elapsed = run_start.elapsed(); tracing::info!( path = path_str, elapsed_ms = elapsed.as_millis() as u64, "export complete" ); Ok(()) }