RivaCube/update_economics.py

# RivaCube/RivDroite/update_economics.py

import os
import logging
from datetime import datetime, timedelta
from dotenv import load_dotenv
from typing import Optional, Dict, List, Any
from utils.economics.collector import EconomicsCollector
from utils.economics.utils import (
    DatabaseManager,
    save_treasury_rates,
    save_economic_indicator,
    save_economic_calendar,
    save_market_risk_premium
)

# Define indicator mappings with their API endpoints
INDICATOR_MAPPINGS = {
    'GDP': 'GDP',
    'INFLATION': 'inflationRate',
    'UNEMPLOYMENT': 'unemploymentRate',
    'INTEREST_RATE': 'federalFunds',
    'RETAIL_SALES': 'retailSales',
    'INDUSTRIAL_PRODUCTION': 'industrialProductionTotalIndex'
}

def setup_logging() -> None:
    """Configure logging with both file and console output."""
    log_dir = 'logs'
    os.makedirs(log_dir, exist_ok=True)
    
    log_file = os.path.join(log_dir, f'economics_update_{datetime.now():%Y%m%d_%H%M%S}.log')
    
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler()
        ]
    )

def init_database(db_manager: DatabaseManager) -> bool:
    """Initialize database schema if not exists."""
    conn = None
    try:
        conn = db_manager.get_connection()
        if not conn:
            raise ValueError("Failed to get database connection")
            
        cursor = conn.cursor()
        
        # Check if tables exist
        cursor.execute("""
            SELECT EXISTS (
                SELECT FROM information_schema.tables 
                WHERE table_schema = 'public' 
                AND table_name = 'treasury_rates'
            );
        """)
        
        tables_exist = cursor.fetchone()[0]
        
        if not tables_exist:
            schema_sql = """
            CREATE TABLE IF NOT EXISTS treasury_rates (
                id SERIAL PRIMARY KEY,
                date DATE NOT NULL,
                maturity VARCHAR(50) NOT NULL,
                rate NUMERIC,
                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                last_updated TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                CONSTRAINT uk_treasury_rates UNIQUE(date, maturity)
            );

            CREATE TABLE IF NOT EXISTS economic_indicators (
                id SERIAL PRIMARY KEY,
                date DATE NOT NULL,
                indicator_name VARCHAR(100) NOT NULL,
                value NUMERIC NOT NULL,
                unit VARCHAR(50),
                frequency VARCHAR(50),
                country VARCHAR(50) DEFAULT 'US',
                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                last_updated TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                CONSTRAINT uk_economic_indicators UNIQUE(date, indicator_name)
            );

            CREATE TABLE IF NOT EXISTS economic_calendar (
                id SERIAL PRIMARY KEY,
                date DATE NOT NULL,
                event VARCHAR(255) NOT NULL,
                country VARCHAR(50) DEFAULT 'US',
                impact VARCHAR(50),
                actual VARCHAR(100),
                previous VARCHAR(100),
                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                last_updated TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                CONSTRAINT uk_economic_calendar UNIQUE(date, event, country)
            );

            CREATE TABLE IF NOT EXISTS market_risk_premium (
                id SERIAL PRIMARY KEY,
                date DATE NOT NULL,
                country VARCHAR(100) NOT NULL,
                country_risk_premium NUMERIC,
                total_equity_risk_premium NUMERIC,
                created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                last_updated TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
                CONSTRAINT uk_market_risk_premium UNIQUE(date, country)
            );

            CREATE INDEX IF NOT EXISTS idx_treasury_rates_date ON treasury_rates(date);
            CREATE INDEX IF NOT EXISTS idx_economic_indicators_date ON economic_indicators(date);
            CREATE INDEX IF NOT EXISTS idx_economic_calendar_date ON economic_calendar(date);
            CREATE INDEX IF NOT EXISTS idx_market_risk_premium_date ON market_risk_premium(date);
            """
            
            cursor.execute(schema_sql)
            conn.commit()
            logging.info("Database schema initialized successfully")
        else:
            logging.info("Database schema already exists")
            
        return True
        
    except Exception as e:
        if conn:
            conn.rollback()
        logging.error(f"Database initialization failed: {e}")
        return False
    finally:
        if cursor:
            cursor.close()
        if conn:
            db_manager.return_connection(conn)

