from datetime import datetime, timezone from sqlalchemy import create_engine, Column, Integer, String, Float, DateTime from sqlalchemy.orm import sessionmaker, declarative_base from config import DATA_DIR, logger import os import threading import time import shutil # Setup database path DB_PATH = DATA_DIR / "alpha_results.db" # Restore database from Hugging Face Dataset if applicable HF_TOKEN = os.getenv("HF_TOKEN") SPACE_ID = os.getenv("SPACE_ID") DEFAULT_DATASET_ID = f"{SPACE_ID}-backup" if SPACE_ID else None HF_DATASET_ID = os.getenv("HF_DATASET_ID") or DEFAULT_DATASET_ID def restore_db_from_hf(): if not HF_TOKEN: logger.warning("HF_TOKEN environment variable is not set. Database restore from Hugging Face Dataset skipped.") return if not HF_DATASET_ID: logger.warning("HF_DATASET_ID and SPACE_ID environment variables are not set. Database restore skipped.") return logger.info(f"Checking for database backup in Hugging Face Dataset: {HF_DATASET_ID}...") try: from huggingface_hub import hf_hub_download cached_path = hf_hub_download( repo_id=HF_DATASET_ID, filename="alpha_results.db", repo_type="dataset", token=HF_TOKEN ) DATA_DIR.mkdir(parents=True, exist_ok=True) shutil.copy(cached_path, DB_PATH) logger.info(f"Successfully restored database from Hugging Face Dataset to {DB_PATH}.") except ImportError: logger.warning("huggingface_hub package is not installed. Database restore skipped.") except Exception as e: logger.warning(f"Could not restore database from Hugging Face Dataset (this is normal if it is the first run): {e}") try: from huggingface_hub import hf_hub_download settings_path = DATA_DIR / "miner_settings.json" cached_settings_path = hf_hub_download( repo_id=HF_DATASET_ID, filename="miner_settings.json", repo_type="dataset", token=HF_TOKEN ) shutil.copy(cached_settings_path, settings_path) logger.info("Successfully restored miner_settings.json from Hugging Face Dataset.") except Exception as e: logger.warning(f"Could not restore miner_settings.json from Hugging Face Dataset: {e}") try: from huggingface_hub import hf_hub_download cache_path = DATA_DIR / "data_fields_cache.json" cached_cache_path = hf_hub_download( repo_id=HF_DATASET_ID, filename="data_fields_cache.json", repo_type="dataset", token=HF_TOKEN ) shutil.copy(cached_cache_path, cache_path) logger.info("Successfully restored data_fields_cache.json from Hugging Face Dataset.") except Exception as e: logger.warning(f"Could not restore data_fields_cache.json from Hugging Face Dataset: {e}") # Try to restore backup before SQLAlchemy engine is initialized restore_db_from_hf() DATABASE_URL = f"sqlite:///{DB_PATH}" def backup_db_to_hf(): if not HF_TOKEN: logger.warning("HF_TOKEN environment variable is not set. Database backup to Hugging Face Dataset skipped.") return if not HF_DATASET_ID: logger.warning("HF_DATASET_ID and SPACE_ID environment variables are not set. Database backup skipped.") return try: from huggingface_hub import HfApi except ImportError: logger.warning("huggingface_hub package is not installed. Background backup thread will not start.") return logger.info(f"Starting background database backup thread (interval: 15 mins) to: {HF_DATASET_ID}...") api = HfApi(token=HF_TOKEN) # Create the dataset repo if it doesn't exist try: api.create_repo( repo_id=HF_DATASET_ID, repo_type="dataset", private=True, exist_ok=True ) except Exception as e: logger.warning(f"Could not verify/create Hugging Face Dataset repository: {e}") is_first_backup = True while True: try: # Sleep for 60 seconds on the first loop, and 15 minutes thereafter if is_first_backup: time.sleep(60) is_first_backup = False else: time.sleep(15 * 60) if not DB_PATH.exists(): logger.debug("Database file does not exist yet. Skipping backup cycle.") continue logger.info(f"Backing up database to Hugging Face Dataset: {HF_DATASET_ID}...") temp_db_path = DB_PATH.parent / "alpha_results_backup_temp.db" try: shutil.copy2(DB_PATH, temp_db_path) api.upload_file( path_or_fileobj=str(temp_db_path), path_in_repo="alpha_results.db", repo_id=HF_DATASET_ID, repo_type="dataset", ) logger.info("Database backup successfully uploaded to Hugging Face Dataset.") except Exception as upload_error: logger.error(f"Failed to upload database file to Hugging Face: {upload_error}") finally: if temp_db_path.exists(): os.remove(temp_db_path) # Also backup miner_settings.json if it exists settings_path = DATA_DIR / "miner_settings.json" if settings_path.exists(): logger.info("Backing up miner_settings.json to Hugging Face Dataset...") try: api.upload_file( path_or_fileobj=str(settings_path), path_in_repo="miner_settings.json", repo_id=HF_DATASET_ID, repo_type="dataset", ) logger.info("miner_settings.json backup successfully uploaded to Hugging Face Dataset.") except Exception as upload_error: logger.error(f"Failed to upload miner_settings.json to Hugging Face: {upload_error}") # Also backup data_fields_cache.json if it exists cache_path = DATA_DIR / "data_fields_cache.json" if cache_path.exists(): logger.info("Backing up data_fields_cache.json to Hugging Face Dataset...") try: api.upload_file( path_or_fileobj=str(cache_path), path_in_repo="data_fields_cache.json", repo_id=HF_DATASET_ID, repo_type="dataset", ) logger.info("data_fields_cache.json backup successfully uploaded to Hugging Face Dataset.") except Exception as upload_error: logger.error(f"Failed to upload data_fields_cache.json to Hugging Face: {upload_error}") # Also backup miner_production.log if it exists prod_log_path = DATA_DIR / "miner_production.log" if prod_log_path.exists(): logger.info("Backing up miner_production.log to Hugging Face Dataset...") try: temp_log_path = DATA_DIR / "miner_production_backup_temp.log" shutil.copy2(prod_log_path, temp_log_path) api.upload_file( path_or_fileobj=str(temp_log_path), path_in_repo="miner_production.log", repo_id=HF_DATASET_ID, repo_type="dataset", ) logger.info("miner_production.log backup successfully uploaded to Hugging Face Dataset.") except Exception as upload_error: logger.error(f"Failed to upload miner_production.log to Hugging Face: {upload_error}") finally: if temp_log_path.exists(): try: os.remove(temp_log_path) except Exception: pass except Exception as e: logger.error(f"Error during background database backup to Hugging Face: {e}") def start_backup_thread(): if HF_TOKEN and HF_DATASET_ID: t = threading.Thread(target=backup_db_to_hf, daemon=True, name="HF_DB_Backup") t.start() logger.info("Database backup thread started successfully.") else: logger.warning("HF_TOKEN or HF_DATASET_ID/SPACE_ID not configured. Backup thread not started.") # Start the periodic background backup thread start_backup_thread() # Initialize SQLAlchemy Engine and Base engine = create_engine(DATABASE_URL, echo=False, connect_args={"timeout": 30}) Base = declarative_base() SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) class AlphaRecord(Base): __tablename__ = "alpha_records" id = Column(Integer, primary_key=True, index=True) expression = Column(String, unique=True, nullable=False) region = Column(String, default='USA') universe = Column(String, default='TOP3000') delay = Column(Integer, default=1) decay = Column(Integer, default=10) truncation = Column(Float, default=0.01) neutralization = Column(String, default='SUBINDUSTRY') # Results sharpe = Column(Float, nullable=True) fitness = Column(Float, nullable=True) turnover = Column(Float, nullable=True) returns = Column(Float, nullable=True) # Metadata status = Column(String, default='GENERATED') # e.g., 'GENERATED', 'SIMULATING', 'FAILED', 'PASSED', 'SUBMITTED' error_log = Column(String, nullable=True) alpha_id = Column(String, nullable=True) date_created = Column(DateTime, default=lambda: datetime.now(timezone.utc)) date_submitted = Column(DateTime, nullable=True) retry_count = Column(Integer, default=0) def __repr__(self): return f"" def init_db(): """Initializes the SQLite database and creates all tables.""" logger.info(f"Initializing database at {DB_PATH}...") Base.metadata.create_all(bind=engine) # Enable WAL mode for better concurrency try: import sqlite3 conn = sqlite3.connect(str(DB_PATH)) cursor = conn.cursor() cursor.execute("PRAGMA journal_mode=WAL") cursor.execute("PRAGMA synchronous=NORMAL") conn.commit() conn.close() logger.info("SQLite WAL mode and synchronous settings configured successfully.") except Exception as e: logger.error(f"Failed to configure SQLite WAL mode: {e}") # Check if alpha_id or date_submitted columns exist, if not, add them safely try: import sqlite3 conn = sqlite3.connect(str(DB_PATH)) cursor = conn.cursor() cursor.execute("PRAGMA table_info(alpha_records)") columns = [col[1] for col in cursor.fetchall()] if "alpha_id" not in columns: logger.info("Adding alpha_id column to alpha_records table...") cursor.execute("ALTER TABLE alpha_records ADD COLUMN alpha_id VARCHAR") conn.commit() logger.info("alpha_id column added successfully.") if "date_submitted" not in columns: logger.info("Adding date_submitted column to alpha_records table...") cursor.execute("ALTER TABLE alpha_records ADD COLUMN date_submitted TIMESTAMP") conn.commit() logger.info("date_submitted column added successfully.") if "retry_count" not in columns: logger.info("Adding retry_count column to alpha_records table...") cursor.execute("ALTER TABLE alpha_records ADD COLUMN retry_count INTEGER DEFAULT 0") conn.commit() logger.info("retry_count column added successfully.") # Migrate SP500 to TOPSP500 in existing records logger.info("Running database migration to replace SP500 with TOPSP500...") cursor.execute("UPDATE alpha_records SET universe = 'TOPSP500' WHERE universe = 'SP500'") conn.commit() logger.info(f"Universe migration query completed. Rows modified: {cursor.rowcount}") conn.close() except Exception as e: logger.error(f"Error checking/adding columns to database: {e}") logger.info("Database initialized successfully.") # Automatically initialize the database and apply migrations on import init_db()