com.amazonaws.services.sqs.model.Message Java Examples

private void setSqsAttributesOnRecord(Message message, Record record, String queueUrl, String queueNamePrefix) {
  final Record.Header header = record.getHeader();

  switch (sqsAttributesOption) {
    case ALL:
      header.setAttribute(SQS_QUEUE_URL_ATTRIBUTE, queueUrl);
      Optional.of(message.getMessageAttributes()).ifPresent(attrs -> {
        attrs.forEach((name, val) -> {
          final String stringValue = val.getStringValue();
          if (stringValue != null) {
            header.setAttribute(SQS_MESSAGE_ATTRIBUTE_PREFIX + name, stringValue);
          }
        });
      });
      final String body = message.getBody();
      if (body != null) {
        header.setAttribute(SQS_MESSAGE_BODY_ATTRIBUTE, body);
      }
      final String bodyMd5 = message.getMD5OfBody();
      if (bodyMd5 != null) {
        header.setAttribute(SQS_MESSAGE_BODY_MD5_ATTRIBUTE, bodyMd5);
      }
      final String attrsMd5 = message.getMD5OfMessageAttributes();
      if (attrsMd5 != null) {
        header.setAttribute(SQS_MESSAGE_ATTRIBUTE_MD5_ATTRIBUTE, attrsMd5);
      }
      // fall through
    case BASIC:
      header.setAttribute(SQS_MESSAGE_ID_ATTRIBUTE, message.getMessageId());
      header.setAttribute(SQS_QUEUE_NAME_PREFIX_ATTRIBUTE, queueNamePrefix);
      header.setAttribute(SQS_REGION_ATTRIBUTE, awsRegionLabel);
      break;
    case NONE:
      // empty block
      break;
  }
}
 

public void receive(String queue, Predicate<List<Message>> predicate) {
    final String queueUrl = getQueue(queue);

    LOG.debug("Consuming messages from {}", queueUrl);

    final ReceiveMessageRequest request = new ReceiveMessageRequest(queueUrl)
            .withWaitTimeSeconds(maxWaitTime)
            .withMaxNumberOfMessages(maxNumberOfMessages);

    while (true) {
        ReceiveMessageResult result = sqs.receiveMessage(request);

        List<Message> messages = result.getMessages();

        if (!predicate.test(messages)) {
            return;
        }
    }
}
 

/**
 * Gets BusinessObjectDataKey from SQS message
 *
 * @param client the AWS SQS client
 * @param queueUrl the AWS SQS queue url
 *
 * @return BusinessObjectDataKey
 * @throws IOException if fails to retrieve BusinessObjectDataKey from SQS message
 * @throws ApiException if fails to make API call
 */
@Retryable(value = ApiException.class, backoff = @Backoff(delay = 2000, multiplier = 2))
BusinessObjectDataKey getBdataKeySqs(AmazonSQS client, String queueUrl) throws IOException, ApiException
{
    ReceiveMessageRequest receiveMessageRequest =
        new ReceiveMessageRequest().withMaxNumberOfMessages(MAX_NUM_MESSAGES).withQueueUrl(queueUrl).withWaitTimeSeconds(WAIT_TIME_SECS)
            .withVisibilityTimeout(VISIBILITY_TIMEOUT_SECS);

    LOGGER.info("Checking queue");
    ReceiveMessageResult receiveMessageResult = client.receiveMessage(receiveMessageRequest);
    if (receiveMessageResult != null && receiveMessageResult.getMessages() != null && receiveMessageResult.getMessages().size() > 0)
    {
        List<Message> sqsMessageList = receiveMessageResult.getMessages();

        LOGGER.info("Scanning {} messages for {} and {}", sqsMessageList.size(), SEARCH_KEYWORD_1, SEARCH_KEYWORD_2);
        // Get message type BUS_OBJCT_DATA_STTS_CHG
        for (Message sqsMessage : sqsMessageList)
        {
            String receivedMessageBody = sqsMessage.getBody();
            if (receivedMessageBody.contains(SEARCH_KEYWORD_1) && receivedMessageBody.contains(SEARCH_KEYWORD_2))
            {
                LOGGER.info("Received Message: {}", receivedMessageBody);
                return mapJsontoBdataKey(receivedMessageBody).getBusinessObjectDataKey();
            }
        }
    }
    throw new ApiException("No SQS message found in queue: " + queueUrl);
}
 

