From c7f1e86d40bbc754fdecbf2495d4273432ecd333 Mon Sep 17 00:00:00 2001
From: Marcos Gonzalez Mayedo <alemayed@amazon.com>
Date: Tue, 11 Jul 2023 11:12:41 -0700
Subject: [PATCH] Added methods to Buffer and InMemoryBuffer for batch fetching

Signed-off-by: Marcos Gonzalez Mayedo <alemayed@amazon.com>
---
 .../dataprepper/plugins/sink/buffer/Buffer.java       |  6 ++++++
 .../plugins/sink/buffer/InMemoryBuffer.java           | 11 +++++++++++
 2 files changed, 17 insertions(+)

diff --git a/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/buffer/Buffer.java b/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/buffer/Buffer.java
index 634566bb91..bfdfb0d825 100644
--- a/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/buffer/Buffer.java
+++ b/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/buffer/Buffer.java
@@ -5,6 +5,8 @@
 
 package org.opensearch.dataprepper.plugins.sink.buffer;
 
+import java.util.ArrayList;
+
 /**
  * Buffer that handles the temporary storage of
  * events. It isolates the implementation of system storage.
@@ -28,4 +30,8 @@ public interface Buffer {
     void writeEvent(byte[] event);
 
     byte[] popEvent();
+
+    ArrayList<byte[]> getBufferedData();
+
+    void clearBuffer();
 }
\ No newline at end of file
diff --git a/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/buffer/InMemoryBuffer.java b/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/buffer/InMemoryBuffer.java
index f670fa177c..25de30ea98 100644
--- a/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/buffer/InMemoryBuffer.java
+++ b/data-prepper-plugins/cloudwatch-logs/src/main/java/org/opensearch/dataprepper/plugins/sink/buffer/InMemoryBuffer.java
@@ -36,4 +36,15 @@ public byte[] popEvent() {
         bufferSize -= eventsBuffered.get(0).length;
         return eventsBuffered.remove(0);
     }
+
+    @Override
+    public ArrayList<byte[]> getBufferedData() {
+        return eventsBuffered;
+    }
+
+    @Override
+    public void clearBuffer() {
+        bufferSize = 0;
+        eventsBuffered.clear();
+    }
 }
\ No newline at end of file