Images can now download from a URL

MrMatAP · Jan 14, 2024 · d14beeb · d14beeb
1 parent abea460
commit d14beeb
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 9 deletions.
diff --git a/src/kaso_mashin/common/generics/images.py b/src/kaso_mashin/common/generics/images.py
@@ -3,6 +3,8 @@
 import pathlib
 import typing
 import uuid
+import httpx
+import aiofiles
 
 from sqlalchemy import String, Integer, Enum, ForeignKey
 from sqlalchemy.orm import Mapped, mapped_column, relationship
@@ -106,7 +108,7 @@ def min_disk(self) -> BinarySizedValue:
     def os_disks(self) -> typing.List[OSDiskValueObject]:
         return self._os_disks
 
-    def __eq__(self, other: 'ImageEntity') -> bool:     # type: ignore[override]
+    def __eq__(self, other: 'ImageEntity') -> bool:  # type: ignore[override]
         return all([
             super().__eq__(other),
             self._name == other.name,
@@ -119,20 +121,32 @@ def __eq__(self, other: 'ImageEntity') -> bool:     # type: ignore[override]
     async def create(owner: 'AggregateRoot',
                      name: str,
                      path: pathlib.Path,
+                     url: str,
                      min_vcpu: int = 0,
                      min_ram: BinarySizedValue = BinarySizedValue(0, BinaryScale.G),
                      min_disk: BinarySizedValue = BinarySizedValue(0, BinaryScale.G)) -> 'ImageEntity':
         if path.exists():
             raise ImageException(code=400, msg=f'Image at {path} already exists')
         path.parent.mkdir(parents=True, exist_ok=True)
+        image = None
         try:
-            image = ImageEntity(owner=owner,
-                                name=name,
-                                path=path,
-                                min_vcpu=min_vcpu,
-                                min_ram=min_ram,
-                                min_disk=min_disk)
-            return await owner.create(image)
+            image = await owner.create(ImageEntity(owner=owner,
+                                                   name=name,
+                                                   path=path,
+                                                   min_vcpu=min_vcpu,
+                                                   min_ram=min_ram,
+                                                   min_disk=min_disk))
+
+            resp = httpx.head(url, follow_redirects=True, timeout=60)
+            image_size = int(resp.headers.get('Content-Length'))
+            current = 0
+            client = httpx.AsyncClient(follow_redirects=True, timeout=60)
+            async with client.stream(method='GET', url=url) as resp, aiofiles.open(path, mode='wb') as i:
+                async for chunk in resp.aiter_bytes(chunk_size=32768):
+                    await i.write(chunk)
+                    current += 32768
+                    print(f'{int(current / image_size * 100)} % complete')
+            return image
         finally:
             pass
 

diff --git a/tests/test_generics.py b/tests/test_generics.py
@@ -36,7 +36,8 @@ async def test_async_images(generics_async_session_maker):
         image_aggregate_root = ImageAggregateRoot(model=ImageModel, session_maker=generics_async_session_maker)
         image = await ImageEntity.create(owner=image_aggregate_root,
                                          name='Test Image',
-                                         path=pathlib.Path(__file__).parent / 'build' / 'test-async-images.qcow2')
+                                         path=pathlib.Path(__file__).parent / 'build' / 'test-async-images.qcow2',
+                                         url='https://cloud-images.ubuntu.com/jammy/current/jammy-server-cloudimg-arm64.img')
         loaded = await image_aggregate_root.get(image.uid)
         assert image == loaded
         await image.set_min_vcpu(10)