@Override
public List<ScanRangeComplete> claimCompleteScanRanges(Duration ttl) {
    List<Message> messages = _sqs.receiveMessage(new ReceiveMessageRequest()
            .withQueueUrl(getQueueUrl(_completeScanRangeQueue))
            .withMaxNumberOfMessages(10)
            .withVisibilityTimeout(toSeconds(ttl))
    ).getMessages();

    return FluentIterable.from(messages)
            .transform(new Function<Message, ScanRangeComplete>() {
                @Override
                public ScanRangeComplete apply(Message message) {
                    QueueScanRangeComplete completion = JsonHelper.fromJson(message.getBody(), QueueScanRangeComplete.class);
                    completion.setMessageId(message.getReceiptHandle());
                    return completion;
                }
            })
            .toList();
}
 

@Test
public void testParse() throws Exception {
    final Message message = new Message()
            .withBody("{\n" +
                    "  \"Type\" : \"Notification\",\n" +
                    "  \"MessageId\" : \"55508fe9-870b-590c-960f-c34960b669f0\",\n" +
                    "  \"TopicArn\" : \"arn:aws:sns:eu-west-1:459220251735:cloudtrail-write\",\n" +
                    "  \"Message\" : \"{\\\"s3Bucket\\\":\\\"cloudtrailbucket\\\",\\\"s3ObjectKey\\\":[\\\"example/AWSLogs/459220251735/CloudTrail/eu-west-1/2014/09/27/459220251735_CloudTrail_eu-west-1_20140927T1625Z_UPwzr7ft2mf0Q1SS.json.gz\\\"]}\",\n" +
                    "  \"Timestamp\" : \"2014-09-27T16:27:41.258Z\",\n" +
                    "  \"SignatureVersion\" : \"1\",\n" +
                    "  \"Signature\" : \"O05joR97NvGHqMJQwsSNXzeSHrtbLqbRcqsXB7xmqARyaCGXjaVh2duwTUL93s4YvoNENnOEMzkILKI5PwmQQPha5/cmj6FSjblwRMMga6Xzf6cMnurT9TphQO7z35foHG49IejW05IkzIwD/DW0GvafJLah+fQI3EFySnShzXLFESGQuumdS8bxnM5r96ne8t+MEAHfBCVyQ/QrduO9tTtfXAz6OeWg1IEwV3TeZ5c5SS5vRxxhsD4hOJSmXAUM99CeQfcG9s7saBcvyyGPZrhPEh8S1uhiTmLvr6h1voM9vgiCbCCUujExvg+bnqsXWTZBmnatF1iOyxFfYcZ6kw==\",\n" +
                    "  \"SigningCertURL\" : \"https://sns.eu-west-1.amazonaws.com/SimpleNotificationService-d6d679a1d18e95c2f9ffcf11f4f9e198.pem\",\n" +
                    "  \"UnsubscribeURL\" : \"https://sns.eu-west-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:eu-west-1:459220251735:cloudtrail-write:9a3a4e76-4173-4c8c-b488-0126315ba643\"\n" +
                    "}");

    CloudtrailSNSNotificationParser parser = new CloudtrailSNSNotificationParser(objectMapper);

    List<CloudtrailSNSNotification> notifications = parser.parse(message);
    assertEquals(1, notifications.size());

    CloudtrailSNSNotification notification = notifications.get(0);

    assertEquals(notification.getS3Bucket(), "cloudtrailbucket");
    assertEquals(notification.getS3ObjectKey(), "example/AWSLogs/459220251735/CloudTrail/eu-west-1/2014/09/27/459220251735_CloudTrail_eu-west-1_20140927T1625Z_UPwzr7ft2mf0Q1SS.json.gz");
}
 

