TimeTracker/app/utils/backup.py

import os
import io
import json
import shutil
import tempfile
import subprocess
from datetime import datetime
from zipfile import ZipFile, ZIP_DEFLATED
from urllib.parse import urlparse


def _get_backup_root_dir(app):
    """Compute the absolute backups directory path (project_root/backups)."""
    project_root = os.path.abspath(os.path.join(app.root_path, ".."))
    backups_dir = os.path.join(project_root, "backups")
    os.makedirs(backups_dir, exist_ok=True)
    return backups_dir


def _now_timestamp():
    """Return a human-readable local timestamp for file names."""
    # Respect user's preference to use local time across the project
    # rather than UTC for user-facing timestamps.
    return datetime.now().strftime("%Y%m%d_%H%M%S")


def _detect_db_type_and_path(app):
    """Detect database type and return a tuple (type, uri, sqlite_path).

    type: 'sqlite' or 'postgresql'
    uri: full SQLAlchemy database URI
    sqlite_path: file path if sqlite, otherwise None
    """
    uri = app.config.get("SQLALCHEMY_DATABASE_URI", "") or ""
    if isinstance(uri, str) and uri.startswith("sqlite:///"):
        return "sqlite", uri, uri.replace("sqlite:///", "")
    if isinstance(uri, str) and (uri.startswith("postgresql") or uri.startswith("postgres")):
        return "postgresql", uri, None
    # Default/fallback
    return "unknown", uri, None


def _get_alembic_revision(db_session):
    """Return current alembic revision string or None if unavailable."""
    try:
        from sqlalchemy import text

        result = db_session.execute(text("SELECT version_num FROM alembic_version"))
        row = result.first()
        return row[0] if row else None
    except Exception:
        return None


def _write_manifest(zf, manifest: dict):
    data = json.dumps(manifest, indent=2, sort_keys=True).encode("utf-8")
    zf.writestr("manifest.json", data)


def _add_directory_to_zip(zf, source_dir: str, arc_prefix: str):
    if not source_dir or not os.path.isdir(source_dir):
        return
    for root, _, files in os.walk(source_dir):
        for file_name in files:
            abs_path = os.path.join(root, file_name)
            rel_path = os.path.relpath(abs_path, start=source_dir)
            zf.write(abs_path, os.path.join(arc_prefix, rel_path))


