Examples with QueryStreamArgs io.confluent.ksql.rest.entity.QueryStreamArgs used on opensource projects

Example 1 with QueryStreamArgs

use of io.confluent.ksql.rest.entity.QueryStreamArgs in project ksql by confluentinc.

the class RestApiTest method shouldExecutePullQuery_allTypes.

@Test
public void shouldExecutePullQuery_allTypes() {
    ImmutableList<String> formats = ImmutableList.of("application/vnd.ksqlapi.delimited.v1", KsqlMediaType.KSQL_V1_JSON.mediaType());
    ImmutableList<HttpVersion> httpVersions = ImmutableList.of(HTTP_1_1, HTTP_2);
    ImmutableList<String> endpoints = ImmutableList.of("/query-stream", "/query");
    ImmutableList<Boolean> migrationFlags = ImmutableList.of(true, false);
    final String query = "SELECT COUNT, USERID from " + AGG_TABLE + " WHERE USERID='" + AN_AGG_KEY + "';";
    for (String format : formats) {
        for (HttpVersion version : httpVersions) {
            for (String endpoint : endpoints) {
                for (Boolean migrationEnabled : migrationFlags) {
                    boolean routedToQueryStream = endpoint.equals("/query") && format.equals("application/vnd.ksqlapi.delimited.v1");
                    boolean migrated = routedToQueryStream || migrationEnabled;
                    if (!migrated && endpoint.equals("/query") && version == HTTP_2) {
                        LOG.info("Skipping pull query combination {} {} {} {}", format, version, endpoint, migrationEnabled);
                        continue;
                    }
                    LOG.info("Trying pull query combination {} {} {} {}", format, version, endpoint, migrationEnabled);
                    Object requestBody;
                    ImmutableMap<String, Object> overrides = ImmutableMap.of(KsqlConfig.KSQL_ENDPOINT_MIGRATE_QUERY_CONFIG, migrationEnabled);
                    if (endpoint.equals("/query-stream")) {
                        requestBody = new QueryStreamArgs(query, overrides, Collections.emptyMap(), Collections.emptyMap());
                    } else if (endpoint.equals("/query")) {
                        requestBody = new KsqlRequest(query, overrides, Collections.emptyMap(), Collections.emptyMap(), null);
                    } else {
                        fail("Unknown endpoint " + endpoint);
                        return;
                    }
                    // this is somewhat hard since they have different data, so we don't try.
                    if (format.equals("application/vnd.ksqlapi.delimited.v1")) {
                        QueryResponse[] queryResponse = new QueryResponse[1];
                        assertThatEventually(() -> {
                            try {
                                HttpResponse<Buffer> resp = RestIntegrationTestUtil.rawRestRequest(REST_APP, version, POST, endpoint, requestBody, "application/vnd.ksqlapi.delimited.v1", Optional.empty(), Optional.empty());
                                queryResponse[0] = new QueryResponse(resp.body().toString());
                                return queryResponse[0].rows.size();
                            } catch (Throwable t) {
                                return Integer.MAX_VALUE;
                            }
                        }, is(1));
                        assertThat(queryResponse[0].rows.get(0).getList(), is(ImmutableList.of(1, "USER_1")));
                    } else if (format.equals(KsqlMediaType.KSQL_V1_JSON.mediaType())) {
                        final Supplier<List<String>> call = () -> {
                            HttpResponse<Buffer> resp = RestIntegrationTestUtil.rawRestRequest(REST_APP, version, POST, endpoint, requestBody, KsqlMediaType.KSQL_V1_JSON.mediaType(), Optional.empty(), Optional.empty());
                            final String response = resp.body().toString();
                            return Arrays.asList(response.split(System.lineSeparator()));
                        };
                        // When:
                        final List<String> messages = assertThatEventually(call, hasSize(HEADER + 1));
                        // Then:
                        assertThat(messages, hasSize(HEADER + 1));
                        assertThat(messages.get(0), startsWith("[{\"header\":{\"queryId\":\""));
                        assertThat(messages.get(0), endsWith("\",\"schema\":\"`COUNT` BIGINT, `USERID` STRING KEY\"}},"));
                        assertThat(messages.get(1), is("{\"row\":{\"columns\":[1,\"USER_1\"]}}]"));
                    } else {
                        fail("Unknown format " + format);
                        return;
                    }
                }
            }
        }
    }
}

Example 2 with QueryStreamArgs

use of io.confluent.ksql.rest.entity.QueryStreamArgs in project ksql by confluentinc.

the class KsqlClientTest method shouldPostQueryRequestStreamedHttp2.

@Test
public void shouldPostQueryRequestStreamedHttp2() throws Exception {
    // Given:
    int numRows = 10;
    List<List<?>> expectedResponse = setQueryStreamDelimitedResponse(numRows);
    String sql = "whateva";
    // When:
    KsqlTarget target = ksqlClient.targetHttp2(serverUri);
    CompletableFuture<RestResponse<StreamPublisher<StreamedRow>>> future = target.postQueryRequestStreamedAsync(sql, ImmutableMap.of());
    RestResponse<StreamPublisher<StreamedRow>> response = future.get();
    // Then:
    assertThat(server.getHttpMethod(), is(HttpMethod.POST));
    assertThat(server.getPath(), is("/query-stream"));
    assertThat(server.getHeaders().get("Accept"), is("application/vnd.ksqlapi.delimited.v1"));
    assertThat(getQueryStreamArgs(), is(new QueryStreamArgs(sql, properties, Collections.emptyMap(), Collections.emptyMap())));
    List<StreamedRow> rows = getElementsFromPublisher(numRows + 1, response.getResponse());
    int i = 0;
    for (StreamedRow row : rows) {
        assertThat(row.getRow().isPresent(), is(true));
        assertThat(row.getRow().get().getColumns(), is(expectedResponse.get(i++)));
    }
}

