diff --git a/app/api_snapshot/snapshot/manager.py b/app/api_snapshot/snapshot/manager.py new file mode 100644 index 0000000..1613973 --- /dev/null +++ b/app/api_snapshot/snapshot/manager.py @@ -0,0 +1,223 @@ +"""Snapshot management module.""" + +import json +import os +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, Dict, List, Optional + +from api_snapshot.recorder.recorder import RequestResponsePair + +SNAPSHOT_VERSION = "1.0" + + +@dataclass +class SnapshotMetadata: + """Metadata for a snapshot.""" + version: str = SNAPSHOT_VERSION + timestamp: str = field(default_factory=lambda: datetime.utcnow().isoformat()) + description: str = "" + source_url: Optional[str] = None + latency_mode: str = "original" + custom_latency_ms: Optional[int] = None + tags: List[str] = field(default_factory=list) + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary.""" + return { + "version": self.version, + "timestamp": self.timestamp, + "description": self.description, + "source_url": self.source_url, + "latency_mode": self.latency_mode, + "custom_latency_ms": self.custom_latency_ms, + "tags": self.tags + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "SnapshotMetadata": + """Create from dictionary.""" + return cls( + version=data.get("version", SNAPSHOT_VERSION), + timestamp=data.get("timestamp", datetime.utcnow().isoformat()), + description=data.get("description", ""), + source_url=data.get("source_url"), + latency_mode=data.get("latency_mode", "original"), + custom_latency_ms=data.get("custom_latency_ms"), + tags=data.get("tags", []) + ) + + +@dataclass +class Snapshot: + """A complete snapshot containing recorded traffic.""" + metadata: SnapshotMetadata + requests: List[RequestResponsePair] + + def to_dict(self) -> Dict[str, Any]: + """Convert to dictionary.""" + return { + "metadata": self.metadata.to_dict(), + "requests": [pair.to_dict() for pair in self.requests] + } + + @classmethod + def from_dict(cls, data: Dict[str, Any]) -> "Snapshot": + """Create from dictionary.""" + metadata = SnapshotMetadata.from_dict(data.get("metadata", {})) + requests = [ + RequestResponsePair.from_dict(req) + for req in data.get("requests", []) + ] + return cls(metadata=metadata, requests=requests) + + +class SnapshotManager: + """Manager for snapshot files.""" + + SNAPSHOT_EXTENSION = ".json" + + def __init__(self, snapshot_dir: str = "./snapshots"): + """Initialize the snapshot manager. + + Args: + snapshot_dir: Directory to store snapshots + """ + self.snapshot_dir = snapshot_dir + os.makedirs(snapshot_dir, exist_ok=True) + + def _get_path(self, name: str) -> str: + """Get the full path for a snapshot.""" + if not name.endswith(self.SNAPSHOT_EXTENSION): + name = name + self.SNAPSHOT_EXTENSION + return os.path.join(self.snapshot_dir, name) + + def save_snapshot( + self, + name: str, + requests: List[RequestResponsePair], + description: str = "", + source_url: Optional[str] = None, + tags: Optional[List[str]] = None + ) -> str: + """Save a snapshot to a file. + + Args: + name: Name of the snapshot (without extension) + requests: List of request-response pairs + description: Optional description + source_url: Optional source URL that was recorded + tags: Optional list of tags + + Returns: + The path to the saved snapshot + """ + path = self._get_path(name) + + metadata = SnapshotMetadata( + version=SNAPSHOT_VERSION, + timestamp=datetime.utcnow().isoformat(), + description=description, + source_url=source_url, + tags=tags or [] + ) + + snapshot = Snapshot(metadata=metadata, requests=requests) + + with open(path, "w", encoding="utf-8") as f: + json.dump(snapshot.to_dict(), f, indent=2, ensure_ascii=False) + + return path + + def load_snapshot(self, name: str) -> Snapshot: + """Load a snapshot from a file. + + Args: + name: Name of the snapshot (with or without extension) + + Returns: + The loaded Snapshot object + + Raises: + FileNotFoundError: If snapshot doesn't exist + ValueError: If snapshot format is invalid + """ + path = self._get_path(name) + + if not os.path.exists(path): + raise FileNotFoundError(f"Snapshot not found: {name}") + + with open(path, "r", encoding="utf-8") as f: + data = json.load(f) + + if "metadata" not in data or "requests" not in data: + raise ValueError(f"Invalid snapshot format: {name}") + + return Snapshot.from_dict(data) + + def delete_snapshot(self, name: str) -> None: + """Delete a snapshot file. + + Args: + name: Name of the snapshot (with or without extension) + + Raises: + FileNotFoundError: If snapshot doesn't exist + """ + path = self._get_path(name) + + if not os.path.exists(path): + raise FileNotFoundError(f"Snapshot not found: {name}") + + os.remove(path) + + def list_snapshots(self) -> List[Dict[str, Any]]: + """List all available snapshots. + + Returns: + List of snapshot info dictionaries + """ + snapshots: List[Dict[str, Any]] = [] + + if not os.path.exists(self.snapshot_dir): + return snapshots + + for filename in os.listdir(self.snapshot_dir): + if filename.endswith(self.SNAPSHOT_EXTENSION): + path = os.path.join(self.snapshot_dir, filename) + + try: + stat = os.stat(path) + snapshot = self.load_snapshot(filename) + + snapshots.append({ + "name": filename[:-len(self.SNAPSHOT_EXTENSION)], + "path": path, + "created": snapshot.metadata.timestamp, + "endpoint_count": len(snapshot.requests), + "size": f"{stat.st_size / 1024:.1f} KB", + "description": snapshot.metadata.description + }) + except Exception: + snapshots.append({ + "name": filename[:-len(self.SNAPSHOT_EXTENSION)], + "path": path, + "created": "Unknown", + "endpoint_count": 0, + "size": f"{os.stat(path).st_size / 1024:.1f} KB", + "description": "Error loading" + }) + + return sorted(snapshots, key=lambda s: s["name"]) + + def snapshot_exists(self, name: str) -> bool: + """Check if a snapshot exists. + + Args: + name: Name of the snapshot + + Returns: + True if snapshot exists + """ + path = self._get_path(name) + return os.path.exists(path)