def create_backup(app) -> str:
    """Create a comprehensive backup archive and return its absolute path.

    Contents:
    - manifest.json (metadata: created_at, db_type, alembic_revision)
    - db.sqlite (if sqlite) OR db.dump (if postgresql, custom format)
    - settings.json (serialized Settings row for quick inspection)
    - uploads/ (logos and other uploaded assets if present)
    """
    # Late imports to avoid circular dependencies
    from app import db
    from app.models.settings import Settings

    backups_dir = _get_backup_root_dir(app)
    timestamp = _now_timestamp()
    archive_name = f"timetracker_backup_{timestamp}.zip"
    archive_path = os.path.join(backups_dir, archive_name)

    db_type, db_uri, sqlite_path = _detect_db_type_and_path(app)

    # Prepare temporary directory for DB dumps if needed
    tmp_dir = tempfile.mkdtemp(prefix="tt_backup_")
    tmp_db_artifact = None

    try:
        # Create DB artifact
        if db_type == "sqlite" and sqlite_path and os.path.exists(sqlite_path):
            tmp_db_artifact = os.path.join(tmp_dir, "db.sqlite")
            shutil.copy2(sqlite_path, tmp_db_artifact)
        elif db_type == "postgresql":
            # Use parsed connection parameters (avoid SQLAlchemy driver suffix in URI)
            database_url = os.getenv("DATABASE_URL", db_uri)
            parsed = urlparse(database_url) if database_url else None
            host = parsed.hostname if parsed and parsed.hostname else os.getenv("POSTGRES_HOST", "db")
            port = parsed.port if parsed and parsed.port else int(os.getenv("POSTGRES_PORT", "5432"))
            user = parsed.username if parsed and parsed.username else os.getenv("POSTGRES_USER", "timetracker")
            password = parsed.password if parsed and parsed.password else os.getenv("POSTGRES_PASSWORD", "timetracker")
            dbname = parsed.path.lstrip("/") if parsed and parsed.path else os.getenv("POSTGRES_DB", "timetracker")

            tmp_db_artifact = os.path.join(tmp_dir, "db.dump")
            pg_dump_cmd = [
                "pg_dump",
                "--format=custom",
                "-h",
                host,
                "-p",
                str(port),
                "-U",
                user,
                "-d",
                dbname,
                f"--file={tmp_db_artifact}",
            ]
            env = os.environ.copy()
            if password:
                env["PGPASSWORD"] = str(password)
            try:
                completed = subprocess.run(
                    pg_dump_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env
                )
            except FileNotFoundError:
                raise RuntimeError("pg_dump not found. Please ensure PostgreSQL client tools are installed.")
            except subprocess.CalledProcessError as e:
                stderr = e.stderr.decode("utf-8", errors="ignore") if e.stderr else ""
                raise RuntimeError(f"pg_dump failed: {stderr.strip() or e}")
        else:
            # Best effort: we continue without DB artifact
            tmp_db_artifact = None

        # Gather metadata
        alembic_rev = _get_alembic_revision(db.session)
        manifest = {
            "created_at": datetime.now().isoformat(timespec="seconds"),
            "db_type": db_type,
            "alembic_revision": alembic_rev,
            "app_version": None,
        }

        # Serialize settings for convenience (DB backup still authoritative)
        settings_obj = Settings.get_settings()
        settings_json = json.dumps(settings_obj.to_dict(), indent=2, sort_keys=True)

        # Write the zip
        with ZipFile(archive_path, mode="w", compression=ZIP_DEFLATED) as zf:
            _write_manifest(zf, manifest)
            if tmp_db_artifact and os.path.exists(tmp_db_artifact):
                arc_name = "db.sqlite" if db_type == "sqlite" else "db.dump"
                zf.write(tmp_db_artifact, arc_name)

            zf.writestr("settings.json", settings_json.encode("utf-8"))

            # Include uploads (e.g., logos)
            uploads_root = os.path.join(app.root_path, "static", "uploads")
            _add_directory_to_zip(zf, uploads_root, "uploads")

        return archive_path
    finally:
        try:
            shutil.rmtree(tmp_dir, ignore_errors=True)
        except Exception:
            pass