Example 3 with QueryStreamArgs

use of io.confluent.ksql.rest.entity.QueryStreamArgs in project ksql by confluentinc.

the class KsqlTarget method executeQueryStreamRequest.

private <T> CompletableFuture<RestResponse<StreamPublisher<T>>> executeQueryStreamRequest(final String ksql, final Map<String, ?> requestProperties, final Function<Buffer, T> mapper) {
    final Map<String, Object> requestPropertiesObject = requestProperties.entrySet().stream().collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    final QueryStreamArgs queryStreamArgs = new QueryStreamArgs(ksql, localProperties.toMap(), Collections.emptyMap(), requestPropertiesObject);
    final AtomicReference<StreamPublisher<T>> pubRef = new AtomicReference<>();
    return executeAsync(HttpMethod.POST, QUERY_STREAM_PATH, Optional.of("application/vnd.ksqlapi.delimited.v1"), queryStreamArgs, resp -> pubRef.get(), (resp, vcf) -> {
        if (resp.statusCode() == 200) {
            pubRef.set(new StreamPublisher<>(Vertx.currentContext(), resp, mapper, vcf, false));
            vcf.complete(new ResponseWithBody(resp));
        } else {
            resp.bodyHandler(body -> vcf.complete(new ResponseWithBody(resp, body)));
        }
    });
}

Example 4 with QueryStreamArgs

use of io.confluent.ksql.rest.entity.QueryStreamArgs in project ksql by confluentinc.

the class RestApiTest method shouldExecuteScalablePushQueryOverHttp2QueryStream.

@Test
public void shouldExecuteScalablePushQueryOverHttp2QueryStream() throws InterruptedException, ExecutionException {
    QueryStreamArgs queryStreamArgs = new QueryStreamArgs("SELECT USERID, PAGEID, VIEWTIME from " + PAGE_VIEW_CSAS + " EMIT CHANGES LIMIT " + LIMIT + ";", ImmutableMap.of("auto.offset.reset", "latest"), Collections.emptyMap(), Collections.emptyMap());
    List<String> messages = new ArrayList<>();
    Semaphore start = new Semaphore(0);
    VertxCompletableFuture<Void> future = RestIntegrationTestUtil.rawRestRequest(REST_APP, HTTP_2, POST, "/query-stream", queryStreamArgs, "application/vnd.ksqlapi.delimited.v1", buffer -> {
        if (buffer == null || buffer.length() == 0) {
            return;
        }
        String str = buffer.toString();
        if (str.startsWith("{") && str.contains("\"queryId\"")) {
            start.release();
        }
        String[] parts = str.split("\n");
        messages.addAll(Arrays.asList(parts));
    }, Optional.empty());
    // Wait to get the metadata so we know we've started.
    start.acquire();
    assertExpectedScalablePushQueries(1);
    // Write some new rows
    TEST_HARNESS.produceRows(PAGE_VIEW2_TOPIC, PAGE_VIEWS2_PROVIDER, FormatFactory.KAFKA, FormatFactory.JSON);
    future.get();
    assertThat(messages.size(), is(HEADER + LIMIT));
    assertThat(messages.get(0), startsWith("{\"queryId\":\"SCALABLE_PUSH_QUERY_"));
    assertThat(messages.get(0), endsWith(",\"columnNames\":[\"USERID\",\"PAGEID\",\"VIEWTIME\"]," + "\"columnTypes\":[\"STRING\",\"STRING\",\"BIGINT\"]}"));
    assertThat(messages.get(1), is("[\"USER_4\",\"PAGE_1\",10]"));
    assertThat(messages.get(2), is("[\"USER_0\",\"PAGE_5\",11]"));
}

Example 5 with QueryStreamArgs

use of io.confluent.ksql.rest.entity.QueryStreamArgs in project ksql by confluentinc.

the class RestApiTest method shouldExecutePullQueryOverHttp2QueryStream.

@Test
public void shouldExecutePullQueryOverHttp2QueryStream() {
    QueryStreamArgs queryStreamArgs = new QueryStreamArgs("SELECT COUNT, USERID from " + AGG_TABLE + " WHERE USERID='" + AN_AGG_KEY + "';", Collections.emptyMap(), Collections.emptyMap(), Collections.emptyMap());
    QueryResponse[] queryResponse = new QueryResponse[1];
    assertThatEventually(() -> {
        try {
            HttpResponse<Buffer> resp = RestIntegrationTestUtil.rawRestRequest(REST_APP, HTTP_2, POST, "/query-stream", queryStreamArgs, "application/vnd.ksqlapi.delimited.v1", Optional.empty(), Optional.empty());
            queryResponse[0] = new QueryResponse(resp.body().toString());
            return queryResponse[0].rows.size();
        } catch (Throwable t) {
            return Integer.MAX_VALUE;
        }
    }, is(1));
    assertThat(queryResponse[0].rows.get(0).getList(), is(ImmutableList.of(1, "USER_1")));
}