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