Lookup tables (#177)

* add generic lookup table support * Use the new LookupResult return type in the lookup function (#176)
Graylog2 · Apr 28, 2017 · 5aa2255 · 5aa2255
1 parent b3cbc50
commit 5aa2255
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 0 deletions.
diff --git a/...c/main/java/org/graylog/plugins/pipelineprocessor/functions/ProcessorFunctionsModule.java b/...c/main/java/org/graylog/plugins/pipelineprocessor/functions/ProcessorFunctionsModule.java
@@ -52,6 +52,7 @@
 import org.graylog.plugins.pipelineprocessor.functions.ips.IpAddressConversion;
 import org.graylog.plugins.pipelineprocessor.functions.json.JsonParse;
 import org.graylog.plugins.pipelineprocessor.functions.json.SelectJsonPath;
+import org.graylog.plugins.pipelineprocessor.functions.lookup.Lookup;
 import org.graylog.plugins.pipelineprocessor.functions.messages.CloneMessage;
 import org.graylog.plugins.pipelineprocessor.functions.messages.CreateMessage;
 import org.graylog.plugins.pipelineprocessor.functions.messages.DropMessage;
@@ -171,6 +172,9 @@ protected void configure() {
         addMessageProcessorFunction(SyslogLevelConversion.NAME, SyslogLevelConversion.class);
         addMessageProcessorFunction(SyslogPriorityConversion.NAME, SyslogPriorityConversion.class);
         addMessageProcessorFunction(SyslogPriorityToStringConversion.NAME, SyslogPriorityToStringConversion.class);
+
+        // Lookup tables
+        addMessageProcessorFunction(Lookup.NAME, Lookup.class);
     }
 
     protected void addMessageProcessorFunction(String name, Class<? extends Function<?>> functionClass) {

diff --git a/plugin/src/main/java/org/graylog/plugins/pipelineprocessor/functions/lookup/Lookup.java b/plugin/src/main/java/org/graylog/plugins/pipelineprocessor/functions/lookup/Lookup.java
@@ -0,0 +1,63 @@
+package org.graylog.plugins.pipelineprocessor.functions.lookup;
+
+import com.google.inject.Inject;
+import com.google.inject.TypeLiteral;
+import org.graylog.plugins.pipelineprocessor.EvaluationContext;
+import org.graylog.plugins.pipelineprocessor.ast.functions.AbstractFunction;
+import org.graylog.plugins.pipelineprocessor.ast.functions.FunctionArgs;
+import org.graylog.plugins.pipelineprocessor.ast.functions.FunctionDescriptor;
+import org.graylog.plugins.pipelineprocessor.ast.functions.ParameterDescriptor;
+import org.graylog2.lookup.LookupTableService;
+import org.graylog2.plugin.lookup.LookupResult;
+
+import java.util.Collections;
+import java.util.Map;
+
+import static org.graylog.plugins.pipelineprocessor.ast.functions.ParameterDescriptor.object;
+import static org.graylog.plugins.pipelineprocessor.ast.functions.ParameterDescriptor.string;
+
+public class Lookup extends AbstractFunction<Map<Object, Object>> {
+
+    public static final String NAME = "lookup";
+
+    private final ParameterDescriptor<String, LookupTableService.Function> lookupTableParam;
+    private final ParameterDescriptor<Object, Object> keyParam;
+    private final ParameterDescriptor<Object, Object> defaultParam;
+
+    @Inject
+    public Lookup(LookupTableService lookupTableService) {
+        lookupTableParam = string("lookup_table", LookupTableService.Function.class)
+                .transform(tableName -> lookupTableService.newBuilder().lookupTable(tableName).build())
+                .build();
+        keyParam = object("key").build();
+        defaultParam = object("default").optional().build();
+    }
+
+    @Override
+    public Map<Object, Object> evaluate(FunctionArgs args, EvaluationContext context) {
+        Object key = keyParam.required(args, context);
+        if (key == null) {
+            return Collections.singletonMap(key, defaultParam.optional(args, context));
+        }
+        LookupTableService.Function table = lookupTableParam.required(args, context);
+        if (table == null) {
+            return Collections.singletonMap(key, defaultParam.optional(args, context));
+        }
+        LookupResult result = table.lookup(key);
+        if (result == null || result.isEmpty()) {
+            return Collections.singletonMap(key, defaultParam.optional(args, context));
+        }
+        return result.asMap();
+    }
+
+    @Override
+    public FunctionDescriptor<Map<Object, Object>> descriptor() {
+        //noinspection unchecked
+        return FunctionDescriptor.<Map<Object, Object>>builder()
+                .name(NAME)
+                .description("Looks a value up in the named lookup table.")
+                .params(lookupTableParam, keyParam)
+                .returnType((Class<? extends Map<Object, Object>>) new TypeLiteral<Map<Object, Object>>() {}.getRawType())
+                .build();
+    }
+}