added event enrichment

src/corporate/update_companies_enrich.rs

@@ -1,22 +1,22 @@
-// src/corporate/update_companies_enrich.rs
-use super::{helpers::*, types::*};
+// src/corporate/update_companies_enrich_events.rs
+use super::{types::*};
 use crate::config::Config;
 use crate::util::directories::DataPaths;
 use crate::util::logger;
 use crate::scraper::yahoo::{YahooClientPool, QuoteSummaryModule};
 
 use std::result::Result::Ok;
-use chrono::{Local, Utc};
-use std::collections::HashMap;
+use chrono::Utc;
+use std::collections::{HashSet};
 use std::sync::Arc;
 use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
-use tokio::fs::{File, OpenOptions};
-use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
+use tokio::fs::{OpenOptions};
+use tokio::io::{AsyncWriteExt};
 use futures::stream::{FuturesUnordered, StreamExt};
 use serde_json::json;
 use tokio::sync::mpsc;
 
-/// Yahoo enriching data per corporate
+/// Yahoo Event enrichment per corporate company
 /// 
 /// # Features
 /// - Graceful shutdown (abort-safe)
@@ -28,7 +28,515 @@ use tokio::sync::mpsc;
 /// 
 /// # Persistence Strategy
 /// - Checkpoint: companies_yahoo_cleaned.jsonl (atomic state)
