mt5-ai-xauusd-trader/scripts/data_cleanup.py

"""
MT5 AI/ML Trading Bot - Data Cleanup Script
Automates the purging of old operational data based on the Data Retention Policy.
"""

import argparse
import logging
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path

from sqlalchemy import create_engine, delete, select
from sqlalchemy.orm import sessionmaker

# Add src to path to import models
sys.path.append(str(Path(__file__).resolve().parents[1]))

from src.core.audit_log import AuditEntry
from src.core.config import get_config
from src.core.trade_logger import ModelSignal, PerformanceMetric, RiskEvent, Trade

# -- Setup Logging ---------------------------------------------------------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger("data_cleanup")

# -- Configuration Constants -----------------------------------------------
# Retention windows in days
RETENTION_LOGS = 90
RETENTION_UNLINKED_SIGNALS = 90
RETENTION_RISK_EVENTS = 2 * 365
RETENTION_PERFORMANCE_METRICS = 2 * 365
RETENTION_TRADES = 7 * 365
RETENTION_AUDIT_LOG = 7 * 365
RETENTION_BACKTESTS = 365


def cleanup_logs(logs_dir: Path, dry_run: bool = False) -> int:
    """Delete log files older than RETENTION_LOGS days."""
    if not logs_dir.exists():
        logger.info(f"Logs directory {logs_dir} does not exist. Skipping.")
        return 0

    count = 0
    cutoff = datetime.now() - timedelta(days=RETENTION_LOGS)

    logger.info(f"Cleaning up logs in {logs_dir} older than {cutoff.date()}...")

    for log_file in logs_dir.glob("*.log*"):
        try:
            mtime = datetime.fromtimestamp(log_file.stat().st_mtime)
            if mtime < cutoff:
                logger.info(f"{'[DRY RUN] ' if dry_run else ''}Deleting old log file: {log_file.name} (mtime: {mtime})")
                if not dry_run:
                    log_file.unlink()
                count += 1
        except Exception as e:
            logger.error(f"Failed to delete log file {log_file}: {e}")

    return count


def cleanup_backtests(backtest_dir: Path, dry_run: bool = False) -> int:
    """Delete backtest results older than RETENTION_BACKTESTS days."""
    if not backtest_dir.exists():
        logger.info(f"Backtest directory {backtest_dir} does not exist. Skipping.")
        return 0

    count = 0
    cutoff = datetime.now() - timedelta(days=RETENTION_BACKTESTS)

    logger.info(f"Cleaning up backtest results in {backtest_dir} older than {cutoff.date()}...")

    # Recursively check for files and directories in backtest_dir
    for item in backtest_dir.rglob("*"):
        if item.is_file():
            try:
                mtime = datetime.fromtimestamp(item.stat().st_mtime)
                if mtime < cutoff:
                    logger.info(f"{'[DRY RUN] ' if dry_run else ''}Deleting old backtest file: {item.relative_to(backtest_dir)} (mtime: {mtime})")
                    if not dry_run:
                        item.unlink()
                    count += 1
            except Exception as e:
                logger.error(f"Failed to delete backtest file {item}: {e}")

    # After deleting files, attempt to delete empty directories
    for item in sorted(backtest_dir.rglob("*"), reverse=True):
        if item.is_dir() and item != backtest_dir and not any(item.iterdir()):
            try:
                logger.info(f"{'[DRY RUN] ' if dry_run else ''}Deleting empty backtest directory: {item.relative_to(backtest_dir)}")
                if not dry_run:
                    item.rmdir()
            except Exception as e:
                logger.error(f"Failed to delete directory {item}: {e}")

    return count


