Search in sources :

Example 1 with SendStream

use of io.confluent.ksql.api.utils.SendStream in project ksql by confluentinc.

the class ApiTest method shouldStreamInserts.

@Test
@CoreApiTest
public void shouldStreamInserts() throws Exception {
    LOG.info("Starting shouldStreamInserts");
    // Given
    JsonObject params = new JsonObject().put("target", "test-stream");
    // Stream for piping the HTTP request body
    SendStream readStream = new SendStream(vertx);
    // Stream for receiving the HTTP response body
    ReceiveStream writeStream = new ReceiveStream(vertx);
    VertxCompletableFuture<HttpResponse<Void>> fut = new VertxCompletableFuture<>();
    List<JsonObject> rows = DEFAULT_INSERT_ROWS;
    // When
    // Make an HTTP request but keep the request body and response streams open
    sendPostRequest("/inserts-stream", (request) -> request.as(BodyCodec.pipe(writeStream)).sendStream(readStream, fut));
    // Write the initial params Json object to the request body
    readStream.acceptBuffer(params.toBuffer().appendString("\n"));
    // Asynchronously on a timer write inserts to the request body
    AtomicInteger rowIndex = new AtomicInteger();
    vertx.setPeriodic(100, tid -> {
        readStream.acceptBuffer(rows.get(rowIndex.getAndIncrement()).toBuffer().appendString("\n"));
        if (rowIndex.get() == rows.size()) {
            vertx.cancelTimer(tid);
            // End the inserts stream and request when we've written all the rows to the stream
            readStream.end();
        }
    });
    // Wait for the response to complete
    LOG.info("Awaiting response from inserts");
    HttpResponse<Void> response = fut.get();
    LOG.info("Got response from inserts");
    // Then
    assertThat(response.statusCode(), is(200));
    assertThat(response.statusMessage(), is("OK"));
    // Verify we got acks for all our inserts
    InsertsResponse insertsResponse = new InsertsResponse(writeStream.getBody().toString());
    assertThat(insertsResponse.acks, hasSize(rows.size()));
    for (int i = 0; i < insertsResponse.acks.size(); i++) {
        final JsonObject ackLine = new JsonObject().put("status", "ok").put("seq", i);
        assertThat(insertsResponse.acks.get(i), is(ackLine));
    }
    // Make sure all inserts made it to the server
    TestInsertsSubscriber insertsSubscriber = testEndpoints.getInsertsSubscriber();
    LOG.info("Checking all rows inserted");
    assertThatEventually(insertsSubscriber::getRowsInserted, is(rows));
    LOG.info("Checking got completion marker");
    assertThatEventually(insertsSubscriber::isCompleted, is(true));
    // Ensure we received at least some of the response before all the request body was written
    // Yay HTTP2!
    assertThat(readStream.getLastSentTime() > writeStream.getFirstReceivedTime(), is(true));
    LOG.info("ShouldStreamInserts complete");
}
Also used : InsertsResponse(io.confluent.ksql.api.utils.InsertsResponse) ReceiveStream(io.confluent.ksql.api.utils.ReceiveStream) JsonObject(io.vertx.core.json.JsonObject) HttpResponse(io.vertx.ext.web.client.HttpResponse) VertxCompletableFuture(io.confluent.ksql.util.VertxCompletableFuture) SendStream(io.confluent.ksql.api.utils.SendStream) AtomicInteger(java.util.concurrent.atomic.AtomicInteger) Test(org.junit.Test)

Aggregations

InsertsResponse (io.confluent.ksql.api.utils.InsertsResponse)1 ReceiveStream (io.confluent.ksql.api.utils.ReceiveStream)1 SendStream (io.confluent.ksql.api.utils.SendStream)1 VertxCompletableFuture (io.confluent.ksql.util.VertxCompletableFuture)1 JsonObject (io.vertx.core.json.JsonObject)1 HttpResponse (io.vertx.ext.web.client.HttpResponse)1 AtomicInteger (java.util.concurrent.atomic.AtomicInteger)1 Test (org.junit.Test)1