change behaviour of SPARQLProtocolWorker for streaming request bodies

dice-group · Sep 21, 2023 · a302807 · a302807
1 parent abe29cc
commit a302807
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 5 deletions.
diff --git a/src/main/java/org/aksw/iguana/cc/query/handler/QueryHandler.java b/src/main/java/org/aksw/iguana/cc/query/handler/QueryHandler.java
@@ -156,7 +156,6 @@ public String value() {
 
     public record QueryStringWrapper(int index, String query) {}
     public record QueryStreamWrapper(int index, InputStream queryInputStream) {}
-    public record QueryStreamSupplierWrapper(int index, Supplier<InputStream> queryStreamSupplier) {}
 
 
     protected final Logger LOGGER = LoggerFactory.getLogger(QueryHandler.class);

diff --git a/src/main/java/org/aksw/iguana/cc/worker/impl/SPARQLProtocolWorker.java b/src/main/java/org/aksw/iguana/cc/worker/impl/SPARQLProtocolWorker.java
@@ -26,6 +26,7 @@
 import java.util.*;
 import java.util.concurrent.*;
 import java.util.function.BiFunction;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 public class SPARQLProtocolWorker extends HttpWorker {
@@ -68,6 +69,21 @@ public HttpRequest buildHttpRequest(InputStream queryStream,
                                             ConnectionConfig connection,
                                             String requestHeader) throws URISyntaxException, IOException {
             HttpRequest.Builder request = HttpRequest.newBuilder().timeout(timeout);
+
+            class CustomStreamSupplier {
+                boolean used = false; // assume, that the stream will only be used again, if the first request failed, because of the client
+                public Supplier<InputStream> getStreamSupplier() {
+                    if (!used) {
+                        used = true;
+                        return () -> queryStream;
+                    }
+                    else
+                        return () -> null;
+                }
+            }
+
+            final var streamSupplier = new CustomStreamSupplier();
+
             if (requestHeader != null)
                 request.header("Accept", requestHeader);
             if (connection.user() != null)
@@ -93,7 +109,7 @@ public HttpRequest buildHttpRequest(InputStream queryStream,
                 case POST_QUERY -> {
                     request.uri(connection.endpoint())
                             .header("Content-Type", "application/sparql-query")
-                            .POST(HttpRequest.BodyPublishers.ofByteArray(queryStream.readAllBytes())); // InputStream BodyPublisher won't work
+                            .POST(HttpRequest.BodyPublishers.ofInputStream(streamSupplier.getStreamSupplier()));
                 }
                 case POST_URL_ENC_UPDATE -> {
                     request.uri(connection.endpoint())
@@ -106,7 +122,7 @@ public HttpRequest buildHttpRequest(InputStream queryStream,
                 case POST_UPDATE -> {
                     request.uri(connection.endpoint())
                             .header("Content-Type", "application/sparql-update")
-                            .POST(HttpRequest.BodyPublishers.ofByteArray(queryStream.readAllBytes()));
+                            .POST(HttpRequest.BodyPublishers.ofInputStream(streamSupplier.getStreamSupplier()));
                 }
             }
             return request.build();
@@ -264,12 +280,27 @@ private ExecutionStats executeQuery(Duration timeout, boolean discardOnFailure)
 
 
     private HttpExecutionResult executeHttpRequest(Duration timeout) {
-        final var queryHandle = config().queries().getNextQueryStreamSupplier();
+        final QueryHandler.QueryStreamWrapper queryHandle;
+        try {
+            queryHandle = config().queries().getNextQueryStream();
+        } catch (IOException e) {
+            return new HttpExecutionResult(
+                config().queries().getCurrentQueryID(),
+                Optional.empty(),
+                Instant.now(),
+                Duration.ZERO,
+                Optional.empty(),
+                OptionalLong.empty(),
+                OptionalLong.empty(),
+                Optional.of(e)
+            );
+        }
+
         final HttpRequest request;
 
         try {
             request = requestFactory.buildHttpRequest(
-                    queryHandle.queryStreamSupplier(),
+                    queryHandle.queryInputStream(),
                     timeout,
                     config().connection(),
                     config().acceptHeader()