Examples with ReceiveStream io.confluent.ksql.api.utils.ReceiveStream used on opensource projects

Example 1 with ReceiveStream

use of io.confluent.ksql.api.utils.ReceiveStream 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");
}

Example 2 with ReceiveStream

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

the class ApiIntegrationTest method shouldExecutePushQueryNoLimit.

@Test
public void shouldExecutePushQueryNoLimit() throws Exception {
    KsqlEngine engine = (KsqlEngine) REST_APP.getEngine();
    // One persistent query for the agg table
    assertThatEventually(engine::numberOfLiveQueries, is(1));
    // Given:
    String sql = "SELECT * from " + TEST_STREAM + " EMIT CHANGES;";
    // Create a write stream to capture the incomplete response
    ReceiveStream writeStream = new ReceiveStream(vertx);
    // Make the request to stream a query
    JsonObject properties = new JsonObject();
    JsonObject requestBody = new JsonObject().put("sql", sql).put("properties", properties);
    VertxCompletableFuture<HttpResponse<Void>> responseFuture = new VertxCompletableFuture<>();
    client.post("/query-stream").as(BodyCodec.pipe(writeStream)).sendJsonObject(requestBody, responseFuture);
    assertThatEventually(engine::numberOfLiveQueries, is(2));
    // Wait for all rows in the response to arrive
    assertThatEventually(() -> {
        try {
            Buffer buff = writeStream.getBody();
            QueryResponse queryResponse = new QueryResponse(buff.toString());
            return queryResponse.rows.size();
        } catch (Throwable t) {
            return -1;
        }
    }, greaterThanOrEqualTo(6));
    // The response shouldn't have ended yet
    assertThat(writeStream.isEnded(), is(false));
    QueryResponse queryResponse = new QueryResponse(writeStream.getBody().toString());
    String queryId = queryResponse.responseObject.getString("queryId");
    // Now send another request to close the query
    JsonObject closeQueryRequestBody = new JsonObject().put("queryId", queryId);
    HttpResponse<Buffer> closeQueryResponse = sendRequest(client, "/close-query", closeQueryRequestBody.toBuffer());
    assertThat(closeQueryResponse.statusCode(), is(200));
    // The response should now be ended
    assertThatEventually(writeStream::isEnded, is(true));
    HttpResponse<Void> response = responseFuture.get();
    assertThat(response.statusCode(), is(200));
    // Make sure it's cleaned up on the server
    assertThatEventually(engine::numberOfLiveQueries, is(1));
}

Example 3 with ReceiveStream

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

the class ApiIntegrationTest method shouldExecutePushQueryFromLatestOffset.

@Test
public void shouldExecutePushQueryFromLatestOffset() {
    KsqlEngine engine = (KsqlEngine) REST_APP.getEngine();
    // One persistent query for the agg table
    assertThatEventually(engine::numberOfLiveQueries, is(1));
    // Given:
    String sql = "SELECT * from " + TEST_STREAM + " EMIT CHANGES LIMIT 1;";
    // Create a write stream to capture the incomplete response
    ReceiveStream writeStream = new ReceiveStream(vertx);
    // Make the request to stream a query
    JsonObject queryProperties = new JsonObject().put("auto.offset.reset", "latest");
    JsonObject queryRequestBody = new JsonObject().put("sql", sql).put("properties", queryProperties);
    VertxCompletableFuture<HttpResponse<Void>> responseFuture = new VertxCompletableFuture<>();
    client.post("/query-stream").as(BodyCodec.pipe(writeStream)).sendJsonObject(queryRequestBody, responseFuture);
    assertThatEventually(engine::numberOfLiveQueries, is(2));
    // New row to insert
    JsonObject row = new JsonObject().put("K", new JsonObject().put("F1", new JsonArray().add("my_key_shouldExecutePushQueryFromLatestOffset"))).put("STR", "Value_shouldExecutePushQueryFromLatestOffset").put("LONG", 2000L).put("DEC", // JsonObject does not accept BigDecimal
    12.34).put("BYTES_", new byte[] { 0, 1, 2 }).put("ARRAY", new JsonArray().add("a_shouldExecutePushQueryFromLatestOffset")).put("MAP", new JsonObject().put("k1", "v1_shouldExecutePushQueryFromLatestOffset")).put("STRUCT", new JsonObject().put("F1", 3)).put("COMPLEX", COMPLEX_FIELD_VALUE);
    // Insert a new row and wait for it to arrive
    assertThatEventually(() -> {
        try {
            // Attempt the insert multiple times, in case the query hasn't started yet
            shouldInsert(row);
            Buffer buff = writeStream.getBody();
            QueryResponse queryResponse = new QueryResponse(buff.toString());
            return queryResponse.rows.size();
        } catch (Throwable t) {
            return Integer.MAX_VALUE;
        }
    }, is(1));
    // Verify that the received row is the expected one
    Buffer buff = writeStream.getBody();
    QueryResponse queryResponse = new QueryResponse(buff.toString());
    assertThat(queryResponse.rows.get(0).getJsonObject(0), is(new JsonObject().put("F1", new JsonArray().add("my_key_shouldExecutePushQueryFromLatestOffset"))));
    assertThat(queryResponse.rows.get(0).getString(1), is("Value_shouldExecutePushQueryFromLatestOffset"));
    assertThat(queryResponse.rows.get(0).getLong(2), is(2000L));
    assertThat(queryResponse.rows.get(0).getDouble(3), is(12.34));
    assertThat(queryResponse.rows.get(0).getBinary(4), is(new byte[] { 0, 1, 2 }));
    assertThat(queryResponse.rows.get(0).getJsonArray(5), is(new JsonArray().add("a_shouldExecutePushQueryFromLatestOffset")));
    assertThat(queryResponse.rows.get(0).getJsonObject(6), is(new JsonObject().put("k1", "v1_shouldExecutePushQueryFromLatestOffset")));
    assertThat(queryResponse.rows.get(0).getJsonObject(7), is(new JsonObject().put("F1", 3)));
    assertThat(queryResponse.rows.get(0).getJsonObject(8), is(COMPLEX_FIELD_VALUE));
    // Check that query is cleaned up on the server
    assertThatEventually(engine::numberOfLiveQueries, is(1));
}

