server: select root disk based on user input during vm import

Signed-off-by: Abhishek Kumar <abhishek.mrt22@gmail.com>

Author: shwstppr

server/src/main/java/org/apache/cloudstack/vm/VmImportManagerImpl.java

@@ -500,6 +500,35 @@ private StoragePool getStoragePool(final UnmanagedInstanceTO.Disk disk, final Da
         return storagePool;
     }
 
+    private Pair<UnmanagedInstanceTO.Disk, List<UnmanagedInstanceTO.Disk>> getRootAndDataDisks(List<UnmanagedInstanceTO.Disk> disks, final Map<String, Long> dataDiskOfferingMap) {
+        UnmanagedInstanceTO.Disk rootDisk = null;
+        List<UnmanagedInstanceTO.Disk> dataDisks = new ArrayList<>();
+        if (disks.size() == 1) {
+            rootDisk = disks.get(0);
+            return new Pair<>(rootDisk, dataDisks);
+        }
+        Set<String> callerDiskIds = dataDiskOfferingMap.keySet();
+        if (callerDiskIds.size() != disks.size() - 1) {
+            String msg = String.format("VM has total %d disks for which %d disk offering mappings provided. %d disks need a disk offering for import", disks.size(), callerDiskIds.size(), disks.size()-1);
+            LOGGER.error(String.format("%s. %s parameter can be used to provide disk offerings for the disks", msg, ApiConstants.DATADISK_OFFERING_LIST));
+            throw new ServerApiException(ApiErrorCode.INTERNAL_ERROR, msg);
+        }
+        List<String> diskIdsWithoutOffering = new ArrayList<>();
+        for (UnmanagedInstanceTO.Disk disk : disks) {
+            String diskId = disk.getDiskId();
+            if (!callerDiskIds.contains(diskId)) {
+                diskIdsWithoutOffering.add(diskId);
+                rootDisk = disk;
+            } else {
+                dataDisks.add(disk);
+            }
+        }
+        if (diskIdsWithoutOffering.size() > 1) {
+            throw new ServerApiException(ApiErrorCode.INTERNAL_ERROR, String.format("VM has total %d disks, disk offering mapping not provided for %d disks. Disk IDs that may need a disk offering - %s", disks.size(), diskIdsWithoutOffering.size()-1, String.join(", ", diskIdsWithoutOffering)));
+        }
+        return new Pair<>(rootDisk, dataDisks);
+    }
+
     private void checkUnmanagedDiskAndOfferingForImport(UnmanagedInstanceTO.Disk disk, DiskOffering diskOffering, ServiceOffering serviceOffering, final Account owner, final DataCenter zone, final Cluster cluster, final boolean migrateAllowed)
             throws ServerApiException, PermissionDeniedException, ResourceAllocationException {
         if (serviceOffering == null && diskOffering == null) {
@@ -910,17 +939,16 @@ private UserVm importVirtualMachineInternal(final UnmanagedInstanceTO unmanagedI
         if (CollectionUtils.isEmpty(unmanagedInstanceDisks)) {
             throw new ServerApiException(ApiErrorCode.INTERNAL_ERROR, String.format("No attached disks found for the unmanaged VM: %s", instanceName));
         }
-        final UnmanagedInstanceTO.Disk rootDisk = unmanagedInstance.getDisks().get(0);
+        Pair<UnmanagedInstanceTO.Disk, List<UnmanagedInstanceTO.Disk>> rootAndDataDisksPair = getRootAndDataDisks(unmanagedInstanceDisks, dataDiskOfferingMap);
+        final UnmanagedInstanceTO.Disk rootDisk = rootAndDataDisksPair.first();
+        final List<UnmanagedInstanceTO.Disk> dataDisks = rootAndDataDisksPair.second();
         if (rootDisk == null || Strings.isNullOrEmpty(rootDisk.getController())) {
             throw new ServerApiException(ApiErrorCode.INTERNAL_ERROR, String.format("VM import failed. Unable to retrieve root disk details for VM: %s ", instanceName));
         }
         allDetails.put(VmDetailConstants.ROOT_DISK_CONTROLLER, rootDisk.getController());
-        List<UnmanagedInstanceTO.Disk> dataDisks = new ArrayList<>();
         try {
             checkUnmanagedDiskAndOfferingForImport(rootDisk, null, validatedServiceOffering, owner, zone, cluster, migrateAllowed);
-            if (unmanagedInstanceDisks.size() > 1) { // Data disk(s) present
-                dataDisks.addAll(unmanagedInstanceDisks);
-                dataDisks.remove(0);
+            if (CollectionUtils.isNotEmpty(dataDisks)) { // Data disk(s) present
                 checkUnmanagedDiskAndOfferingForImport(dataDisks, dataDiskOfferingMap, owner, zone, cluster, migrateAllowed);
                 allDetails.put(VmDetailConstants.DATA_DISK_CONTROLLER, dataDisks.get(0).getController());
             }