Search in sources :

Example 1 with ListJobMessagesResponse

use of com.google.api.services.dataflow.model.ListJobMessagesResponse in project beam by apache.

the class MonitoringUtilTest method testGetJobMessages.

@Test
public void testGetJobMessages() throws IOException {
    DataflowClient dataflowClient = mock(DataflowClient.class);
    ListJobMessagesResponse firstResponse = new ListJobMessagesResponse();
    firstResponse.setJobMessages(new ArrayList<JobMessage>());
    for (int i = 0; i < 100; ++i) {
        JobMessage message = new JobMessage();
        message.setId("message_" + i);
        message.setTime(TimeUtil.toCloudTime(new Instant(i)));
        firstResponse.getJobMessages().add(message);
    }
    String pageToken = "page_token";
    firstResponse.setNextPageToken(pageToken);
    ListJobMessagesResponse secondResponse = new ListJobMessagesResponse();
    secondResponse.setJobMessages(new ArrayList<JobMessage>());
    for (int i = 100; i < 150; ++i) {
        JobMessage message = new JobMessage();
        message.setId("message_" + i);
        message.setTime(TimeUtil.toCloudTime(new Instant(i)));
        secondResponse.getJobMessages().add(message);
    }
    when(dataflowClient.listJobMessages(JOB_ID, null)).thenReturn(firstResponse);
    when(dataflowClient.listJobMessages(JOB_ID, pageToken)).thenReturn(secondResponse);
    MonitoringUtil util = new MonitoringUtil(dataflowClient);
    List<JobMessage> messages = util.getJobMessages(JOB_ID, -1);
    assertEquals(150, messages.size());
}
Also used : DataflowClient(org.apache.beam.runners.dataflow.DataflowClient) JobMessage(com.google.api.services.dataflow.model.JobMessage) Instant(org.joda.time.Instant) ListJobMessagesResponse(com.google.api.services.dataflow.model.ListJobMessagesResponse) Test(org.junit.Test)

Example 2 with ListJobMessagesResponse

use of com.google.api.services.dataflow.model.ListJobMessagesResponse in project beam by apache.

the class MonitoringUtil method getJobMessages.

/**
   * Return job messages sorted in ascending order by timestamp.
   * @param jobId The id of the job to get the messages for.
   * @param startTimestampMs Return only those messages with a
   *   timestamp greater than this value.
   * @return collection of messages
   */
public List<JobMessage> getJobMessages(String jobId, long startTimestampMs) throws IOException {
    // TODO: Allow filtering messages by importance
    Instant startTimestamp = new Instant(startTimestampMs);
    ArrayList<JobMessage> allMessages = new ArrayList<>();
    String pageToken = null;
    while (true) {
        ListJobMessagesResponse response = dataflowClient.listJobMessages(jobId, pageToken);
        if (response == null || response.getJobMessages() == null) {
            return allMessages;
        }
        for (JobMessage m : response.getJobMessages()) {
            @Nullable Instant timestamp = fromCloudTime(m.getTime());
            if (timestamp == null) {
                continue;
            }
            if (timestamp.isAfter(startTimestamp)) {
                allMessages.add(m);
            }
        }
        if (response.getNextPageToken() == null) {
            break;
        } else {
            pageToken = response.getNextPageToken();
        }
    }
    Collections.sort(allMessages, new TimeStampComparator());
    return allMessages;
}
Also used : Instant(org.joda.time.Instant) JobMessage(com.google.api.services.dataflow.model.JobMessage) ArrayList(java.util.ArrayList) ListJobMessagesResponse(com.google.api.services.dataflow.model.ListJobMessagesResponse) Nullable(javax.annotation.Nullable)

Aggregations

JobMessage (com.google.api.services.dataflow.model.JobMessage)2 ListJobMessagesResponse (com.google.api.services.dataflow.model.ListJobMessagesResponse)2 Instant (org.joda.time.Instant)2 ArrayList (java.util.ArrayList)1 Nullable (javax.annotation.Nullable)1 DataflowClient (org.apache.beam.runners.dataflow.DataflowClient)1 Test (org.junit.Test)1