-/// - Log: companies_update.log (append-only updates)
+/// - 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)
+/// - Batched fsync (every 10 writes or 10 seconds)
+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; // Limit parallel enrichment tasks
+    
+    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");
+    let state_path = data_path.join("state.jsonl");
+    
+    // Check input exists
+    if !input_path.exists() {
+        logger::log_warn("  companies_yahoo_cleaned.jsonl not found, skipping event enrichment").await;
+        return Ok(0);
+    }
+    
+    // Check if already completed
+    if state_path.exists() {
+        let state_content = tokio::fs::read_to_string(&state_path).await?;
+        
+        for line in state_content.lines() {
+            if line.trim().is_empty() {
+                continue;
+            }
+            
+            if let Ok(state) = serde_json::from_str::<serde_json::Value>(line) {
+                if state.get("yahoo_events_enrichment_complete").and_then(|v| v.as_bool()).unwrap_or(false) {
+                    logger::log_info("  Yahoo events enrichment already completed").await;
+                    
+                    // Count enriched companies
+                    let count = count_enriched_companies(paths).await?;
+                    logger::log_info(&format!("  ✓ Found {} companies with event data", count)).await;
+                    return Ok(count);
+                }
+            }
+        }
+    }
+    
+    // === RECOVERY PHASE: Track enriched companies ===
+    let mut enriched_companies: HashSet<String> = HashSet::new();
+    
+    if log_path.exists() {
+        logger::log_info("Loading enrichment progress from log...").await;
+        let log_content = tokio::fs::read_to_string(&log_path).await?;
+        
+        for line in log_content.lines() {
+            if line.trim().is_empty() || !line.ends_with('}') {
+                continue; // Skip incomplete lines
+            }
+            
+            match serde_json::from_str::<serde_json::Value>(line) {
+                Ok(entry) => {
+                    if let Some(name) = entry.get("company_name").and_then(|v| v.as_str()) {
+                        if entry.get("status").and_then(|v| v.as_str()) == Some("enriched") {
+                            enriched_companies.insert(name.to_string());
+                        }
+                    }
+                }
+                Err(e) => {
+                    logger::log_warn(&format!("Skipping invalid log line: {}", e)).await;
+                }
+            }
+        }
+        logger::log_info(&format!("Loaded {} enriched companies from log", enriched_companies.len())).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<CompanyCrossPlatformInfo> = 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;
+        mark_enrichment_complete(&state_path).await?;
+        return Ok(enriched_companies.len());
+    }
+    
+    // === PROCESSING PHASE: Enrich companies with events ===
+    
+    // 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, mut log_rx) = mpsc::channel::<LogCommand>(1000);
+    
+    // Spawn log writer task
+    let log_writer_handle = {
+        let log_path = log_path.clone();
+        let processed_count = Arc::clone(&processed_count);
+        let total_companies = total_companies;
+        
+        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;
+                    }
+                }
+            }
+        })
+    };
+    
+    // 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),
+            );
+            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 detected, stopping new enrichment tasks...").await;
+            break;
+        }
+        
+        // Spawn next task
+        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),
+            );
+            tasks.push(task);
+        }
+        
+        // Periodic checkpoint
+        checkpoint_counter += 1;
+        if checkpoint_counter % CHECKPOINT_INTERVAL == 0 {
+            let _ = log_tx.send(LogCommand::Checkpoint).await;
+        }
+    }
+    
+    // Shutdown log writer
+    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) {
+        mark_enrichment_complete(&state_path).await?;
+        logger::log_info("  ✓ Event enrichment marked as complete").await;
+    }
+    
+    Ok(final_success)
+}
+
+/// Spawn a single enrichment task with panic isolation
+fn spawn_enrichment_task(
+    company: CompanyCrossPlatformInfo,
+    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>,
+) -> 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;
+        }
+        
+        // Process company
+        let result = enrich_company_with_events(
+            &company,
+            &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;
+    })
+}
+
+/// Enrich a single company with event data
+async fn enrich_company_with_events(
+    company: &CompanyCrossPlatformInfo,
+    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(())
+}
+
+/// Extract first valid Yahoo ticker from company
+fn extract_first_yahoo_ticker(company: &CompanyCrossPlatformInfo) -> Option<String> {
+    for tickers in company.isin_tickers_map.values() {
+        for ticker in tickers {
+            if ticker.starts_with("YAHOO:") 
+                && ticker != "YAHOO:NO_RESULTS" 
+                && ticker != "YAHOO:ERROR" 
+            {
+                return Some(ticker.trim_start_matches("YAHOO:").to_string());
+            }
+        }
+    }
+    None
+}
+
+/// Sanitize company name for file system
+fn sanitize_company_name(name: &str) -> String {
+    name.replace("/", "_")
+        .replace("\\", "_")
+        .replace(":", "_")
+        .replace("*", "_")
+        .replace("?", "_")
+        .replace("\"", "_")
+        .replace("<", "_")
+        .replace(">", "_")
+        .replace("|", "_")
+}
+
+/// Load companies from JSONL file
+async fn load_companies_from_jsonl(path: &std::path::Path) -> anyhow::Result<Vec<CompanyCrossPlatformInfo>> {
+    let content = tokio::fs::read_to_string(path).await?;
+    let mut companies = Vec::new();
+    
+    for line in content.lines() {
+        if line.trim().is_empty() {
+            continue;
+        }
+        if let Ok(company) = serde_json::from_str::<CompanyCrossPlatformInfo>(line) {
+            companies.push(company);
+        }
+    }
+    
+    Ok(companies)
+}
+
+/// Count enriched companies (companies with event data)
+async fn count_enriched_companies(paths: &DataPaths) -> anyhow::Result<usize> {
+    let corporate_dir = paths.corporate_dir();
+    
+    if !corporate_dir.exists() {
+        return Ok(0);
+    }
+    
+    let mut count = 0;
+    let mut entries = tokio::fs::read_dir(&corporate_dir).await?;
+    
+    while let Some(entry) = entries.next_entry().await? {
+        let path = entry.path();
+        if path.is_dir() {
+            let events_dir = path.join("events");
+            let events_file = events_dir.join("data.jsonl");
+            
+            if events_file.exists() {
+                count += 1;
+            }
+        }
+    }
+    
+    Ok(count)
+}
+
+/// Mark enrichment as complete in state file
+async fn mark_enrichment_complete(state_path: &std::path::Path) -> anyhow::Result<()> {
+    let enrichment_complete = json!({
+        "yahoo_events_enrichment_complete": true,
+        "completed_at": Utc::now().to_rfc3339(),
+    });
+    
+    let mut state_file = OpenOptions::new()
+        .create(true)
+        .append(true)
+        .open(state_path)
+        .await?;
+    
+    let state_line = serde_json::to_string(&enrichment_complete)?;
+    state_file.write_all(state_line.as_bytes()).await?;
+    state_file.write_all(b"\n").await?;
+    state_file.flush().await?;
+    state_file.sync_all().await?;
+    
+    Ok(())
+}
+
+/// Log command enum
+enum LogCommand {
+    Write(serde_json::Value),
+    Checkpoint,
+    Shutdown,
+}