public ReceiveMessageResult receiveMessage(ReceiveMessageRequest request) {
    Integer waitTimeSeconds = request.getWaitTimeSeconds();
    long timeout = waitTimeSeconds != null ? waitTimeSeconds : 0; 
    
    try {
        MessageContent messageContent = visibleMessages.poll(timeout, TimeUnit.SECONDS);
        ReceiveMessageResult result = new ReceiveMessageResult();
        if (messageContent != null) {
            Message message = messageContent.toMessage();
            String receiptHandle = UUID.randomUUID().toString();
            inflight.put(receiptHandle, messageContent);
            message.withReceiptHandle(receiptHandle);
            result.withMessages(message);
        }
        return result;
    } catch (InterruptedException e) {
        // Imitate what the real SDK does
        throw new AmazonClientException(e);
    }
}
 

@Override
public void run() {
  while (pollForMessages) {
    try {
      List<Message> messages = temporarySQSQueue.fetchMessages();

      if (messages.isEmpty()) {
        continue;
      }

      messages.forEach(
          message -> {
            S3Event s3Event = unmarshall(objectMapper, message.getBody());
            if (s3Event != null) {
              tick(s3Event);
            }
            temporarySQSQueue.markMessageAsHandled(message.getReceiptHandle());
          });
    } catch (Exception e) {
      log.error("Failed to poll for messages", e);
      registry.counter("s3.eventing.pollErrors").increment();
    }
  }
}
 

private SQSMessageConsumerPrefetch.MessageManager createFifoMessageManager(String queueUrl, String groupId, String messageId, String receiptHandle) throws JMSException {
    Message message = new Message();
    message.setBody("body");
    message.setMessageId(messageId);
    message.setReceiptHandle(receiptHandle);
    Map<String, String> attributes = new HashMap<String, String>();
    attributes.put(SQSMessagingClientConstants.SEQUENCE_NUMBER, "728374687246872364");
    attributes.put(SQSMessagingClientConstants.MESSAGE_DEDUPLICATION_ID, messageId);
    attributes.put(SQSMessagingClientConstants.MESSAGE_GROUP_ID, groupId);
    attributes.put(SQSMessagingClientConstants.APPROXIMATE_RECEIVE_COUNT, "0");
    message.setAttributes(attributes);
    SQSMessage sqsMessage = new SQSTextMessage(acknowledger, queueUrl, message);
    PrefetchManager prefetchManager = mock(PrefetchManager.class);
    when(prefetchManager.getMessageConsumer())
        .thenReturn(consumer);
    SQSMessageConsumerPrefetch.MessageManager msgManager = new SQSMessageConsumerPrefetch.MessageManager(prefetchManager, sqsMessage);
    return msgManager;
}
 

@Test
public void testSimpleDelete() throws IOException {
    // Setup - put one message in queue
    SendMessageResult sendMessageResult = sqsClient.sendMessage(TEST_QUEUE_URL, "Test message");
    assertEquals(200, sendMessageResult.getSdkHttpMetadata().getHttpStatusCode());

    // Setup - receive message to get receipt handle
    ReceiveMessageResult receiveMessageResult = sqsClient.receiveMessage(TEST_QUEUE_URL);
    assertEquals(200, receiveMessageResult.getSdkHttpMetadata().getHttpStatusCode());
    Message deleteMessage = receiveMessageResult.getMessages().get(0);
    String receiptHandle = deleteMessage.getReceiptHandle();

    // Test - delete message with DeleteSQS
    final TestRunner runner = TestRunners.newTestRunner(new DeleteSQS());
    runner.setProperty(DeleteSQS.CREDENTIALS_FILE, CREDENTIALS_FILE);
    runner.setProperty(DeleteSQS.QUEUE_URL, TEST_QUEUE_URL);
    runner.setProperty(DeleteSQS.REGION, TEST_REGION);
    final Map<String, String> ffAttributes = new HashMap<>();
    ffAttributes.put("filename", "1.txt");
    ffAttributes.put("sqs.receipt.handle", receiptHandle);
    runner.enqueue("TestMessageBody", ffAttributes);

    runner.run(1);

    runner.assertAllFlowFilesTransferred(DeleteSQS.REL_SUCCESS, 1);
}
 

