storage: the variant storage engine uses a Solr client with extended …

…timeout for write operations, #TASK-6981
opencb · Sep 27, 2024 · 2757197 · 2757197
1 parent 8f7cec5
commit 2757197
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 4 deletions.
diff --git a/...rage-core/src/main/java/org/opencb/opencga/storage/core/variant/VariantStorageEngine.java b/...rage-core/src/main/java/org/opencb/opencga/storage/core/variant/VariantStorageEngine.java
@@ -34,7 +34,6 @@
 import org.opencb.opencga.core.models.operations.variant.VariantAggregateFamilyParams;
 import org.opencb.opencga.core.models.operations.variant.VariantAggregateParams;
 import org.opencb.opencga.core.models.variant.VariantSetupParams;
-import org.opencb.opencga.storage.core.variant.query.VariantQueryResult;
 import org.opencb.opencga.storage.core.StorageEngine;
 import org.opencb.opencga.storage.core.StoragePipelineResult;
 import org.opencb.opencga.storage.core.exceptions.StorageEngineException;
@@ -60,6 +59,7 @@
 import org.opencb.opencga.storage.core.variant.io.VariantWriterFactory.VariantOutputFormat;
 import org.opencb.opencga.storage.core.variant.query.ParsedVariantQuery;
 import org.opencb.opencga.storage.core.variant.query.VariantQueryParser;
+import org.opencb.opencga.storage.core.variant.query.VariantQueryResult;
 import org.opencb.opencga.storage.core.variant.query.VariantQueryUtils;
 import org.opencb.opencga.storage.core.variant.query.executors.*;
 import org.opencb.opencga.storage.core.variant.score.VariantScoreFormatDescriptor;
@@ -794,7 +794,7 @@ protected void searchIndexLoadedFiles(List<URI> inputFiles, ObjectMap options) t
 
     protected SolrInputDocumentDataWriter newVariantSearchDataWriter(String collection) throws StorageEngineException {
         return new SolrInputDocumentDataWriter(collection,
-                getVariantSearchManager().getSolrClient(),
+                getVariantSearchManager().getSolrManager().newSolrClient(configuration.getSearch().getWriteTimeout()), true,
                 getVariantSearchManager().getInsertBatchSize());
     }
 

diff --git a/...java/org/opencb/opencga/storage/core/variant/search/solr/SolrInputDocumentDataWriter.java b/...java/org/opencb/opencga/storage/core/variant/search/solr/SolrInputDocumentDataWriter.java
@@ -9,12 +9,14 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.io.IOException;
 import java.util.List;
 
 public class SolrInputDocumentDataWriter implements DataWriter<SolrInputDocument> {
 
     private final String collection;
     private final SolrClient solrClient;
+    private boolean closeSolrClient;
     private final int insertBatchSize;
     private int serverBufferSize = 0;
     private int insertedDocuments = 0;
@@ -23,8 +25,13 @@ public class SolrInputDocumentDataWriter implements DataWriter<SolrInputDocument
     private final Logger logger = LoggerFactory.getLogger(SolrInputDocumentDataWriter.class);
 
     public SolrInputDocumentDataWriter(String collection, SolrClient solrClient, int insertBatchSize) {
+        this(collection, solrClient, false, insertBatchSize);
+    }
+
+    public SolrInputDocumentDataWriter(String collection, SolrClient solrClient, boolean closeSolrClient, int insertBatchSize) {
         this.collection = collection;
         this.solrClient = solrClient;
+        this.closeSolrClient = closeSolrClient;
         this.insertBatchSize = insertBatchSize;
     }
 
@@ -54,6 +61,18 @@ public boolean post() {
         return true;
     }
 
+    @Override
+    public boolean close() {
+        if (closeSolrClient) {
+            try {
+                solrClient.close();
+            } catch (IOException e) {
+                Throwables.propagate(e);
+            }
+        }
+        return true;
+    }
+
     protected void add(List<SolrInputDocument> batch) throws Exception {
         UpdateResponse response = solrClient.add(collection, batch);
         addTimeMs += response.getElapsedTime();

diff --git a/...c/test/java/org/opencb/opencga/storage/core/variant/solr/VariantSolrExternalResource.java b/...c/test/java/org/opencb/opencga/storage/core/variant/solr/VariantSolrExternalResource.java
@@ -116,8 +116,7 @@ public VariantSearchManager configure(VariantStorageEngine variantStorageEngine)
         variantStorageEngine.getConfiguration().getSearch().setMode("core");
         variantStorageEngine.getConfiguration().getSearch().setActive(true);
         VariantSearchManager variantSearchManager = variantStorageEngine.getVariantSearchManager();
-        variantSearchManager.setSolrManager(new SolrManager(solrClient, "localhost", "core",
-                variantStorageEngine.getConfiguration().getSearch().getTimeout()));
+        variantSearchManager.setSolrManager(new SolrManager(solrClient, "localhost", "core"));
         variantSearchManager.setSolrClient(solrClient);
         return variantSearchManager;
     }