Support qcow2 format

Exactly the same as vmdk.

1 files changed, 14 insertions, 7 deletions

diff --git a/kpov_judge/create_disk_images.py b/kpov_judge/create_disk_images.py
index 381ba25..6777226 100755
--- a/kpov_judge/create_disk_images.py
+++ b/kpov_judge/create_disk_images.py
@@ -35,6 +35,9 @@ def create_snapshot(class_id, task_id, student_id, disk_name, fmt='vmdk', overwr
     task_path = os.path.join(settings.STUDENT_DISK_PATH, task_dir)
 
     if not os.path.exists(os.path.join(task_path)) or overwrite:
+        if not os.path.exists(os.path.join(settings.DISK_TEMPLATE_PATH, template)):
+            raise Exception('template not found: {}'.format(template))
+
         # ensure task dir exists
         os.makedirs(task_path, exist_ok=True)
 
@@ -46,7 +49,7 @@ def create_snapshot(class_id, task_id, student_id, disk_name, fmt='vmdk', overwr
             else:
                 subprocess.call(['cp', os.path.join(settings.DISK_TEMPLATE_PATH, template), snap])
 
-        elif fmt == 'vmdk':
+        elif fmt in ('qcow2', 'vmdk'):
             # qemu-img create stores backing-file path as given, so link all
             # backing images to task directory where target image will be
             # generated
@@ -173,14 +176,18 @@ if __name__ == '__main__':
             except IOError:
                 continue
 
-            print("Creating {}/{} for {}".format(class_id, task_id, student_id))
             all_disks = collections.defaultdict(dict)
             for fmt in settings.STUDENT_DISK_FORMATS:
-                for computer, disks in prepare_task_disks(class_id, task_id, student_id, fmt, computers).items():
-                    for disk, urls in disks.items():
-                        d = all_disks[computer].setdefault(disk, {'formats': []})
-                        d['formats'] += [fmt]
-                        d[fmt] = urls
+                print("Creating {}/{} for {} [format={}]".format(class_id, task_id, student_id, fmt))
+                try:
+                    for computer, disks in prepare_task_disks(class_id, task_id, student_id, fmt, computers).items():
+                        for disk, urls in disks.items():
+                            d = all_disks[computer].setdefault(disk, {'formats': []})
+                            d['formats'] += [fmt]
+                            d[fmt] = urls
+                except Exception as ex:
+                    print(ex)
+                    continue
 
             lock_fp.write("saving URLs\n")
             for computer in computers: