AMBARI-5134. Request Schedule Batch request, return code, status and message not persisted. (swagle)

ambari-server/src/main/java/org/apache/ambari/server/orm/dao/RequestScheduleBatchRequestDAO.java

@@ -68,8 +68,7 @@ public class RequestScheduleBatchRequestDAO {
   }
 
   @Transactional
-  public RequestScheduleBatchRequestEntity merge
-    (RequestScheduleBatchRequestEntity batchRequestEntity) {
+  public RequestScheduleBatchRequestEntity merge(RequestScheduleBatchRequestEntity batchRequestEntity) {
 
     return entityManagerProvider.get().merge(batchRequestEntity);
   }

ambari-server/src/main/java/org/apache/ambari/server/orm/entities/RequestScheduleBatchRequestEntity.java

@@ -184,4 +184,18 @@ public class RequestScheduleBatchRequestEntity {
     result = 31 * result + batchId.hashCode();
     return result;
   }
+
+  @Override
+  public String toString() {
+    return "RequestScheduleBatchRequestEntity{" +
+      "scheduleId=" + scheduleId +
+      ", batchId=" + batchId +
+      ", requestId=" + requestId +
+      ", requestType='" + requestType + '\'' +
+      ", requestUri='" + requestUri + '\'' +
+      ", requestStatus='" + requestStatus + '\'' +
+      ", returnCode=" + returnCode +
+      ", returnMessage='" + returnMessage + '\'' +
+      '}';
+  }
 }

ambari-server/src/main/java/org/apache/ambari/server/state/scheduler/RequestExecutionImpl.java

@@ -427,24 +427,28 @@ public class RequestExecutionImpl implements RequestExecution {
   public void updateBatchRequest(long batchId,
                                  BatchRequestResponse batchRequestResponse,
                                  boolean statusOnly) {
-    long executionId = requestScheduleEntity.getScheduleId();
 
-    RequestScheduleBatchRequestEntityPK batchRequestEntityPK = new
-        RequestScheduleBatchRequestEntityPK();
-    batchRequestEntityPK.setScheduleId(executionId);
-    batchRequestEntityPK.setBatchId(batchId);
-    RequestScheduleBatchRequestEntity batchRequestEntity =
-        batchRequestDAO.findByPk(batchRequestEntityPK);
+    RequestScheduleBatchRequestEntity batchRequestEntity = null;
 
-    batchRequestEntity.setRequestStatus(batchRequestResponse.getStatus());
-
-    if (!statusOnly) {
-      batchRequestEntity.setReturnCode(batchRequestResponse.getReturnCode());
-      batchRequestEntity.setRequestId(batchRequestResponse.getRequestId());
-      batchRequestEntity.setReturnMessage(batchRequestResponse.getReturnMessage());
+    for (RequestScheduleBatchRequestEntity entity :
+        requestScheduleEntity.getRequestScheduleBatchRequestEntities()) {
+      if (entity.getBatchId() == batchId
+          && entity.getScheduleId() == requestScheduleEntity.getScheduleId()) {
+        batchRequestEntity = entity;
+      }
     }
 
-    batchRequestDAO.merge(batchRequestEntity);
+    if (batchRequestEntity != null) {
+      batchRequestEntity.setRequestStatus(batchRequestResponse.getStatus());
+
+      if (!statusOnly) {
+        batchRequestEntity.setReturnCode(batchRequestResponse.getReturnCode());
+        batchRequestEntity.setRequestId(batchRequestResponse.getRequestId());
+        batchRequestEntity.setReturnMessage(batchRequestResponse.getReturnMessage());
+      }
+
+      batchRequestDAO.merge(batchRequestEntity);
+    }
 
     BatchRequest batchRequest = getBatchRequest(batchId);
 
@@ -456,6 +460,7 @@ public class RequestExecutionImpl implements RequestExecution {
     }
 
     setLastExecutionStatus(batchRequestResponse.getStatus());
+    requestScheduleDAO.merge(requestScheduleEntity);
   }
 
   @Override

ambari-server/src/test/java/org/apache/ambari/server/state/RequestExecutionTest.java

@@ -31,6 +31,7 @@ import org.apache.ambari.server.orm.entities.RequestScheduleBatchRequestEntity;
 import org.apache.ambari.server.orm.entities.RequestScheduleEntity;
 import org.apache.ambari.server.state.scheduler.Batch;
 import org.apache.ambari.server.state.scheduler.BatchRequest;
+import org.apache.ambari.server.state.scheduler.BatchRequestResponse;
 import org.apache.ambari.server.state.scheduler.BatchSettings;
 import org.apache.ambari.server.state.scheduler.RequestExecution;
 import org.apache.ambari.server.state.scheduler.RequestExecutionFactory;
@@ -341,4 +342,45 @@ public class RequestExecutionTest {
     Assert.assertEquals(RequestExecution.Status.COMPLETED.name(),
       scheduleEntity.getStatus());
   }
+
+  @Test
+  public void testUpdateBatchRequest() throws Exception {
+    RequestExecution requestExecution = createRequestSchedule();
+    Assert.assertNotNull(requestExecution);
+    Assert.assertNotNull(cluster.getAllRequestExecutions().get
+      (requestExecution.getId()));
+
+    RequestScheduleEntity scheduleEntity = requestScheduleDAO.findById(requestExecution.getId());
+    Assert.assertNotNull(scheduleEntity);
+    Assert.assertEquals(RequestExecution.Status.SCHEDULED.name(), scheduleEntity.getStatus());
+
+    Collection<RequestScheduleBatchRequestEntity> batchRequestEntities =
+      scheduleEntity.getRequestScheduleBatchRequestEntities();
+
+    Assert.assertNotNull(batchRequestEntities);
+    Assert.assertEquals(2, batchRequestEntities.size());
+
+    BatchRequestResponse batchRequestResponse = new BatchRequestResponse();
+    batchRequestResponse.setRequestId(1L);
+    batchRequestResponse.setReturnCode(200);
+    batchRequestResponse.setReturnMessage("test");
+    batchRequestResponse.setStatus("IN_PROGRESS");
+
+    requestExecution.updateBatchRequest(10L, batchRequestResponse, false);
+
+    scheduleEntity = requestScheduleDAO.findById(requestExecution.getId());
+    RequestScheduleBatchRequestEntity testEntity = null;
+
+    for (RequestScheduleBatchRequestEntity entity :
+        scheduleEntity.getRequestScheduleBatchRequestEntities()) {
+      if (entity.getBatchId().equals(10L)) {
+        testEntity = entity;
+      }
+    }
+
+    Assert.assertNotNull(testEntity);
+    Assert.assertEquals(200, testEntity.getReturnCode().intValue());
+    Assert.assertEquals("test", testEntity.getReturnMessage());
+    Assert.assertEquals("IN_PROGRESS", testEntity.getRequestStatus());
+  }
 }