allow for output field configuration

xloouis · Aug 15, 2019 · f8f5f49 · f8f5f49
1 parent e0b6d30
commit f8f5f49
Show file tree

Hide file tree

Showing 3 changed files with 81 additions and 11 deletions.
diff --git a/README.md b/README.md
@@ -18,7 +18,8 @@ PUT _ingest/pipeline/custom-pipeline-name
     {
       "regionip": {
         "field": "field_containing_ip_address",
-        "target_field": "target_field_name"
+        "target_field": "target_field_name",
+        "properties": ["country_name", "region_name", "city_name"]
       }
     }
   ]
@@ -29,9 +30,7 @@ PUT _ingest/pipeline/custom-pipeline-name
 target_field result example
 
 "regionip": {
-  "ip": "127.0.0.1",
   "country_name": "中国",
-  "isp_name": "电信",
   "city_name": "上海市",
   "region_name": "上海"
 }
@@ -47,6 +46,7 @@ target_field result example
 | target_field   | no | Field name to write region info to, defaults to `regionip` |
 | ignore_missing | no | If set to true, doc missing specified field will not throw a exception, defaults to `false`. |
 | ip2region_algorithm | no |`BTREE`/`BINARY`/`MEMORY`, defaults to `MEMORY` [[link]](https://github.com/lionsoul2014/ip2region) |
+| properties | no | `ip`, `country_name`, `region_name`, `city_name`, `isp_name`, default to all properties |
 
 ## Build
 

diff --git a/build.gradle b/build.gradle
@@ -11,7 +11,7 @@ buildscript {
 }
 
 group = 'org.ltx.elasticsearch.plugin'
-version = '1.0'
+version = '1.1'
 
 apply plugin: 'java'
 apply plugin: 'elasticsearch.esplugin'
@@ -32,7 +32,7 @@ esplugin {
     name 'ingest-regionip'
     description 'Ingest processor to add Chinese region info based on provided ip address'
     classname 'org.ltx.elasticsearch.plugin.ingest.regionip.IngestRegionIpPlugin'
-    version '1.0'
+    version '1.1'
 
     licenseFile rootProject.file('LICENSE.txt')
     noticeFile rootProject.file('NOTICE.txt')

diff --git a/src/main/java/org/ltx/elasticsearch/plugin/ingest/regionip/RegionIpProcessor.java b/src/main/java/org/ltx/elasticsearch/plugin/ingest/regionip/RegionIpProcessor.java
@@ -1,11 +1,17 @@
 package org.ltx.elasticsearch.plugin.ingest.regionip;
 
 import static org.elasticsearch.ingest.ConfigurationUtils.readBooleanProperty;
+import static org.elasticsearch.ingest.ConfigurationUtils.readOptionalList;
 import static org.elasticsearch.ingest.ConfigurationUtils.readStringProperty;
 
 import java.io.IOException;
+import java.util.Arrays;
+import java.util.EnumSet;
 import java.util.HashMap;
+import java.util.List;
+import java.util.Locale;
 import java.util.Map;
+import java.util.Set;
 import org.elasticsearch.ingest.AbstractProcessor;
 import org.elasticsearch.ingest.IngestDocument;
 import org.elasticsearch.ingest.Processor;
@@ -19,6 +25,7 @@ public class RegionIpProcessor extends AbstractProcessor {
   private final String targetField;
   private final String algorithm;
   private final boolean ignoreMissing;
+  private final Set<Property> properties;
   private final Ip2RegionSearcher ip2RegionSearcher;
 
   RegionIpProcessor(
@@ -27,6 +34,7 @@ public class RegionIpProcessor extends AbstractProcessor {
       String targetField,
       String algorithm,
       boolean ignoreMissing,
+      Set<Property> properties,
       Ip2RegionSearcher ip2RegionSearcher)
       throws IOException {
 
@@ -35,6 +43,7 @@ public class RegionIpProcessor extends AbstractProcessor {
     this.targetField = targetField;
     this.algorithm = algorithm;
     this.ignoreMissing = ignoreMissing;
+    this.properties = properties;
     this.ip2RegionSearcher = ip2RegionSearcher;
   }
 
@@ -52,11 +61,29 @@ public void execute(IngestDocument ingestDocument) throws Exception {
     RegionIp regionIp = ip2RegionSearcher.searchIp(ip, algorithm);
     if (regionIp != null) {
       HashMap<String, Object> regionData = new HashMap<>(8);
-      regionData.put("ip", regionIp.getIp());
-      regionData.put("country_name", regionIp.getCountryName());
-      regionData.put("region_name", regionIp.getRegionName());
-      regionData.put("city_name", regionIp.getCityName());
-      regionData.put("isp_name", regionIp.getIspName());
+
+      for (Property property : properties) {
+        switch (property) {
+          case IP:
+            regionData.put("ip", regionIp.getIp());
+            break;
+          case COUNTRY_NAME:
+            regionData.put("country_name", regionIp.getCountryName());
+            break;
+          case REGION_NAME:
+            regionData.put("region_name", regionIp.getRegionName());
+            break;
+          case CITY_NAME:
+            regionData.put("city_name", regionIp.getCityName());
+            break;
+          case ISP_NAME:
+            regionData.put("isp_name", regionIp.getIspName());
+            break;
+          default:
+            // do nothing
+        }
+      }
+
       ingestDocument.setFieldValue(targetField, regionData);
     }
   }
@@ -77,9 +104,52 @@ public Processor create(
       String targetField = readStringProperty(TYPE, tag, config, "target_field", "regionip");
       Boolean ignoreMissing = readBooleanProperty(TYPE, tag, config, "ignore_missing", false);
       String algorithm = readStringProperty(TYPE, tag, config, "ip2region_algorithm", "MEMORY");
+      List<String> propertyNames = readOptionalList(TYPE, tag, config, "properties");
+
+      final Set<Property> properties;
+      if (propertyNames != null) {
+        properties = EnumSet.noneOf(Property.class);
+        for (String fieldName : propertyNames) {
+          try {
+            properties.add(Property.parseProperty(fieldName));
+          } catch (IllegalArgumentException e) {
+            throw new RegionIpException(e.getMessage());
+          }
+        }
+      } else {
+        properties = Property.ALL_PROPERTIES;
+      }
 
       return new RegionIpProcessor(
-          tag, field, targetField, algorithm, ignoreMissing, new Ip2RegionSearcher());
+          tag, field, targetField, algorithm, ignoreMissing, properties, new Ip2RegionSearcher());
+    }
+  }
+
+  enum Property {
+    /** output property */
+    IP,
+    COUNTRY_NAME,
+    REGION_NAME,
+    CITY_NAME,
+    ISP_NAME;
+
+    static final EnumSet<Property> ALL_PROPERTIES = EnumSet.allOf(Property.class);
+
+    public static Property parseProperty(String value) {
+      try {
+        Property property = valueOf(value.toUpperCase(Locale.ROOT));
+        if (!ALL_PROPERTIES.contains(property)) {
+          throw new IllegalArgumentException("invalid");
+        }
+
+        return property;
+      } catch (IllegalArgumentException e) {
+        throw new IllegalArgumentException(
+            "illegal property value ["
+                + value
+                + "]. valid values are "
+                + Arrays.toString(ALL_PROPERTIES.toArray()));
+      }
     }
   }