/**
 * As long as there is at least one CloudTrail log object:
 * <p>
 *     <li>Add the CloudTrail log object key to the list.</li>
 *     <li>Add <code>accountId</code> extracted from log object key to <code>sqsMessage</code>.</li>
 *     <li>Add {@link SourceType#CloudTrailLog} to the <code>sqsMessage</code>.</li>
 * </p>
 *
 * If there is no CloudTrail log object and it is a valid S3 message, CPL adds only {@link SourceType#Other}
 * to the <code>sqsMessage</code>.
 *
 */
private void addCloudTrailLogsAndMessageAttributes(Message sqsMessage, JsonNode s3RecordsNode, List<CloudTrailLog> cloudTrailLogs) {
    SourceType sourceType = SourceType.Other;

    for (JsonNode s3Record: s3RecordsNode) {
        String bucketName = s3Record.at(S3_BUCKET_NAME).textValue();
        String objectKey = s3Record.at(S3_OBJECT_KEY).textValue();
        String eventName = s3Record.get(EVENT_NAME).textValue();

        SourceType currSourceType = sourceIdentifier.identifyWithEventName(objectKey, eventName);
        if (currSourceType == SourceType.CloudTrailLog) {
            cloudTrailLogs.add(new CloudTrailLog(bucketName, objectKey));
            sourceType = currSourceType;
            LibraryUtils.setMessageAccountId(sqsMessage, objectKey);
        }
    }

    sqsMessage.addAttributesEntry(SourceAttributeKeys.SOURCE_TYPE.getAttributeKey(), sourceType.name());
}
 

private void handleMessage(Message message) {
  Exception caught = null;
  try {
    messageHandler.handleMessage(message);
    getSuccessCounter().increment();
  } catch (Exception e) {
    log.error("failed to process message {}", message, e);
    getErrorCounter(e).increment();
    caught = e;
  }

  if (caught == null) {
    messageAcknowledger.ack(subscriptionInfo, message);
  } else {
    messageAcknowledger.nack(subscriptionInfo, message);
  }
}
 

private void process(final List<Message> messages) {
  if (messages.size() == 0) return;

  final List<DeleteMessageBatchRequestEntry> toDelete = new ArrayList<>();
  int count = 0;
  for (Message message : messages) {
    final String deleteId = String.valueOf(count++);
    try {
      String stringBody = message.getBody();
      if (stringBody.isEmpty() || stringBody.equals("[]")) continue;
      // allow plain-text json, but permit base64 encoded thrift or json
      byte[] serialized =
          stringBody.charAt(0) == '[' ? stringBody.getBytes(UTF_8) : Base64.decode(stringBody);
      metrics.incrementMessages();
      metrics.incrementBytes(serialized.length);
      collector.acceptSpans(
          serialized,
          new Callback<Void>() {
            @Override
            public void onSuccess(Void value) {
              toDelete.add(
                  new DeleteMessageBatchRequestEntry(deleteId, message.getReceiptHandle()));
            }

            @Override
            public void onError(Throwable t) {
              logger.log(Level.WARNING, "collector accept failed", t);
              // for cases that are not recoverable just discard the message,
              // otherwise ignore so processing can be retried.
              if (t instanceof IllegalArgumentException) {
                toDelete.add(
                    new DeleteMessageBatchRequestEntry(deleteId, message.getReceiptHandle()));
              }
            }
          });
    } catch (RuntimeException | Error e) {
      logger.log(Level.WARNING, "message decoding failed", e);
      toDelete.add(new DeleteMessageBatchRequestEntry(deleteId, message.getReceiptHandle()));
    }
  }

  if (!toDelete.isEmpty()) {
    delete(toDelete);
  }
}
 

