import os import shutil import datetime import sqlite3 import logging from typing import List from sqlalchemy import text from sqlalchemy.orm import Session from .database import DATA_DIR, engine from . import models, schemas logger = logging.getLogger("ainventory") BACKUP_DIR = os.path.join(DATA_DIR, "backups") os.makedirs(BACKUP_DIR, exist_ok=True) class DbManager: @staticmethod def get_backup_list() -> List[schemas.BackupInfo]: """List all available backup files with metadata.""" backups = [] if not os.path.exists(BACKUP_DIR): return [] for filename in os.listdir(BACKUP_DIR): if filename.endswith(".db"): path = os.path.join(BACKUP_DIR, filename) stats = os.stat(path) backups.append(schemas.BackupInfo( filename=filename, size_bytes=stats.st_size, created_at=datetime.datetime.fromtimestamp(stats.st_mtime) )) # Sort by creation time descending backups.sort(key=lambda x: x.created_at, reverse=True) return backups @staticmethod def get_stats() -> schemas.DatabaseStats: """Calculate storage statistics for backups.""" backups = DbManager.get_backup_list() total_size = sum(b.size_bytes for b in backups) return schemas.DatabaseStats( backup_count=len(backups), total_size_bytes=total_size ) @staticmethod def create_backup(db: Session, label: str = "manual", user_id: int = None) -> str: """ Creates a consistent snapshot of the active database using VACUUM INTO. Records the action in AuditLog. """ timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") filename = f"inventory_{label}_{timestamp}.db" target_path = os.path.join(BACKUP_DIR, filename) try: # Use SQLite's VACUUM INTO for a safe, non-blocking backup # We need a clean string path without potential SQL injection (filenames are generated here) db.execute(text(f"VACUUM INTO '{target_path}'")) logger.info(f"[DB] Backup created successfully: {filename}") # Audit Log audit = models.AuditLog( user_id=user_id, action="DB_BACKUP", details=f"Backup type: {label}. Created file: {filename}" ) db.add(audit) db.commit() # Enforce retention policy DbManager.enforce_retention(db) return filename except Exception as e: logger.error(f"[DB] Backup failed: {str(e)}") db.rollback() raise Exception(f"Backup failed: {str(e)}") @staticmethod def enforce_retention(db: Session): """Deletes oldest backups if count exceeds retention limit.""" try: # Get retention limit from settings limit_setting = db.query(models.SystemSetting).filter(models.SystemSetting.key == "backup_retention_count").first() limit = int(limit_setting.value) if limit_setting else 10 # Default to 10 backups = DbManager.get_backup_list() if len(backups) > limit: to_delete = backups[limit:] for b in to_delete: path = os.path.join(BACKUP_DIR, b.filename) if os.path.exists(path): os.remove(path) logger.info(f"[DB] Retention policy: deleted old backup {b.filename}") except Exception as e: logger.error(f"[DB] Retention enforcement failed: {str(e)}") @staticmethod def restore_backup(filename: str, db: Session, user_id: int) -> bool: """ Restores a database from a backup file. IMPORTANT: This replaces the primary inventory.db file. """ source_path = os.path.join(BACKUP_DIR, filename) active_db_path = os.path.join(DATA_DIR, "inventory.db") if not os.path.exists(source_path): raise Exception("Backup file not found") try: # 1. Create a safety rollback backup of current state logger.info("[DB] Creating safety rollback backup before restore...") DbManager.create_backup(db, label="rollback", user_id=user_id) # 2. Close all connections (or as many as possible via pool dispose) # engine.dispose() drops the current connection pool engine.dispose() # 3. Physically swap files # Note: shutil.copy2 handles file metadata preservation shutil.copy2(source_path, active_db_path) logger.warning(f"[DB] RESTORE COMPLETED from {filename} by user_id={user_id}") # Record the restore in the NEW database (since we just swapped it) # Re-initializing session logic might be needed but usually next request handles it. # However, for the current request, the 'db' session is still bound to the OLD file mapping possibly? # Actually, the file on disk is changed. Next commit might fail or work depending on how sqlite handles it. # It's safest to return success and let the frontend trigger a reload. return True except Exception as e: logger.error(f"[DB] Restore failed: {str(e)}") raise Exception(f"Restore failed: {str(e)}") @staticmethod def export_db() -> str: """Returns the path to the active database file for download.""" active_db_path = os.path.join(DATA_DIR, "inventory.db") if not os.path.exists(active_db_path): raise Exception("Database file not found") return active_db_path @staticmethod def import_db(file_bytes: bytes, db: Session, user_id: int) -> bool: """ Overwrites the active database with the provided file bytes. Creates a safety rollback backup first. """ active_db_path = os.path.join(DATA_DIR, "inventory.db") try: # 1. Safety backup DbManager.create_backup(db, label="import_rollback", user_id=user_id) # 2. Close pool connections engine.dispose() # 3. Write new file with open(active_db_path, "wb") as f: f.write(file_bytes) logger.warning(f"[DB] IMPORT COMPLETED by user_id={user_id}") return True except Exception as e: logger.error(f"[DB] Import failed: {str(e)}") raise Exception(f"Import failed: {str(e)}")