1058 lines
37 KiB
Rust
1058 lines
37 KiB
Rust
// src/corporate/update_companies_enrich.rs - MERGED VERSION WITH GENERIC ENRICHMENT
|
|
use super::{types::*, helpers::*};
|
|
use crate::config::Config;
|
|
use crate::corporate::checkpoint_helpers;
|
|
use crate::util::directories::DataPaths;
|
|
use crate::util::integrity::{StateManager, directory_reference, DataStage};
|
|
use crate::util::logger;
|
|
use crate::scraper::yahoo::{YahooClientPool, QuoteSummaryModule};
|
|
|
|
use std::result::Result::Ok;
|
|
use chrono::{TimeZone, Utc};
|
|
use std::collections::{HashSet};
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
|
|
use tokio::fs::{OpenOptions};
|
|
use tokio::io::{AsyncWriteExt};
|
|
use futures::stream::{FuturesUnordered, StreamExt};
|
|
use serde_json::json;
|
|
use tokio::sync::mpsc;
|
|
use std::future::Future;
|
|
use std::pin::Pin;
|
|
|
|
/// Log command enum (shared across all enrichment types)
|
|
enum LogCommand {
|
|
Write(serde_json::Value),
|
|
Checkpoint,
|
|
Shutdown,
|
|
}
|
|
|
|
/// Type alias for enrichment function
|
|
type EnrichmentFn = Arc<
|
|
dyn Fn(CompanyData, Arc<YahooClientPool>, DataPaths)
|
|
-> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send>>
|
|
+ Send
|
|
+ Sync
|
|
>;
|
|
|
|
// ============================================================================
|
|
// EVENTS ENRICHMENT
|
|
// ============================================================================
|
|
|
|
/// Yahoo Event enrichment per corporate company
|
|
///
|
|
/// # Features
|
|
/// - Graceful shutdown (abort-safe)
|
|
/// - Task panic isolation (tasks fail independently)
|
|
/// - Crash-safe persistence (checkpoint + log with fsync)
|
|
/// - Smart skip logic (only process incomplete data)
|
|
/// - Uses pending queue instead of retry mechanism
|
|
/// - Content integrity validation with hash tracking
|
|
///
|
|
/// # Persistence Strategy
|
|
/// - Checkpoint: companies_yahoo_cleaned.jsonl (atomic state)
|
|
/// - Log: companies_events_updates.log (append-only updates)
|
|
/// - On restart: Load checkpoint + replay log
|
|
/// - Periodic checkpoints (every 50 companies)
|
|
/// - Batched fsync (every 10 writes or 10 seconds)
|
|
/// - Hash validation of all event data directories
|
|
pub async fn enrich_companies_with_events(
|
|
paths: &DataPaths,
|
|
_config: &Config,
|
|
yahoo_pool: Arc<YahooClientPool>,
|
|
shutdown_flag: &Arc<AtomicBool>,
|
|
) -> anyhow::Result<usize> {
|
|
// Configuration constants
|
|
const CHECKPOINT_INTERVAL: usize = 50;
|
|
const FSYNC_BATCH_SIZE: usize = 10;
|
|
const FSYNC_INTERVAL_SECS: u64 = 10;
|
|
const CONCURRENCY_LIMIT: usize = 50;
|
|
|
|
let data_path = paths.data_dir();
|
|
|
|
// File paths
|
|
let input_path = data_path.join("companies_yahoo_cleaned.jsonl");
|
|
let log_path = data_path.join("companies_events_updates.log");
|
|
|
|
// Check input exists
|
|
if !input_path.exists() {
|
|
logger::log_warn(" companies_yahoo_cleaned.jsonl not found, skipping event enrichment").await;
|
|
return Ok(0);
|
|
}
|
|
|
|
let manager = StateManager::new(paths.integrity_dir()).await?;
|
|
let step_name = "yahoo_events_enrichment_complete";
|
|
|
|
if manager.is_step_valid(step_name).await? {
|
|
logger::log_info(" Yahoo events enrichment already completed and valid").await;
|
|
|
|
// Count enriched companies
|
|
let count = checkpoint_helpers::count_enriched_companies(paths, "events").await?;
|
|
logger::log_info(&format!(" ✓ Found {} companies with valid event data", count)).await;
|
|
return Ok(count);
|
|
}
|
|
|
|
logger::log_info(" Event data needs refresh - starting enrichment").await;
|
|
|
|
// === RECOVERY PHASE: Track enriched companies ===
|
|
let enriched_companies: HashSet<String> = checkpoint_helpers::load_enrichment_progress(&log_path).await?;
|
|
|
|
// Load all companies from input
|
|
logger::log_info("Loading companies from companies_yahoo_cleaned.jsonl...").await;
|
|
let companies = load_companies_from_jsonl(&input_path).await?;
|
|
let total_companies = companies.len();
|
|
logger::log_info(&format!("Found {} companies to process", total_companies)).await;
|
|
|
|
// Filter companies that need enrichment
|
|
let pending_companies: Vec<CompanyData> = companies
|
|
.into_iter()
|
|
.filter(|company| !enriched_companies.contains(&company.name))
|
|
.collect();
|
|
|
|
let pending_count = pending_companies.len();
|
|
logger::log_info(&format!(
|
|
" {} already enriched, {} pending",
|
|
enriched_companies.len(),
|
|
pending_count
|
|
)).await;
|
|
|
|
if pending_count == 0 {
|
|
logger::log_info(" ✓ All companies already enriched").await;
|
|
track_events_completion(&manager, paths, step_name).await?;
|
|
return Ok(enriched_companies.len());
|
|
}
|
|
|
|
// === PROCESSING PHASE: Enrich companies with events ===
|
|
|
|
// Create enrichment function
|
|
let enrichment_fn: EnrichmentFn = Arc::new(move |company, pool, paths| {
|
|
let company = company.clone();
|
|
let pool = Arc::clone(&pool);
|
|
let paths = paths.clone();
|
|
Box::pin(async move {
|
|
enrich_company_with_events(&company, &pool, &paths).await
|
|
})
|
|
});
|
|
|
|
// Shared counters
|
|
let processed_count = Arc::new(AtomicUsize::new(enriched_companies.len()));
|
|
let success_count = Arc::new(AtomicUsize::new(enriched_companies.len()));
|
|
let failed_count = Arc::new(AtomicUsize::new(0));
|
|
|
|
// Log writer channel with batching and fsync
|
|
let (log_tx, log_rx) = mpsc::channel::<LogCommand>(1000);
|
|
|
|
// Spawn log writer task
|
|
let log_writer_handle = spawn_log_writer(
|
|
log_path,
|
|
log_rx,
|
|
Arc::clone(&processed_count),
|
|
total_companies,
|
|
FSYNC_BATCH_SIZE,
|
|
FSYNC_INTERVAL_SECS,
|
|
);
|
|
|
|
// Process companies concurrently with task panic isolation
|
|
let mut tasks = FuturesUnordered::new();
|
|
let mut pending_iter = pending_companies.into_iter();
|
|
let semaphore = Arc::new(tokio::sync::Semaphore::new(CONCURRENCY_LIMIT));
|
|
|
|
// Initial batch of tasks
|
|
for _ in 0..CONCURRENCY_LIMIT.min(pending_count) {
|
|
if let Some(company) = pending_iter.next() {
|
|
let task = spawn_enrichment_task(
|
|
company,
|
|
Arc::clone(&yahoo_pool),
|
|
paths.clone(),
|
|
Arc::clone(&processed_count),
|
|
Arc::clone(&success_count),
|
|
Arc::clone(&failed_count),
|
|
log_tx.clone(),
|
|
Arc::clone(&semaphore),
|
|
Arc::clone(shutdown_flag),
|
|
Arc::clone(&enrichment_fn),
|
|
);
|
|
tasks.push(task);
|
|
}
|
|
}
|
|
|
|
// Process results and spawn new tasks
|
|
let mut checkpoint_counter = enriched_companies.len();
|
|
|
|
while let Some(result) = tasks.next().await {
|
|
// Handle task result (even if panicked)
|
|
match result {
|
|
Ok(_) => {
|
|
// Task completed successfully
|
|
}
|
|
Err(e) => {
|
|
logger::log_warn(&format!("Task panicked: {}", e)).await;
|
|
failed_count.fetch_add(1, Ordering::SeqCst);
|
|
}
|
|
}
|
|
|
|
// Check for shutdown
|
|
if shutdown_flag.load(Ordering::SeqCst) {
|
|
logger::log_warn("Shutdown signal received, stopping event enrichment").await;
|
|
break;
|
|
}
|
|
|
|
// Checkpoint periodically
|
|
checkpoint_counter += 1;
|
|
if checkpoint_counter % CHECKPOINT_INTERVAL == 0 {
|
|
let _ = log_tx.send(LogCommand::Checkpoint).await;
|
|
}
|
|
|
|
// Spawn next task if available
|
|
if let Some(company) = pending_iter.next() {
|
|
let task = spawn_enrichment_task(
|
|
company,
|
|
Arc::clone(&yahoo_pool),
|
|
paths.clone(),
|
|
Arc::clone(&processed_count),
|
|
Arc::clone(&success_count),
|
|
Arc::clone(&failed_count),
|
|
log_tx.clone(),
|
|
Arc::clone(&semaphore),
|
|
Arc::clone(shutdown_flag),
|
|
Arc::clone(&enrichment_fn),
|
|
);
|
|
tasks.push(task);
|
|
}
|
|
}
|
|
|
|
// Signal log writer to shutdown
|
|
let _ = log_tx.send(LogCommand::Shutdown).await;
|
|
drop(log_tx);
|
|
|
|
// Wait for log writer to finish
|
|
let _ = log_writer_handle.await;
|
|
|
|
let final_processed = processed_count.load(Ordering::SeqCst);
|
|
let final_success = success_count.load(Ordering::SeqCst);
|
|
let final_failed = failed_count.load(Ordering::SeqCst);
|
|
|
|
logger::log_info(&format!(
|
|
" Event enrichment summary: {} total, {} success, {} failed",
|
|
final_processed, final_success, final_failed
|
|
)).await;
|
|
|
|
// Mark as complete if all companies processed
|
|
if final_processed >= total_companies && !shutdown_flag.load(Ordering::SeqCst) {
|
|
track_events_completion(&manager, paths, step_name).await?;
|
|
logger::log_info(" ✓ Event enrichment marked as complete with integrity tracking").await;
|
|
}
|
|
|
|
Ok(final_success)
|
|
}
|
|
|
|
/// Track event enrichment completion with content hash verification
|
|
async fn track_events_completion(
|
|
manager: &StateManager,
|
|
paths: &DataPaths,
|
|
step_name: &str,
|
|
) -> anyhow::Result<()> {
|
|
// Create content reference for all event data
|
|
// This will hash ALL files matching the pattern: {company}/events/data.jsonl
|
|
let content_reference = directory_reference(
|
|
paths.corporate_dir(),
|
|
Some(vec![
|
|
"*/events/*.jsonl".to_string(), // Main pattern for events data
|
|
"*/events/data.jsonl".to_string(), // Specific pattern (more precise)
|
|
]),
|
|
Some(vec![
|
|
"*.log".to_string(), // Exclude log files
|
|
"*.tmp".to_string(), // Exclude temp files
|
|
"*.bak".to_string(), // Exclude backup files
|
|
]),
|
|
);
|
|
|
|
// Track completion with:
|
|
// - Content reference: All event directories
|
|
// - Data stage: Data (7-day TTL by default)
|
|
// - Dependencies: Depends on cleaned companies data
|
|
manager.update_entry(
|
|
step_name.to_string(),
|
|
content_reference,
|
|
DataStage::Data,
|
|
None, // Use default TTL (7 days for Data stage)
|
|
).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Enrich a single company with event data
|
|
async fn enrich_company_with_events(
|
|
company: &CompanyData,
|
|
yahoo_pool: &Arc<YahooClientPool>,
|
|
paths: &DataPaths,
|
|
) -> anyhow::Result<()> {
|
|
use std::collections::HashMap;
|
|
|
|
let ticker = match extract_first_yahoo_ticker(company) {
|
|
Some(t) => t,
|
|
None => {
|
|
return Err(anyhow::anyhow!("No valid Yahoo ticker found"));
|
|
}
|
|
};
|
|
|
|
// Combined summary to accumulate data from all available modules
|
|
let mut combined_modules: HashMap<String, serde_json::Value> = HashMap::new();
|
|
let timestamp = chrono::Utc::now().timestamp();
|
|
|
|
// Try each event module individually
|
|
let event_modules = QuoteSummaryModule::event_modules();
|
|
|
|
for module in event_modules {
|
|
match yahoo_pool.get_quote_summary(&ticker, &[module]).await {
|
|
Ok(summary) => {
|
|
// Merge this module's data into combined summary
|
|
for (key, value) in summary.modules {
|
|
combined_modules.insert(key, value);
|
|
}
|
|
}
|
|
Err(e) => {
|
|
// Module not available - silently continue for expected errors
|
|
let err_str = e.to_string();
|
|
if err_str.contains("500") || err_str.contains("404") || err_str.contains("Not Found") {
|
|
// Expected for securities without this data - continue silently
|
|
continue;
|
|
} else {
|
|
// Unexpected error - log but continue trying other modules
|
|
logger::log_warn(&format!(
|
|
" Unexpected error fetching event module for {}: {}",
|
|
ticker, e
|
|
)).await;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Only save if we got at least some data
|
|
if combined_modules.is_empty() {
|
|
return Err(anyhow::anyhow!("No event data available for any module"));
|
|
}
|
|
|
|
// Create combined summary with all available modules
|
|
let combined_summary = crate::scraper::yahoo::QuoteSummary {
|
|
symbol: ticker.clone(),
|
|
modules: combined_modules,
|
|
timestamp,
|
|
};
|
|
|
|
// Save the combined event data
|
|
save_company_event_data(paths, &company.name, &combined_summary).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Save event data to company directory
|
|
async fn save_company_event_data(
|
|
paths: &DataPaths,
|
|
company_name: &str,
|
|
summary: &crate::scraper::yahoo::QuoteSummary,
|
|
) -> anyhow::Result<()> {
|
|
use tokio::fs;
|
|
|
|
let safe_name = sanitize_company_name(company_name);
|
|
|
|
let company_dir = paths.corporate_dir().join(&safe_name).join("events");
|
|
fs::create_dir_all(&company_dir).await?;
|
|
|
|
let data_path = company_dir.join("data.jsonl");
|
|
let json_line = serde_json::to_string(summary)?;
|
|
|
|
let mut file = fs::File::create(&data_path).await?;
|
|
file.write_all(json_line.as_bytes()).await?;
|
|
file.write_all(b"\n").await?;
|
|
file.flush().await?;
|
|
file.sync_all().await?; // Ensure data is persisted
|
|
|
|
Ok(())
|
|
}
|
|
|
|
// ============================================================================
|
|
// OPTION ENRICHMENT
|
|
// ============================================================================
|
|
|
|
/// Yahoo Option enrichment per corporate company
|
|
///
|
|
/// # Features
|
|
/// - Graceful shutdown (abort-safe)
|
|
/// - Task panic isolation (tasks fail independently)
|
|
/// - Crash-safe persistence (checkpoint + log with fsync)
|
|
/// - Smart skip logic (only process incomplete data)
|
|
/// - Uses pending queue instead of retry mechanism
|
|
/// - Content integrity validation with hash tracking
|
|
///
|
|
/// # Persistence Strategy
|
|
/// - Checkpoint: companies_yahoo_cleaned.jsonl (atomic state)
|
|
/// - Log: companies_option_updates.log (append-only updates)
|
|
/// - On restart: Load checkpoint + replay log
|
|
/// - Periodic checkpoints (every 50 companies)
|
|
/// - Batched fsync (every 10 writes or 10 seconds)
|
|
/// - Hash validation of all option data directories
|
|
pub async fn enrich_companies_with_option(
|
|
paths: &DataPaths,
|
|
_config: &Config,
|
|
yahoo_pool: Arc<YahooClientPool>,
|
|
shutdown_flag: &Arc<AtomicBool>,
|
|
) -> anyhow::Result<usize> {
|
|
// Configuration constants
|
|
const CHECKPOINT_INTERVAL: usize = 50;
|
|
const FSYNC_BATCH_SIZE: usize = 10;
|
|
const FSYNC_INTERVAL_SECS: u64 = 10;
|
|
const CONCURRENCY_LIMIT: usize = 50;
|
|
|
|
let data_path = paths.data_dir();
|
|
|
|
// File paths
|
|
let input_path = data_path.join("companies_yahoo_cleaned.jsonl");
|
|
let log_path = data_path.join("companies_option_updates.log");
|
|
|
|
// Check input exists
|
|
if !input_path.exists() {
|
|
logger::log_warn(" companies_yahoo_cleaned.jsonl not found, skipping option enrichment").await;
|
|
return Ok(0);
|
|
}
|
|
|
|
let manager = StateManager::new(paths.integrity_dir()).await?;
|
|
let step_name = "yahoo_option_enrichment_complete";
|
|
|
|
if manager.is_step_valid(step_name).await? {
|
|
logger::log_info(" Yahoo option enrichment already completed and valid").await;
|
|
let count = checkpoint_helpers::count_enriched_companies(paths, "option").await?;
|
|
logger::log_info(&format!(" ✓ Found {} companies with valid option data", count)).await;
|
|
return Ok(count);
|
|
}
|
|
|
|
logger::log_info(" Option data needs refresh - starting enrichment").await;
|
|
|
|
// === RECOVERY PHASE: Track enriched companies ===
|
|
let enriched_companies = checkpoint_helpers::load_enrichment_progress(&log_path).await?;
|
|
|
|
// Load all companies from input
|
|
logger::log_info("Loading companies from companies_yahoo_cleaned.jsonl...").await;
|
|
let companies = load_companies_from_jsonl(&input_path).await?;
|
|
let total_companies = companies.len();
|
|
logger::log_info(&format!("Found {} companies to process", total_companies)).await;
|
|
|
|
// Filter companies that need enrichment
|
|
let pending_companies: Vec<CompanyData> = companies
|
|
.into_iter()
|
|
.filter(|company| !enriched_companies.contains(&company.name))
|
|
.collect();
|
|
|
|
let pending_count = pending_companies.len();
|
|
logger::log_info(&format!(
|
|
" {} already enriched, {} pending",
|
|
enriched_companies.len(),
|
|
pending_count
|
|
)).await;
|
|
|
|
if pending_count == 0 {
|
|
logger::log_info(" ✓ All companies already enriched").await;
|
|
track_option_completion(&manager, paths, step_name).await?;
|
|
return Ok(enriched_companies.len());
|
|
}
|
|
|
|
// === PROCESSING PHASE: Enrich companies with option ===
|
|
|
|
// Create enrichment function
|
|
let enrichment_fn: EnrichmentFn = Arc::new(move |company, pool, paths| {
|
|
let company = company.clone();
|
|
let pool = Arc::clone(&pool);
|
|
let paths = paths.clone();
|
|
Box::pin(async move {
|
|
enrich_company_with_option(&company, &pool, &paths).await
|
|
})
|
|
});
|
|
|
|
// Shared counters
|
|
let processed_count = Arc::new(AtomicUsize::new(enriched_companies.len()));
|
|
let success_count = Arc::new(AtomicUsize::new(enriched_companies.len()));
|
|
let failed_count = Arc::new(AtomicUsize::new(0));
|
|
|
|
// Log writer channel with batching and fsync
|
|
let (log_tx, log_rx) = mpsc::channel::<LogCommand>(1000);
|
|
|
|
// Spawn log writer task
|
|
let log_writer_handle = spawn_log_writer(
|
|
log_path,
|
|
log_rx,
|
|
Arc::clone(&processed_count),
|
|
total_companies,
|
|
FSYNC_BATCH_SIZE,
|
|
FSYNC_INTERVAL_SECS,
|
|
);
|
|
|
|
// Process companies concurrently with task panic isolation
|
|
let mut tasks = FuturesUnordered::new();
|
|
let mut pending_iter = pending_companies.into_iter();
|
|
let semaphore = Arc::new(tokio::sync::Semaphore::new(CONCURRENCY_LIMIT));
|
|
|
|
// Initial batch of tasks
|
|
for _ in 0..CONCURRENCY_LIMIT.min(pending_count) {
|
|
if let Some(company) = pending_iter.next() {
|
|
let task = spawn_enrichment_task(
|
|
company,
|
|
Arc::clone(&yahoo_pool),
|
|
paths.clone(),
|
|
Arc::clone(&processed_count),
|
|
Arc::clone(&success_count),
|
|
Arc::clone(&failed_count),
|
|
log_tx.clone(),
|
|
Arc::clone(&semaphore),
|
|
Arc::clone(shutdown_flag),
|
|
Arc::clone(&enrichment_fn),
|
|
);
|
|
tasks.push(task);
|
|
}
|
|
}
|
|
|
|
// Process tasks as they complete and spawn new ones
|
|
let mut checkpoint_counter = 0;
|
|
while let Some(_result) = tasks.next().await {
|
|
// Check for shutdown
|
|
if shutdown_flag.load(Ordering::SeqCst) {
|
|
logger::log_warn("Shutdown signal received, stopping option enrichment").await;
|
|
break;
|
|
}
|
|
|
|
// Checkpoint periodically
|
|
checkpoint_counter += 1;
|
|
if checkpoint_counter % CHECKPOINT_INTERVAL == 0 {
|
|
let _ = log_tx.send(LogCommand::Checkpoint).await;
|
|
}
|
|
|
|
// Spawn next task if available
|
|
if let Some(company) = pending_iter.next() {
|
|
let task = spawn_enrichment_task(
|
|
company,
|
|
Arc::clone(&yahoo_pool),
|
|
paths.clone(),
|
|
Arc::clone(&processed_count),
|
|
Arc::clone(&success_count),
|
|
Arc::clone(&failed_count),
|
|
log_tx.clone(),
|
|
Arc::clone(&semaphore),
|
|
Arc::clone(shutdown_flag),
|
|
Arc::clone(&enrichment_fn),
|
|
);
|
|
tasks.push(task);
|
|
}
|
|
}
|
|
|
|
// Signal log writer to shutdown
|
|
let _ = log_tx.send(LogCommand::Shutdown).await;
|
|
drop(log_tx);
|
|
|
|
// Wait for log writer to finish
|
|
let _ = log_writer_handle.await;
|
|
|
|
let final_processed = processed_count.load(Ordering::SeqCst);
|
|
let final_success = success_count.load(Ordering::SeqCst);
|
|
let final_failed = failed_count.load(Ordering::SeqCst);
|
|
|
|
logger::log_info(&format!(
|
|
" Option enrichment summary: {} total, {} success, {} failed",
|
|
final_processed, final_success, final_failed
|
|
)).await;
|
|
|
|
// Mark as complete if all companies processed
|
|
if final_processed >= total_companies && !shutdown_flag.load(Ordering::SeqCst) {
|
|
track_option_completion(&manager, paths, step_name).await?;
|
|
logger::log_info(" ✓ Option enrichment marked as complete with integrity tracking").await;
|
|
}
|
|
|
|
Ok(final_success)
|
|
}
|
|
|
|
/// Track option enrichment completion with content hash verification
|
|
async fn track_option_completion(
|
|
manager: &StateManager,
|
|
paths: &DataPaths,
|
|
step_name: &str,
|
|
) -> anyhow::Result<()> {
|
|
// Create content reference for all option data
|
|
// This will hash ALL files matching the pattern: {company}/option/data.jsonl
|
|
let content_reference = directory_reference(
|
|
paths.corporate_dir(),
|
|
Some(vec![
|
|
"*/option/*.jsonl".to_string(), // Main pattern for option data
|
|
"*/option/data.jsonl".to_string(), // Specific pattern (more precise)
|
|
]),
|
|
Some(vec![
|
|
"*.log".to_string(), // Exclude log files
|
|
"*.tmp".to_string(), // Exclude temp files
|
|
"*.bak".to_string(), // Exclude backup files
|
|
]),
|
|
);
|
|
|
|
// Track completion with:
|
|
// - Content reference: All option directories
|
|
// - Data stage: Data (7-day TTL by default)
|
|
// - Dependencies: Depends on cleaned companies data
|
|
manager.update_entry(
|
|
step_name.to_string(),
|
|
content_reference,
|
|
DataStage::Data,
|
|
None, // Use default TTL (7 days for Data stage)
|
|
).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Enrich a single company with option data
|
|
async fn enrich_company_with_option(
|
|
company: &CompanyData,
|
|
yahoo_pool: &Arc<YahooClientPool>,
|
|
paths: &DataPaths,
|
|
) -> anyhow::Result<()> {
|
|
let ticker = match extract_first_yahoo_ticker(company) {
|
|
Some(t) => t,
|
|
None => {
|
|
return Err(anyhow::anyhow!("No valid Yahoo ticker found"));
|
|
}
|
|
};
|
|
|
|
// Get option data for all available expiration dates
|
|
let option_data = yahoo_pool.get_option_data(&ticker, None).await?;
|
|
|
|
// Only save if we got meaningful data
|
|
if option_data.option.is_empty() {
|
|
return Err(anyhow::anyhow!("No option data available"));
|
|
}
|
|
|
|
// Save the option data
|
|
save_company_data(paths, &company.name, &option_data, "option").await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
// ============================================================================
|
|
// CHART ENRICHMENT
|
|
// ============================================================================
|
|
|
|
/// Yahoo Chart enrichment per corporate company
|
|
///
|
|
/// # Features
|
|
/// - Graceful shutdown (abort-safe)
|
|
/// - Task panic isolation (tasks fail independently)
|
|
/// - Crash-safe persistence (checkpoint + log with fsync)
|
|
/// - Smart skip logic (only process incomplete data)
|
|
/// - Uses pending queue instead of retry mechanism
|
|
/// - Content integrity validation with hash tracking
|
|
///
|
|
/// # Persistence Strategy
|
|
/// - Checkpoint: companies_yahoo_cleaned.jsonl (atomic state)
|
|
/// - Log: companies_chart_updates.log (append-only updates)
|
|
/// - On restart: Load checkpoint + replay log
|
|
/// - Periodic checkpoints (every 50 companies)
|
|
/// - Batched fsync (every 10 writes or 10 seconds)
|
|
/// - Hash validation of all chart data directories
|
|
pub async fn enrich_companies_with_chart(
|
|
paths: &DataPaths,
|
|
_config: &Config,
|
|
yahoo_pool: Arc<YahooClientPool>,
|
|
shutdown_flag: &Arc<AtomicBool>,
|
|
) -> anyhow::Result<usize> {
|
|
// Configuration constants
|
|
const CHECKPOINT_INTERVAL: usize = 50;
|
|
const FSYNC_BATCH_SIZE: usize = 10;
|
|
const FSYNC_INTERVAL_SECS: u64 = 10;
|
|
const CONCURRENCY_LIMIT: usize = 50;
|
|
|
|
let data_path = paths.data_dir();
|
|
|
|
// File paths
|
|
let input_path = data_path.join("companies_yahoo_cleaned.jsonl");
|
|
let log_path = data_path.join("companies_chart_updates.log");
|
|
|
|
// Check input exists
|
|
if !input_path.exists() {
|
|
logger::log_warn(" companies_yahoo_cleaned.jsonl not found, skipping chart enrichment").await;
|
|
return Ok(0);
|
|
}
|
|
|
|
let manager = StateManager::new(paths.integrity_dir()).await?;
|
|
let step_name = "yahoo_chart_enrichment_complete";
|
|
|
|
if manager.is_step_valid(step_name).await? {
|
|
logger::log_info(" Yahoo chart enrichment already completed and valid").await;
|
|
let count = checkpoint_helpers::count_enriched_companies(paths, "chart").await?;
|
|
logger::log_info(&format!(" ✓ Found {} companies with valid chart data", count)).await;
|
|
return Ok(count);
|
|
}
|
|
|
|
logger::log_info(" Chart data needs refresh - starting enrichment").await;
|
|
|
|
// === RECOVERY PHASE: Track enriched companies ===
|
|
let enriched_companies = checkpoint_helpers::load_enrichment_progress(&log_path).await?;
|
|
|
|
// Load all companies from input
|
|
logger::log_info("Loading companies from companies_yahoo_cleaned.jsonl...").await;
|
|
let companies = load_companies_from_jsonl(&input_path).await?;
|
|
let total_companies = companies.len();
|
|
logger::log_info(&format!("Found {} companies to process", total_companies)).await;
|
|
|
|
// Filter companies that need enrichment
|
|
let pending_companies: Vec<CompanyData> = companies
|
|
.into_iter()
|
|
.filter(|company| !enriched_companies.contains(&company.name))
|
|
.collect();
|
|
|
|
let pending_count = pending_companies.len();
|
|
logger::log_info(&format!(
|
|
" {} already enriched, {} pending",
|
|
enriched_companies.len(),
|
|
pending_count
|
|
)).await;
|
|
|
|
if pending_count == 0 {
|
|
logger::log_info(" ✓ All companies already enriched").await;
|
|
track_chart_completion(&manager, paths, step_name).await?;
|
|
return Ok(enriched_companies.len());
|
|
}
|
|
|
|
// === PROCESSING PHASE: Enrich companies with chart ===
|
|
|
|
// Create enrichment function
|
|
let enrichment_fn: EnrichmentFn = Arc::new(move |company, pool, paths| {
|
|
let company = company.clone();
|
|
let pool = Arc::clone(&pool);
|
|
let paths = paths.clone();
|
|
Box::pin(async move {
|
|
enrich_company_with_chart(&company, &pool, &paths).await
|
|
})
|
|
});
|
|
|
|
// Shared counters
|
|
let processed_count = Arc::new(AtomicUsize::new(enriched_companies.len()));
|
|
let success_count = Arc::new(AtomicUsize::new(enriched_companies.len()));
|
|
let failed_count = Arc::new(AtomicUsize::new(0));
|
|
|
|
// Log writer channel with batching and fsync
|
|
let (log_tx, log_rx) = mpsc::channel::<LogCommand>(1000);
|
|
|
|
// Spawn log writer task
|
|
let log_writer_handle = spawn_log_writer(
|
|
log_path,
|
|
log_rx,
|
|
Arc::clone(&processed_count),
|
|
total_companies,
|
|
FSYNC_BATCH_SIZE,
|
|
FSYNC_INTERVAL_SECS,
|
|
);
|
|
|
|
// Process companies concurrently with task panic isolation
|
|
let mut tasks = FuturesUnordered::new();
|
|
let mut pending_iter = pending_companies.into_iter();
|
|
let semaphore = Arc::new(tokio::sync::Semaphore::new(CONCURRENCY_LIMIT));
|
|
|
|
// Initial batch of tasks
|
|
for _ in 0..CONCURRENCY_LIMIT.min(pending_count) {
|
|
if let Some(company) = pending_iter.next() {
|
|
let task = spawn_enrichment_task(
|
|
company,
|
|
Arc::clone(&yahoo_pool),
|
|
paths.clone(),
|
|
Arc::clone(&processed_count),
|
|
Arc::clone(&success_count),
|
|
Arc::clone(&failed_count),
|
|
log_tx.clone(),
|
|
Arc::clone(&semaphore),
|
|
Arc::clone(shutdown_flag),
|
|
Arc::clone(&enrichment_fn),
|
|
);
|
|
tasks.push(task);
|
|
}
|
|
}
|
|
|
|
// Process tasks as they complete and spawn new ones
|
|
let mut checkpoint_counter = 0;
|
|
while let Some(_result) = tasks.next().await {
|
|
// Check for shutdown
|
|
if shutdown_flag.load(Ordering::SeqCst) {
|
|
logger::log_warn("Shutdown signal received, stopping chart enrichment").await;
|
|
break;
|
|
}
|
|
|
|
// Checkpoint periodically
|
|
checkpoint_counter += 1;
|
|
if checkpoint_counter % CHECKPOINT_INTERVAL == 0 {
|
|
let _ = log_tx.send(LogCommand::Checkpoint).await;
|
|
}
|
|
|
|
// Spawn next task if available
|
|
if let Some(company) = pending_iter.next() {
|
|
let task = spawn_enrichment_task(
|
|
company,
|
|
Arc::clone(&yahoo_pool),
|
|
paths.clone(),
|
|
Arc::clone(&processed_count),
|
|
Arc::clone(&success_count),
|
|
Arc::clone(&failed_count),
|
|
log_tx.clone(),
|
|
Arc::clone(&semaphore),
|
|
Arc::clone(shutdown_flag),
|
|
Arc::clone(&enrichment_fn),
|
|
);
|
|
tasks.push(task);
|
|
}
|
|
}
|
|
|
|
// Signal log writer to shutdown
|
|
let _ = log_tx.send(LogCommand::Shutdown).await;
|
|
drop(log_tx);
|
|
|
|
// Wait for log writer to finish
|
|
let _ = log_writer_handle.await;
|
|
|
|
let final_processed = processed_count.load(Ordering::SeqCst);
|
|
let final_success = success_count.load(Ordering::SeqCst);
|
|
let final_failed = failed_count.load(Ordering::SeqCst);
|
|
|
|
logger::log_info(&format!(
|
|
" Chart enrichment summary: {} total, {} success, {} failed",
|
|
final_processed, final_success, final_failed
|
|
)).await;
|
|
|
|
// Mark as complete if all companies processed
|
|
if final_processed >= total_companies && !shutdown_flag.load(Ordering::SeqCst) {
|
|
track_chart_completion(&manager, paths, step_name).await?;
|
|
logger::log_info(" ✓ Chart enrichment marked as complete with integrity tracking").await;
|
|
}
|
|
|
|
Ok(final_success)
|
|
}
|
|
|
|
/// Track chart enrichment completion with content hash verification
|
|
async fn track_chart_completion(
|
|
manager: &StateManager,
|
|
paths: &DataPaths,
|
|
step_name: &str,
|
|
) -> anyhow::Result<()> {
|
|
// Create content reference for all chart data
|
|
// This will hash ALL files matching the pattern: {company}/chart/data.jsonl
|
|
let content_reference = directory_reference(
|
|
paths.corporate_dir(),
|
|
Some(vec![
|
|
"*/chart/*.jsonl".to_string(), // Main pattern for chart data
|
|
"*/chart/data.jsonl".to_string(), // Specific pattern (more precise)
|
|
]),
|
|
Some(vec![
|
|
"*.log".to_string(), // Exclude log files
|
|
"*.tmp".to_string(), // Exclude temp files
|
|
"*.bak".to_string(), // Exclude backup files
|
|
]),
|
|
);
|
|
|
|
// Track completion with:
|
|
// - Content reference: All chart directories
|
|
// - Data stage: Data (7-day TTL by default)
|
|
// - Dependencies: Depends on cleaned companies data
|
|
manager.update_entry(
|
|
step_name.to_string(),
|
|
content_reference,
|
|
DataStage::Data,
|
|
None, // Use default TTL (7 days for Data stage)
|
|
).await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Enrich a single company with chart data
|
|
async fn enrich_company_with_chart(
|
|
company: &CompanyData,
|
|
yahoo_pool: &Arc<YahooClientPool>,
|
|
paths: &DataPaths,
|
|
) -> anyhow::Result<()> {
|
|
let ticker = match extract_first_yahoo_ticker(company) {
|
|
Some(t) => t,
|
|
None => {
|
|
return Err(anyhow::anyhow!("No valid Yahoo ticker found"));
|
|
}
|
|
};
|
|
|
|
// Get 1 year of daily chart data
|
|
let now = chrono::Utc::now().timestamp();
|
|
let start = chrono::Utc
|
|
.with_ymd_and_hms(2000, 1, 1, 0, 0, 0)
|
|
.unwrap()
|
|
.timestamp();
|
|
|
|
let chart_data = yahoo_pool.get_chart_data(&ticker, "1d", start, now).await?;
|
|
|
|
// Only save if we got meaningful data
|
|
if chart_data.quotes.is_empty() {
|
|
return Err(anyhow::anyhow!("No chart data available"));
|
|
}
|
|
|
|
// Save the chart data
|
|
save_company_data(paths, &company.name, &chart_data, "chart").await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Save data to company directory (generic version)
|
|
async fn save_company_data<T: serde::Serialize>(
|
|
paths: &DataPaths,
|
|
company_name: &str,
|
|
data: &T,
|
|
data_type: &str,
|
|
) -> anyhow::Result<()> {
|
|
use tokio::fs;
|
|
|
|
let safe_name = sanitize_company_name(company_name);
|
|
|
|
let company_dir = paths.corporate_dir().join(&safe_name).join(data_type);
|
|
fs::create_dir_all(&company_dir).await?;
|
|
|
|
let data_path = company_dir.join("data.jsonl");
|
|
let json_line = serde_json::to_string(data)?;
|
|
|
|
let mut file = fs::File::create(&data_path).await?;
|
|
file.write_all(json_line.as_bytes()).await?;
|
|
file.write_all(b"\n").await?;
|
|
file.flush().await?;
|
|
file.sync_all().await?; // Ensure data is persisted
|
|
|
|
Ok(())
|
|
}
|
|
|
|
// ============================================================================
|
|
// GENERIC SHARED FUNCTIONS
|
|
// ============================================================================
|
|
|
|
/// Spawn log writer task (shared across all enrichment types)
|
|
fn spawn_log_writer(
|
|
log_path: std::path::PathBuf,
|
|
mut log_rx: mpsc::Receiver<LogCommand>,
|
|
processed_count: Arc<AtomicUsize>,
|
|
total_companies: usize,
|
|
fsync_batch_size: usize,
|
|
fsync_interval_secs: u64,
|
|
) -> tokio::task::JoinHandle<()> {
|
|
tokio::spawn(async move {
|
|
let mut log_file = OpenOptions::new()
|
|
.create(true)
|
|
.append(true)
|
|
.open(&log_path)
|
|
.await
|
|
.expect("Failed to open log file");
|
|
|
|
let mut write_count = 0;
|
|
let mut last_fsync = tokio::time::Instant::now();
|
|
|
|
while let Some(cmd) = log_rx.recv().await {
|
|
match cmd {
|
|
LogCommand::Write(entry) => {
|
|
let json_line = serde_json::to_string(&entry).expect("Serialization failed");
|
|
log_file.write_all(json_line.as_bytes()).await.expect("Write failed");
|
|
log_file.write_all(b"\n").await.expect("Write failed");
|
|
|
|
write_count += 1;
|
|
|
|
// Batched fsync
|
|
if write_count >= fsync_batch_size
|
|
|| last_fsync.elapsed().as_secs() >= fsync_interval_secs
|
|
{
|
|
log_file.flush().await.expect("Flush failed");
|
|
log_file.sync_all().await.expect("Fsync failed");
|
|
write_count = 0;
|
|
last_fsync = tokio::time::Instant::now();
|
|
}
|
|
}
|
|
LogCommand::Checkpoint => {
|
|
// Force fsync on checkpoint
|
|
log_file.flush().await.expect("Flush failed");
|
|
log_file.sync_all().await.expect("Fsync failed");
|
|
write_count = 0;
|
|
last_fsync = tokio::time::Instant::now();
|
|
|
|
let current = processed_count.load(Ordering::SeqCst);
|
|
logger::log_info(&format!(
|
|
" Checkpoint: {}/{} companies processed",
|
|
current, total_companies
|
|
)).await;
|
|
}
|
|
LogCommand::Shutdown => {
|
|
// Final fsync before shutdown
|
|
log_file.flush().await.expect("Flush failed");
|
|
log_file.sync_all().await.expect("Fsync failed");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
})
|
|
}
|
|
|
|
/// Spawn a single enrichment task with panic isolation (GENERIC VERSION)
|
|
///
|
|
/// This generic version accepts an enrichment function as a parameter,
|
|
/// allowing it to be reused for events, options, charts, or any other enrichment type.
|
|
///
|
|
/// # Parameters
|
|
/// - `company`: The company to enrich
|
|
/// - `yahoo_pool`: Yahoo API client pool
|
|
/// - `paths`: Data paths
|
|
/// - `processed_count`: Counter for processed companies
|
|
/// - `success_count`: Counter for successful enrichments
|
|
/// - `failed_count`: Counter for failed enrichments
|
|
/// - `log_tx`: Channel to send log commands
|
|
/// - `semaphore`: Semaphore for concurrency control
|
|
/// - `shutdown_flag`: Flag to signal shutdown
|
|
/// - `enrichment_fn`: The specific enrichment function to call (events, option, chart, etc.)
|
|
fn spawn_enrichment_task(
|
|
company: CompanyData,
|
|
yahoo_pool: Arc<YahooClientPool>,
|
|
paths: DataPaths,
|
|
processed_count: Arc<AtomicUsize>,
|
|
success_count: Arc<AtomicUsize>,
|
|
failed_count: Arc<AtomicUsize>,
|
|
log_tx: mpsc::Sender<LogCommand>,
|
|
semaphore: Arc<tokio::sync::Semaphore>,
|
|
shutdown_flag: Arc<AtomicBool>,
|
|
enrichment_fn: EnrichmentFn,
|
|
) -> tokio::task::JoinHandle<()> {
|
|
tokio::spawn(async move {
|
|
// Acquire semaphore permit
|
|
let _permit = semaphore.acquire().await.expect("Semaphore closed");
|
|
|
|
// Check shutdown before processing
|
|
if shutdown_flag.load(Ordering::SeqCst) {
|
|
return;
|
|
}
|
|
|
|
// Call the enrichment function (this is where the type-specific logic happens)
|
|
let result = enrichment_fn(company.clone(), Arc::clone(&yahoo_pool), paths).await;
|
|
|
|
// Update counters
|
|
processed_count.fetch_add(1, Ordering::SeqCst);
|
|
|
|
let status = match result {
|
|
Ok(_) => {
|
|
success_count.fetch_add(1, Ordering::SeqCst);
|
|
"enriched"
|
|
}
|
|
Err(e) => {
|
|
failed_count.fetch_add(1, Ordering::SeqCst);
|
|
logger::log_warn(&format!(
|
|
" Failed to enrich {}: {}",
|
|
company.name, e
|
|
)).await;
|
|
"failed"
|
|
}
|
|
};
|
|
|
|
// Log result
|
|
let log_entry = json!({
|
|
"company_name": company.name,
|
|
"status": status,
|
|
"timestamp": Utc::now().to_rfc3339(),
|
|
});
|
|
|
|
let _ = log_tx.send(LogCommand::Write(log_entry)).await;
|
|
})
|
|
} |