spoolman2/spoolman/api/v1/backup.py

"""Database backup and restore API endpoints."""

import json
import logging
from datetime import datetime
from typing import Annotated, Any, Optional

from fastapi import APIRouter, Depends, File, UploadFile
from fastapi.responses import JSONResponse, Response
from pydantic import BaseModel, Field
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from spoolman.database.database import get_db_session
from spoolman.database.models import (
    Filament,
    FilamentField,
    Setting,
    Spool,
    SpoolAdjustment,
    SpoolField,
    Vendor,
    VendorField,
)

logger = logging.getLogger(__name__)

router = APIRouter(
    prefix="/backup",
    tags=["backup"],
)


class BackupData(BaseModel):
    """Schema for backup data."""

    version: str = Field(default="1.0", description="Backup format version")
    exported_at: str = Field(description="ISO timestamp of export")
    vendors: list[dict[str, Any]] = Field(default_factory=list)
    filaments: list[dict[str, Any]] = Field(default_factory=list)
    spools: list[dict[str, Any]] = Field(default_factory=list)
    spool_adjustments: list[dict[str, Any]] = Field(default_factory=list)
    settings: list[dict[str, Any]] = Field(default_factory=list)


class ImportResult(BaseModel):
    """Result of import operation."""

    success: bool
    vendors_imported: int = 0
    filaments_imported: int = 0
    spools_imported: int = 0
    adjustments_imported: int = 0
    settings_imported: int = 0
    errors: list[str] = Field(default_factory=list)


def _serialize_datetime(obj: Any) -> Any:
    """Convert datetime objects to ISO strings."""
    if isinstance(obj, datetime):
        return obj.isoformat()
    return obj


def _model_to_dict(obj: Any, exclude: Optional[set[str]] = None) -> dict[str, Any]:
    """Convert SQLAlchemy model to dict, handling extra fields."""
    exclude = exclude or set()
    result = {}

    for column in obj.__table__.columns:
        if column.name not in exclude:
            value = getattr(obj, column.name)
            result[column.name] = _serialize_datetime(value)

    # Handle extra fields if present
    if hasattr(obj, "extra") and obj.extra:
        result["extra"] = {field.key: field.value for field in obj.extra}

    return result


@router.get(
    "/export",
    name="Export all data",
    description="Export all Spoolman data as a single JSON file for backup or migration.",
    response_class=Response,
)
async def export_all(
    db: Annotated[AsyncSession, Depends(get_db_session)],
) -> Response:
    """Export all database data as JSON."""

    # Fetch all vendors
    vendors_result = await db.execute(select(Vendor))
    vendors = [_model_to_dict(v) for v in vendors_result.unique().scalars().all()]

    # Fetch all filaments
    filaments_result = await db.execute(select(Filament))
    filaments = [_model_to_dict(f) for f in filaments_result.unique().scalars().all()]

    # Fetch all spools
    spools_result = await db.execute(select(Spool))
    spools = [_model_to_dict(s) for s in spools_result.unique().scalars().all()]

    # Fetch all spool adjustments
    adjustments_result = await db.execute(select(SpoolAdjustment))
    adjustments = [_model_to_dict(a, exclude={"spool"}) for a in adjustments_result.scalars().all()]

    # Fetch all settings
    settings_result = await db.execute(select(Setting))
    settings = [_model_to_dict(s) for s in settings_result.scalars().all()]

    backup = BackupData(
        version="1.0",
        exported_at=datetime.utcnow().isoformat(),
        vendors=vendors,
        filaments=filaments,
        spools=spools,
        spool_adjustments=adjustments,
        settings=settings,
    )

    # Generate filename with timestamp
    filename = f"spoolman_backup_{datetime.utcnow().strftime('%Y%m%d_%H%M%S')}.json"

    return Response(
        content=json.dumps(backup.model_dump(), indent=2, default=str),
        media_type="application/json",
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )


