This is an automated email from the ASF dual-hosted git repository.

oleewere pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/ambari.git

commit c3ccc4eb89370ef0f1983dd07dc1601d18d53d20
Author: Oliver Szabo <[email protected]>
AuthorDate: Wed Jun 20 19:42:21 2018 +0200

    AMBARI-23945. Infra Solr Client - add number of docs to collection dump
---
 .../org/apache/ambari/infra/solr/AmbariSolrCloudCLI.java  |  8 ++++++++
 .../apache/ambari/infra/solr/AmbariSolrCloudClient.java   |  6 ++++++
 .../ambari/infra/solr/AmbariSolrCloudClientBuilder.java   |  6 ++++++
 .../infra/solr/commands/DumpCollectionsCommand.java       | 15 ++++++++++++++-
 .../ambari/infra/solr/domain/json/SolrCollection.java     |  9 +++++++++
 5 files changed, 43 insertions(+), 1 deletion(-)

diff --git 
a/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudCLI.java
 
b/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudCLI.java
index 44db3af..b0c7781 100644
--- 
a/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudCLI.java
+++ 
b/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudCLI.java
@@ -374,6 +374,11 @@ public class AmbariSolrCloudCLI {
       .numberOfArgs(1)
       .build();
 
+    final Option includeDocNumberOption = Option.builder("idn")
+      .longOpt("include-doc-number")
+      .desc("Include the number of docs as well in collection dump")
+      .build();
+
     options.addOption(helpOption);
     options.addOption(retryOption);
     options.addOption(removeAdminHandlerOption);
@@ -421,6 +426,7 @@ public class AmbariSolrCloudCLI {
     options.addOption(setupKerberosPluginOption);
     options.addOption(securityJsonLocationOption);
     options.addOption(outputOption);
+    options.addOption(includeDocNumberOption);
 
     AmbariSolrCloudClient solrCloudClient = null;
 
@@ -520,6 +526,7 @@ public class AmbariSolrCloudCLI {
       String copyDest = cli.hasOption("cpd") ? cli.getOptionValue("cpd") : 
null;
       String transferMode = cli.hasOption("tm") ? cli.getOptionValue("tm") : 
"NONE";
       String output = cli.hasOption("o") ? cli.getOptionValue("o") : null;
+      boolean includeDocNumber = cli.hasOption("idn");
 
       AmbariSolrCloudClientBuilder clientBuilder = new 
AmbariSolrCloudClientBuilder()
         .withZkConnectString(zkConnectString)
@@ -547,6 +554,7 @@ public class AmbariSolrCloudCLI {
         .withCopySrc(copySrc)
         .withCopyDest(copyDest)
         .withOutput(output)
+        .withIncludeDocNumber(includeDocNumber)
         .withSecurityJsonLocation(securityJsonLocation)
         .withZnode(znode)
         .withSecure(isSecure)
diff --git 
a/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudClient.java
 
b/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudClient.java
index 5aecae0..7571c99 100644
--- 
a/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudClient.java
+++ 
b/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudClient.java
@@ -79,6 +79,7 @@ public class AmbariSolrCloudClient {
   private final String copySrc;
   private final String copyDest;
   private final String output;
+  private final boolean includeDocNumber;
 
   public AmbariSolrCloudClient(AmbariSolrCloudClientBuilder builder) {
     this.zkConnectString = builder.zkConnectString;
@@ -106,6 +107,7 @@ public class AmbariSolrCloudClient {
     this.copySrc = builder.copySrc;
     this.copyDest = builder.copyDest;
     this.output = builder.output;
+    this.includeDocNumber = builder.includeDocNumber;
   }
 
   /**
@@ -396,4 +398,8 @@ public class AmbariSolrCloudClient {
   public String getOutput() {
     return output;
   }
+
+  public boolean isIncludeDocNumber() {
+    return includeDocNumber;
+  }
 }
diff --git 
a/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudClientBuilder.java
 
b/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudClientBuilder.java
index 87ad1be..db4396b 100644
--- 
a/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudClientBuilder.java
+++ 
b/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/AmbariSolrCloudClientBuilder.java
@@ -58,6 +58,7 @@ public class AmbariSolrCloudClientBuilder {
   String copySrc;
   String copyDest;
   String output;
+  public boolean includeDocNumber;
 
   public AmbariSolrCloudClient build() {
     return new AmbariSolrCloudClient(this);
@@ -221,6 +222,11 @@ public class AmbariSolrCloudClientBuilder {
     return this;
   }
 
+  public AmbariSolrCloudClientBuilder withIncludeDocNumber(boolean 
includeDocNumber) {
+    this.includeDocNumber = includeDocNumber;
+    return this;
+  }
+
   public AmbariSolrCloudClientBuilder withSecurityJsonLocation(String 
securityJson) {
     this.securityJsonLocation = securityJson;
     return this;
diff --git 
a/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/commands/DumpCollectionsCommand.java
 
b/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/commands/DumpCollectionsCommand.java
index 2b58bfc..708ecac 100644
--- 
a/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/commands/DumpCollectionsCommand.java
+++ 
b/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/commands/DumpCollectionsCommand.java
@@ -24,6 +24,7 @@ import org.apache.ambari.infra.solr.AmbariSolrCloudClient;
 import org.apache.ambari.infra.solr.domain.json.SolrCollection;
 import org.apache.ambari.infra.solr.domain.json.SolrCoreData;
 import org.apache.ambari.infra.solr.domain.json.SolrShard;
+import org.apache.solr.client.solrj.SolrQuery;
 import org.apache.solr.client.solrj.impl.CloudSolrClient;
 import org.apache.solr.common.cloud.DocCollection;
 import org.apache.solr.common.cloud.Replica;
@@ -58,7 +59,12 @@ public class DumpCollectionsCommand extends 
AbstractZookeeperRetryCommand<String
     if (!this.collections.isEmpty()) {
       for (String collection : this.collections) {
         SolrCollection solrCollection = new SolrCollection();
-        Collection<Slice> slices = getSlices(client.getSolrCloudClient(), 
collection);
+        CloudSolrClient solrClient = client.getSolrCloudClient();
+        if (client.isIncludeDocNumber()) {
+          long numberOfDocs = getNumberOfDocs(solrClient, collection);
+          solrCollection.setNumberOfDocs(numberOfDocs);
+        }
+        Collection<Slice> slices = getSlices(solrClient, collection);
         Integer numShards = slices.size();
         Map<String, SolrShard> solrShardMap = new HashMap<>();
         Map<String, List<String>> leaderHostCoreMap = new HashMap<>();
@@ -141,4 +147,11 @@ public class DumpCollectionsCommand extends 
AbstractZookeeperRetryCommand<String
     DocCollection docCollection = 
reader.getClusterState().getCollection(collection);
     return docCollection.getSlices();
   }
+
+  private long getNumberOfDocs(CloudSolrClient solrClient, String collection) 
throws Exception {
+    solrClient.setDefaultCollection(collection);
+    SolrQuery q = new SolrQuery("*:*");
+    q.setRows(0);
+    return solrClient.query(q).getResults().getNumFound();
+  }
 }
diff --git 
a/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/domain/json/SolrCollection.java
 
b/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/domain/json/SolrCollection.java
index 86751e5..21820ec 100644
--- 
a/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/domain/json/SolrCollection.java
+++ 
b/ambari-infra/ambari-infra-solr-client/src/main/java/org/apache/ambari/infra/solr/domain/json/SolrCollection.java
@@ -24,12 +24,21 @@ import java.util.Map;
 
 public class SolrCollection {
   private String name;
+  private long numberOfDocs = -1;
   private Map<String, SolrShard> shards = new HashMap<>();
   private Map<String, List<String>> leaderHostCoreMap = new HashMap<>();
   private Map<String, SolrCoreData> leaderSolrCoreDataMap = new HashMap<>();
   private Map<String, List<String>> leaderShardsMap = new HashMap<>();
   private Map<String, String> leaderCoreHostMap = new HashMap<>();
 
+  public long getNumberOfDocs() {
+    return numberOfDocs;
+  }
+
+  public void setNumberOfDocs(long numberOfDocs) {
+    this.numberOfDocs = numberOfDocs;
+  }
+
   public Map<String, SolrShard> getShards() {
     return shards;
   }

Reply via email to