Merge pull request influxdata#259 from andyflury/chuncking

majst01 · web-flow · commit 08a075cc8e58 · 2016-12-22T13:01:28.000+01:00
InfluxDB-Java Chunking
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@
  - Support writing by UDP protocal.
  - Support gzip compress for http request body.
  - Support setting thread factory for batch processor.
+ - Support chunking
 
 #### Fixes
 
diff --git a/README.md b/README.md
@@ -96,6 +96,16 @@ public void write(final int udpPort, final Point point);
 ```
 note: make sure write content's total size should not > UDP protocol's limit(64K), or you should use http instead of udp.
 
+
+#### chunking support (version 2.5+ required):
+
+influxdb-java client now supports influxdb chunking. The following example uses a chunkSize of 20 and invokes the specified Consumer (e.g. System.out.println) for each received QueryResult
+```
+Query query = new Query("SELECT idle FROM cpu", dbName);
+influxDB.query(query, 20, queryResult -> System.out.println(queryResult));
+```
+
+
 ### Other Usages:
 For additional usage examples have a look at [InfluxDBTest.java](https://github.com/influxdb/influxdb-java/blob/master/src/test/java/org/influxdb/InfluxDBTest.java "InfluxDBTest.java")
 
diff --git a/src/main/java/org/influxdb/InfluxDB.java b/src/main/java/org/influxdb/InfluxDB.java
@@ -3,6 +3,7 @@
 import java.util.List;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
 
 import org.influxdb.dto.BatchPoints;
 import org.influxdb.dto.Point;
@@ -217,6 +218,18 @@ public void write(final String database, final String retentionPolicy,
    */
   public QueryResult query(final Query query);
 
+  /**
+   * Execute a streaming query against a database.
+   *
+   * @param query
+   *            the query to execute.
+   * @param chunkSize
+   *            the number of QueryResults to process in one chunk.
+   * @param consumer
+   *            the consumer to invoke for each received QueryResult
+   */
+    public void query(Query query, int chunkSize, Consumer<QueryResult> consumer);
+
   /**
    * Execute a query against a database.
    *
diff --git a/src/main/java/org/influxdb/impl/InfluxDBImpl.java b/src/main/java/org/influxdb/impl/InfluxDBImpl.java
@@ -7,6 +7,9 @@
 import com.google.common.base.Stopwatch;
 import com.google.common.base.Strings;
 import com.google.common.collect.Lists;
+import com.squareup.moshi.JsonAdapter;
+import com.squareup.moshi.Moshi;
+
 import org.influxdb.InfluxDB;
 import org.influxdb.dto.BatchPoints;
 import org.influxdb.dto.Point;
@@ -24,11 +27,14 @@
 import okhttp3.ResponseBody;
 import okhttp3.logging.HttpLoggingInterceptor;
 import okhttp3.logging.HttpLoggingInterceptor.Level;
+import okio.BufferedSource;
 import retrofit2.Call;
+import retrofit2.Callback;
 import retrofit2.Response;
 import retrofit2.Retrofit;
 import retrofit2.converter.moshi.MoshiConverterFactory;
 
+import java.io.EOFException;
 import java.io.IOException;
 import java.net.DatagramPacket;
 import java.net.DatagramSocket;
@@ -41,6 +47,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
 
 /**
  * Implementation of a InluxDB API.
@@ -66,10 +73,12 @@ public class InfluxDBImpl implements InfluxDB {
   private final HttpLoggingInterceptor loggingInterceptor;
   private final GzipRequestInterceptor gzipRequestInterceptor;
   private LogLevel logLevel = LogLevel.NONE;
+  private JsonAdapter<QueryResult> adapter;
 
   public InfluxDBImpl(final String url, final String username, final String password,
       final OkHttpClient.Builder client) {
     super();
+    Moshi moshi = new Moshi.Builder().build();
     this.hostAddress = parseHostAddress(url);
     this.username = username;
     this.password = password;
@@ -82,6 +91,7 @@ public InfluxDBImpl(final String url, final String username, final String passwo
         .addConverterFactory(MoshiConverterFactory.create())
         .build();
     this.influxDBService = this.retrofit.create(InfluxDBService.class);
+    this.adapter = moshi.adapter(QueryResult.class);
   }
 
   private InetAddress parseHostAddress(final String url) {
@@ -327,6 +337,49 @@ public QueryResult query(final Query query) {
     return execute(call);
   }
 
+  /**
+   * {@inheritDoc}
+   */
+  @Override
+    public void query(final Query query, final int chunkSize, final Consumer<QueryResult> consumer) {
+
+        if (version().startsWith("0.") || version().startsWith("1.0")) {
+            throw new RuntimeException("chunking not supported");
+        }
+
+        Call<ResponseBody> call = this.influxDBService.query(this.username, this.password,
+                query.getDatabase(), query.getCommandWithUrlEncoded(), chunkSize);
+
+        call.enqueue(new Callback<ResponseBody>() {
+            @Override
+            public void onResponse(final Call<ResponseBody> call, final Response<ResponseBody> response) {
+                try {
+                    if (response.isSuccessful()) {
+                        BufferedSource source = response.body().source();
+                        while (true) {
+                            QueryResult result = InfluxDBImpl.this.adapter.fromJson(source);
+                            if (result != null) {
+                                consumer.accept(result);
+                            }
+                        }
+                    }
+                    try (ResponseBody errorBody = response.errorBody()) {
+                        throw new RuntimeException(errorBody.string());
+                    }
+                } catch (EOFException e) {
+                    // do nothing
+                } catch (IOException e) {
+                    throw new RuntimeException(e);
+                }
+            }
+
+            @Override
+            public void onFailure(final Call<ResponseBody> call, final Throwable t) {
+                throw new RuntimeException(t);
+            }
+        });
+  }
+
   /**
    * {@inheritDoc}
    */
diff --git a/src/main/java/org/influxdb/impl/InfluxDBService.java b/src/main/java/org/influxdb/impl/InfluxDBService.java
@@ -9,6 +9,7 @@
 import retrofit2.http.GET;
 import retrofit2.http.POST;
 import retrofit2.http.Query;
+import retrofit2.http.Streaming;
 
 interface InfluxDBService {
 
@@ -20,6 +21,7 @@ interface InfluxDBService {
   public static final String PRECISION = "precision";
   public static final String CONSISTENCY = "consistency";
   public static final String EPOCH = "epoch";
+  public static final String CHUNK_SIZE = "chunk_size";
 
   @GET("/ping")
   public Call<ResponseBody> ping();
@@ -61,4 +63,9 @@ public Call<QueryResult> query(@Query(U) String username, @Query(P) String passw
   public Call<QueryResult> postQuery(@Query(U) String username,
       @Query(P) String password, @Query(value = Q, encoded = true) String query);
 
+  @Streaming
+  @GET("/query?chunked=true")
+  public Call<ResponseBody> query(@Query(U) String username,
+      @Query(P) String password, @Query(DB) String db, @Query(value = Q, encoded = true) String query,
+      @Query(CHUNK_SIZE) int chunkSize);
 }
diff --git a/src/test/java/org/influxdb/InfluxDBTest.java b/src/test/java/org/influxdb/InfluxDBTest.java
@@ -1,19 +1,29 @@
 package org.influxdb;
 
-import org.influxdb.InfluxDB.LogLevel;
-import org.influxdb.dto.*;
-import org.influxdb.impl.InfluxDBImpl;
-import org.junit.Assert;
-import org.junit.Before;
-import org.junit.Test;
-
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Set;
+import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadFactory;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Consumer;
+
+import org.influxdb.InfluxDB.LogLevel;
+import org.influxdb.dto.BatchPoints;
+import org.influxdb.dto.Point;
+import org.influxdb.dto.Pong;
+import org.influxdb.dto.Query;
+import org.influxdb.dto.QueryResult;
+import org.influxdb.impl.InfluxDBImpl;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
 
 import com.google.common.util.concurrent.Uninterruptibles;
 
@@ -29,6 +39,7 @@ public class InfluxDBTest {
 	private final static int UDP_PORT = 8089;
 	private final static String UDP_DATABASE = "udp";
 
+    @Rule public final ExpectedException exception = ExpectedException.none();
 	/**
 	 * Create a influxDB connection before all tests start.
 	 *
@@ -57,7 +68,7 @@ public void setUp() throws InterruptedException, IOException {
 		this.influxDB.createDatabase(UDP_DATABASE);
 		// String logs = CharStreams.toString(new InputStreamReader(containerLogsStream,
 		// Charsets.UTF_8));
-		System.out.println("##################################################################################");
+        System.out.println("################################################################################## ");
 		// System.out.println("Container Logs: \n" + logs);
 		System.out.println("#  Connected to InfluxDB Version: " + this.influxDB.version() + " #");
 		System.out.println("##################################################################################");
@@ -219,8 +230,8 @@ public void testWriteStringDataThroughUDP() {
     @Test
     public void testWriteMultipleStringDataThroughUDP() {
         String measurement = TestUtils.getRandomMeasurement();
-        this.influxDB.write(UDP_PORT, measurement + ",atag=test1 idle=100,usertime=10,system=1\n" + 
-                                      measurement + ",atag=test2 idle=200,usertime=20,system=2\n" + 
+        this.influxDB.write(UDP_PORT, measurement + ",atag=test1 idle=100,usertime=10,system=1\n" +
+                                      measurement + ",atag=test2 idle=200,usertime=20,system=2\n" +
                                       measurement + ",atag=test3 idle=300,usertime=30,system=3");
         Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
         Query query = new Query("SELECT * FROM " + measurement + " GROUP BY *", UDP_DATABASE);
@@ -250,12 +261,12 @@ public void testWriteMultipleStringDataLinesThroughUDP() {
         Assert.assertEquals(3, result.getResults().get(0).getSeries().size());
         Assert.assertEquals(result.getResults().get(0).getSeries().get(0).getTags().get("atag"), "test1");
         Assert.assertEquals(result.getResults().get(0).getSeries().get(1).getTags().get("atag"), "test2");
-        Assert.assertEquals(result.getResults().get(0).getSeries().get(2).getTags().get("atag"), "test3"); 
+        Assert.assertEquals(result.getResults().get(0).getSeries().get(2).getTags().get("atag"), "test3");
     }
 
     /**
      * When batch of points' size is over UDP limit, the expected exception
-     * is java.lang.RuntimeException: java.net.SocketException: 
+     * is java.lang.RuntimeException: java.net.SocketException:
      * The message is larger than the maximum supported by the underlying transport: Datagram send failed
      * @throws Exception
      */
@@ -487,4 +498,93 @@ public void testWriteEnableGzipAndDisableGzip() {
         }
     }
 
+    /**
+     * Test chunking.
+     * @throws InterruptedException
+     */
+    @Test
+    public void testChunking() throws InterruptedException {
+        if (this.influxDB.version().startsWith("0.") || this.influxDB.version().startsWith("1.0")) {
+            // do not test version 0.13 and 1.0
+            return;
+        }
+        String dbName = "write_unittest_" + System.currentTimeMillis();
+        this.influxDB.createDatabase(dbName);
+        String rp = TestUtils.defaultRetentionPolicy(this.influxDB.version());
+        BatchPoints batchPoints = BatchPoints.database(dbName).retentionPolicy(rp).build();
+        Point point1 = Point.measurement("disk").tag("atag", "a").addField("used", 60L).addField("free", 1L).build();
+        Point point2 = Point.measurement("disk").tag("atag", "b").addField("used", 70L).addField("free", 2L).build();
+        Point point3 = Point.measurement("disk").tag("atag", "c").addField("used", 80L).addField("free", 3L).build();
+        batchPoints.point(point1);
+        batchPoints.point(point2);
+        batchPoints.point(point3);
+        this.influxDB.write(batchPoints);
+
+        Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
+        final BlockingQueue<QueryResult> queue = new LinkedBlockingQueue<>();
+        Query query = new Query("SELECT * FROM disk", dbName);
+        this.influxDB.query(query, 2, new Consumer<QueryResult>() {
+            @Override
+            public void accept(QueryResult result) {
+                queue.add(result);
+            }});
+
+        Uninterruptibles.sleepUninterruptibly(2, TimeUnit.SECONDS);
+        this.influxDB.deleteDatabase(dbName);
+
+        QueryResult result = queue.poll(20, TimeUnit.SECONDS);
+        Assert.assertNotNull(result);
+        System.out.println(result);
+        Assert.assertEquals(2, result.getResults().get(0).getSeries().get(0).getValues().size());
+
+        result = queue.poll(20, TimeUnit.SECONDS);
+        Assert.assertNotNull(result);
+        System.out.println(result);
+        Assert.assertEquals(1, result.getResults().get(0).getSeries().get(0).getValues().size());
+    }
+
+    /**
+     * Test chunking edge case.
+     * @throws InterruptedException
+     */
+    @Test
+    public void testChunkingFail() throws InterruptedException {
+        if (this.influxDB.version().startsWith("0.") || this.influxDB.version().startsWith("1.0")) {
+            // do not test version 0.13 and 1.0
+            return;
+        }
+        String dbName = "write_unittest_" + System.currentTimeMillis();
+        this.influxDB.createDatabase(dbName);
+        final CountDownLatch countDownLatch = new CountDownLatch(1);
+        Query query = new Query("XXX", dbName);
+        this.influxDB.query(query, 10, new Consumer<QueryResult>() {
+            @Override
+            public void accept(QueryResult result) {
+                countDownLatch.countDown();
+            }
+        });
+        this.influxDB.deleteDatabase(dbName);
+        Assert.assertFalse(countDownLatch.await(10, TimeUnit.SECONDS));
+    }
+
+    /**
+     * Test chunking on 0.13 and 1.0.
+     * @throws InterruptedException
+     */
+    @Test()
+    public void testChunkingOldVersion() throws InterruptedException {
+
+        if (this.influxDB.version().startsWith("0.") || this.influxDB.version().startsWith("1.0")) {
+
+            this.exception.expect(RuntimeException.class);
+            String dbName = "write_unittest_" + System.currentTimeMillis();
+            Query query = new Query("SELECT * FROM cpu GROUP BY *", dbName);
+            this.influxDB.query(query, 10, new Consumer<QueryResult>() {
+                @Override
+                public void accept(QueryResult result) {
+                }
+            });
+        }
+    }
+
 }