def restore_backup(app, archive_path: str, progress_callback=None) -> tuple[bool, str]:
    """Restore a backup archive.

    Steps:
    - Extract archive to temp dir
    - Restore DB depending on type
    - Copy uploads back
    - Run migrations to head to ensure compatibility with newer code

    Returns: (success, message)
    """
    from app import db
    from time import sleep

    if not archive_path or not os.path.exists(archive_path):
        return False, f"Backup archive not found: {archive_path}"

    db_type, db_uri, sqlite_path = _detect_db_type_and_path(app)
    tmp_dir = tempfile.mkdtemp(prefix="tt_restore_")

    def _progress(label: str, percent: int):
        try:
            if callable(progress_callback):
                progress_callback(label, percent)
        except Exception:
            pass

    try:
        # Extract archive
        with ZipFile(archive_path, mode="r") as zf:
            zf.extractall(tmp_dir)
        _progress("Archive extracted", 10)

        # Read manifest (optional)
        manifest_path = os.path.join(tmp_dir, "manifest.json")
        if os.path.exists(manifest_path):
            try:
                with open(manifest_path, "r", encoding="utf-8") as f:
                    _ = json.load(f)
            except Exception:
                pass

        # Proactively close any open DB connections before modifying files
        try:
            with app.app_context():
                db.session.remove()
                db.engine.dispose()
        except Exception:
            pass

        # Restore DB
        if db_type == "sqlite":
            src_sqlite = os.path.join(tmp_dir, "db.sqlite")
            if not os.path.exists(src_sqlite):
                return False, "Backup does not contain db.sqlite for SQLite restore"

            if not sqlite_path:
                return False, "Current configuration is not SQLite or path not found"

            # Ensure destination directory exists
            dest_dir = os.path.dirname(sqlite_path)
            if dest_dir and not os.path.exists(dest_dir):
                os.makedirs(dest_dir, exist_ok=True)

            # Safety copy of current DB if exists
            if os.path.exists(sqlite_path):
                safety_copy = sqlite_path + f".bak_{_now_timestamp()}"
                shutil.copy2(sqlite_path, safety_copy)

            # Replace DB file
            os.makedirs(os.path.dirname(sqlite_path), exist_ok=True)
            # Retry a few times in case the file is briefly locked
            last_err = None
            for _ in range(3):
                try:
                    shutil.copy2(src_sqlite, sqlite_path)
                    last_err = None
                    break
                except Exception as e:
                    last_err = e
                    sleep(0.2)
            if last_err:
                return False, f"Failed to write SQLite database file: {last_err}"
            _progress("SQLite database restored", 60)

        elif db_type == "postgresql":
            src_dump = os.path.join(tmp_dir, "db.dump")
            if not os.path.exists(src_dump):
                return False, "Backup does not contain db.dump for PostgreSQL restore"

            database_url = os.getenv("DATABASE_URL", db_uri)
            parsed = urlparse(database_url) if database_url else None
            host = parsed.hostname if parsed and parsed.hostname else os.getenv("POSTGRES_HOST", "db")
            port = parsed.port if parsed and parsed.port else int(os.getenv("POSTGRES_PORT", "5432"))
            user = parsed.username if parsed and parsed.username else os.getenv("POSTGRES_USER", "timetracker")
            password = parsed.password if parsed and parsed.password else os.getenv("POSTGRES_PASSWORD", "timetracker")
            dbname = parsed.path.lstrip("/") if parsed and parsed.path else os.getenv("POSTGRES_DB", "timetracker")

            pg_restore_cmd = [
                "pg_restore",
                "--clean",
                "--if-exists",
                "--no-owner",
                "-h",
                host,
                "-p",
                str(port),
                "-U",
                user,
                "-d",
                dbname,
                src_dump,
            ]
            env = os.environ.copy()
            if password:
                env["PGPASSWORD"] = str(password)
            try:
                _progress("Restoring PostgreSQL database", 20)
                subprocess.run(pg_restore_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env)
            except FileNotFoundError:
                return False, "pg_restore not found. Please install PostgreSQL client tools."
            except subprocess.CalledProcessError as e:
                stderr = e.stderr.decode("utf-8", errors="ignore") if e.stderr else ""
                return False, f"pg_restore failed: {stderr.strip() or e}"
            _progress("PostgreSQL database restored", 60)
        else:
            return False, "Unsupported or unknown database type for restore"

        # Restore uploads
        extracted_uploads = os.path.join(tmp_dir, "uploads")
        if os.path.isdir(extracted_uploads):
            target_uploads = os.path.join(app.root_path, "static", "uploads")
            os.makedirs(target_uploads, exist_ok=True)
            # Merge copy
            for root, _, files in os.walk(extracted_uploads):
                rel = os.path.relpath(root, extracted_uploads)
                target_dir = os.path.join(target_uploads, rel) if rel != "." else target_uploads
                os.makedirs(target_dir, exist_ok=True)
                for fn in files:
                    shutil.copy2(os.path.join(root, fn), os.path.join(target_dir, fn))
        _progress("Uploads restored", 80)

        # Run migrations to ensure compatibility with current code
        try:
            from flask_migrate import upgrade as alembic_upgrade

            with app.app_context():
                _progress("Running migrations", 90)
                alembic_upgrade()
        except Exception as e:
            # If migrations fail, report failure to caller for visibility
            return False, f"Restore completed but migration failed: {e}"

        # Dispose connections once more after restore/migrate to ensure clean state
        try:
            with app.app_context():
                db.session.remove()
                db.engine.dispose()
        except Exception:
            pass

        _progress("Restore completed successfully", 100)
        return True, "Restore completed successfully"
    finally:
        try:
            shutil.rmtree(tmp_dir, ignore_errors=True)
        except Exception:
            pass