ATLAS-3674: ZipFileMigationImporter: Migration status display fix. Part 2

repository/src/main/java/org/apache/atlas/repository/migration/DataMigrationStatusService.java

@@ -101,12 +101,12 @@ public class DataMigrationStatusService {
 
     public void savePosition(Long position) {
         this.status.setCurrentIndex(position);
-        this.migrationStatusVertexManagement.updateVertexPartial(this.status);
+        this.migrationStatusVertexManagement.updateVertexPartialPosition(this.status);
     }
 
     public void setStatus(String status) {
         this.status.setOperationStatus(status);
-        this.migrationStatusVertexManagement.updateVertexPartial(this.status);
+        this.migrationStatusVertexManagement.updateVertexPartialStatus(this.status);
     }
 
     private static class MigrationStatusVertexManagement {
@@ -184,7 +184,7 @@ public class DataMigrationStatusService {
             return null;
         }
 
-        public void updateVertexPartial(MigrationImportStatus status) {
+        public void updateVertexPartialPosition(MigrationImportStatus status) {
             try {
                 setEncodedProperty(vertex, PROPERTY_KEY_POSITION, status.getCurrentIndex());
             } catch (Exception e) {
@@ -194,6 +194,16 @@ public class DataMigrationStatusService {
             }
         }
 
+        public void updateVertexPartialStatus(MigrationImportStatus status) {
+            try {
+                setEncodedProperty(vertex, PROPERTY_KEY_STATUS, status.getOperationStatus());
+            } catch (Exception e) {
+                LOG.warn("Error updating status. Please rely on log messages.", e);
+            } finally {
+                atlasGraph.commit();
+            }
+        }
+
         private void updateVertex(AtlasVertex vertex, MigrationImportStatus status) {
             try {
                 setEncodedProperty(vertex, Constants.GUID_PROPERTY_KEY, status.getName());

repository/src/main/java/org/apache/atlas/repository/store/graph/v2/bulkimport/pc/EntityConsumer.java

@@ -71,9 +71,9 @@ public class EntityConsumer extends WorkItemConsumer<AtlasEntity.AtlasEntityWith
 
     @Override
     protected void processItem(AtlasEntity.AtlasEntityWithExtInfo entityWithExtInfo) {
-        int delta = (MapUtils.isEmpty(entityWithExtInfo.getReferredEntities())
+        int delta = MapUtils.isEmpty(entityWithExtInfo.getReferredEntities())
                 ? 1
-                : entityWithExtInfo.getReferredEntities().size()) + 1;
+                : entityWithExtInfo.getReferredEntities().size() + 1;
 
         long currentCount = counter.addAndGet(delta);
         currentBatch.addAndGet(delta);
@@ -83,7 +83,7 @@ public class EntityConsumer extends WorkItemConsumer<AtlasEntity.AtlasEntityWith
             processEntity(entityWithExtInfo, currentCount);
             attemptCommit();
         } catch (Exception e) {
-            LOG.info("Data loss: Please re-submit!", e);
+            LOG.info("Invalid entities. Possible data loss: Please correct and re-submit!", e);
         }
     }
 

repository/src/test/java/org/apache/atlas/repository/impexp/DataMigrationStatusServiceTest.java

@@ -38,7 +38,9 @@ public class DataMigrationStatusServiceTest {
     AtlasGraph atlasGraph;
 
     @Test
-    public void createUpdateDelete() throws AtlasBaseException {
+    public void createUpdateDelete() {
+        final String STATUS_DONE = "DONE";
+
         DataMigrationStatusService dataMigrationStatusService = new DataMigrationStatusService(atlasGraph);
 
         MigrationImportStatus expected = new MigrationImportStatus("/tmp/defg.zip");
@@ -59,6 +61,11 @@ public class DataMigrationStatusServiceTest {
         assertNotNull(dataMigrationStatusService.getStatus().getCurrentIndex(), "100");
         assertNotNull(dataMigrationStatusService.getCreate(expected).getCurrentIndex(), "100");
 
+
+        dataMigrationStatusService.setStatus(STATUS_DONE);
+        assertNotNull(dataMigrationStatusService.getCreate(expected).getOperationStatus());
+        assertEquals(dataMigrationStatusService.getCreate(expected).getOperationStatus(), STATUS_DONE);
+
         dataMigrationStatusService.delete();
         assertNull(dataMigrationStatusService.getStatus());
         assertNull(dataMigrationStatusService.getByName(ret.getName()));