Account for multivalue lookups in doRefresh and getPreview (#23436)

* Account for multivalue lookups in doRefresh and getPreview

* unit tests

* refresh unit test

* add pr to changelog

Author: kingzacko1

changelog/unreleased/issue-4203.toml

@@ -2,4 +2,4 @@ type = "added"
 message = "Added option for multi-value lookups to CSV File data adapters."
 
 issues = ["4203"]
-pulls = ["23346"]
+pulls = ["23346", "23436"]

graylog2-server/src/main/java/org/graylog2/lookup/adapters/CSVFileDataAdapter.java

@@ -163,7 +163,11 @@ protected void doRefresh(LookupCachePurge cachePurge) throws Exception {
             }
 
             LOG.debug("CSV file {} has changed, updating data", config.path());
-            setLookupRefFromCSV();
+            if (!config.isMultiValueLookup()) {
+                setLookupRefFromCSV();
+            } else {
+                setMultiValueLookupRefFromCSV();
+            }
             cachePurge.purgeAll();
             // If the file has been moved, then moved back, the fileInfo might have been disconnected.
             // In this case, create a new fileInfo.
@@ -417,14 +421,27 @@ public LookupPreview getPreview(int size) {
             return cidrLookupRef.get().getPreview(size);
         } else {
             final Map<Object, Object> result = new HashMap<>();
-            final Map<String, String> lookup = lookupRef.get();
-            for (Map.Entry<String, String> entries : lookup.entrySet()) {
-                if (result.size() == size) {
-                    break;
+            final long totalSize;
+            if (!config.isMultiValueLookup()) {
+                final Map<String, String> lookup = lookupRef.get();
+                totalSize = lookup.size();
+                for (Map.Entry<String, String> entries : lookup.entrySet()) {
+                    if (result.size() == size) {
+                        break;
+                    }
+                    result.put(entries.getKey(), entries.getValue());
+                }
+            } else {
+                final Map<String, Map<Object, Object>> multiLookup = multiValueLookupRef.get();
+                totalSize = multiLookup.size();
+                for (Map.Entry<String, Map<Object, Object>> entries : multiLookup.entrySet()) {
+                    if (result.size() == size) {
+                        break;
+                    }
+                    result.put(entries.getKey(), toSingleValue(entries.getValue()));
                 }
-                result.put(entries.getKey(), entries.getValue());
             }
-            return new LookupPreview(lookup.size(), result);
+            return new LookupPreview(totalSize, result);
         }
     }
 

graylog2-server/src/test/java/org/graylog2/lookup/adapters/CSVFileDataAdapterTest.java

@@ -21,6 +21,7 @@
 import com.google.common.io.Resources;
 import org.graylog2.lookup.AllowedAuxiliaryPathChecker;
 import org.graylog2.plugin.lookup.LookupCachePurge;
+import org.graylog2.plugin.lookup.LookupPreview;
 import org.graylog2.plugin.lookup.LookupResult;
 import org.junit.Rule;
 import org.junit.Test;
@@ -236,6 +237,30 @@ public void testConfigValidationFailsIfValueColumnMissingForNonMultiValueLookup(
                 .contains("Value column must be set unless multi-value lookup is enabled.");
     }
 
+    @Test
+    public void testGetPreview() throws Exception {
+        final Config config = baseConfig();
+        csvFileDataAdapter = spy(new CSVFileDataAdapter("id", "name", config, new MetricRegistry(), pathChecker));
+        when(pathChecker.fileIsInAllowedPath(isA(Path.class))).thenReturn(true);
+        csvFileDataAdapter.doStart();
+
+        LookupPreview previewResult = csvFileDataAdapter.getPreview(1);
+        assertThat(previewResult).satisfies(r -> {
+            assertThat(r.supported()).isTrue();
+            assertThat(r.total()).isEqualTo(2);
+            assertThat(r.results()).hasSize(1);
+        });
+
+        previewResult = csvFileDataAdapter.getPreview(3);
+        assertThat(previewResult).satisfies(r -> {
+            assertThat(r.supported()).isTrue();
+            assertThat(r.total()).isEqualTo(2);
+            assertThat(r.results()).hasSize(2);
+            assertThat(r.results().get("foo")).isEqualTo("23");
+            assertThat(r.results().get("bar")).isEqualTo("42");
+        });
+    }
+
     @Test
     public void testCIDRLookups() throws Exception {
         final Config config = cidrLookupConfig();
@@ -292,6 +317,23 @@ public void testMultiValueLookups() throws Exception {
         assertThat(csvFileDataAdapter.doGet("999999")).isEqualTo(LookupResult.empty());
     }
 
+    @Test
+    public void refreshMultiValueSuccess() throws Exception {
+        csvFileDataAdapter = spy(new CSVFileDataAdapter("id", "name", multiValueConfig(), new MetricRegistry(), pathChecker));
+        when(pathChecker.fileIsInAllowedPath(isA(Path.class))).thenReturn(true);
+        csvFileDataAdapter.doStart();
+        csvFileDataAdapter.doRefresh(cachePurge);
+        assertFalse(csvFileDataAdapter.getError().isPresent());
+
+        LookupResult result1 = csvFileDataAdapter.doGet("000001");
+        assertThat(result1.multiValue())
+                .containsEntry("first_name", "Adam")
+                .containsEntry("last_name", "Alpha")
+                .containsEntry("username", "aalpha")
+                .containsEntry("phone", "123-4567")
+                .containsEntry("address", "123 Sleepy Hollow Lane");
+    }
+
     @Test
     public void testMultiValueCIDRLookups() throws Exception {
         csvFileDataAdapter = spy(new CSVFileDataAdapter("id", "name", multiValueCidrConfig(), new MetricRegistry(), pathChecker));
@@ -325,6 +367,27 @@ public void testMultiValueCIDRLookups() throws Exception {
         assertThat(csvFileDataAdapter.doGet("8.8.8.8")).isEqualTo(LookupResult.empty());
     }
 
+    @Test
+    public void testMultiValuePreview() throws Exception {
+        csvFileDataAdapter = spy(new CSVFileDataAdapter("id", "name", multiValueConfig(), new MetricRegistry(), pathChecker));
+        when(pathChecker.fileIsInAllowedPath(isA(Path.class))).thenReturn(true);
+        csvFileDataAdapter.doStart();
+
+        LookupPreview previewResult = csvFileDataAdapter.getPreview(3);
+        assertThat(previewResult).satisfies(r -> {
+            assertThat(r.supported()).isTrue();
+            assertThat(r.total()).isEqualTo(10);
+            assertThat(r.results()).hasSize(3);
+        });
+
+        previewResult = csvFileDataAdapter.getPreview(15);
+        assertThat(previewResult).satisfies(r -> {
+            assertThat(r.supported()).isTrue();
+            assertThat(r.total()).isEqualTo(10);
+            assertThat(r.results()).hasSize(10);
+        });
+    }
+
     private Config baseConfig() {
         return Config.builder()
                 .type(NAME)