Refactor handler into smaller methods

Author: KernelPanicAUS

src/main/java/com/github/kernelpanicaus/ddb_loader/LambdaFunctionHandler.java

@@ -6,11 +6,14 @@
 import com.jayway.jsonpath.JsonPath;
 import org.json.simple.JSONObject;
 import org.json.simple.parser.JSONParser;
+import org.json.simple.parser.ParseException;
 import software.amazon.awssdk.core.ResponseInputStream;
 import software.amazon.awssdk.core.sync.ResponseTransformer;
+import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.model.BatchWriteItemRequest;
+import software.amazon.awssdk.services.dynamodb.model.BatchWriteItemResponse;
 import software.amazon.awssdk.services.dynamodb.model.PutRequest;
 import software.amazon.awssdk.services.dynamodb.model.WriteRequest;
 import software.amazon.awssdk.services.s3.S3Client;
@@ -22,7 +25,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.util.ArrayList;
-import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 import java.util.Scanner;
 import java.util.zip.GZIPInputStream;
@@ -36,28 +39,30 @@ public class LambdaFunctionHandler implements RequestStreamHandler {
     /**
      * Provide the AWS region which your DynamoDB table is hosted.
      */
-    static final Region AWS_REGION = Region.of(System.getenv("AWS_REGION"));
+    private static final Region AWS_REGION = Region.of(System.getenv("AWS_REGION"));
 
     /**
      * The DynamoDB table name.
      */
     // TODO: Make this dynamic, from the S3 event.
-    static final String DYNAMO_TABLE_NAME = System.getenv("DYNAMO_TABLE_NAME");
+    private static final String DYNAMO_TABLE_NAME = System.getenv("DYNAMO_TABLE_NAME");
 
     /**
      * Configurable batch size
      */
-    static final int BATCH_SIZE = Integer.parseInt(System.getenv().getOrDefault("BATCH_SIZE", "25"));
+    private static final int BATCH_SIZE = Integer.parseInt(System.getenv().getOrDefault("BATCH_SIZE", "25"));
 //
 //    static final ClientConfiguration config = new ClientConfiguration()
 //            .withMaxConnections(ClientConfiguration.DEFAULT_MAX_CONNECTIONS * 2);
 
-    final S3Client s3Client = S3Client.builder()
+    private static final S3Client s3Client = S3Client.builder()
             .region(AWS_REGION)
+            .httpClientBuilder(UrlConnectionHttpClient.builder())
             .build();
 
-    final DynamoDbClient dynamoDBClient = DynamoDbClient.builder()
+    private static final DynamoDbClient dynamoDBClient = DynamoDbClient.builder()
             .region(AWS_REGION)
+            .httpClientBuilder(UrlConnectionHttpClient.builder())
             .build();
 
     public void handleRequest(InputStream inputStream, OutputStream outputStream, Context context) {
@@ -81,73 +86,53 @@ public void handleRequest(InputStream inputStream, OutputStream outputStream, Co
             String srcKey = JsonPath.read(event, "$.Records[0].s3.object.key");
             String srcBucket = JsonPath.read(event, "$.Records[0].s3.bucket.name");
 
-
             logger.log("Bucket name: " + srcBucket + "\n");
             logger.log("Key name: " + srcKey + "\n");
             logger.log("S3 Object: " + srcBucket + "/" + srcKey + "\n");
 
             logger.log("S3 Event Received: " + srcBucket + "/" + srcKey + "\n");
 
-            ResponseInputStream<GetObjectResponse> responseInputStream = s3Client.getObject(
-                    GetObjectRequest.builder()
-                            .bucket(srcBucket)
-                            .key(srcKey)
-                            .build()
-                    , ResponseTransformer.toInputStream()
-            );
+            ResponseInputStream<GetObjectResponse> responseInputStream = getS3ClientObject(srcKey, srcBucket);
 
             logger.log("Reading input stream \n");
 
             GZIPInputStream gis = new GZIPInputStream(responseInputStream);
             Scanner fileIn = new Scanner(gis);
             var parser = new JSONParser();
 
-            Collection<WriteRequest> itemList = new ArrayList<>();
-
             int counter = 0;
             int batchCounter = 0;
+            List<WriteRequest> itemList = new ArrayList<>();
+
             while (fileIn.hasNext()) {
-                var line = fileIn.nextLine();
-                JSONObject jsonLine = (JSONObject) parser.parse(line);
-                JSONObject jsonItem = (JSONObject) jsonLine.get("Item");
 
-                WriteRequest item = WriteRequest.builder()
-                        .putRequest(PutRequest.builder().item(jsonItem).build())
-                        .build();
+                JSONObject jsonItem = getWriteItemRequest(fileIn, parser);
 
-                itemList.add(item);
+                itemList.add(getWriteItemRequest(jsonItem));
 
                 logger.log("[" + batchCounter + "/" + counter + "] Adding item to itemlist \n");
                 counter++;
 
                 if (counter == BATCH_SIZE) {
-                    batchCounter++;
-
-                    var batchItemRequest = BatchWriteItemRequest.builder()
-                            .requestItems(Map.of(DYNAMO_TABLE_NAME, itemList))
-                            .build();
 
                     logger.log("Sending Batch " + batchCounter + " \n");
-
-                    var outcome = dynamoDBClient.batchWriteItem(batchItemRequest);
+                    BatchWriteItemResponse outcome = getBatchWriteItemResponse(Map.of(DYNAMO_TABLE_NAME, itemList));
 
                     do {
-                        var unprocessedItemsRequest = BatchWriteItemRequest.builder()
-                                .requestItems(outcome.unprocessedItems())
-                                .build();
+                        BatchWriteItemRequest unprocessedItemsRequest = getBatchWriteItemRequest(outcome.unprocessedItems());
 
                         if (outcome.unprocessedItems().size() > 0) {
                             logger.log("Retrieving the unprocessed " + outcome.unprocessedItems().size() + " items, batch [" + batchCounter + "].");
-                            outcome = dynamoDBClient.batchWriteItem(unprocessedItemsRequest);
+                            outcome = batchWrite(unprocessedItemsRequest);
                         }
 
                     } while (outcome.unprocessedItems().size() > 0);
                     itemList.clear();
+                    batchCounter++;
                     counter = 0;
                 }
             }
 
-
             logger.log("Load finish in " + (System.currentTimeMillis() - startTime) + "ms");
             fileIn.close();
             gis.close();
@@ -158,4 +143,90 @@ public void handleRequest(InputStream inputStream, OutputStream outputStream, Co
         }
         statusReport.setExecutiongTime(System.currentTimeMillis() - startTime);
     }
+
+    /**
+     * <p>
+     * Builds and returns a <code>WriteRequest</code> for JSON object.
+     * </p>
+     *
+     * @param jsonItem
+     * @return WriteRequest object
+     */
+    private WriteRequest getWriteItemRequest(JSONObject jsonItem) {
+        return WriteRequest.builder()
+                .putRequest(PutRequest.builder().item(jsonItem).build())
+                .build();
+    }
+
+    /**
+     * <p>
+     * Takes a Map of Dynamo Table name to a List of Write requests and executes a bulk write.
+     * </p>
+     *
+     * @param items Mapping of table name to collection of write requests.
+     * @return BatchWriteItemResponse
+     */
+    private BatchWriteItemResponse getBatchWriteItemResponse(Map<String, List<WriteRequest>> items) {
+        return batchWrite(getBatchWriteItemRequest(items));
+    }
+
+    /**
+     * <p>
+     * Builds and return a BatchWriteItemRequest object.
+     * </p>
+     *
+     * @param items Mapping of table name to collection of write requests.
+     * @return BatchWriteItemRequest
+     */
+    private BatchWriteItemRequest getBatchWriteItemRequest(Map<String, List<WriteRequest>> items) {
+        return BatchWriteItemRequest.builder()
+                .requestItems(items)
+                .build();
+    }
+
+    /**
+     * <p>
+     * Executes BatchWriteItem operation against a Dynamo Table.
+     * </p>
+     *
+     * @param `batchItemRequest`
+     * @return <code>BatchWriteItemResponse</code>
+     */
+    private BatchWriteItemResponse batchWrite(BatchWriteItemRequest batchItemRequest) {
+        return dynamoDBClient.batchWriteItem(batchItemRequest);
+    }
+
+    /**
+     * <p>
+     * Returns inner JSON object.
+     * </p>
+     *
+     * @param fileIn
+     * @param parser
+     * @return <code>JSONObject</code>
+     * @throws ParseException
+     */
+    private JSONObject getWriteItemRequest(Scanner fileIn, JSONParser parser) throws ParseException {
+        var line = fileIn.nextLine();
+        JSONObject jsonLine = (JSONObject) parser.parse(line);
+        return (JSONObject) jsonLine.get("Item");
+    }
+
+    /**
+     * <p>
+     * getS3ClientObject and returns an S3 Object as a stream.
+     * </p>
+     *
+     * @param srcKey
+     * @param srcBucket
+     * @return
+     */
+    private ResponseInputStream<GetObjectResponse> getS3ClientObject(String srcKey, String srcBucket) {
+        GetObjectRequest objectRequest = GetObjectRequest.builder()
+                .bucket(srcBucket)
+                .key(srcKey)
+                .build();
+
+        return s3Client.getObject(objectRequest, ResponseTransformer.toInputStream());
+    }
 }