Merge pull request #576 from lonvia/extra-attributes

Add new import option '-extra-tags'
komoot · May 7, 2021 · 1c52419 · 1c52419
2 parents e7f37bb + 53228a1
commit 1c52419
Show file tree

Hide file tree

Showing 13 changed files with 202 additions and 23 deletions.
diff --git a/src/main/java/de/komoot/photon/App.java b/src/main/java/de/komoot/photon/App.java
@@ -107,7 +107,7 @@ private static void startRecreatingIndex(Server esServer) {
     private static void startJsonDump(CommandLineArgs args) {
         try {
             final String filename = args.getJsonDump();
-            final JsonDumper jsonDumper = new JsonDumper(filename, args.getLanguages());
+            final JsonDumper jsonDumper = new JsonDumper(filename, args.getLanguages(), args.getExtraTags());
             NominatimConnector nominatimConnector = new NominatimConnector(args.getHost(), args.getPort(), args.getDatabase(), args.getUser(), args.getPassword());
             nominatimConnector.setImporter(jsonDumper);
             nominatimConnector.readEntireDatabase(args.getCountryCodes().split(","));
@@ -133,7 +133,7 @@ private static void startNominatimImport(CommandLineArgs args, Server esServer,
         }
 
         log.info("starting import from nominatim to photon with languages: " + args.getLanguages());
-        de.komoot.photon.elasticsearch.Importer importer = new de.komoot.photon.elasticsearch.Importer(esNodeClient, args.getLanguages());
+        de.komoot.photon.elasticsearch.Importer importer = new de.komoot.photon.elasticsearch.Importer(esNodeClient, args.getLanguages(), args.getExtraTags());
         NominatimConnector nominatimConnector = new NominatimConnector(args.getHost(), args.getPort(), args.getDatabase(), args.getUser(), args.getPassword());
         nominatimConnector.setImporter(importer);
         nominatimConnector.readEntireDatabase(args.getCountryCodes().split(","));
@@ -149,7 +149,7 @@ private static void startNominatimImport(CommandLineArgs args, Server esServer,
      */
     private static NominatimUpdater setupNominatimUpdater(CommandLineArgs args, Client esNodeClient) {
         NominatimUpdater nominatimUpdater = new NominatimUpdater(args.getHost(), args.getPort(), args.getDatabase(), args.getUser(), args.getPassword());
-        Updater updater = new de.komoot.photon.elasticsearch.Updater(esNodeClient, args.getLanguages());
+        Updater updater = new de.komoot.photon.elasticsearch.Updater(esNodeClient, args.getLanguages(), args.getExtraTags());
         nominatimUpdater.setUpdater(updater);
         return nominatimUpdater;
     }

diff --git a/src/main/java/de/komoot/photon/CommandLineArgs.java b/src/main/java/de/komoot/photon/CommandLineArgs.java
@@ -32,6 +32,9 @@ public class CommandLineArgs {
     @Parameter(names = "-country-codes", description = "country codes filter that nominatim importer should import, comma separated. If empty full planet is done")
     private String countryCodes = "";
 
+    @Parameter(names = "-extra-tags", description = "additional tags to save for each place")
+    private String extraTags = "";
+
     @Parameter(names = "-json", description = "import nominatim database and dump it to a json like files in (useful for developing)")
     private String jsonDump = null;
 

diff --git a/src/main/java/de/komoot/photon/JsonDumper.java b/src/main/java/de/komoot/photon/JsonDumper.java
@@ -15,17 +15,19 @@
 public class JsonDumper implements Importer {
     private PrintWriter writer = null;
     private final String[] languages;
+    private final String[] extraTags;
 
-    public JsonDumper(String filename, String languages) throws FileNotFoundException {
+    public JsonDumper(String filename, String languages, String extraTags) throws FileNotFoundException {
         this.writer = new PrintWriter(filename);
         this.languages = languages.split(",");
+        this.extraTags = extraTags.split(",");
     }
 
     @Override
     public void add(PhotonDoc doc) {
         try {
             writer.println("{\"index\": {}}");
-            writer.println(Utils.convert(doc, this.languages).string());
+            writer.println(Utils.convert(doc, languages, extraTags).string());
         } catch (IOException e) {
             log.error("error writing json file", e);
         }

diff --git a/src/main/java/de/komoot/photon/Utils.java b/src/main/java/de/komoot/photon/Utils.java
@@ -23,7 +23,7 @@
 public class Utils {
     private static final Joiner commaJoiner = Joiner.on(", ").skipNulls();
 
-    public static XContentBuilder convert(PhotonDoc doc, String[] languages) throws IOException {
+    public static XContentBuilder convert(PhotonDoc doc, String[] languages, String[] extraTags) throws IOException {
         final AddressType atype = doc.getAddressType();
         XContentBuilder builder = XContentFactory.jsonBuilder().startObject()
                 .field(Constants.OSM_ID, doc.getOsmId())
@@ -56,6 +56,7 @@ public static XContentBuilder convert(PhotonDoc doc, String[] languages) throws
         if (countryCode != null)
             builder.field(Constants.COUNTRYCODE, countryCode.getAlpha2());
         writeContext(builder, doc.getContext(), languages);
+        writeExtraTags(builder, doc.getExtratags(), extraTags);
         writeExtent(builder, doc.getBbox());
 
         builder.endObject();
@@ -64,6 +65,25 @@ public static XContentBuilder convert(PhotonDoc doc, String[] languages) throws
         return builder;
     }
 
+    private static void writeExtraTags(XContentBuilder builder, Map<String, String> docTags, String[] extraTags) throws IOException {
+        boolean foundTag = false;
+
+        for (String tag: extraTags) {
+            String value = docTags.get(tag);
+            if (value != null) {
+                if (!foundTag) {
+                    builder.startObject("extra");
+                    foundTag = true;
+                }
+                builder.field(tag, value);
+            }
+        }
+
+        if (foundTag) {
+            builder.endObject();
+        }
+    }
+
     private static void writeExtent(XContentBuilder builder, Envelope bbox) throws IOException {
         if (bbox == null) return;
 

diff --git a/src/main/java/de/komoot/photon/elasticsearch/Importer.java b/src/main/java/de/komoot/photon/elasticsearch/Importer.java
@@ -21,18 +21,20 @@ public class Importer implements de.komoot.photon.Importer {
     private final Client esClient;
     private BulkRequestBuilder bulkRequest;
     private final String[] languages;
+    private final String[] extraTags;
 
-    public Importer(Client esClient, String languages) {
+    public Importer(Client esClient, String languages, String extraTags) {
         this.esClient = esClient;
         this.bulkRequest = esClient.prepareBulk();
         this.languages = languages.split(",");
+        this.extraTags = extraTags.split(",");
     }
 
     @Override
     public void add(PhotonDoc doc) {
         try {
             this.bulkRequest.add(this.esClient.prepareIndex(PhotonIndex.NAME, PhotonIndex.TYPE).
-                    setSource(Utils.convert(doc, languages)).setId(doc.getUid()));
+                    setSource(Utils.convert(doc, languages, extraTags)).setId(doc.getUid()));
         } catch (IOException e) {
             log.error("could not bulk add document " + doc.getUid(), e);
             return;

diff --git a/src/main/java/de/komoot/photon/elasticsearch/Server.java b/src/main/java/de/komoot/photon/elasticsearch/Server.java
@@ -49,6 +49,7 @@ public class Server {
     private File esDirectory;
 
     private final String[] languages;
+    private String[] extraTags = new String[0];
 
     private String transportAddresses;
 
@@ -62,6 +63,9 @@ public MyNode(Settings preparedSettings, Collection<Class<? extends Plugin>> cla
 
     public Server(CommandLineArgs args) {
         this(args.getCluster(), args.getDataDirectory(), args.getLanguages(), args.getTransportAddresses());
+        if (args.getExtraTags().length() > 0) {
+            this.extraTags = args.getExtraTags().split(",");
+        }
     }
 
     public Server(String clusterName, String mainDirectory, String languages, String transportAddresses) {

diff --git a/src/main/java/de/komoot/photon/elasticsearch/Updater.java b/src/main/java/de/komoot/photon/elasticsearch/Updater.java
@@ -19,11 +19,13 @@ public class Updater implements de.komoot.photon.Updater {
     private final Client esClient;
     private BulkRequestBuilder bulkRequest;
     private final String[] languages;
+    private final String[] extraTags;
 
-    public Updater(Client esClient, String languages) {
+    public Updater(Client esClient, String languages, String extraTags) {
         this.esClient = esClient;
         this.bulkRequest = esClient.prepareBulk();
         this.languages = languages.split(",");
+        this.extraTags = extraTags.split(",");
     }
 
     public void finish() {
@@ -33,7 +35,7 @@ public void finish() {
     @Override
     public void create(PhotonDoc doc) {
         try {
-            this.bulkRequest.add(this.esClient.prepareIndex(PhotonIndex.NAME, PhotonIndex.TYPE).setSource(Utils.convert(doc, this.languages)).setId(String.valueOf(doc.getPlaceId())));
+            bulkRequest.add(esClient.prepareIndex(PhotonIndex.NAME, PhotonIndex.TYPE).setSource(Utils.convert(doc, languages, extraTags)).setId(String.valueOf(doc.getPlaceId())));
         } catch (IOException e) {
             log.error(String.format("creation of new doc [%s] failed", doc), e);
         }

diff --git a/src/main/java/de/komoot/photon/utils/ConvertToJson.java b/src/main/java/de/komoot/photon/utils/ConvertToJson.java
@@ -64,6 +64,11 @@ public List<JSONObject> convert(SearchResponse searchResponse) {
                 properties.put("extent", new JSONArray(Lists.newArrayList(nw.get(0), nw.get(1), se.get(0), se.get(1))));
             }
 
+            final Map<String, String> extraTags = (Map<String, String>) source.get("extra");
+            if (extraTags != null) {
+                properties.put("extra", extraTags);
+            }
+
             feature.put(Constants.PROPERTIES, properties);
 
             list.add(feature);

diff --git a/src/test/java/de/komoot/photon/ApiIntegrationTest.java b/src/test/java/de/komoot/photon/ApiIntegrationTest.java
@@ -26,7 +26,7 @@ public class ApiIntegrationTest extends ESBaseTester {
     @Before
     public void setUp() throws Exception {
         setUpES();
-        Importer instance = new Importer(getClient(), "en");
+        Importer instance = new Importer(getClient(), "en", "");
         instance.add(createDoc(13.38886, 52.51704, 1000, 1000, "place", "city"));
         instance.add(createDoc(13.39026, 52.54714, 1001, 1001, "place", "town"));
         instance.finish();

diff --git a/src/test/java/de/komoot/photon/elasticsearch/ImporterTest.java b/src/test/java/de/komoot/photon/elasticsearch/ImporterTest.java
@@ -7,6 +7,9 @@
 import org.junit.Test;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
 
 import static org.junit.Assert.*;
 
@@ -19,15 +22,57 @@ public void tearDown() {
     }
 
     @Test
-    public void addSimpleDoc() throws IOException {
+    public void testAddSimpleDoc() throws IOException {
         setUpES();
-        Importer instance = new Importer(getClient(), "en");
-        instance.add(new PhotonDoc(1234, "N", 1000, "place", "city"));
+        Importer instance = new Importer(getClient(), "en", "");
+        instance.add(new PhotonDoc(1234, "N", 1000, "place", "city")
+                .extraTags(Collections.singletonMap("maxspeed", "100")));
         instance.finish();
         refresh();
 
         GetResponse response = getById(1234);
 
         assertTrue(response.isExists());
+
+        Map<String, Object> source = response.getSource();
+
+        assertEquals("N", source.get("osm_type"));
+        assertEquals(1000, source.get("osm_id"));
+        assertEquals("place", source.get("osm_key"));
+        assertEquals("city", source.get("osm_value"));
+
+        assertNull(source.get("extra"));
+    }
+
+    @Test
+    public void testSelectedExtraTagsCanBeIncluded() throws IOException {
+        setUpES();
+        Importer instance = new Importer(getClient(), "en", "maxspeed,website");
+
+        Map<String, String> extratags = new HashMap<>();
+        extratags.put("website", "foo");
+        extratags.put("maxspeed", "100 mph");
+        extratags.put("source", "survey");
+
+        instance.add(new PhotonDoc(1234, "N", 1000, "place", "city").extraTags(extratags));
+        instance.add(new PhotonDoc(1235, "N", 1001, "place", "city")
+                .extraTags(Collections.singletonMap("wikidata", "100")));
+        instance.finish();
+        refresh();
+
+        GetResponse response = getById(1234);
+        assertTrue(response.isExists());
+
+        Map<String, String> extra = (Map<String, String>) response.getSource().get("extra");
+        assertNotNull(extra);
+
+        assertEquals(2, extra.size());
+        assertEquals("100 mph", extra.get("maxspeed"));
+        assertEquals("foo", extra.get("website"));
+
+        response = getById(1235);
+        assertTrue(response.isExists());
+
+        assertNull(response.getSource().get("extra"));
     }
 }
diff --git a/src/test/java/de/komoot/photon/elasticsearch/UpdaterTest.java b/src/test/java/de/komoot/photon/elasticsearch/UpdaterTest.java
@@ -7,6 +7,7 @@
 import org.junit.Test;
 
 import java.io.IOException;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -28,13 +29,13 @@ public void addNameToDoc() throws IOException {
         PhotonDoc doc = new PhotonDoc(1234, "N", 1000, "place", "city").names(names);
 
         setUpES();
-        Importer instance = new Importer(getClient(), "en");
+        Importer instance = new Importer(getClient(), "en", "");
         instance.add(doc);
         instance.finish();
         refresh();
 
         names.put("name:en", "Enfoo");
-        Updater updater = new Updater(getClient(), "en,de");
+        Updater updater = new Updater(getClient(), "en,de", "");
         updater.create(doc);
         updater.finish();
         refresh();
@@ -55,13 +56,13 @@ public void removeNameFromDoc() throws IOException {
         PhotonDoc doc = new PhotonDoc(1234, "N", 1000, "place", "city").names(names);
 
         setUpES();
-        Importer instance = new Importer(getClient(), "en");
+        Importer instance = new Importer(getClient(), "en", "");
         instance.add(doc);
         instance.finish();
         refresh();
 
         names.remove("name");
-        Updater updater = new Updater(getClient(), "en,de");
+        Updater updater = new Updater(getClient(), "en,de", "");
         updater.create(doc);
         updater.finish();
         refresh();
@@ -73,4 +74,36 @@ public void removeNameFromDoc() throws IOException {
         assertFalse(out_names.containsKey("default"));
         assertEquals("Enfoo", out_names.get("en"));
     }
+
+    @Test
+    public void addExtraTagsToDoc() throws IOException {
+        Map<String, String> names = new HashMap<>();
+        names.put("name", "Foo");
+        PhotonDoc doc = new PhotonDoc(1234, "N", 1000, "place", "city").names(names);
+
+        setUpES();
+        Importer instance = new Importer(getClient(), "en", "website");
+        instance.add(doc);
+        instance.finish();
+        refresh();
+
+        GetResponse response = getById(1234);
+        assertTrue(response.isExists());
+
+        assertNull(response.getSource().get("extra"));
+
+        doc.extraTags(Collections.singletonMap("website", "http://site.foo"));
+        Updater updater = new Updater(getClient(), "en,de", "website");
+        updater.create(doc);
+        updater.finish();
+        refresh();
+
+        response = getById(1234);
+        assertTrue(response.isExists());
+
+        Map<String, String> extra = (Map<String, String>) response.getSource().get("extra");
+
+        assertNotNull(extra);
+        assertEquals(Collections.singletonMap("website", "http://site.foo"), extra);
+    }
 }
diff --git a/src/test/java/de/komoot/photon/query/PhotonQueryBuilderSearchTest.java b/src/test/java/de/komoot/photon/query/PhotonQueryBuilderSearchTest.java
@@ -19,10 +19,6 @@
 import org.junit.Before;
 import org.junit.Test;
 
-import java.io.IOException;
-import java.util.Set;
-
-import static org.junit.Assert.assertEquals;
 
 /**
  * Created by Sachin Dole on 2/20/2015.
@@ -34,7 +30,7 @@ public void setUp() throws Exception {
         setUpES();
         ImmutableList<String> tags = ImmutableList.of("tourism", "attraction", "tourism", "hotel", "tourism", "museum", "tourism", "information", "amenity",
                 "parking", "amenity", "restaurant", "amenity", "information", "food", "information", "railway", "station");
-        Importer instance = new Importer(getClient(), "en");
+        Importer instance = new Importer(getClient(), "en", "");
         double lon = 13.38886;
         double lat = 52.51704;
         for (int i = 0; i < tags.size(); i++) {