Merge pull request #30 from DarioBalinzo/additional-fields

Additional fields

Author: DarioBalinzo

src/main/java/com/github/dariobalinzo/elastic/ElasticRepository.java

@@ -88,9 +88,7 @@ public PageResult searchAfter(String index, Cursor cursor) throws IOException, I
 
         SearchResponse response = executeSearch(searchRequest);
 
-        List<Map<String, Object>> documents = Arrays.stream(response.getHits().getHits())
-                .map(SearchHit::getSourceAsMap)
-                .collect(Collectors.toList());
+        List<Map<String, Object>> documents = extractDocuments(response);
 
         Cursor lastCursor;
         if (documents.isEmpty()) {
@@ -102,6 +100,16 @@ public PageResult searchAfter(String index, Cursor cursor) throws IOException, I
         return new PageResult(index, documents, lastCursor);
     }
 
+    private List<Map<String, Object>> extractDocuments(SearchResponse response) {
+        return Arrays.stream(response.getHits().getHits())
+                .map(hit -> {
+                    Map<String, Object> sourceMap = hit.getSourceAsMap();
+                    sourceMap.put("es-id", hit.getId());
+                    sourceMap.put("es-index", hit.getIndex());
+                    return sourceMap;
+                }).collect(Collectors.toList());
+    }
+
     public PageResult searchAfterWithSecondarySort(String index, Cursor cursor) throws IOException, InterruptedException {
         Objects.requireNonNull(secondaryCursorField);
         String primaryCursor = cursor.getPrimaryCursor();
@@ -122,9 +130,7 @@ public PageResult searchAfterWithSecondarySort(String index, Cursor cursor) thro
 
         SearchResponse response = executeSearch(searchRequest);
 
-        List<Map<String, Object>> documents = Arrays.stream(response.getHits().getHits())
-                .map(SearchHit::getSourceAsMap)
-                .collect(Collectors.toList());
+        List<Map<String, Object>> documents = extractDocuments(response);
 
         Cursor lastCursor;
         if (documents.isEmpty()) {

src/test/java/com/github/dariobalinzo/elastic/ElasticRepositoryTest.java

@@ -24,8 +24,7 @@
 import java.io.IOException;
 import java.util.Collections;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNull;
+import static org.junit.Assert.*;
 
 public class ElasticRepositoryTest extends TestContainersContext {
 
@@ -94,4 +93,19 @@ public void shouldFetchDataUsingSecondarySortField() throws IOException, Interru
         assertNull(emptyPage.getLastCursor().getSecondaryCursor());
     }
 
+    @Test
+    public void shouldFetchDataWithAdditionalField() throws IOException, InterruptedException {
+        deleteTestIndex();
+
+        insertMockData(110, "customerA", TEST_INDEX);
+        insertMockData(111, "customerB", TEST_INDEX);
+        refreshIndex();
+
+        PageResult firstPage = repository.searchAfter(TEST_INDEX, Cursor.empty());
+        firstPage.getDocuments().forEach(item -> {
+            assertNotNull(item.get((String) "es-index"));
+            assertNotNull(item.get((String) "es-id"));
+        });
+    }
+
 }