Associate images when creating or updating a template.

gns3server/db/repositories/templates.py

@@ -16,9 +16,10 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from uuid import UUID
-from typing import List, Union
-from sqlalchemy import select, update, delete
+from typing import List, Union, Optional
+from sqlalchemy import select, delete
 from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import selectinload
 from sqlalchemy.orm.session import make_transient
 
 from .base import BaseRepository
@@ -41,19 +42,22 @@ TEMPLATE_TYPE_TO_MODEL = {
 
 
 class TemplatesRepository(BaseRepository):
+
     def __init__(self, db_session: AsyncSession) -> None:
 
         super().__init__(db_session)
 
     async def get_template(self, template_id: UUID) -> Union[None, models.Template]:
 
-        query = select(models.Template).where(models.Template.template_id == template_id)
+        query = select(models.Template).\
+            options(selectinload(models.Template.images)).\
+            where(models.Template.template_id == template_id)
         result = await self._db_session.execute(query)
         return result.scalars().first()
 
     async def get_templates(self) -> List[models.Template]:
 
-        query = select(models.Template)
+        query = select(models.Template).options(selectinload(models.Template.images))
         result = await self._db_session.execute(query)
         return result.scalars().all()
 
@@ -66,20 +70,14 @@ class TemplatesRepository(BaseRepository):
         await self._db_session.refresh(db_template)
         return db_template
 
-    async def update_template(self, template_id: UUID, template_update: schemas.TemplateUpdate) -> schemas.Template:
+    async def update_template(self, db_template: models.Template, template_settings: dict) -> schemas.Template:
 
-        update_values = template_update.dict(exclude_unset=True)
-
-        query = update(models.Template). \
-            where(models.Template.template_id == template_id). \
-            values(update_values)
-
-        await self._db_session.execute(query)
+        # update the fields directly because update() query couldn't work
+        for key, value in template_settings.items():
+            setattr(db_template, key, value)
         await self._db_session.commit()
-        template_db = await self.get_template(template_id)
-        if template_db:
-            await self._db_session.refresh(template_db)  # force refresh of updated_at value
-        return template_db
+        await self._db_session.refresh(db_template)  # force refresh of updated_at value
+        return db_template
 
     async def delete_template(self, template_id: UUID) -> bool:
 
@@ -88,18 +86,72 @@ class TemplatesRepository(BaseRepository):
         await self._db_session.commit()
         return result.rowcount > 0
 
-    async def duplicate_template(self, template_id: UUID) -> schemas.Template:
+    async def duplicate_template(self, template_id: UUID) -> Optional[schemas.Template]:
 
-        query = select(models.Template).where(models.Template.template_id == template_id)
+        query = select(models.Template).\
+            options(selectinload(models.Template.images)).\
+            where(models.Template.template_id == template_id)
         db_template = (await self._db_session.execute(query)).scalars().first()
-        if not db_template:
-            return db_template
-
-        # duplicate db object with new primary key (template_id)
-        self._db_session.expunge(db_template)
-        make_transient(db_template)
-        db_template.template_id = None
-        self._db_session.add(db_template)
-        await self._db_session.commit()
-        await self._db_session.refresh(db_template)
+        if db_template:
+            # duplicate db object with new primary key (template_id)
+            self._db_session.expunge(db_template)
+            make_transient(db_template)
+            db_template.template_id = None
+            self._db_session.add(db_template)
+            await self._db_session.commit()
+            await self._db_session.refresh(db_template)
         return db_template
+
+    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 add_image_to_template(
+            self,
+            template_id: UUID,
+            image: models.Image
+    ) -> Union[None, models.Template]:
+        """
+        Add an image to template.
+        """
+
+        query = select(models.Template).\
+            options(selectinload(models.Template.images)).\
+            where(models.Template.template_id == template_id)
+        result = await self._db_session.execute(query)
+        template_in_db = result.scalars().first()
+        if not template_in_db:
+            return None
+
+        template_in_db.images.append(image)
+        await self._db_session.commit()
+        await self._db_session.refresh(template_in_db)
+        return template_in_db
+
+    async def remove_image_from_template(
+            self,
+            template_id: UUID,
+            image: models.Image
+    ) -> Union[None, models.Template]:
+        """
+        Remove an image from a template.
+        """
+
+        query = select(models.Template).\
+            options(selectinload(models.Template.images)).\
+            where(models.Template.template_id == template_id)
+        result = await self._db_session.execute(query)
+        template_in_db = result.scalars().first()
+        if not template_in_db:
+            return None
+
+        if image in template_in_db.images:
+            template_in_db.images.remove(image)
+            await self._db_session.commit()
+            await self._db_session.refresh(template_in_db)
+        return template_in_db

gns3server/services/templates.py

@@ -14,6 +14,7 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
+import os
 import uuid
 import pydantic
 
@@ -22,6 +23,7 @@ from fastapi.encoders import jsonable_encoder
 from typing import List
 
 from gns3server import schemas
+import gns3server.db.models as models
 from gns3server.db.repositories.templates import TemplatesRepository
 from gns3server.controller import Controller
 from gns3server.controller.controller_error import (
@@ -131,6 +133,7 @@ BUILTIN_TEMPLATES = [
 
 
 class TemplatesService:
+
     def __init__(self, templates_repo: TemplatesRepository):
 
         self._templates_repo = templates_repo
@@ -152,6 +155,44 @@ class TemplatesService:
             templates.append(jsonable_encoder(builtin_template))
         return templates
 
+    async def _find_image(self, image_name):
+
+        image = await self._templates_repo.get_image(image_name)
+        if not image or not os.path.exists(image.path):
+            raise ControllerNotFoundError(f"Image {image_name} could not be found")
+        return image
+
+    async def _find_images(self, template_type: str, settings: dict) -> List[models.Image]:
+
+        images_to_add_to_template = []
+        if template_type == "dynamips":
+            if settings["image"]:
+                image = await self._find_image(settings["image"])
+                if image.image_type != "ios":
+                    raise ControllerBadRequestError(
+                        f"Image '{image.filename}' type is not 'ios' but '{image.image_type}'"
+                    )
+                images_to_add_to_template.append(image)
+        elif template_type == "iou":
+            if settings["path"]:
+                image = await self._find_image(settings["path"])
+                if image.image_type != "iou":
+                    raise ControllerBadRequestError(
+                        f"Image '{image.filename}' type is not 'iou' but '{image.image_type}'"
+                    )
+                images_to_add_to_template.append(image)
+        elif template_type == "qemu":
+            for key, value in settings.items():
+                if key.endswith("_image") and value:
+                    image = await self._find_image(value)
+                    if image.image_type != "qemu":
+                        raise ControllerBadRequestError(
+                            f"Image '{image.filename}' type is not 'qemu' but '{image.image_type}'"
+                        )
+                    if image not in images_to_add_to_template:
+                        images_to_add_to_template.append(image)
+        return images_to_add_to_template
+
     async def create_template(self, template_create: schemas.TemplateCreate) -> dict:
 
         try:
@@ -167,7 +208,11 @@ class TemplatesService:
                 settings = dynamips_template_settings_with_defaults.dict()
         except pydantic.ValidationError as e:
             raise ControllerBadRequestError(f"JSON schema error received while creating new template: {e}")
+
+        images_to_add_to_template = await self._find_images(template_create.template_type, settings)
         db_template = await self._templates_repo.create_template(template_create.template_type, settings)
+        for image in images_to_add_to_template:
+            await self._templates_repo.add_image_to_template(db_template.template_id, image)
         template = db_template.asjson()
         self._controller.notification.controller_emit("template.created", template)
         return template
@@ -183,13 +228,34 @@ class TemplatesService:
             raise ControllerNotFoundError(f"Template '{template_id}' not found")
         return template
 
+    async def _remove_image(self, template_id: UUID, image:str) -> None:
+
+        image = await self._templates_repo.get_image(image)
+        await self._templates_repo.remove_image_from_template(template_id, image)
+
     async def update_template(self, template_id: UUID, template_update: schemas.TemplateUpdate) -> dict:
 
         if self.get_builtin_template(template_id):
             raise ControllerForbiddenError(f"Template '{template_id}' cannot be updated because it is built-in")
-        db_template = await self._templates_repo.update_template(template_id, template_update)
+        template_settings = jsonable_encoder(template_update, exclude_unset=True)
+
+        db_template = await self._templates_repo.get_template(template_id)
         if not db_template:
             raise ControllerNotFoundError(f"Template '{template_id}' not found")
+
+        images_to_add_to_template = await self._find_images(db_template.template_type, template_settings)
+        if db_template.template_type == "dynamips" and "image" in template_settings:
+            await self._remove_image(db_template.template_id, db_template.image)
+        elif db_template.template_type == "iou" and "path" in template_settings:
+            await self._remove_image(db_template.template_id, db_template.path)
+        elif db_template.template_type == "qemu":
+            for key in template_update.dict().keys():
+                if key.endswith("_image") and key in template_settings:
+                    await self._remove_image(db_template.template_id, db_template.__dict__[key])
+
+        db_template = await self._templates_repo.update_template(db_template, template_settings)
+        for image in images_to_add_to_template:
+            await self._templates_repo.add_image_to_template(db_template.template_id, image)
         template = db_template.asjson()
         self._controller.notification.controller_emit("template.updated", template)
         return template