tfm_ainventory/backend/db_manager.py

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)}")