@Test
public void testGetMessageAndAutoDelete() {
    runner.setProperty(GetSQS.QUEUE_URL, "https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000");
    runner.setProperty(GetSQS.AUTO_DELETE, "true");

    Message message1 = new Message();
    message1.setBody("TestMessage1");
    message1.setMessageId("test-message-id-1");
    message1.setReceiptHandle("test-receipt-handle-1");
    Message message2 = new Message();
    message2.setBody("TestMessage2");
    message2.setMessageId("test-message-id-2");
    message2.setReceiptHandle("test-receipt-handle-2");
    ReceiveMessageResult receiveMessageResult = new ReceiveMessageResult()
            .withMessages(message1, message2);
    Mockito.when(mockSQSClient.receiveMessage(Mockito.any(ReceiveMessageRequest.class))).thenReturn(receiveMessageResult);

    runner.run(1);

    ArgumentCaptor<ReceiveMessageRequest> captureReceiveRequest = ArgumentCaptor.forClass(ReceiveMessageRequest.class);
    Mockito.verify(mockSQSClient, Mockito.times(1)).receiveMessage(captureReceiveRequest.capture());
    ReceiveMessageRequest receiveRequest = captureReceiveRequest.getValue();
    assertEquals("https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000", receiveRequest.getQueueUrl());

    ArgumentCaptor<DeleteMessageBatchRequest> captureDeleteRequest = ArgumentCaptor.forClass(DeleteMessageBatchRequest.class);
    Mockito.verify(mockSQSClient, Mockito.times(1)).deleteMessageBatch(captureDeleteRequest.capture());
    DeleteMessageBatchRequest deleteRequest = captureDeleteRequest.getValue();
    assertEquals("https://sqs.us-west-2.amazonaws.com/123456789012/test-queue-000000000", deleteRequest.getQueueUrl());
    assertEquals("test-message-id-1", deleteRequest.getEntries().get(0).getId());
    assertEquals("test-message-id-2", deleteRequest.getEntries().get(1).getId());

    runner.assertAllFlowFilesTransferred(GetSQS.REL_SUCCESS, 2);
    List<MockFlowFile> flowFiles = runner.getFlowFilesForRelationship(GetSQS.REL_SUCCESS);
    MockFlowFile ff0 = flowFiles.get(0);
    ff0.assertAttributeEquals("sqs.message.id", "test-message-id-1");
    MockFlowFile ff1 = flowFiles.get(1);
    ff1.assertAttributeEquals("sqs.message.id", "test-message-id-2");
}
 

@Override
public CloudTrailSource getSource(Message sqsMessage) throws IOException {
    List<CloudTrailLog> cloudTrailLogs = new ArrayList<>();
    JsonNode messageNode = messageExtractor.getMessageBody(sqsMessage);

    addCloudTrailLogsAndMessageAttributes(sqsMessage, cloudTrailLogs, messageNode);
    addRestMessageAttributes(sqsMessage, messageNode);

    return new SQSBasedSource(sqsMessage, cloudTrailLogs);
}
 

@Test(expected = ReceiptHandleIsInvalidException.class)
public void deleteMessageFailsWithInvalidReceiptHandle() {
  final String queueName = "bizo";
  final String messageBody = "hi everybody";

  final CreateQueueRequest createQueueRequest = new CreateQueueRequest().withQueueName(queueName);
  sqs.createQueue(createQueueRequest);

  final GetQueueUrlRequest getQueueUrlRequest = new GetQueueUrlRequest().withQueueName(queueName);
  final GetQueueUrlResult getQueueUrlResult = sqs.getQueueUrl(getQueueUrlRequest);
  final String queueUrl = getQueueUrlResult.getQueueUrl();

  final SendMessageRequest sendMessageRequest =
    new SendMessageRequest().withQueueUrl(queueUrl).withMessageBody(messageBody);
  sqs.sendMessage(sendMessageRequest);

  final ReceiveMessageRequest receiveMessageRequest = new ReceiveMessageRequest().withQueueUrl(queueUrl);
  final ReceiveMessageResult receiveMessageResult = sqs.receiveMessage(receiveMessageRequest);
  final List<Message> messages = receiveMessageResult.getMessages();
  assertThat(messages.size(), equalTo(1));

  final String receiptHandle = "bizo";
  final DeleteMessageRequest deleteMessageRequest =
    new DeleteMessageRequest().withQueueUrl(queueUrl).withReceiptHandle(receiptHandle);

  sqs.deleteMessage(deleteMessageRequest);
}
 

public CloudTrailSource getCloudTrailSource(Message sqsMessage, JsonNode s3MessageNode) throws IOException {
    JsonNode s3RecordsNode = s3MessageNode.get(RECORDS);
    List<CloudTrailLog> cloudTrailLogs = new ArrayList<>();

    addCloudTrailLogsAndMessageAttributes(sqsMessage, s3RecordsNode, cloudTrailLogs);

    return new SQSBasedSource(sqsMessage, cloudTrailLogs);
}
 

