""" Data Archiving System for AegisLM SaaS Backend. Production-ready data archiving with automated cleanup, compression, and retrieval capabilities. """ import asyncio import json import gzip import shutil from datetime import datetime, timedelta from pathlib import Path from typing import List, Dict, Optional, Any, Tuple from sqlalchemy import text, select, and_, func from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.orm import Session import logging import csv import io from .database import async_engine, get_sync_db from .config import settings from .backup_manager import backup_manager logger = logging.getLogger(__name__) class ArchivePolicy: """Data archiving policy configuration.""" def __init__(self, table_name: str, retention_days: int, archive_condition: Optional[str] = None, compress: bool = True, batch_size: int = 1000): self.table_name = table_name self.retention_days = retention_days self.archive_condition = archive_condition self.compress = compress self.batch_size = batch_size def get_cutoff_date(self) -> datetime: """Get cutoff date for archiving.""" return datetime.utcnow() - timedelta(days=self.retention_days) class ArchivedRecord: """Archived record metadata.""" def __init__(self, archive_id: str, table_name: str, record_count: int, archive_date: datetime, file_path: str, compressed: bool): self.archive_id = archive_id self.table_name = table_name self.record_count = record_count self.archive_date = archive_date self.file_path = file_path self.compressed = compressed def to_dict(self) -> Dict[str, Any]: """Convert to dictionary.""" return { "archive_id": self.archive_id, "table_name": self.table_name, "record_count": self.record_count, "archive_date": self.archive_date.isoformat(), "file_path": str(self.file_path), "compressed": self.compressed } class DataArchiver: """Data archiving management system.""" def __init__(self): self.archive_dir = Path(getattr(settings, 'ARCHIVE_DIR', './archives')) self.archive_dir.mkdir(exist_ok=True) # Default archiving policies self.policies = { 'evaluations': ArchivePolicy( table_name='evaluations', retention_days=getattr(settings, 'EVALUATION_RETENTION_DAYS', 365), archive_condition="status IN ('completed', 'failed', 'cancelled')", compress=True, batch_size=500 ), 'api_keys': ArchivePolicy( table_name='api_keys', retention_days=getattr(settings, 'API_KEY_RETENTION_DAYS', 730), archive_condition="is_active = false AND expires_at < NOW()", compress=True, batch_size=1000 ), } self.redis_client = None async def get_redis(self): """Get Redis client.""" if not self.redis_client: from .database import get_redis self.redis_client = await get_redis() return self.redis_client def _generate_archive_id(self, table_name: str) -> str: """Generate unique archive ID.""" timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S") return f"archive_{table_name}_{timestamp}" async def archive_table_data(self, table_name: str, policy: Optional[ArchivePolicy] = None) -> Optional[ArchivedRecord]: """Archive data from a specific table.""" try: # Use provided policy or default if not policy: policy = self.policies.get(table_name) if not policy: logger.error(f"No archiving policy found for table: {table_name}") return None # Get cutoff date cutoff_date = policy.get_cutoff_date() # Build archive query base_query = f"FROM {table_name} WHERE created_at < :cutoff_date" if policy.archive_condition: base_query += f" AND {policy.archive_condition}" # Get count of records to archive count_query = f"SELECT COUNT(*) {base_query}" async with async_engine.begin() as conn: result = await conn.execute(text(count_query), {"cutoff_date": cutoff_date}) record_count = result.scalar() if record_count == 0: logger.info(f"No records to archive in {table_name}") return None logger.info(f"Archiving {record_count} records from {table_name}") # Generate archive ID and file path archive_id = self._generate_archive_id(table_name) file_extension = '.csv.gz' if policy.compress else '.csv' file_path = self.archive_dir / f"{archive_id}{file_extension}" # Export data in batches await self._export_data_in_batches( conn, table_name, base_query, cutoff_date, file_path, policy.batch_size, policy.compress ) # Delete archived records delete_query = f"DELETE {base_query}" await conn.execute(text(delete_query), {"cutoff_date": cutoff_date}) # Create archive record archived_record = ArchivedRecord( archive_id=archive_id, table_name=table_name, record_count=record_count, archive_date=datetime.utcnow(), file_path=file_path, compressed=policy.compress ) # Record archive in Redis await self._record_archive(archived_record) logger.info(f"Successfully archived {record_count} records from {table_name}") return archived_record except Exception as e: logger.error(f"Failed to archive data from {table_name}: {e}") return None async def _export_data_in_batches(self, conn, table_name: str, base_query: str, cutoff_date: datetime, file_path: Path, batch_size: int, compress: bool): """Export data in batches to avoid memory issues.""" offset = 0 # Get column names columns_query = f""" SELECT column_name FROM information_schema.columns WHERE table_name = :table_name ORDER BY ordinal_position """ result = await conn.execute(text(columns_query), {"table_name": table_name}) columns = [row.column_name for row in result.fetchall()] # Open file for writing if compress: output_file = gzip.open(file_path, 'wt', newline='', encoding='utf-8') else: output_file = open(file_path, 'w', newline='', encoding='utf-8') try: writer = csv.DictWriter(output_file, fieldnames=columns) writer.writeheader() while True: # Get batch of data select_query = f"SELECT * {base_query} ORDER BY id LIMIT :batch_size OFFSET :offset" result = await conn.execute( text(select_query), {"cutoff_date": cutoff_date, "batch_size": batch_size, "offset": offset} ) rows = result.fetchall() if not rows: break # Write batch to CSV for row in rows: row_dict = dict(row._mapping) # Handle datetime serialization for key, value in row_dict.items(): if isinstance(value, datetime): row_dict[key] = value.isoformat() elif value is None: row_dict[key] = '' writer.writerow(row_dict) offset += batch_size logger.info(f"Exported batch {offset // batch_size} for {table_name}") finally: output_file.close() async def restore_archived_data(self, archive_id: str) -> bool: """Restore data from archive.""" try: # Get archive metadata archive_record = await self._get_archive_record(archive_id) if not archive_record: logger.error(f"Archive record not found: {archive_id}") return False file_path = Path(archive_record['file_path']) if not file_path.exists(): logger.error(f"Archive file not found: {file_path}") return False # Open archive file if archive_record['compressed']: input_file = gzip.open(file_path, 'rt', newline='', encoding='utf-8') else: input_file = open(file_path, 'r', newline='', encoding='utf-8') try: # Read CSV data reader = csv.DictReader(input_file) rows = list(reader) if not rows: logger.info(f"No data to restore from archive {archive_id}") return True # Restore data in batches table_name = archive_record['table_name'] await self._restore_data_in_batches(table_name, rows, batch_size=500) logger.info(f"Restored {len(rows)} records from archive {archive_id}") return True finally: input_file.close() except Exception as e: logger.error(f"Failed to restore archive {archive_id}: {e}") return False async def _restore_data_in_batches(self, table_name: str, rows: List[Dict], batch_size: int): """Restore data in batches.""" async with async_engine.begin() as conn: for i in range(0, len(rows), batch_size): batch = rows[i:i + batch_size] for row in batch: # Convert empty strings back to None for key, value in row.items(): if value == '': row[key] = None elif key in ['created_at', 'updated_at', 'started_at', 'completed_at', 'last_login_at', 'expires_at', 'last_used_at'] and value: row[key] = datetime.fromisoformat(value) # Build INSERT query columns = list(batch[0].keys()) placeholders = [f":{col}" for col in columns] insert_query = f""" INSERT INTO {table_name} ({', '.join(columns)}) VALUES ({', '.join(placeholders)}) """ await conn.execute(text(insert_query), batch) logger.info(f"Restored batch {i // batch_size + 1} for {table_name}") async def _record_archive(self, archive: ArchivedRecord): """Record archive metadata in Redis.""" try: redis_client = await self.get_redis() await redis_client.setex( f"archive:{archive.archive_id}", 86400 * 30, # 30 days json.dumps(archive.to_dict()) ) except Exception as e: logger.error(f"Failed to record archive: {e}") async def _get_archive_record(self, archive_id: str) -> Optional[Dict[str, Any]]: """Get archive record from Redis.""" try: redis_client = await self.get_redis() record = await redis_client.get(f"archive:{archive_id}") return json.loads(record) if record else None except Exception as e: logger.error(f"Failed to get archive record: {e}") return None async def list_archives(self, table_name: Optional[str] = None) -> List[Dict[str, Any]]: """List available archives.""" try: redis_client = await self.get_redis() # Get all archive keys cursor = 0 archives = [] while True: cursor, keys = await redis_client.scan(cursor, match="archive:*", count=100) for key in keys: archive_id = key.decode().split(':')[1] record = await self._get_archive_record(archive_id) if record and (not table_name or record['table_name'] == table_name): archives.append(record) if cursor == 0: break # Sort by archive date (newest first) archives.sort(key=lambda x: x['archive_date'], reverse=True) return archives except Exception as e: logger.error(f"Failed to list archives: {e}") return [] async def get_archive_statistics(self) -> Dict[str, Any]: """Get archiving statistics.""" try: archives = await self.list_archives() # Calculate statistics total_archives = len(archives) total_records = sum(archive['record_count'] for archive in archives) # Group by table table_stats = {} for archive in archives: table = archive['table_name'] if table not in table_stats: table_stats[table] = { 'archive_count': 0, 'total_records': 0, 'total_size_mb': 0 } table_stats[table]['archive_count'] += 1 table_stats[table]['total_records'] += archive['record_count'] # Calculate file size file_path = Path(archive['file_path']) if file_path.exists(): size_mb = file_path.stat().st_size / (1024 * 1024) table_stats[table]['total_size_mb'] += size_mb # Get current table sizes current_sizes = await self._get_current_table_sizes() return { "total_archives": total_archives, "total_records_archived": total_records, "table_statistics": table_stats, "current_table_sizes": current_sizes, "archive_directory": str(self.archive_dir), "policies": {name: { 'retention_days': policy.retention_days, 'archive_condition': policy.archive_condition } for name, policy in self.policies.items()} } except Exception as e: logger.error(f"Failed to get archive statistics: {e}") return {"error": str(e)} async def _get_current_table_sizes(self) -> Dict[str, Any]: """Get current table sizes.""" try: async with async_engine.begin() as conn: result = await conn.execute(text(""" SELECT tablename, pg_size_pretty(pg_total_relation_size(schemaname||'.'||tablename)) as size, pg_total_relation_size(schemaname||'.'||tablename) as size_bytes FROM pg_tables WHERE schemaname = 'public' ORDER BY pg_total_relation_size(schemaname||'.'||tablename) DESC """)) sizes = {} for row in result.fetchall(): sizes[row.tablename] = { "size_pretty": row.size, "size_bytes": row.size_bytes } return sizes except Exception as e: logger.error(f"Failed to get table sizes: {e}") return {} async def run_scheduled_archiving(self) -> Dict[str, Any]: """Run scheduled data archiving for all policies.""" results = { "archived": [], "failed": [], "skipped": [] } try: for table_name, policy in self.policies.items(): logger.info(f"Running scheduled archiving for {table_name}") archive_record = await self.archive_table_data(table_name, policy) if archive_record: results["archived"].append({ "table": table_name, "archive_id": archive_record.archive_id, "record_count": archive_record.record_count }) else: results["skipped"].append({ "table": table_name, "reason": "No data to archive" }) # Clean up old archive records await self._cleanup_old_archive_records() logger.info(f"Scheduled archiving completed: {len(results['archived'])} archived, {len(results['failed'])} failed, {len(results['skipped'])} skipped") except Exception as e: logger.error(f"Scheduled archiving failed: {e}") results["error"] = str(e) return results async def _cleanup_old_archive_records(self): """Clean up old archive records from Redis.""" try: redis_client = await self.get_redis() # Get all archive keys cursor = 0 old_keys = [] while True: cursor, keys = await redis_client.scan(cursor, match="archive:*", count=100) for key in keys: archive_id = key.decode().split(':')[1] record = await self._get_archive_record(archive_id) if record: archive_date = datetime.fromisoformat(record['archive_date']) if archive_date < datetime.utcnow() - timedelta(days=90): old_keys.append(key) if cursor == 0: break # Delete old keys if old_keys: await redis_client.delete(*old_keys) logger.info(f"Cleaned up {len(old_keys)} old archive records") except Exception as e: logger.error(f"Failed to cleanup old archive records: {e}") async def verify_archive_integrity(self, archive_id: str) -> Dict[str, Any]: """Verify archive file integrity.""" try: archive_record = await self._get_archive_record(archive_id) if not archive_record: return {"valid": False, "error": "Archive record not found"} file_path = Path(archive_record['file_path']) if not file_path.exists(): return {"valid": False, "error": "Archive file not found"} # Check file size file_size = file_path.stat().st_size # For compressed files, verify gzip integrity if archive_record['compressed']: try: with gzip.open(file_path, 'rb') as f: f.read(1024) # Try to read first 1KB except Exception as e: return {"valid": False, "error": f"Corrupted gzip file: {e}"} # Count records in file if archive_record['compressed']: input_file = gzip.open(file_path, 'rt', newline='', encoding='utf-8') else: input_file = open(file_path, 'r', newline='', encoding='utf-8') try: reader = csv.DictReader(input_file) record_count = sum(1 for _ in reader) finally: input_file.close() # Verify record count matches count_matches = record_count == archive_record['record_count'] return { "valid": True, "file_size_bytes": file_size, "actual_record_count": record_count, "expected_record_count": archive_record['record_count'], "count_matches": count_matches, "archive_date": archive_record['archive_date'] } except Exception as e: return {"valid": False, "error": str(e)} # Global data archiver instance data_archiver = DataArchiver() # Scheduled archiving task async def scheduled_archiving_task(): """Run scheduled data archiving.""" logger.info("Starting scheduled data archiving...") try: results = await data_archiver.run_scheduled_archiving() # Store results in Redis redis_client = await data_archiver.get_redis() await redis_client.setex( "archiving_results", 3600, # 1 hour json.dumps(results, default=str) ) logger.info("Scheduled data archiving completed") except Exception as e: logger.error(f"Scheduled data archiving failed: {e}") if __name__ == "__main__": import sys async def main(): command = sys.argv[1] if len(sys.argv) > 1 else "help" if command == "archive": table_name = sys.argv[2] if len(sys.argv) > 2 else None if not table_name: print("Error: archive requires table name") sys.exit(1) archive_record = await data_archiver.archive_table_data(table_name) if archive_record: print(f"✅ Archived {archive_record.record_count} records: {archive_record.archive_id}") else: print("❌ Archiving failed") elif command == "restore": archive_id = sys.argv[2] if len(sys.argv) > 2 else None if not archive_id: print("Error: restore requires archive ID") sys.exit(1) success = await data_archiver.restore_archived_data(archive_id) if success: print(f"✅ Archive {archive_id} restored successfully") else: print(f"❌ Archive {archive_id} restore failed") elif command == "list": table_name = sys.argv[2] if len(sys.argv) > 2 else None archives = await data_archiver.list_archives(table_name) print(f"Archives: {len(archives)}") for archive in archives: print(f" - {archive['archive_id']}: {archive['table_name']} ({archive['record_count']} records, {archive['archive_date']})") elif command == "stats": stats = await data_archiver.get_archive_statistics() print(json.dumps(stats, indent=2, default=str)) elif command == "verify": archive_id = sys.argv[2] if len(sys.argv) > 2 else None if not archive_id: print("Error: verify requires archive ID") sys.exit(1) verification = await data_archiver.verify_archive_integrity(archive_id) if verification["valid"]: print(f"✅ Archive {archive_id} is valid") else: print(f"❌ Archive {archive_id} verification failed") print(f"Error: {verification.get('error', 'Unknown error')}") elif command == "scheduled": results = await data_archiver.run_scheduled_archiving() print(json.dumps(results, indent=2, default=str)) else: print("Usage: python data_archiver.py [args]") print("Commands: archive , restore , list [table], stats, verify , scheduled") asyncio.run(main())