Skip to content

Commit

Permalink
datastore: add a method to SolrManager to create a Solr client, #TASK…
Browse files Browse the repository at this point in the history
…-6981
  • Loading branch information
jtarraga committed Sep 27, 2024
1 parent 9c1bb4a commit b49926c
Showing 1 changed file with 58 additions and 46 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -23,20 +23,24 @@
import org.apache.solr.client.solrj.impl.CloudSolrClient;
import org.apache.solr.client.solrj.impl.HttpSolrClient;
import org.apache.solr.client.solrj.impl.LBHttpSolrClient;
import org.apache.solr.client.solrj.request.*;
import org.apache.solr.client.solrj.request.CollectionAdminRequest;
import org.apache.solr.client.solrj.request.CoreAdminRequest;
import org.apache.solr.client.solrj.request.CoreStatus;
import org.apache.solr.client.solrj.request.SolrPing;
import org.apache.solr.client.solrj.response.SolrPingResponse;
import org.apache.solr.common.SolrException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.stream.Collectors;

public class SolrManager {

private String host;
private List<String> hosts;
private String mode;
private SolrClient solrClient;

Expand Down Expand Up @@ -65,51 +69,23 @@ public SolrManager(String host, String mode, int timeout) {
* @param timeout Read timeout
*/
public SolrManager(List<String> hosts, String mode, int timeout) {
hosts = hosts.stream().flatMap(s -> Arrays.stream(s.split(","))).collect(Collectors.toList());

this.host = String.join(",", hosts);
this.hosts = hosts.stream().flatMap(s -> Arrays.stream(s.split(","))).collect(Collectors.toList());
this.mode = mode;

if (hosts.get(0).startsWith("http")) {
if (hosts.size() == 1) {
// Single HTTP endpoint.
this.solrClient = new HttpSolrClient.Builder(host).build();
((HttpSolrClient) this.solrClient).setRequestWriter(new BinaryRequestWriter());
((HttpSolrClient) this.solrClient).setSoTimeout(timeout);
} else {
// Use a LoadBalancer if there are multiple http hosts
this.solrClient = new LBHttpSolrClient.Builder().withBaseSolrUrls(hosts.toArray(new String[0])).build();

((LBHttpSolrClient) this.solrClient).setRequestWriter(new BinaryRequestWriter());
((LBHttpSolrClient) this.solrClient).setSoTimeout(timeout);
}
} else {
// If the provided hosts are not http, assume zookeeper hosts like HOST:PORT
// This client will use Zookeeper to discover Solr endpoints for SolrCloud collections, and then use the
// LBHttpSolrClient to issue requests.
if (isCloud()) {
this.solrClient = new CloudSolrClient.Builder().withZkHost(hosts).build();

((CloudSolrClient) this.solrClient).setRequestWriter(new BinaryRequestWriter());
((CloudSolrClient) this.solrClient).setSoTimeout(timeout);
} else {
throw new IllegalArgumentException("Can not initialize SolrManager from Zookeeper host not in Cloud mode");
}
}
this.solrClient = newSolrClient(timeout);
}

public SolrManager(SolrClient solrClient, String host, String mode) {
this.solrClient = solrClient;
this.host = host;
this.hosts = Collections.singletonList(host);
this.mode = mode;
}

@Deprecated
public SolrManager(SolrClient solrClient, String host, String mode, int timeout) {
this.solrClient = solrClient;
this.host = host;
this.mode = mode;
}
// @Deprecated
// public SolrManager(SolrClient solrClient, String host, String mode, int timeout) {
// this.solrClient = solrClient;
// this.hosts = Collectionshost;
// this.mode = mode;
// }

public SolrCollection getCollection(String collection) throws SolrException {
checkIsAlive();
Expand Down Expand Up @@ -188,7 +164,7 @@ public void create(String dbName, String configSet) throws SolrException {
*/
public void createCore(String coreName, String configSet) throws SolrException {
try {
logger.debug("Creating core: host={}, core={}, configSet={}", host, coreName, configSet);
logger.debug("Creating core: host={}, core={}, configSet={}", StringUtils.join(",", hosts), coreName, configSet);
CoreAdminRequest.Create request = new CoreAdminRequest.Create();
request.setCoreName(coreName);
request.setConfigSet(configSet);
Expand All @@ -210,7 +186,7 @@ public void createCore(String coreName, String configSet) throws SolrException {
*/
public void createCollection(String collectionName, String configSet) throws SolrException {
logger.debug("Creating collection: host={}, collection={}, config={}, numShards={}, numReplicas={}",
host, collectionName, configSet, 1, 1);
StringUtils.join(",", hosts), collectionName, configSet, 1, 1);
try {
CollectionAdminRequest request = CollectionAdminRequest.createCollection(collectionName, configSet, 1, 1);
request.process(solrClient);
Expand Down Expand Up @@ -353,6 +329,42 @@ public void close() throws IOException {
}
}

public SolrClient newSolrClient(int timeout) {
SolrClient solrClient;
if (hosts.get(0).startsWith("http")) {
if (hosts.size() == 1) {
// Single HTTP endpoint.
solrClient = new HttpSolrClient.Builder(hosts.get(0)).build();
((HttpSolrClient) solrClient).setRequestWriter(new BinaryRequestWriter());
((HttpSolrClient) solrClient).setSoTimeout(timeout);
} else {
// Use a LoadBalancer if there are multiple http hosts
solrClient = new LBHttpSolrClient.Builder().withBaseSolrUrls(hosts.toArray(new String[0])).build();

((LBHttpSolrClient) solrClient).setRequestWriter(new BinaryRequestWriter());
((LBHttpSolrClient) solrClient).setSoTimeout(timeout);
}
} else {
// If the provided hosts are not http, assume zookeeper hosts like HOST:PORT
// This client will use Zookeeper to discover Solr endpoints for SolrCloud collections, and then use the
// LBHttpSolrClient to issue requests.
if (isCloud()) {
solrClient = new CloudSolrClient.Builder().withZkHost(hosts).build();

((CloudSolrClient) solrClient).setRequestWriter(new BinaryRequestWriter());
((CloudSolrClient) solrClient).setSoTimeout(timeout);
} else {
throw new IllegalArgumentException("Can not initialize SolrManager from Zookeeper host not in Cloud mode");
}
}

// Sanity check
if (solrClient == null) {
throw new IllegalArgumentException("Can not build a Solr client, please, check Solr configuration");
}
return solrClient;
}

private boolean isCloud() {
if (StringUtils.isEmpty(mode)) {
logger.warn("Solr 'mode' is empty, setting default 'cloud'");
Expand All @@ -376,19 +388,19 @@ private boolean isCloud() {
@Override
public String toString() {
final StringBuilder sb = new StringBuilder("SolrManager{");
sb.append("host='").append(host).append('\'');
sb.append("hosts='").append(hosts).append('\'');
sb.append(", mode='").append(mode).append('\'');
sb.append(", solrClient=").append(solrClient);
sb.append('}');
return sb.toString();
}

public String getHost() {
return host;
public List<String> getHosts() {
return hosts;
}

public SolrManager setHost(String host) {
this.host = host;
public SolrManager setHosts(List<String> hosts) {
this.hosts = hosts;
return this;
}

Expand Down

0 comments on commit b49926c

Please sign in to comment.