SQSMessageConsumer(AmazonSQS sqs, String queueUrl, Consumer<Message> consumer,
                   Runnable shutdownHook, Consumer<Exception> exceptionHandler,
                   int maxWaitTimeSeconds, int pollingThreadCount) {
    this.sqs = Objects.requireNonNull(sqs);
    this.queueUrl = Objects.requireNonNull(queueUrl);
    this.consumer = Objects.requireNonNull(consumer);
    this.shutdownHook = Objects.requireNonNull(shutdownHook);
    this.exceptionHandler = Objects.requireNonNull(exceptionHandler);
    this.maxWaitTimeSeconds = maxWaitTimeSeconds;
    this.pollingThreadCount = pollingThreadCount;
}
 

@Test
public void testParseWithTwoS3Objects() throws Exception {
    final Message doubleMessage = new Message()
            .withBody("{\n" +
                    "  \"Type\" : \"Notification\",\n" +
                    "  \"MessageId\" : \"11a04c4a-094e-5395-b297-00eaefda2893\",\n" +
                    "  \"TopicArn\" : \"arn:aws:sns:eu-west-1:459220251735:cloudtrail-write\",\n" +
                    "  \"Message\" : \"{\\\"s3Bucket\\\":\\\"cloudtrailbucket\\\",\\\"s3ObjectKey\\\":[\\\"example/AWSLogs/459220251735/CloudTrail/eu-west-1/2014/09/27/459220251735_CloudTrail_eu-west-1_20140927T1620Z_Nk2SdmlEzA0gDpPr.json.gz\\\", \\\"example/AWSLogs/459220251735/CloudTrail/eu-west-1/2014/09/27/459220251999_CloudTrail2_eu-west-1_20140927T1620Z_Nk2SdmlEzA0gDpPr.json.gz\\\"]}\",\n" +
                    "  \"Timestamp\" : \"2014-09-27T16:22:44.011Z\",\n" +
                    "  \"SignatureVersion\" : \"1\",\n" +
                    "  \"Signature\" : \"q9xmJZ8nJR5iaAYMLN3M8v9HyLbUqbLjGGFlmmvIK9UDQiQO0wmvlYeo5/lQqvANW/v+NVXZxxOoWx06p6Rv5BwXIa2ASVh7RlXc2y+U2pQgLaQlJ671cA33iBi/iH1al/7lTLrlIkUb9m2gAdEyulbhZfBfAQOm7GN1PHR/nW+CtT61g4KvMSonNzj23jglLTb0r6pxxQ5EmXz6Jo5DOsbXvuFt0BSyVP/8xRXT1ap0S7BuUOstz8+FMqdUyOQSR9RA9r61yUsJ4nnq0KfK5/1gjTTDPmE4OkGvk6AuV9YTME7FWTY/wU4LPg5/+g/rUo2UDGrxnGoJ3OUW5yrtyQ==\",\n" +
                    "  \"SigningCertURL\" : \"https://sns.eu-west-1.amazonaws.com/SimpleNotificationService-d6d679a1d18e95c2f9ffcf11f4f9e198.pem\",\n" +
                    "  \"UnsubscribeURL\" : \"https://sns.eu-west-1.amazonaws.com/?Action=Unsubscribe&SubscriptionArn=arn:aws:sns:eu-west-1:459220251735:cloudtrail-write:9a3a4e76-4173-4c8c-b488-0126315ba643\"\n" +
                    "}");

    CloudtrailSNSNotificationParser parser = new CloudtrailSNSNotificationParser(objectMapper);

    List<CloudtrailSNSNotification> notifications = parser.parse(doubleMessage);
    assertEquals(2, notifications.size());

    CloudtrailSNSNotification notification1 = notifications.get(0);
    CloudtrailSNSNotification notification2 = notifications.get(1);

    assertEquals(notification1.getS3Bucket(), "cloudtrailbucket");
    assertEquals(notification1.getS3ObjectKey(), "example/AWSLogs/459220251735/CloudTrail/eu-west-1/2014/09/27/459220251735_CloudTrail_eu-west-1_20140927T1620Z_Nk2SdmlEzA0gDpPr.json.gz");

    assertEquals(notification2.getS3Bucket(), "cloudtrailbucket");
    assertEquals(notification2.getS3ObjectKey(), "example/AWSLogs/459220251735/CloudTrail/eu-west-1/2014/09/27/459220251999_CloudTrail2_eu-west-1_20140927T1620Z_Nk2SdmlEzA0gDpPr.json.gz");
}
 

