224 lines
6.9 KiB
Python
224 lines
6.9 KiB
Python
"""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)
|