303 lines
8.8 KiB
Rust
303 lines
8.8 KiB
Rust
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<tracing_appender::non_blocking::WorkerGuard>,
|
|
}
|
|
|
|
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<String> = 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(())
|
|
}
|