@Test
public void testWrite() {
  final AmazonSQS client = embeddedSqsRestServer.getClient();
  final String queueUrl = embeddedSqsRestServer.getQueueUrl();

  List<SendMessageRequest> messages = new ArrayList<>();
  for (int i = 0; i < 100; i++) {
    final SendMessageRequest request = new SendMessageRequest(queueUrl, "This is a test " + i);
    messages.add(request);
  }
  pipeline.apply(Create.of(messages)).apply(SqsIO.write());
  pipeline.run().waitUntilFinish();

  List<String> received = new ArrayList<>();
  while (received.size() < 100) {
    final ReceiveMessageResult receiveMessageResult = client.receiveMessage(queueUrl);

    if (receiveMessageResult.getMessages() != null) {
      for (Message message : receiveMessageResult.getMessages()) {
        received.add(message.getBody());
      }
    }
  }
  assertEquals(100, received.size());
  for (int i = 0; i < 100; i++) {
    received.contains("This is a test " + i);
  }
}
 

@Override
public void run() {
    ReceiveMessageRequest receiveMessageRequest = new ReceiveMessageRequest()
        .withQueueUrl(queueUrl)
        .withMaxNumberOfMessages(MAX_NUMBER_OF_MESSAGES)
        .withVisibilityTimeout(VISIBILITY_TIMEOUT)
        .withWaitTimeSeconds(WAIT_TIME_SECONDS);
    
    while (!shutdown) {
        // Long pole the SQS queue.  This will return as soon as a message
        // is received, or when WAIT_TIME_SECONDS has elapsed.
        List<Message> messages = amazonSqs.receiveMessage(receiveMessageRequest).getMessages();
        if (messages == null) {
            // If there were no messages during this poll period, SQS
            // will return this list as null.  Continue polling.
            continue;
        }
        
        synchronized(handlers) {
            for (Message message : messages) {
                try {
                    // Parse notification and call handlers.
                    JobStatusNotification notification = parseNotification(message);
                    for (JobStatusNotificationHandler handler : handlers) {
                        handler.handle(notification);
                    }
                } catch (IOException e) {
                    System.out.println("Failed to convert notification: " + e.getMessage());
                }
                
                // Delete the message from the queue.
                amazonSqs.deleteMessage(new DeleteMessageRequest().withQueueUrl(queueUrl).withReceiptHandle(message.getReceiptHandle()));
            }
        }
    }
}
 

public void submit(Executor executor, Callable<List<Message>> callable, String queueUrl, Integer visibilityTimeoutSecondsOverride) {
    long visibilityTimeoutNanos;
    if (visibilityTimeoutSecondsOverride == null) {
        visibilityTimeoutNanos = defaultVisibilityTimeoutNanos;
    } else {
        visibilityTimeoutNanos = TimeUnit.SECONDS.toNanos(visibilityTimeoutSecondsOverride);
    }
    ReceiveMessageBatchTask task = new ReceiveMessageBatchTask(callable, queueUrl, visibilityTimeoutNanos);
    executor.execute(task);
}
 

/**
 * Constructs a receive task waiting the specified time before calling SQS.
 * 
 * @param waitTimeMs
 *            the time to wait before calling SQS
 */
ReceiveMessageBatchTask(Callable<List<Message>> callable, String sourceQueueUrl, long visibilityTimeoutNanos) {
    super(callable);
    this.sourceQueueUrl = sourceQueueUrl;
    this.visibilityTimeoutNanos = visibilityTimeoutNanos;
    messages = EMPTY_DEQUE;
}
 

