Provide functionality for the readBatch and SendJob

gchq · Nov 20, 2024 · 3a41f30 · 3a41f30
1 parent cb3317a
commit 3a41f30
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 6 deletions.
diff --git a/...bda/src/main/java/sleeper/compaction/job/creation/lambda/CompactionJobDispatchLambda.java b/...bda/src/main/java/sleeper/compaction/job/creation/lambda/CompactionJobDispatchLambda.java
@@ -17,9 +17,13 @@
 
 import com.amazonaws.services.dynamodbv2.AmazonDynamoDB;
 import com.amazonaws.services.s3.AmazonS3;
+import com.amazonaws.services.sqs.AmazonSQS;
 import org.apache.hadoop.conf.Configuration;
 
+import sleeper.compaction.core.job.CompactionJobSerDe;
 import sleeper.compaction.core.job.dispatch.CompactionJobDispatcher;
+import sleeper.compaction.core.job.dispatch.CompactionJobDispatcher.ReadBatch;
+import sleeper.compaction.core.job.dispatch.CompactionJobDispatcher.SendJob;
 import sleeper.configuration.properties.S3InstanceProperties;
 import sleeper.configuration.properties.S3TableProperties;
 import sleeper.core.properties.instance.InstanceProperties;
@@ -28,18 +32,31 @@
 import java.time.Instant;
 import java.util.function.Supplier;
 
+import static sleeper.core.properties.instance.CdkDefinedInstanceProperty.COMPACTION_JOB_QUEUE_URL;
+
 public class CompactionJobDispatchLambda {
 
     private CompactionJobDispatchLambda() {
     }
 
     public static CompactionJobDispatcher dispatcher(
-            AmazonS3 s3, AmazonDynamoDB dynamoDB, Configuration conf, String configBucket, Supplier<Instant> timeSupplier) {
+            AmazonS3 s3, AmazonDynamoDB dynamoDB, AmazonSQS sqs, Configuration conf, String configBucket, Supplier<Instant> timeSupplier) {
         InstanceProperties instanceProperties = S3InstanceProperties.loadFromBucket(s3, configBucket);
+        CompactionJobSerDe compactionJobSerDe = new CompactionJobSerDe();
+
         return new CompactionJobDispatcher(instanceProperties,
                 S3TableProperties.createProvider(instanceProperties, s3, dynamoDB),
                 StateStoreFactory.createProvider(instanceProperties, s3, dynamoDB, conf),
-                null, null, null, timeSupplier);
+                readBatch(s3, compactionJobSerDe), sendJob(instanceProperties, sqs, compactionJobSerDe), null, timeSupplier);
     }
 
+    private static SendJob sendJob(InstanceProperties instanceProperties, AmazonSQS sqsClient, CompactionJobSerDe compactionJobSerDe) {
+        return compactionJob -> sqsClient.sendMessage(
+                instanceProperties.get(COMPACTION_JOB_QUEUE_URL),
+                compactionJobSerDe.toJson(compactionJob));
+    }
+
+    private static ReadBatch readBatch(AmazonS3 s3, CompactionJobSerDe compactionJobSerDe) {
+        return (bucketName, key) -> compactionJobSerDe.batchFromJson(s3.getObjectAsString(bucketName, key));
+    }
 }
diff --git a/.../src/test/java/sleeper/compaction/job/creation/lambda/CompactionJobsDispatchLambdaIT.java b/.../src/test/java/sleeper/compaction/job/creation/lambda/CompactionJobsDispatchLambdaIT.java
@@ -21,8 +21,10 @@
 import com.amazonaws.services.s3.AmazonS3ClientBuilder;
 import com.amazonaws.services.sqs.AmazonSQS;
 import com.amazonaws.services.sqs.AmazonSQSClientBuilder;
+import com.amazonaws.services.sqs.model.Message;
+import com.amazonaws.services.sqs.model.ReceiveMessageRequest;
+import com.amazonaws.services.sqs.model.ReceiveMessageResult;
 import org.apache.hadoop.conf.Configuration;
-import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
 import org.testcontainers.containers.localstack.LocalStackContainer;
 import org.testcontainers.junit.jupiter.Container;
@@ -86,7 +88,6 @@ public class CompactionJobsDispatchLambdaIT {
     StateStore stateStore = stateStoreProvider.getStateStore(tableProperties);
 
     @Test
-    @Disabled("TODO")
     void shouldSendCompactionJobsInABatchWhenAllFilesAreAssigned() {
 
         // Given
@@ -156,10 +157,15 @@ private void dispatchWithTimeAtRetryCheck(CompactionJobDispatchRequest request,
     }
 
     private CompactionJobDispatcher dispatcher(List<Instant> times) {
-        return CompactionJobDispatchLambda.dispatcher(s3, dynamoDB, conf, instanceProperties.get(CONFIG_BUCKET), times.iterator()::next);
+        return CompactionJobDispatchLambda.dispatcher(s3, dynamoDB, sqs, conf, instanceProperties.get(CONFIG_BUCKET), times.iterator()::next);
     }
 
     private List<CompactionJob> receiveCompactionJobs() {
-        return null;
+        ReceiveMessageResult result = sqs.receiveMessage(new ReceiveMessageRequest()
+                .withQueueUrl(instanceProperties.get(COMPACTION_JOB_QUEUE_URL))
+                .withMaxNumberOfMessages(10));
+        return result.getMessages().stream()
+                .map(Message::getBody)
+                .map(new CompactionJobSerDe()::fromJson).toList();
     }
 }