def process_market_risk_premium(collector: EconomicsCollector, conn: Any) -> bool:
    """Process market risk premium data with proper error handling."""
    try:
        logging.info("Updating market risk premium...")
        data = collector.get_market_risk_premium()
        
        if not data:
            logging.info("No market risk premium data received")
            return True
            
        # Filter out entries with None values for both premiums
        valid_data = [
            item for item in data
            if (item.get('countryRiskPremium') is not None or 
                item.get('totalEquityRiskPremium') is not None)
        ]
        
        if valid_data:
            save_market_risk_premium(conn, valid_data)
            return True
        else:
            logging.info("No valid market risk premium data after filtering")
            return True  # Not considering this as an error
            
    except Exception as e:
        logging.error(f"Error processing market risk premium: {e}")
        return False

def process_economic_data(
    collector: EconomicsCollector,
    db_name: str,
    api_name: str,
    conn: Any
) -> bool:
    """Process a single economic indicator."""
    try:
        logging.info(f"Updating {db_name} (API: {api_name})...")
        data = collector.get_economic_indicator(api_name)
        
        if not data:
            logging.info(f"No data received for {db_name}")
            return True
            
        save_economic_indicator(conn, db_name, data)
        return True
        
    except Exception as e:
        logging.error(f"Error processing {db_name}: {e}")
        return False

def update_all_economics(db_manager: DatabaseManager, collector: EconomicsCollector) -> bool:
    """Update all economic data in the database."""
    success = True
    conn = None
    
    try:
        conn = db_manager.get_connection()
        if not conn:
            raise ValueError("Failed to get database connection")

        # Get date range for data collection
        to_date = datetime.now()
        days_to_fetch = int(os.getenv('DAYS_TO_FETCH', '2'))
        from_date = to_date - timedelta(days=days_to_fetch)
        
        # Update treasury rates
        try:
            logging.info("Updating treasury rates...")
            rates = collector.get_treasury_rates(
                from_date.strftime('%Y-%m-%d'),
                to_date.strftime('%Y-%m-%d')
            )
            if rates:
                save_treasury_rates(conn, rates)
        except Exception as e:
            logging.error(f"Error updating treasury rates: {e}")
            success = False

        # Update economic indicators
        for db_name, api_name in INDICATOR_MAPPINGS.items():
            if not process_economic_data(collector, db_name, api_name, conn):
                success = False

        # Update economic calendar
        try:
            logging.info("Updating economic calendar...")
            events = collector.get_economic_calendar(
                from_date.strftime('%Y-%m-%d'),
                to_date.strftime('%Y-%m-%d')
            )
            if events:
                save_economic_calendar(conn, events)
        except Exception as e:
            logging.error(f"Error updating economic calendar: {e}")
            success = False

        # Update market risk premium
        if not process_market_risk_premium(collector, conn):
            success = False

        return success
        
    except Exception as e:
        logging.error(f"Error in update_all_economics: {e}")
        return False
    finally:
        if conn:
            db_manager.return_connection(conn)

def main() -> None:
    """Main function to run the economic data update process."""
    load_dotenv()
    setup_logging()
    
    start_time = datetime.now()
    logging.info(f"Starting economic data update at {start_time}")
    
    db_manager = None
    success = False
    
    try:
        api_key = os.getenv('FMP_API_KEY')
        if not api_key:
            raise ValueError("FMP_API_KEY not found in environment variables")
        
        db_manager = DatabaseManager()
        if not init_database(db_manager):
            raise ValueError("Failed to initialize database")
            
        collector = EconomicsCollector(api_key)
        success = update_all_economics(db_manager, collector)
        
    except Exception as e:
        logging.error(f"Fatal error in main execution: {e}")
    finally:
        if db_manager:
            db_manager.close_all()
            logging.info("Database connections closed")
        
        end_time = datetime.now()
        duration = end_time - start_time
        status = "successfully" if success else "with errors"
        logging.info(f"Economic data update completed {status} in {duration}")

if __name__ == "__main__":
    main()