Spaces:
Sleeping
Sleeping
| import os | |
| import sys | |
| from datetime import datetime, timedelta | |
| import logging | |
| # Add project root to sys.path | |
| sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) | |
| from src.database.models import SessionLocal, RawNews, VerifiedNews, SavedArticle, TopicTracking, ReadHistory, BreakingNews, TrackNotification | |
| logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') | |
| logger = logging.getLogger("DataRetention") | |
| def prune_old_data(days=10): | |
| """ | |
| Deletes news articles older than the specified number of days (default 10) to save space and egress. | |
| Articles saved by users or tracked are preserved. | |
| """ | |
| db = SessionLocal() | |
| threshold = datetime.utcnow() - timedelta(days=days) | |
| try: | |
| logger.info(f"Starting data pruning (Threshold: articles older than {threshold})") | |
| # 1. Get IDs of articles that should NOT be deleted | |
| logger.info("Identifying protected articles...") | |
| protected_ids = set() | |
| # Saved by users | |
| saved_ids = db.query(SavedArticle.news_id).all() | |
| protected_ids.update([sid[0] for sid in saved_ids if sid[0]]) | |
| # Tracked by users | |
| tracked_ids = db.query(TopicTracking.news_id).all() | |
| protected_ids.update([tid[0] for tid in tracked_ids if tid[0]]) | |
| # Notifications sent | |
| notify_ids = db.query(TrackNotification.news_id).all() | |
| protected_ids.update([nid[0] for nid in notify_ids if nid[0]]) | |
| # Breaking news references | |
| breaking_ids = db.query(BreakingNews.verified_news_id).all() | |
| protected_ids.update([bid[0] for bid in breaking_ids if bid[0]]) | |
| logger.info(f"Protected VerifiedNews count: {len(protected_ids)}") | |
| # 2. Delete old VerifiedNews | |
| # Filter: older than threshold AND not in protected list | |
| verified_query = db.query(VerifiedNews).filter( | |
| VerifiedNews.published_at < threshold, | |
| ~VerifiedNews.id.in_(list(protected_ids)) if protected_ids else True | |
| ) | |
| verified_to_delete = verified_query.count() | |
| if verified_to_delete > 0: | |
| logger.info(f"Deleting {verified_to_delete} old VerifiedNews records...") | |
| verified_query.delete(synchronize_session=False) | |
| else: | |
| logger.info("No old VerifiedNews to delete.") | |
| # 3. Delete old RawNews | |
| # Filter: older than threshold AND NOT referenced by any VerifiedNews | |
| logger.info("Identifying RawNews referenced by any VerifiedNews...") | |
| referenced_raw_ids = db.query(VerifiedNews.raw_news_id).filter(VerifiedNews.raw_news_id != None).all() | |
| protected_raw_ids = set([r[0] for r in referenced_raw_ids]) | |
| raw_query = db.query(RawNews).filter( | |
| RawNews.published_at < threshold, | |
| ~RawNews.id.in_(list(protected_raw_ids)) if protected_raw_ids else True | |
| ) | |
| raw_to_delete = raw_query.count() | |
| if raw_to_delete > 0: | |
| logger.info(f"Deleting {raw_to_delete} old RawNews records...") | |
| raw_query.delete(synchronize_session=False) | |
| else: | |
| logger.info("No old RawNews to delete.") | |
| # 4. Cleanup old ReadHistory (optional but good for DB size) | |
| # Keep 30 days of history? | |
| history_threshold = datetime.utcnow() - timedelta(days=30) | |
| history_query = db.query(ReadHistory).filter(ReadHistory.read_at < history_threshold) | |
| history_to_delete = history_query.count() | |
| if history_to_delete > 0: | |
| logger.info(f"Deleting {history_to_delete} old ReadHistory records...") | |
| history_query.delete(synchronize_session=False) | |
| db.commit() | |
| logger.info("Pruning completed successfully.") | |
| except Exception as e: | |
| db.rollback() | |
| logger.error(f"Error during pruning: {e}") | |
| finally: | |
| db.close() | |
| if __name__ == "__main__": | |
| # Default to 10 days as requested | |
| prune_old_data(days=10) | |