def cleanup_database(db_url: str, audit_db_url: str = None, dry_run: bool = False) -> dict:
    """Purge old records from the database according to the retention policy."""
    engine = create_engine(db_url)
    Session = sessionmaker(bind=engine)
    results = {
        "model_signals": 0,
        "risk_events": 0,
        "performance_metrics": 0,
        "trades": 0,
        "audit_log": 0,
    }

    now = datetime.now(timezone.utc)

    with Session() as session:
        # 1. Cleanup Unlinked Model Signals (older than 90 days)
        # We must exclude signals that are linked to trades
        signal_cutoff = now - timedelta(days=RETENTION_UNLINKED_SIGNALS)

        # Subquery for linked signals
        linked_signal_ids = select(Trade.signal_id).where(Trade.signal_id.is_not(None))

        unlinked_signals_query = (
            select(ModelSignal.id)
            .where(ModelSignal.created_at < signal_cutoff)
            .where(ModelSignal.id.not_in(linked_signal_ids))
        )

        unlinked_ids = session.execute(unlinked_signals_query).scalars().all()
        results["model_signals"] = len(unlinked_ids)

        if unlinked_ids:
            logger.info(f"{'[DRY RUN] ' if dry_run else ''}Purging {len(unlinked_ids)} unlinked signals older than {signal_cutoff.date()}")
            if not dry_run:
                session.execute(delete(ModelSignal).where(ModelSignal.id.in_(unlinked_ids)))

        # 2. Cleanup Risk Events (older than 2 years)
        risk_cutoff = now - timedelta(days=RETENTION_RISK_EVENTS)
        risk_query = select(RiskEvent.id).where(RiskEvent.created_at < risk_cutoff)
        risk_ids = session.execute(risk_query).scalars().all()
        results["risk_events"] = len(risk_ids)

        if risk_ids:
            logger.info(f"{'[DRY RUN] ' if dry_run else ''}Purging {len(risk_ids)} risk events older than {risk_cutoff.date()}")
            if not dry_run:
                session.execute(delete(RiskEvent).where(RiskEvent.id.in_(risk_ids)))

        # 3. Cleanup Performance Metrics (older than 2 years)
        perf_cutoff = now - timedelta(days=RETENTION_PERFORMANCE_METRICS)
        perf_query = select(PerformanceMetric.id).where(PerformanceMetric.created_at < perf_cutoff)
        perf_ids = session.execute(perf_query).scalars().all()
        results["performance_metrics"] = len(perf_ids)

        if perf_ids:
            logger.info(f"{'[DRY RUN] ' if dry_run else ''}Purging {len(perf_ids)} performance metrics older than {perf_cutoff.date()}")
            if not dry_run:
                session.execute(delete(PerformanceMetric).where(PerformanceMetric.id.in_(perf_ids)))

        # 4. Cleanup Trades (older than 7 years)
        # Note: This also enables cleanup of linked signals in the next run after 7 years
        trade_cutoff = now - timedelta(days=RETENTION_TRADES)
        trade_query = select(Trade.id).where(Trade.created_at < trade_cutoff)
        trade_ids = session.execute(trade_query).scalars().all()
        results["trades"] = len(trade_ids)

        if trade_ids:
            logger.info(f"{'[DRY RUN] ' if dry_run else ''}Purging {len(trade_ids)} trade records older than {trade_cutoff.date()}")
            if not dry_run:
                session.execute(delete(Trade).where(Trade.id.in_(trade_ids)))

        if not dry_run:
            session.commit()

    # 5. Cleanup Audit Log (older than 7 years)
    # Handle potentially separate audit database
    audit_engine = create_engine(audit_db_url) if audit_db_url and audit_db_url != db_url else engine
    AuditSession = sessionmaker(bind=audit_engine)

    with AuditSession() as session:
        audit_cutoff = now - timedelta(days=RETENTION_AUDIT_LOG)
        audit_query = select(AuditEntry.id).where(AuditEntry.created_at < audit_cutoff)
        audit_ids = session.execute(audit_query).scalars().all()
        results["audit_log"] = len(audit_ids)

        if audit_ids:
            logger.info(f"{'[DRY RUN] ' if dry_run else ''}Purging {len(audit_ids)} audit log entries older than {audit_cutoff.date()}")
            if not dry_run:
                session.execute(delete(AuditEntry).where(AuditEntry.id.in_(audit_ids)))

        if not dry_run:
            session.commit()

    return results


def main():
    parser = argparse.ArgumentParser(description="MT5 AI/ML Trading Bot - Data Cleanup Utility")
    parser.add_argument("--dry-run", action="store_true", help="Perform a dry run without deleting any data.")
    parser.add_argument("--db-url", help="Override the primary database URL from config.")
    parser.add_argument("--audit-db-url", help="Override the audit database URL.")
    parser.add_argument("--logs-dir", help="Override the logs directory from config.")
    parser.add_argument("--backtest-dir", help="Override the backtest results directory.")

    args = parser.parse_args()
    cfg = get_config()

    # Determine DB URLs mirroring main.py logic
    if args.db_url:
        db_url = args.db_url
    else:
        db_url = cfg.database_url if "sqlite" in cfg.database_url else "sqlite:///trades.db"

    if args.audit_db_url:
        audit_db_url = args.audit_db_url
    else:
        audit_db_url = cfg.database_url if "sqlite" in cfg.database_url else "sqlite:///audit.db"

    # Ensure we don't accidentally wipe a production PG DB unless intended
    if "sqlite" not in db_url:
        logger.warning(f"Using production-like primary DB URL: {db_url}")

    if "sqlite" not in audit_db_url:
        logger.warning(f"Using production-like audit DB URL: {audit_db_url}")

    logs_dir = Path(args.logs_dir) if args.logs_dir else cfg.logs_dir
    backtest_dir = Path(args.backtest_dir) if args.backtest_dir else Path(__file__).resolve().parents[1] / "backtest_results"

    logger.info(f"Starting data cleanup (dry_run={args.dry_run})")

    # Filesystem cleanup - Logs
    log_count = cleanup_logs(logs_dir, dry_run=args.dry_run)
    logger.info(f"Log cleanup complete. Total files processed: {log_count}")

    # Filesystem cleanup - Backtests
    backtest_count = cleanup_backtests(backtest_dir, dry_run=args.dry_run)
    logger.info(f"Backtest cleanup complete. Total files processed: {backtest_count}")

    # Database cleanup
    db_results = cleanup_database(db_url, audit_db_url=audit_db_url, dry_run=args.dry_run)
    logger.info("Database cleanup complete.")
    for table, count in db_results.items():
        logger.info(f"  - {table}: {count} records {'identified' if args.dry_run else 'purged'}")

    logger.info("Cleanup process finished successfully.")


if __name__ == "__main__":
    main()