102 lines
4.2 KiB
Python
102 lines
4.2 KiB
Python
from typing import Any, Callable, Dict, List
|
|
|
|
from fastapi import HTTPException
|
|
from sqlmodel import Session, select
|
|
|
|
from models.bot import BotInstance, NanobotImage
|
|
|
|
|
|
class ImageService:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
cache: Any,
|
|
cache_key_images: Callable[[], str],
|
|
invalidate_images_cache: Callable[[], None],
|
|
reconcile_image_registry: Callable[[Session], None],
|
|
docker_manager: Any,
|
|
) -> None:
|
|
self._cache = cache
|
|
self._cache_key_images = cache_key_images
|
|
self._invalidate_images_cache = invalidate_images_cache
|
|
self._reconcile_image_registry = reconcile_image_registry
|
|
self._docker_manager = docker_manager
|
|
|
|
def list_images(self, *, session: Session) -> List[Dict[str, Any]]:
|
|
cached = self._cache.get_json(self._cache_key_images())
|
|
if isinstance(cached, list) and all(isinstance(row, dict) for row in cached):
|
|
return cached
|
|
if isinstance(cached, list):
|
|
self._invalidate_images_cache()
|
|
self._reconcile_image_registry(session)
|
|
rows = session.exec(select(NanobotImage)).all()
|
|
payload = [row.model_dump() for row in rows]
|
|
self._cache.set_json(self._cache_key_images(), payload, ttl=60)
|
|
return payload
|
|
|
|
def delete_image(self, *, session: Session, tag: str) -> Dict[str, Any]:
|
|
image = session.get(NanobotImage, tag)
|
|
if not image:
|
|
raise HTTPException(status_code=404, detail="Image not found")
|
|
|
|
bots_using = session.exec(select(BotInstance).where(BotInstance.image_tag == tag)).all()
|
|
if bots_using:
|
|
raise HTTPException(status_code=400, detail=f"Cannot delete image: {len(bots_using)} bots are using it.")
|
|
|
|
session.delete(image)
|
|
session.commit()
|
|
self._invalidate_images_cache()
|
|
return {"status": "deleted"}
|
|
|
|
def list_docker_images(self, *, repository: str = "nanobot-base") -> List[Dict[str, Any]]:
|
|
return self._docker_manager.list_images_by_repo(repository)
|
|
|
|
def register_image(self, *, session: Session, payload: Dict[str, Any]) -> NanobotImage:
|
|
tag = str(payload.get("tag") or "").strip()
|
|
source_dir = str(payload.get("source_dir") or "manual").strip() or "manual"
|
|
if not tag:
|
|
raise HTTPException(status_code=400, detail="tag is required")
|
|
|
|
if not self._docker_manager.has_image(tag):
|
|
raise HTTPException(status_code=404, detail=f"Docker image not found: {tag}")
|
|
|
|
version = tag.split(":")[-1].removeprefix("v") if ":" in tag else tag
|
|
try:
|
|
docker_img = self._docker_manager.client.images.get(tag) if self._docker_manager.client else None
|
|
image_id = docker_img.id if docker_img else None
|
|
except Exception:
|
|
image_id = None
|
|
|
|
row = session.get(NanobotImage, tag)
|
|
if not row:
|
|
row = NanobotImage(
|
|
tag=tag,
|
|
version=version,
|
|
status="READY",
|
|
source_dir=source_dir,
|
|
image_id=image_id,
|
|
)
|
|
else:
|
|
row.version = version
|
|
row.status = "READY"
|
|
row.source_dir = source_dir
|
|
row.image_id = image_id
|
|
session.add(row)
|
|
session.commit()
|
|
session.refresh(row)
|
|
self._invalidate_images_cache()
|
|
return row
|
|
|
|
def require_ready_image(self, session: Session, image_tag: str, *, require_local_image: bool) -> NanobotImage:
|
|
normalized_tag = str(image_tag or "").strip()
|
|
if not normalized_tag:
|
|
raise HTTPException(status_code=400, detail="image_tag is required")
|
|
image_row = session.get(NanobotImage, normalized_tag)
|
|
if not image_row:
|
|
raise HTTPException(status_code=400, detail=f"Image not registered in DB: {normalized_tag}")
|
|
if image_row.status != "READY":
|
|
raise HTTPException(status_code=400, detail=f"Image status is not READY: {normalized_tag} ({image_row.status})")
|
|
if require_local_image and not self._docker_manager.has_image(normalized_tag):
|
|
raise HTTPException(status_code=400, detail=f"Docker image not found locally: {normalized_tag}")
|
|
return image_row
|