Nexus/server/application/services/sync_engine_v2.py

"""Nexus — Sync Engine v2 (Step 4 Upgrade)

S1: Existing rsync push retained as Sync file mode
S5: Parallel execution + result aggregation with Semaphore
"""

import asyncio
import json
import logging
from datetime import datetime, timezone
from typing import List, Optional

from sqlalchemy.ext.asyncio import AsyncSession

from server.domain.models import Server, SyncLog, AuditLog
from server.infrastructure.database.server_repo import ServerRepositoryImpl
from server.infrastructure.database.sync_log_repo import SyncLogRepositoryImpl
from server.infrastructure.database.audit_log_repo import AuditLogRepositoryImpl
from server.infrastructure.database.push_schedule_repo import PushRetryJobRepositoryImpl
from server.infrastructure.redis.client import get_redis
from server.infrastructure.ssh.asyncssh_pool import exec_ssh_command, ssh_pool

logger = logging.getLogger("nexus.sync_v2")

# S5: Semaphore for concurrency control
MAX_CONCURRENT = 10


class SyncEngineV2:
    """Unified sync engine: file sync + preview, with parallel execution (S5)"""

    def __init__(
        self,
        server_repo: ServerRepositoryImpl,
        sync_log_repo: SyncLogRepositoryImpl,
        audit_repo: AuditLogRepositoryImpl,
        retry_repo: PushRetryJobRepositoryImpl,
    ):
        self.server_repo = server_repo
        self.sync_log_repo = sync_log_repo
        self.audit_repo = audit_repo
        self.retry_repo = retry_repo

    # ── S1: Sync File Mode (rsync-based) ──

    async def sync_files(
        self,
        server_ids: List[int],
        source_path: str,
        target_path: Optional[str] = None,
        sync_mode: str = "incremental",
        operator: str = "admin",
        batch_size: int = 50,
        concurrency: int = 10,
        trigger_type: str = "manual",
    ) -> dict:
        """S1: File sync using rsync over SSH"""
        concurrency = min(concurrency, MAX_CONCURRENT)
        sem = asyncio.Semaphore(concurrency)

        results = {}
        # Build server list
        servers = []
        for sid in server_ids:
            server = await self.server_repo.get_by_id(sid)
            if server:
                servers.append(server)

        total = len(servers)
        completed = 0
        failed = 0
        _counter_lock = asyncio.Lock()

        async def _sync_one(server: Server):
            nonlocal completed, failed
            async with sem:
                # Create sync log
                sync_log = SyncLog(
                    server_id=server.id,
                    source_path=source_path,
                    target_path=target_path or server.target_path or "/tmp/sync",
                    trigger_type=trigger_type,
                    operator=operator,
                    status="running",
                    sync_mode=sync_mode,
                    started_at=datetime.now(timezone.utc),
                )
                sync_log = await self.sync_log_repo.create(sync_log)

                # Execute rsync
                import shlex
                rsync_cmd = f"rsync -az --delete {shlex.quote(source_path)}/ {shlex.quote(target_path or server.target_path or '/tmp/sync')}/"
                result = await exec_ssh_command(server, rsync_cmd, timeout=300)

                # Update sync log
                sync_log.status = "success" if result["exit_code"] == 0 else "failed"
                sync_log.finished_at = datetime.now(timezone.utc)
                sync_log.duration_seconds = int(
                    (sync_log.finished_at - sync_log.started_at).total_seconds()
                ) if sync_log.started_at else 0
                sync_log.error_message = result["stderr"][:1000] if result["exit_code"] != 0 else None
                sync_log = await self.sync_log_repo.update(sync_log)

                async with _counter_lock:
                    if sync_log.status == "success":
                        completed += 1
                    else:
                        failed += 1

                results[server.id] = sync_log
                return sync_log

        # S5: Parallel execution with concurrency control
        tasks = [asyncio.create_task(_sync_one(s)) for s in servers]
        await asyncio.gather(*tasks, return_exceptions=True)

        # Audit
        await self._audit(
            "sync_files", "sync", 0,
            f"File sync: {source_path} → {total} servers ({completed} ok, {failed} failed)",
            operator,
        )

        return {
            "total": total,
            "completed": completed,
            "failed": failed,
            "results": {sid: _log_to_dict(log) for sid, log in results.items()},
        }

    async def _audit(self, action: str, target_type: str, target_id: int, detail: str, operator: str = "admin"):
        """Create audit log entry"""
        try:
            audit_log = AuditLog(
                admin_username=operator,
                action=action,
                target_type=target_type,
                target_id=target_id,
                detail=detail,
            )
            await self.audit_repo.create(audit_log)
        except Exception as e:
            logger.error(f"Audit log failed: {e}")


    # ── Preview (dry-run) ──

    async def preview_sync(
        self,
        server_id: int,
        source_path: str,
        target_path: Optional[str] = None,
        sync_mode: str = "incremental",
        verbose: bool = False,
    ) -> dict:
        """Run rsync --dry-run --stats against one server; no data is transferred.

        Returns parsed stats and optional file list (verbose, capped at 200 lines).
        """
        import re
        server = await self.server_repo.get_by_id(server_id)
        if not server:
            return {"error": "Server not found", "exit_code": -1}

        target = target_path or server.target_path or "/tmp/sync"

        # Build flags matching the actual sync mode
        mode_flags = {
            "incremental": "-az",
            "full":        "-az --delete",
            "overwrite":   "-az --inplace",
            "checksum":    "-az --checksum",
        }
        flags = mode_flags.get(sync_mode, "-az")
        dry_flags = f"{flags} --dry-run --stats"
        if verbose:
            dry_flags += " -v"

        cmd = (
            f"rsync {dry_flags} "
            f"{shlex.quote(source_path)}/ "
            f"{shlex.quote(target)}/"
        )

        result = await exec_ssh_command(server, cmd, timeout=60)
        stdout = result.get("stdout", "")
        stderr = result.get("stderr", "")
        exit_code = result.get("exit_code", -1)

        # exit_code 23 = partial transfer (some files not transferred) is OK for dry-run
        if exit_code not in (0, 23):
            return {
                "server_id": server_id,
                "server_name": server.name,
                "error": stderr[:500] or f"rsync exited with code {exit_code}",
                "exit_code": exit_code,
            }

        stats = _parse_rsync_stats(stdout)

        files: list[str] = []
        files_truncated = False
        if verbose:
            all_lines = [
                l for l in stdout.split("\n")
                if l.strip()
                and not l.startswith("Number of")
                and not l.startswith("Total")
                and not l.startswith("Literal")
                and not l.startswith("Matched")
                and not l.startswith("File list")
                and not l.startswith("sent ")
                and not l.startswith("speedup")
                and not l.startswith("sending")
            ]
            files = all_lines[:200]
            files_truncated = len(all_lines) > 200

        return {
            "server_id": server_id,
            "server_name": server.name,
            "source_path": source_path,
            "target_path": target,
            "sync_mode": sync_mode,
            "dry_run": True,
            "stats": stats,
            "files": files,
            "files_truncated": files_truncated,
            "exit_code": exit_code,
        }


def _parse_rsync_stats(output: str) -> dict:
    """Extract key numbers from rsync --stats output."""
    import re
    stats: dict = {}
    patterns = {
        "files_total":       r"Number of files:\s+([\d,]+)",
        "files_created":     r"Number of created files:\s+([\d,]+)",
        "files_deleted":     r"Number of deleted files:\s+([\d,]+)",
        "files_transferred": r"Number of regular files transferred:\s+([\d,]+)",
        "total_size_bytes":  r"Total file size:\s+([\d,]+)\s+bytes",
        "transfer_size_bytes": r"Total transferred file size:\s+([\d,]+)\s+bytes",
    }
    for key, pattern in patterns.items():
        m = re.search(pattern, output)
        if m:
            try:
                stats[key] = int(m.group(1).replace(",", ""))
            except ValueError:
                stats[key] = m.group(1)
    return stats


def _log_to_dict(log: SyncLog) -> dict:
    return {
        "id": log.id, "server_id": log.server_id,
        "status": log.status, "sync_mode": log.sync_mode,
        "error_message": log.error_message,
        "started_at": str(log.started_at) if log.started_at else None,
        "finished_at": str(log.finished_at) if log.finished_at else None,
        "duration_seconds": log.duration_seconds,
    }