public void populateMessage(int populateMessageSize) throws JMSException {
    String queueUrl = baseQueueUrl + 0;
    for (int i = 0; i < populateMessageSize; i++) {
        // Change queueUrl depending on how many messages there are.
        if (i == 11) {
            queueUrl = baseQueueUrl + 1;
        } else if (i == 22) {
            queueUrl = baseQueueUrl + 2;
        } else if (i == 33) {
            queueUrl = baseQueueUrl + 3;
        } else if (i == 44) {
            queueUrl = baseQueueUrl + 4;
        }
        
        Message sqsMessage = mock(Message.class);
        when(sqsMessage.getReceiptHandle()).thenReturn("ReceiptHandle" + i);
        when(sqsMessage.getMessageId()).thenReturn("MessageId" + i);
        // Add mock Attributes
        Map<String, String> mockAttributes = new HashMap<String, String>();
        mockAttributes.put(SQSMessagingClientConstants.APPROXIMATE_RECEIVE_COUNT, "2");
        when(sqsMessage.getAttributes()).thenReturn(mockAttributes);
        
        SQSMessage message = (SQSMessage) new SQSTextMessage(acknowledger, queueUrl, sqsMessage);
        
        populatedMessages.add(message);
        acknowledger.notifyMessageReceived(message);
    }
    Assert.assertEquals(populateMessageSize, acknowledger.getUnAckMessages().size());
}
 

@Override
public CloudTrailSource getSource(Message sqsMessage) throws IOException {
    if (messageExtractor.getMessageText(sqsMessage).equals(CLOUD_TRAIL_VALIDATION_MESSAGE)) {
        sqsMessage.addAttributesEntry(SourceAttributeKeys.SOURCE_TYPE.getAttributeKey(), SourceType.CloudTrailValidationMessage.name());
        return new SQSBasedSource(sqsMessage, null);
    }

    return null;
}
 

private void processMessages() {
    if (this.isShutDown) {
        return;
    }

    final List<Message> messages = this.channel.getMessages();
    List<Message> proceedMessages = notifyListeners(messages);
    log.debug("Received %d messages, proceed %d messages", messages.size(), proceedMessages.size());
    this.channel.deleteMessages(messages);
}
 

@Override
public List<Message> handleMessages(final List<Message> messages) {
    List<Message> proceedMessages = new ArrayList<>();
    for (final Message message : messages) {
        if (this.handleMessage(message)) {
            proceedMessages.add(message);
        }
    }
    return proceedMessages;
}
 

public void changeMessageVisibility(Message msg, int value) {
  logger.info("Change visibility to {} seconds", value);
  if (value > 36000) {
    value = 36000;
  }
  ChangeMessageVisibilityRequest
      request =
      new ChangeMessageVisibilityRequest()
          .withQueueUrl(this.queueUrl)
          .withReceiptHandle(msg.getReceiptHandle()).withVisibilityTimeout(value);
  this.getClient().changeMessageVisibility(request);
}
 

private MessageExecutor(String logicalQueueName, Message message,
		QueueAttributes queueAttributes) {
	this.logicalQueueName = logicalQueueName;
	this.message = message;
	this.queueUrl = queueAttributes.getReceiveMessageRequest().getQueueUrl();
	this.hasRedrivePolicy = queueAttributes.hasRedrivePolicy();
	this.deletionPolicy = queueAttributes.getDeletionPolicy();
}
 

/**
 * Deletes all messages in the given queue.
 * 
 * @param queue 
 */
public void clearAll(String queue) {
    do {
        List<Message> messages = pull(queue, batchCount);
 
        if(messages.isEmpty()) {
            break;
        }
        
        delete(queue, messages);
    } while(true);
}
 

private static void mockReceiveMessage(AmazonSQSAsync sqs, String queueUrl,
		String messageContent, String receiptHandle) {
	when(sqs.receiveMessage(new ReceiveMessageRequest(queueUrl)
			.withAttributeNames("All").withMessageAttributeNames("All")
			.withMaxNumberOfMessages(10).withWaitTimeSeconds(20)))
					.thenReturn(
							new ReceiveMessageResult().withMessages(
									new Message().withBody(messageContent)
											.withReceiptHandle(receiptHandle)),
							new ReceiveMessageResult());
}