From 4d9e4e1059496767e15af61aadaa40023f160d06 Mon Sep 17 00:00:00 2001 From: grossmj Date: Fri, 20 Aug 2021 15:58:41 +0930 Subject: [PATCH] Add prune images endpoint. Use many-to-many relationship between images and templates. --- gns3server/api/routes/controller/images.py | 14 +++++- gns3server/db/models/images.py | 14 ++++-- gns3server/db/models/templates.py | 5 +- gns3server/db/repositories/images.py | 58 +++++++++++++++++++--- tests/api/routes/controller/test_images.py | 12 +++++ 5 files changed, 90 insertions(+), 13 deletions(-) diff --git a/gns3server/api/routes/controller/images.py b/gns3server/api/routes/controller/images.py index 84a6f00e..73a6cd4a 100644 --- a/gns3server/api/routes/controller/images.py +++ b/gns3server/api/routes/controller/images.py @@ -22,7 +22,7 @@ import os import logging import urllib.parse -from fastapi import APIRouter, Request, Depends, status +from fastapi import APIRouter, Request, Response, Depends, status from typing import List from gns3server import schemas @@ -125,3 +125,15 @@ async def delete_image( success = await images_repo.delete_image(image_name) if not success: raise ControllerError(f"Image '{image_name}' could not be deleted") + + +@router.post("/prune", status_code=status.HTTP_204_NO_CONTENT) +async def prune_images( + images_repo: ImagesRepository = Depends(get_repository(ImagesRepository)), +) -> Response: + """ + Prune images not attached to any template. + """ + + await images_repo.prune_images() + return Response(status_code=status.HTTP_204_NO_CONTENT) diff --git a/gns3server/db/models/images.py b/gns3server/db/models/images.py index ed41564d..aba30387 100644 --- a/gns3server/db/models/images.py +++ b/gns3server/db/models/images.py @@ -15,10 +15,18 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . -from sqlalchemy import Column, String, Integer, BigInteger +from sqlalchemy import Table, Column, String, Integer, ForeignKey, BigInteger from sqlalchemy.orm import relationship -from .base import BaseTable +from .base import Base, BaseTable, GUID + + +image_template_link = Table( + "images_templates_link", + Base.metadata, + Column("image_id", Integer, ForeignKey("images.id", ondelete="CASCADE")), + Column("template_id", GUID, ForeignKey("templates.template_id", ondelete="CASCADE")) +) class Image(BaseTable): @@ -32,4 +40,4 @@ class Image(BaseTable): path = Column(String) checksum = Column(String) checksum_algorithm = Column(String) - templates = relationship("Template") + templates = relationship("Template", secondary=image_template_link, back_populates="images") diff --git a/gns3server/db/models/templates.py b/gns3server/db/models/templates.py index 75795039..71515271 100644 --- a/gns3server/db/models/templates.py +++ b/gns3server/db/models/templates.py @@ -17,8 +17,10 @@ from sqlalchemy import Boolean, Column, String, Integer, ForeignKey, PickleType +from sqlalchemy.orm import relationship from .base import BaseTable, generate_uuid, GUID +from .images import image_template_link class Template(BaseTable): @@ -34,8 +36,7 @@ class Template(BaseTable): compute_id = Column(String) usage = Column(String) template_type = Column(String) - - image_id = Column(Integer, ForeignKey('images.id', ondelete="CASCADE")) + images = relationship("Image", secondary=image_template_link, back_populates="templates") __mapper_args__ = { "polymorphic_identity": "templates", diff --git a/gns3server/db/repositories/images.py b/gns3server/db/repositories/images.py index 1d2f802e..02b77a5c 100644 --- a/gns3server/db/repositories/images.py +++ b/gns3server/db/repositories/images.py @@ -15,6 +15,8 @@ # You should have received a copy of the GNU General Public License # along with this program. If not, see . +import os + from typing import Optional, List from sqlalchemy import select, delete from sqlalchemy.ext.asyncio import AsyncSession @@ -23,6 +25,10 @@ from .base import BaseRepository import gns3server.db.models as models +import logging + +log = logging.getLogger(__name__) + class ImagesRepository(BaseRepository): @@ -31,33 +37,48 @@ class ImagesRepository(BaseRepository): super().__init__(db_session) async def get_image(self, image_name: str) -> Optional[models.Image]: + """ + Get an image by its name (filename). + """ query = select(models.Image).where(models.Image.filename == image_name) result = await self._db_session.execute(query) return result.scalars().first() + async def get_image_by_checksum(self, checksum: str) -> Optional[models.Image]: + """ + Get an image by its checksum. + """ + + query = select(models.Image).where(models.Image.checksum == checksum) + result = await self._db_session.execute(query) + return result.scalars().first() + async def get_images(self) -> List[models.Image]: + """ + Get all images. + """ query = select(models.Image) result = await self._db_session.execute(query) return result.scalars().all() async def get_image_templates(self, image_id: int) -> Optional[List[models.Template]]: + """ + Get all templates that an image belongs to. + """ query = select(models.Template).\ - join(models.Image.templates). \ + join(models.Template.images).\ filter(models.Image.id == image_id) result = await self._db_session.execute(query) return result.scalars().all() - async def get_image_by_checksum(self, checksum: str) -> Optional[models.Image]: - - query = select(models.Image).where(models.Image.checksum == checksum) - result = await self._db_session.execute(query) - return result.scalars().first() - async def add_image(self, image_name, image_type, image_size, path, checksum, checksum_algorithm) -> models.Image: + """ + Create a new image. + """ db_image = models.Image( id=None, @@ -75,8 +96,31 @@ class ImagesRepository(BaseRepository): return db_image async def delete_image(self, image_name: str) -> bool: + """ + Delete an image. + """ query = delete(models.Image).where(models.Image.filename == image_name) result = await self._db_session.execute(query) await self._db_session.commit() return result.rowcount > 0 + + async def prune_images(self) -> int: + """ + Prune images not attached to any template. + """ + + query = select(models.Image).\ + filter(~models.Image.templates.any()) + result = await self._db_session.execute(query) + images = result.scalars().all() + images_deleted = 0 + for image in images: + try: + os.remove(image.path) + except OSError: + log.warning(f"Could not delete image file {image.path}") + if await self.delete_image(image.filename): + images_deleted += 1 + log.info(f"{images_deleted} image have been deleted") + return images_deleted diff --git a/tests/api/routes/controller/test_images.py b/tests/api/routes/controller/test_images.py index dd81b732..f81ee1e1 100644 --- a/tests/api/routes/controller/test_images.py +++ b/tests/api/routes/controller/test_images.py @@ -19,9 +19,12 @@ import os import pytest import hashlib +from sqlalchemy.ext.asyncio import AsyncSession from fastapi import FastAPI, status from httpx import AsyncClient +from gns3server.db.repositories.images import ImagesRepository + pytestmark = pytest.mark.asyncio @@ -193,3 +196,12 @@ class TestImageRoutes: response = await client.delete(app.url_path_for("delete_image", image_name=image_name)) assert response.status_code == status.HTTP_204_NO_CONTENT assert not os.path.exists(os.path.join(images_dir, "QEMU", image_name)) + + async def test_prune_images(self, app: FastAPI, client: AsyncClient, db_session: AsyncSession) -> None: + + response = await client.post(app.url_path_for("prune_images")) + assert response.status_code == status.HTTP_204_NO_CONTENT + + images_repo = ImagesRepository(db_session) + images_in_db = await images_repo.get_images() + assert len(images_in_db) == 0