ATLAS-2920: Update JanusGraph Solr clients to use all zookeeper entries

graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusGraphDatabase.java

@@ -18,6 +18,7 @@
 
 package org.apache.atlas.repository.graphdb.janus;
 
+import com.google.common.collect.ImmutableMap;
 import org.apache.atlas.ApplicationProperties;
 import org.apache.atlas.AtlasException;
 import org.apache.atlas.repository.graphdb.AtlasGraph;
@@ -30,19 +31,24 @@ import org.apache.atlas.typesystem.types.DataTypes.TypeCategory;
 import org.apache.atlas.utils.AtlasPerfTracer;
 import org.apache.commons.configuration.Configuration;
 import org.apache.tinkerpop.gremlin.structure.io.graphson.GraphSONMapper;
+import org.janusgraph.core.JanusGraph;
 import org.janusgraph.core.JanusGraphException;
-import org.janusgraph.graphdb.database.serialize.attribute.SerializableSerializer;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
 import org.janusgraph.core.JanusGraphFactory;
-import org.janusgraph.core.JanusGraph;
 import org.janusgraph.core.schema.JanusGraphManagement;
+import org.janusgraph.diskstorage.StandardIndexProvider;
+import org.janusgraph.diskstorage.solr.Solr6Index;
+import org.janusgraph.graphdb.database.serialize.attribute.SerializableSerializer;
 import org.janusgraph.graphdb.tinkerpop.JanusGraphIoRegistry;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
+import java.lang.reflect.Field;
+import java.lang.reflect.Modifier;
 import java.math.BigDecimal;
 import java.math.BigInteger;
 import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.Map;
 
 /**
  * Default implementation for Graph Provider that doles out JanusGraph.
@@ -59,6 +65,7 @@ public class AtlasJanusGraphDatabase implements GraphDatabase<AtlasJanusVertex,
     public static final String GRAPH_PREFIX         = "atlas.graph";
     public static final String INDEX_BACKEND_CONF   = "index.search.backend";
     public static final String SOLR_ZOOKEEPER_URL   = "atlas.graph.index.search.solr.zookeeper-url";
+    public static final String SOLR_ZOOKEEPER_URLS  = "atlas.graph.index.search.solr.zookeeper-urls";
     public static final String INDEX_BACKEND_LUCENE = "lucene";
     public static final String INDEX_BACKEND_ES     = "elasticsearch";
 
@@ -74,7 +81,10 @@ public class AtlasJanusGraphDatabase implements GraphDatabase<AtlasJanusVertex,
         startLocalSolr();
 
         Configuration configProperties = ApplicationProperties.get();
-        Configuration janusConfig      = ApplicationProperties.getSubsetConfiguration(configProperties, GRAPH_PREFIX);
+
+        configProperties.setProperty(SOLR_ZOOKEEPER_URLS, configProperties.getStringArray(SOLR_ZOOKEEPER_URL));
+
+        Configuration janusConfig = ApplicationProperties.getSubsetConfiguration(configProperties, GRAPH_PREFIX);
 
         //add serializers for non-standard property value types that Atlas uses
         janusConfig.addProperty("attributes.custom.attribute1.attribute-class", TypeCategory.class.getName());
@@ -93,6 +103,30 @@ public class AtlasJanusGraphDatabase implements GraphDatabase<AtlasJanusVertex,
         return janusConfig;
     }
 
+    static {
+        addSolr6Index();
+    }
+
+    private static void addSolr6Index() {
+        try {
+            Field field = StandardIndexProvider.class.getDeclaredField("ALL_MANAGER_CLASSES");
+            field.setAccessible(true);
+
+            Field modifiersField = Field.class.getDeclaredField("modifiers");
+            modifiersField.setAccessible(true);
+            modifiersField.setInt(field, field.getModifiers() & ~Modifier.FINAL);
+
+            Map<String, String> customMap = new HashMap<>(StandardIndexProvider.getAllProviderClasses());
+            customMap.put("solr", Solr6Index.class.getName());
+            ImmutableMap<String, String> immap = ImmutableMap.copyOf(customMap);
+            field.set(null, immap);
+
+            LOG.debug("Injected solr6 index - {}", Solr6Index.class.getName());
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     public static JanusGraph getGraphInstance() {
         if (graphInstance == null) {
             synchronized (AtlasJanusGraphDatabase.class) {