support custom writer for streaming

ClickHouse · Jun 23, 2022 · 3b48e07 · 3b48e07
1 parent a1a617c
commit 3b48e07
Show file tree

Hide file tree

Showing 2 changed files with 126 additions and 0 deletions.
diff --git a/clickhouse-client/src/main/java/com/clickhouse/client/ClickHouseRequest.java b/clickhouse-client/src/main/java/com/clickhouse/client/ClickHouseRequest.java
@@ -1,8 +1,10 @@
 package com.clickhouse.client;
 
+import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.io.Serializable;
+import java.io.UncheckedIOException;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
@@ -14,8 +16,10 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Map.Entry;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.Optional;
 import java.util.Properties;
@@ -39,6 +43,8 @@ public class ClickHouseRequest<SelfT extends ClickHouseRequest<SelfT>> implement
      * Mutation request.
      */
     public static class Mutation extends ClickHouseRequest<Mutation> {
+        private ClickHouseWriter writer;
+
         protected Mutation(ClickHouseRequest<?> request, boolean sealed) {
             super(request.getClient(), request.server, request.serverRef, request.options, sealed);
             this.settings.putAll(request.settings);
@@ -105,6 +111,20 @@ public Mutation format(ClickHouseFormat format) {
             return super.format(format);
         }
 
+        /**
+         * Sets custom writer for streaming. This will create a piped stream between the
+         * writer and ClickHouse server.
+         *
+         * @param writer writer
+         * @return mutation request
+         */
+        public Mutation data(ClickHouseWriter writer) {
+            checkSealed();
+
+            this.writer = changeProperty(PROP_WRITER, this.writer, writer);
+            return this;
+        }
+
         /**
          * Loads data from given file which may or may not be compressed.
          *
@@ -197,6 +217,70 @@ public Mutation data(ClickHouseDeferredValue<ClickHouseInputStream> input) {
             return this;
         }
 
+        @Override
+        public CompletableFuture<ClickHouseResponse> execute() {
+            if (writer != null) {
+                ClickHouseConfig c = getConfig();
+                ClickHousePipedOutputStream stream = ClickHouseDataStreamFactory.getInstance()
+                        .createPipedOutputStream(c, null);
+                data(stream.getInputStream());
+                CompletableFuture<ClickHouseResponse> future = null;
+                if (c.isAsync()) {
+                    future = getClient().execute(isSealed() ? this : seal());
+                }
+                try (ClickHouseOutputStream out = stream) {
+                    writer.write(out);
+                } catch (IOException e) {
+                    throw new CompletionException(e);
+                }
+                if (future != null) {
+                    return future;
+                }
+            }
+
+            return getClient().execute(isSealed() ? this : seal());
+        }
+
+        @Override
+        public ClickHouseResponse executeAndWait() throws ClickHouseException {
+            if (writer != null) {
+                ClickHouseConfig c = getConfig();
+                ClickHousePipedOutputStream stream = ClickHouseDataStreamFactory.getInstance()
+                        .createPipedOutputStream(c, null);
+                data(stream.getInputStream());
+                CompletableFuture<ClickHouseResponse> future = null;
+                if (c.isAsync()) {
+                    future = getClient().execute(isSealed() ? this : seal());
+                }
+                try (ClickHouseOutputStream out = stream) {
+                    writer.write(out);
+                } catch (IOException e) {
+                    throw ClickHouseException.of(e, getServer());
+                }
+                if (future != null) {
+                    try {
+                        return future.get();
+                    } catch (InterruptedException e) {
+                        Thread.currentThread().interrupt();
+                        throw ClickHouseException.forCancellation(e, getServer());
+                    } catch (CancellationException e) {
+                        throw ClickHouseException.forCancellation(e, getServer());
+                    } catch (ExecutionException | UncheckedIOException e) {
+                        Throwable cause = e.getCause();
+                        if (cause == null) {
+                            cause = e;
+                        }
+                        throw cause instanceof ClickHouseException ? (ClickHouseException) cause
+                                : ClickHouseException.of(cause, getServer());
+                    } catch (RuntimeException e) { // unexpected
+                        throw ClickHouseException.of(e, getServer());
+                    }
+                }
+            }
+
+            return getClient().executeAndWait(isSealed() ? this : seal());
+        }
+
         /**
          * Sends mutation requets for execution. Same as
          * {@code client.execute(request.seal())}.
@@ -256,6 +340,7 @@ public Mutation seal() {
     static final String PROP_PREPARED_QUERY = "preparedQuery";
     static final String PROP_QUERY = "query";
     static final String PROP_QUERY_ID = "queryId";
+    static final String PROP_WRITER = "writer";
 
     private final boolean sealed;
 

diff --git a/clickhouse-client/src/test/java/com/clickhouse/client/ClientIntegrationTest.java b/clickhouse-client/src/test/java/com/clickhouse/client/ClientIntegrationTest.java
@@ -21,6 +21,7 @@
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
 
 import com.clickhouse.client.ClickHouseClientBuilder.Agent;
@@ -909,6 +910,46 @@ public void testCustomRead() throws Exception {
         Assert.assertEquals(count, 1000L);
     }
 
+    @Test(groups = { "integration" })
+    public void testCustomWriter() throws Exception {
+        ClickHouseNode server = getServer();
+        ClickHouseClient.send(server, "drop table if exists test_custom_writer",
+                "create table test_custom_writer(a Int8) engine=Memory")
+                .get();
+
+        try (ClickHouseClient client = getClient()) {
+            AtomicInteger i = new AtomicInteger(1);
+            ClickHouseRequest.Mutation req = client.connect(server).write().format(ClickHouseFormat.RowBinary)
+                    .table("test_custom_writer").data(o -> {
+                        o.write(i.getAndIncrement());
+                    });
+            for (boolean b : new boolean[] { true, false }) {
+                req.option(ClickHouseClientOption.ASYNC, b);
+
+                try (ClickHouseResponse resp = req.send().get()) {
+                    Assert.assertNotNull(resp);
+                }
+
+                try (ClickHouseResponse resp = req.sendAndWait()) {
+                    Assert.assertNotNull(resp);
+                }
+
+                try (ClickHouseResponse resp = req.execute().get()) {
+                    Assert.assertNotNull(resp);
+                }
+
+                try (ClickHouseResponse resp = req.executeAndWait()) {
+                    Assert.assertNotNull(resp);
+                }
+            }
+
+            try (ClickHouseResponse resp = client.connect(server).query("select count(1) from test_custom_writer")
+                    .executeAndWait()) {
+                Assert.assertEquals(resp.firstRecord().getValue(0).asInteger(), i.get() - 1);
+            }
+        }
+    }
+
     @Test(groups = { "integration" })
     public void testDumpAndLoadFile() throws Exception {
         // super.testLoadRawData();