Tag virt queues

src/main/java/com/amazonaws/services/sqs/AmazonSQSRequesterClient.java

@@ -29,21 +29,23 @@ class AmazonSQSRequesterClient implements AmazonSQSRequester {
     private final AmazonSQS sqs;
     private final String queuePrefix;
     private final Map<String, String> queueAttributes;
+    private final Map<String, String> queueTags;
     private final Consumer<Exception> exceptionHandler;
 
     private final Set<SQSMessageConsumer> responseConsumers = Collections.newSetFromMap(new ConcurrentHashMap<>());
 
     private Runnable shutdownHook;
 
-    AmazonSQSRequesterClient(AmazonSQS sqs, String queuePrefix, Map<String, String> queueAttributes) {
-        this(sqs, queuePrefix, queueAttributes, SQSQueueUtils.DEFAULT_EXCEPTION_HANDLER);
+    AmazonSQSRequesterClient(AmazonSQS sqs, String queuePrefix, Map<String, String> queueAttributes, Map<String, String> queueTags) {
+        this(sqs, queuePrefix, queueAttributes, queueTags, SQSQueueUtils.DEFAULT_EXCEPTION_HANDLER);
     }
 
     AmazonSQSRequesterClient(AmazonSQS sqs, String queuePrefix, Map<String, String> queueAttributes,
-                Consumer<Exception> exceptionHandler) {
+                             Map<String, String> queueTags, Consumer<Exception> exceptionHandler) {
         this.sqs = sqs;
         this.queuePrefix = queuePrefix;
         this.queueAttributes = new HashMap<>(queueAttributes);
+        this.queueTags = new HashMap<>(queueTags);
         this.exceptionHandler = exceptionHandler;
     }
 
@@ -66,7 +68,8 @@ public CompletableFuture<Message> sendMessageAndGetResponseAsync(SendMessageRequ
         String queueName = queuePrefix + UUID.randomUUID().toString();
         CreateQueueRequest createQueueRequest = new CreateQueueRequest()
                 .withQueueName(queueName)
-                .withAttributes(queueAttributes);
+                .withAttributes(queueAttributes)
+                .withTags(queueTags);
         String responseQueueUrl = sqs.createQueue(createQueueRequest).getQueueUrl();
 
         SendMessageRequest requestWithResponseUrl = SQSQueueUtils.copyWithExtraAttributes(request,

src/main/java/com/amazonaws/services/sqs/AmazonSQSRequesterClientBuilder.java

@@ -13,6 +13,7 @@ public class AmazonSQSRequesterClientBuilder {
     private String internalQueuePrefix = "__RequesterClientQueues__";
 
     private Map<String, String> queueAttributes = Collections.emptyMap();
+    private Map<String, String> queueTags = Collections.emptyMap();
 
     private int idleQueueSweepingPeriod = 5;
     private long queueHeartbeatInterval = AmazonSQSIdleQueueDeletingClient.HEARTBEAT_INTERVAL_SECONDS_DEFAULT;
@@ -72,6 +73,19 @@ public AmazonSQSRequesterClientBuilder withQueueAttributes(Map<String, String> q
         return this;
     }
 
+    public Map<String, String> getQueueTags() {
+        return Collections.unmodifiableMap(queueTags);
+    }
+
+    public void setQueueTags(Map<String, String> queueTags) {
+        this.queueTags = new HashMap<>(queueTags);
+    }
+
+    public AmazonSQSRequesterClientBuilder withQueueTags(Map<String, String> queueTags) {
+        setQueueTags(queueTags);
+        return this;
+    }
+
     public int getIdleQueueSweepingPeriod() {
         return idleQueueSweepingPeriod;
     }

src/main/java/com/amazonaws/services/sqs/AmazonSQSTemporaryQueuesClient.java

@@ -53,14 +53,16 @@ class AmazonSQSTemporaryQueuesClient extends AbstractAmazonSQSClientWrapper {
 
     private final String prefix;
     private final long idleQueueRetentionPeriodSeconds;
+    private final Map<String, String> queueTags;
 
     private AmazonSQSRequester requester;
 
-    private AmazonSQSTemporaryQueuesClient(AmazonSQS virtualizer, AmazonSQSIdleQueueDeletingClient deleter, String queueNamePrefix, Long idleQueueRetentionPeriodSeconds) {
+    private AmazonSQSTemporaryQueuesClient(AmazonSQS virtualizer, AmazonSQSIdleQueueDeletingClient deleter, String queueNamePrefix, Long idleQueueRetentionPeriodSeconds, Map<String, String> queueTags) {
         super(virtualizer);
         this.virtualizer = virtualizer;
         this.deleter = deleter;
         this.prefix = queueNamePrefix + UUID.randomUUID().toString();
+        this.queueTags = queueTags;
 
         if (idleQueueRetentionPeriodSeconds != null) {
             AmazonSQSIdleQueueDeletingClient.checkQueueRetentionPeriodBounds(idleQueueRetentionPeriodSeconds);
@@ -69,9 +71,13 @@ private AmazonSQSTemporaryQueuesClient(AmazonSQS virtualizer, AmazonSQSIdleQueue
             this.idleQueueRetentionPeriodSeconds = AmazonSQSTemporaryQueuesClientBuilder.IDLE_QUEUE_RETENTION_PERIOD_SECONDS_DEFAULT;
         }
     }
+
+    private AmazonSQSTemporaryQueuesClient(AmazonSQS virtualizer, AmazonSQSIdleQueueDeletingClient deleter, String queueNamePrefix, Long idleQueueRetentionPeriodSeconds) {
+        this(virtualizer, deleter, queueNamePrefix, idleQueueRetentionPeriodSeconds, null);
+    }
 
     private AmazonSQSTemporaryQueuesClient(AmazonSQS virtualizer, AmazonSQSIdleQueueDeletingClient deleter, String queueNamePrefix) {
-        this(virtualizer, deleter, queueNamePrefix, null);
+        this(virtualizer, deleter, queueNamePrefix, null, null);
     }
 
     public static AmazonSQSTemporaryQueuesClient make(AmazonSQSRequesterClientBuilder builder) {
@@ -81,8 +87,8 @@ public static AmazonSQSTemporaryQueuesClient make(AmazonSQSRequesterClientBuilde
                 .withAmazonSQS(deleter)
                 .withHeartbeatIntervalSeconds(builder.getQueueHeartbeatInterval())
                 .build();
-        AmazonSQSTemporaryQueuesClient temporaryQueuesClient = new AmazonSQSTemporaryQueuesClient(virtualizer, deleter, builder.getInternalQueuePrefix(), builder.getIdleQueueRetentionPeriodSeconds());
-        AmazonSQSRequesterClient requester = new AmazonSQSRequesterClient(temporaryQueuesClient, builder.getInternalQueuePrefix(), builder.getQueueAttributes());
+        AmazonSQSTemporaryQueuesClient temporaryQueuesClient = new AmazonSQSTemporaryQueuesClient(virtualizer, deleter, builder.getInternalQueuePrefix(), builder.getIdleQueueRetentionPeriodSeconds(), builder.getQueueTags());
+        AmazonSQSRequesterClient requester = new AmazonSQSRequesterClient(temporaryQueuesClient, builder.getInternalQueuePrefix(), builder.getQueueAttributes(), builder.getQueueTags());
         AmazonSQSResponderClient responder = new AmazonSQSResponderClient(temporaryQueuesClient);
         temporaryQueuesClient.startIdleQueueSweeper(requester, responder,
                 builder.getIdleQueueSweepingPeriod(), builder.getIdleQueueSweepingTimeUnit());
@@ -122,6 +128,11 @@ public CreateQueueResult createQueue(CreateQueueRequest request) {
                     + String.join(", ", unsupportedQueueAttributes));
         }
 
+        // Add the tags specified for each Virtual queue
+        Map<String, String> tags = request.getTags();
+        queueTags.forEach(tags::putIfAbsent);
+        request.withTags(tags);
+
         Map<String, String> extraQueueAttributes = new HashMap<>();
         // Add the retention period to both the host queue and each virtual queue
         extraQueueAttributes.put(AmazonSQSIdleQueueDeletingClient.IDLE_QUEUE_RETENTION_PERIOD, Long.toString(idleQueueRetentionPeriodSeconds));
@@ -136,7 +147,8 @@ public CreateQueueResult createQueue(CreateQueueRequest request) {
         // queue or else the client may think we're trying to set them independently!
         CreateQueueRequest createVirtualQueueRequest = new CreateQueueRequest()
                 .withQueueName(request.getQueueName())
-                .withAttributes(extraQueueAttributes);
+                .withAttributes(extraQueueAttributes)
+                .withTags(tags);
         return amazonSqsToBeExtended.createQueue(createVirtualQueueRequest);
     }
 

src/main/java/com/amazonaws/services/sqs/AmazonSQSTemporaryQueuesClientBuilder.java

@@ -1,5 +1,6 @@
 package com.amazonaws.services.sqs;
 
+import java.util.Map;
 import java.util.Optional;
 import java.util.concurrent.TimeUnit;
 
@@ -67,6 +68,19 @@ public AmazonSQSTemporaryQueuesClientBuilder withIdleQueueSweepingPeriod(int per
         setIdleQueueSweepingPeriod(period, timeUnit);
         return this;
     }
+
+    public Map<String, String> getQueueTags() {
+        return requesterBuilder.getQueueTags();
+    }
+
+    public void setQueueTags(Map<String, String> queueTags) {
+        requesterBuilder.setQueueTags(queueTags);
+    }
+
+    public AmazonSQSTemporaryQueuesClientBuilder withQueueTags(Map<String, String> queueTags) {
+        setQueueTags(queueTags);
+        return this;
+    }
 
     /**
      * @return Create new instance of builder with all defaults set.

src/main/java/com/amazonaws/services/sqs/AmazonSQSVirtualQueuesClient.java

@@ -181,6 +181,10 @@ public CreateQueueResult createQueue(CreateQueueRequest request) {
             throw new IllegalStateException("Cannot create virtual queue: the number of virtual queues would exceed the maximum of "
                     + MAXIMUM_VIRTUAL_QUEUES_COUNT);
         }
+
+        TagQueueRequest tagQueueRequest = new TagQueueRequest().withTags(request.getTags());
+        virtualQueue.tagQueue(tagQueueRequest);
+
         virtualQueues.put(virtualQueue.getID().getVirtualQueueName(), virtualQueue);
 
         if (LOG.isDebugEnabled()) {

src/test/java/com/amazonaws/services/sqs/AmazonSQSIdleQueueDeletingIT.java

@@ -31,7 +31,7 @@ public class AmazonSQSIdleQueueDeletingIT extends IntegrationTest {
     public void setup() {
         client = new AmazonSQSIdleQueueDeletingClient(sqs, queueNamePrefix);
         requester = new AmazonSQSRequesterClient(sqs, queueNamePrefix,
-                Collections.emptyMap(), exceptionHandler);
+                Collections.emptyMap(), Collections.emptyMap(), exceptionHandler);
         responder = new AmazonSQSResponderClient(sqs);
     }
 

src/test/java/com/amazonaws/services/sqs/AmazonSQSRequesterClientTest.java

@@ -12,11 +12,11 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
+import com.amazonaws.services.sqs.model.ListQueueTagsResult;
 import org.junit.After;
 import org.junit.Test;
 
 import com.amazonaws.services.sqs.model.Message;
-import com.amazonaws.services.sqs.model.QueueDoesNotExistException;
 import com.amazonaws.services.sqs.model.SendMessageRequest;
 import com.amazonaws.services.sqs.util.ExceptionAsserter;
 import com.amazonaws.services.sqs.util.MockSQS;
@@ -35,6 +35,7 @@ public AmazonSQSRequesterClientTest() {
         this.accountPrefix = "http://queue.amazon.com/123456789012/";
         this.sqs = new MockSQS(accountPrefix);
         this.requesterClient = new AmazonSQSRequesterClient(sqs, "RequesterClientQueues",
+                                                            Collections.emptyMap(),
                                                             Collections.emptyMap(),
                                                             exceptionHandler);
         this.responderClient = new AmazonSQSResponderClient(sqs);
@@ -61,7 +62,8 @@ public void happyPath() throws TimeoutException, InterruptedException, Execution
         assertEquals(requestMessageBody, requestMessage.getBody());
         String responseQueueUrl = requestMessage.getMessageAttributes().get(AmazonSQSRequesterClient.RESPONSE_QUEUE_URL_ATTRIBUTE_NAME).getStringValue();
         assertNotNull(responseQueueUrl);
-
+        ListQueueTagsResult listQueueTagsResult = sqs.listQueueTags(responseQueueUrl);
+
         responderClient.sendResponseMessage(MessageContent.fromMessage(requestMessage), new MessageContent(responseMessageBody));
 
         Message response = future.get(5, TimeUnit.SECONDS);
@@ -80,13 +82,12 @@ public void timeout() throws TimeoutException, InterruptedException, ExecutionEx
         SendMessageRequest request = new SendMessageRequest()
                 .withQueueUrl(queueUrl)
                 .withMessageBody(requestMessageBody);
-        Future<Message> future = requesterClient.sendMessageAndGetResponseAsync(request, 1, TimeUnit.SECONDS);     
+        Future<Message> future = requesterClient.sendMessageAndGetResponseAsync(request, 1, TimeUnit.SECONDS);
 
         Message requestMessage = sqs.receiveMessage(queueUrl).getMessages().get(0);
         assertEquals(requestMessageBody, requestMessage.getBody());
         String responseQueueUrl = requestMessage.getMessageAttributes().get(AmazonSQSRequesterClient.RESPONSE_QUEUE_URL_ATTRIBUTE_NAME).getStringValue();
         assertNotNull(responseQueueUrl);
-
         // TODO-RS: Junit 5
         try {
             future.get();

src/test/java/com/amazonaws/services/sqs/AmazonSQSResponsesClientCrossAccountIT.java

@@ -34,7 +34,7 @@ public void setup() {
         String policyString = allowSendMessagePolicy(getBuddyRoleARN()).toJson();
         sqsRequester = new AmazonSQSRequesterClient(sqs, queueNamePrefix,
                 Collections.singletonMap(QueueAttributeName.Policy.toString(), policyString),
-                exceptionHandler);
+                Collections.emptyMap(), exceptionHandler);
 
         requestQueueUrl = sqs.createQueue("RequestQueue-" + UUID.randomUUID().toString()).getQueueUrl();
     }

src/test/java/com/amazonaws/services/sqs/AmazonSQSResponsesClientIT.java

@@ -25,7 +25,7 @@ public class AmazonSQSResponsesClientIT extends IntegrationTest {
     @Before
     public void setup() {
         sqsRequester = new AmazonSQSRequesterClient(sqs, queueNamePrefix,
-                Collections.emptyMap(), exceptionHandler);
+                Collections.emptyMap(), Collections.emptyMap(), exceptionHandler);
         sqsResponder = new AmazonSQSResponderClient(sqs);
         requestQueueUrl = sqs.createQueue("RequestQueue-" + UUID.randomUUID().toString()).getQueueUrl();
     }

src/test/java/com/amazonaws/services/sqs/AmazonSQSTemporaryQueuesClientIT.java

@@ -2,9 +2,12 @@
 
 import static com.amazonaws.services.sqs.AmazonSQSIdleQueueDeletingClient.IDLE_QUEUE_RETENTION_PERIOD;
 import static com.amazonaws.services.sqs.AmazonSQSVirtualQueuesClient.VIRTUAL_QUEUE_HOST_QUEUE_ATTRIBUTE;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertTrue;
 
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.Map;
 
 import com.amazonaws.services.sqs.model.CreateQueueRequest;
@@ -36,10 +39,18 @@ public void createQueueAddsAttributes() {
         createQueueShouldSetRetentionPeriod(null);
     }
 
+    @Test
+    public void createQueueAddsTags() {
+        Map<String, String> tags = new HashMap<>();
+        tags.put("Tag1", "Value1");
+        tags.put("Tag2", "Value2");
+        createQueueShouldSetTags(tags);
+    }
+
     @Test
     public void createQueueWithUnsupportedAttributes() {
         Assertions.assertThrows(IllegalArgumentException.class, () -> {
-            setupClient(null);
+            setupClient(null, new HashMap<>());
             client.createQueue(new CreateQueueRequest()
                     .withQueueName(queueNamePrefix + "InvalidQueue")
                     .withAttributes(Collections.singletonMap(QueueAttributeName.FifoQueue.name(), "true")));
@@ -54,17 +65,18 @@ public void createQueueConfigurableIdleQueueRetentionPeriod() {
     @Test
     public void createQueueWithUnsupportedIdleQueueRetentionPeriod() {
         Assertions.assertThrows(IllegalArgumentException.class, () -> {
-            setupClient(-10L);
+            setupClient(-10L, new HashMap<>());
         });
     }
 
-    private void setupClient(Long idleQueueRetentionPeriod) {
+    private void setupClient(Long idleQueueRetentionPeriod, Map<String, String> tags) {
         AmazonSQSRequesterClientBuilder requesterBuilder;
 
         requesterBuilder =
                 AmazonSQSRequesterClientBuilder.standard()
                         .withAmazonSQS(sqs)
-                        .withInternalQueuePrefix(queueNamePrefix);
+                        .withInternalQueuePrefix(queueNamePrefix)
+                        .withQueueTags(tags);
         if (idleQueueRetentionPeriod != null) {
             requesterBuilder = requesterBuilder.withIdleQueueRetentionPeriodSeconds(idleQueueRetentionPeriod);
         }
@@ -73,7 +85,7 @@ private void setupClient(Long idleQueueRetentionPeriod) {
     }
 
     private void createQueueShouldSetRetentionPeriod(Long idleQueueRetentionPeriod) {
-        setupClient(idleQueueRetentionPeriod);
+        setupClient(idleQueueRetentionPeriod, new HashMap<>());
         idleQueueRetentionPeriod = (idleQueueRetentionPeriod != null) ? idleQueueRetentionPeriod : 300L;
         queueUrl = client.createQueue(queueNamePrefix + "TestQueue").getQueueUrl();
         Map<String, String> attributes = client.getQueueAttributes(queueUrl, Collections.singletonList("All")).getAttributes();
@@ -84,4 +96,15 @@ private void createQueueShouldSetRetentionPeriod(Long idleQueueRetentionPeriod)
         Map<String, String> hostQueueAttributes = client.getQueueAttributes(queueUrl, Collections.singletonList("All")).getAttributes();
         Assert.assertEquals(Long.toString(idleQueueRetentionPeriod), hostQueueAttributes.get(AmazonSQSIdleQueueDeletingClient.IDLE_QUEUE_RETENTION_PERIOD));
     }
+
+    private void createQueueShouldSetTags(Map<String, String> tags) {
+        setupClient(null, tags);
+        CreateQueueRequest createQueueRequest = new CreateQueueRequest().withQueueName(queueNamePrefix + "TestQueue");
+        queueUrl = client.createQueue(createQueueRequest).getQueueUrl();
+        Map<String, String> actualTags = client.listQueueTags(queueUrl).getTags();
+        tags.forEach((k, v) -> {
+            assertTrue(actualTags.containsKey(k));
+            assertEquals(v, actualTags.get(k));
+        });
+    }
 }

src/test/java/com/amazonaws/services/sqs/SQSExecutorServiceIT.java

@@ -74,7 +74,7 @@ protected Object writeReplace() throws ObjectStreamException {
     @Before
     public void setup() {
         requester = new AmazonSQSRequesterClient(sqs, queueNamePrefix,
-                Collections.emptyMap(), exceptionHandler);
+                Collections.emptyMap(), Collections.emptyMap(), exceptionHandler);
         responder = new AmazonSQSResponderClient(sqs);
         queueUrl = sqs.createQueue(queueNamePrefix + "-RequestQueue").getQueueUrl();
         tasksCompletedLatch = new CountDownLatch(1);

src/test/java/com/amazonaws/services/sqs/SQSScheduledExecutorServiceIT.java

@@ -64,7 +64,7 @@ protected Object writeReplace() throws ObjectStreamException {
     @Before
     public void setup() {
         requester = new AmazonSQSRequesterClient(sqs, queueNamePrefix,
-                Collections.emptyMap(), exceptionHandler);
+                Collections.emptyMap(), Collections.emptyMap(), exceptionHandler);
         responder = new AmazonSQSResponderClient(sqs);
         queueUrl = sqs.createQueue(queueNamePrefix + "-RequestQueue").getQueueUrl();
         tasksRemaining = new AtomicInteger(1);