@router.post(
    "/import",
    name="Import data",
    description="Import Spoolman data from a backup JSON file.",
    response_model=ImportResult,
)
async def import_all(
    db: Annotated[AsyncSession, Depends(get_db_session)],
    file: Annotated[UploadFile, File(description="Backup JSON file to import")],
    replace: bool = False,
) -> ImportResult:
    """Import database data from JSON backup.

    Args:
        file: The backup JSON file to import
        replace: If True, clear existing data before import. If False, merge/skip conflicts.
    """
    result = ImportResult(success=False)

    try:
        content = await file.read()
        data = json.loads(content)

        # Validate version
        version = data.get("version", "1.0")
        if version != "1.0":
            result.errors.append(f"Unsupported backup version: {version}")
            return result

        # If replace mode, clear existing data (order matters due to foreign keys)
        if replace:
            await db.execute(SpoolAdjustment.__table__.delete())
            await db.execute(SpoolField.__table__.delete())
            await db.execute(Spool.__table__.delete())
            await db.execute(FilamentField.__table__.delete())
            await db.execute(Filament.__table__.delete())
            await db.execute(VendorField.__table__.delete())
            await db.execute(Vendor.__table__.delete())
            await db.commit()

        # Track ID mappings for relationships (old_id -> new_id)
        vendor_id_map: dict[int, int] = {}
        filament_id_map: dict[int, int] = {}
        spool_id_map: dict[int, int] = {}

        # Import vendors
        for vendor_data in data.get("vendors", []):
            old_id = vendor_data.pop("id", None)
            extra_fields = vendor_data.pop("extra", {})

            # Parse datetime
            if "registered" in vendor_data and isinstance(vendor_data["registered"], str):
                vendor_data["registered"] = datetime.fromisoformat(vendor_data["registered"])

            vendor = Vendor(**vendor_data)
            db.add(vendor)
            await db.flush()  # Get the new ID

            if old_id is not None:
                vendor_id_map[old_id] = vendor.id

            # Add extra fields
            for key, value in extra_fields.items():
                db.add(VendorField(vendor_id=vendor.id, key=key, value=value))

            result.vendors_imported += 1

        # Import filaments
        for filament_data in data.get("filaments", []):
            old_id = filament_data.pop("id", None)
            extra_fields = filament_data.pop("extra", {})

            # Remap vendor_id
            if filament_data.get("vendor_id") and filament_data["vendor_id"] in vendor_id_map:
                filament_data["vendor_id"] = vendor_id_map[filament_data["vendor_id"]]

            # Parse datetime
            if "registered" in filament_data and isinstance(filament_data["registered"], str):
                filament_data["registered"] = datetime.fromisoformat(filament_data["registered"])

            filament = Filament(**filament_data)
            db.add(filament)
            await db.flush()

            if old_id is not None:
                filament_id_map[old_id] = filament.id

            # Add extra fields
            for key, value in extra_fields.items():
                db.add(FilamentField(filament_id=filament.id, key=key, value=value))

            result.filaments_imported += 1

        # Import spools
        for spool_data in data.get("spools", []):
            old_id = spool_data.pop("id", None)
            extra_fields = spool_data.pop("extra", {})

            # Remap filament_id
            if spool_data.get("filament_id") and spool_data["filament_id"] in filament_id_map:
                spool_data["filament_id"] = filament_id_map[spool_data["filament_id"]]

            # Parse datetimes
            for field in ["registered", "first_used", "last_used"]:
                if field in spool_data and isinstance(spool_data[field], str):
                    spool_data[field] = datetime.fromisoformat(spool_data[field])

            spool = Spool(**spool_data)
            db.add(spool)
            await db.flush()

            if old_id is not None:
                spool_id_map[old_id] = spool.id

            # Add extra fields
            for key, value in extra_fields.items():
                db.add(SpoolField(spool_id=spool.id, key=key, value=value))

            result.spools_imported += 1

        # Import spool adjustments
        for adj_data in data.get("spool_adjustments", []):
            adj_data.pop("id", None)  # Remove old ID

            # Remap spool_id
            if adj_data.get("spool_id") and adj_data["spool_id"] in spool_id_map:
                adj_data["spool_id"] = spool_id_map[adj_data["spool_id"]]
            else:
                # Skip if spool doesn't exist
                continue

            # Parse timestamp
            if "timestamp" in adj_data and isinstance(adj_data["timestamp"], str):
                adj_data["timestamp"] = datetime.fromisoformat(adj_data["timestamp"])

            adjustment = SpoolAdjustment(**adj_data)
            db.add(adjustment)
            result.adjustments_imported += 1

        # Import settings
        for setting_data in data.get("settings", []):
            # Parse datetime
            if "last_updated" in setting_data and isinstance(setting_data["last_updated"], str):
                setting_data["last_updated"] = datetime.fromisoformat(setting_data["last_updated"])

            # Use merge to handle existing settings
            setting = Setting(**setting_data)
            await db.merge(setting)
            result.settings_imported += 1

        await db.commit()
        result.success = True

    except json.JSONDecodeError as e:
        result.errors.append(f"Invalid JSON: {e}")
    except Exception as e:
        logger.exception("Import failed")
        result.errors.append(str(e))
        await db.rollback()

    return result