Example 4 with ReceiveStream

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

the class ApiTest method shouldCloseQuery.

@Test
@CoreApiTest
public void shouldCloseQuery() throws Exception {
    // Create a write stream to capture the incomplete response
    ReceiveStream writeStream = new ReceiveStream(vertx);
    VertxCompletableFuture<HttpResponse<Void>> responseFuture = new VertxCompletableFuture<>();
    // Make the request to stream a query
    sendPostRequest("/query-stream", (request) -> request.as(BodyCodec.pipe(writeStream)).sendJsonObject(DEFAULT_PUSH_QUERY_REQUEST_BODY, responseFuture));
    // Wait for all rows in the response to arrive
    assertThatEventually(() -> {
        try {
            Buffer buff = writeStream.getBody();
            QueryResponse queryResponse = new QueryResponse(buff.toString());
            return queryResponse.rows.size();
        } catch (Throwable t) {
            return Integer.MAX_VALUE;
        }
    }, is(DEFAULT_JSON_ROWS.size()));
    // The response shouldn't have ended yet
    assertThat(writeStream.isEnded(), is(false));
    // Assert the query is still live on the server
    QueryResponse queryResponse = new QueryResponse(writeStream.getBody().toString());
    String queryId = queryResponse.responseObject.getString("queryId");
    assertThat(server.getQueryIDs().contains(new PushQueryId(queryId)), is(true));
    assertThat(server.getQueryIDs(), hasSize(1));
    assertThat(testEndpoints.getQueryPublishers(), hasSize(1));
    // Now send another request to close the query
    JsonObject closeQueryRequestBody = new JsonObject().put("queryId", queryId);
    HttpResponse<Buffer> closeQueryResponse = sendPostRequest("/close-query", closeQueryRequestBody.toBuffer());
    assertThat(closeQueryResponse.statusCode(), is(200));
    // Assert the query no longer exists on the server
    assertThat(server.getQueryIDs(), not(hasItem(new PushQueryId(queryId))));
    assertThat(server.getQueryIDs(), hasSize(0));
    // The response should now be ended
    assertThatEventually(writeStream::isEnded, is(true));
    HttpResponse<Void> response = responseFuture.get();
    assertThat(response.statusCode(), is(200));
}

Example 5 with ReceiveStream

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

the class BaseApiTest method executePushQueryAndWaitForRows.

protected QueryResponse executePushQueryAndWaitForRows(final WebClient client, final JsonObject requestBody) throws Exception {
    ReceiveStream writeStream = new ReceiveStream(vertx);
    sendPostRequest(client, "/query-stream", (request) -> request.as(BodyCodec.pipe(writeStream)).sendJsonObject(requestBody, ar -> {
    }));
    // Wait for all rows to arrive
    assertThatEventually(() -> {
        try {
            Buffer buff = writeStream.getBody();
            QueryResponse queryResponse = new QueryResponse(buff.toString());
            return queryResponse.rows.size();
        } catch (Throwable t) {
            return Integer.MAX_VALUE;
        }
    }, is(DEFAULT_JSON_ROWS.size()));
    // Note, the response hasn't ended at this point
    assertThat(writeStream.isEnded(), is(false));
    return new QueryResponse(writeStream.getBody().toString());
}