Examples with SqlQuery org.apache.druid.sql.http.SqlQuery used on opensource projects

Example 1 with SqlQuery

use of org.apache.druid.sql.http.SqlQuery in project druid by druid-io.

the class SqlQueryTest method testSerde.

@Test
public void testSerde() throws Exception {
    final ObjectMapper jsonMapper = TestHelper.makeJsonMapper();
    final SqlQuery query = new SqlQuery("SELECT ?", ResultFormat.ARRAY, true, true, true, ImmutableMap.of("useCache", false), ImmutableList.of(new SqlParameter(SqlType.INTEGER, 1)));
    Assert.assertEquals(query, jsonMapper.readValue(jsonMapper.writeValueAsString(query), SqlQuery.class));
}

Example 2 with SqlQuery

use of org.apache.druid.sql.http.SqlQuery in project druid by druid-io.

the class AsyncQueryForwardingServletTest method testSqlQueryProxy.

@Test
public void testSqlQueryProxy() throws Exception {
    final SqlQuery query = new SqlQuery("SELECT * FROM foo", ResultFormat.ARRAY, false, false, false, null, null);
    final QueryHostFinder hostFinder = EasyMock.createMock(QueryHostFinder.class);
    EasyMock.expect(hostFinder.findServerSql(query)).andReturn(new TestServer("http", "1.2.3.4", 9999)).once();
    EasyMock.replay(hostFinder);
    Properties properties = new Properties();
    properties.setProperty("druid.router.sql.enable", "true");
    verifyServletCallsForQuery(query, true, hostFinder, properties);
}

Example 3 with SqlQuery

use of org.apache.druid.sql.http.SqlQuery in project druid by druid-io.

the class AsyncQueryForwardingServlet method service.

@Override
protected void service(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {
    final boolean isSmile = SmileMediaTypes.APPLICATION_JACKSON_SMILE.equals(request.getContentType()) || APPLICATION_SMILE.equals(request.getContentType());
    final ObjectMapper objectMapper = isSmile ? smileMapper : jsonMapper;
    request.setAttribute(OBJECTMAPPER_ATTRIBUTE, objectMapper);
    final String requestURI = request.getRequestURI();
    final String method = request.getMethod();
    final Server targetServer;
    // The Router does not have the ability to look inside SQL queries and route them intelligently, so just treat
    // them as a generic request.
    final boolean isNativeQueryEndpoint = requestURI.startsWith("/druid/v2") && !requestURI.startsWith("/druid/v2/sql");
    final boolean isSqlQueryEndpoint = requestURI.startsWith("/druid/v2/sql");
    final boolean isAvaticaJson = requestURI.startsWith("/druid/v2/sql/avatica");
    final boolean isAvaticaPb = requestURI.startsWith("/druid/v2/sql/avatica-protobuf");
    if (isAvaticaPb) {
        byte[] requestBytes = IOUtils.toByteArray(request.getInputStream());
        Service.Request protobufRequest = this.protobufTranslation.parseRequest(requestBytes);
        String connectionId = getAvaticaProtobufConnectionId(protobufRequest);
        targetServer = hostFinder.findServerAvatica(connectionId);
        request.setAttribute(AVATICA_QUERY_ATTRIBUTE, requestBytes);
        LOG.debug("Forwarding protobuf JDBC connection [%s] to broker [%s]", connectionId, targetServer);
    } else if (isAvaticaJson) {
        Map<String, Object> requestMap = objectMapper.readValue(request.getInputStream(), JacksonUtils.TYPE_REFERENCE_MAP_STRING_OBJECT);
        String connectionId = getAvaticaConnectionId(requestMap);
        targetServer = hostFinder.findServerAvatica(connectionId);
        byte[] requestBytes = objectMapper.writeValueAsBytes(requestMap);
        request.setAttribute(AVATICA_QUERY_ATTRIBUTE, requestBytes);
        LOG.debug("Forwarding JDBC connection [%s] to broker [%s]", connectionId, targetServer.getHost());
    } else if (HttpMethod.DELETE.is(method)) {
        // query cancellation request
        targetServer = hostFinder.pickDefaultServer();
        broadcastQueryCancelRequest(request, targetServer);
        LOG.debug("Broadcasting cancellation request to all brokers");
    } else if (isNativeQueryEndpoint && HttpMethod.POST.is(method)) {
        // query request
        try {
            Query inputQuery = objectMapper.readValue(request.getInputStream(), Query.class);
            if (inputQuery != null) {
                targetServer = hostFinder.pickServer(inputQuery);
                if (inputQuery.getId() == null) {
                    inputQuery = inputQuery.withId(UUID.randomUUID().toString());
                }
                LOG.debug("Forwarding JSON query [%s] to broker [%s]", inputQuery.getId(), targetServer.getHost());
            } else {
                targetServer = hostFinder.pickDefaultServer();
                LOG.debug("Forwarding JSON request to broker [%s]", targetServer.getHost());
            }
            request.setAttribute(QUERY_ATTRIBUTE, inputQuery);
        } catch (IOException e) {
            handleQueryParseException(request, response, objectMapper, e, true);
            return;
        } catch (Exception e) {
            handleException(response, objectMapper, e);
            return;
        }
    } else if (routeSqlByStrategy && isSqlQueryEndpoint && HttpMethod.POST.is(method)) {
        try {
            SqlQuery inputSqlQuery = objectMapper.readValue(request.getInputStream(), SqlQuery.class);
            request.setAttribute(SQL_QUERY_ATTRIBUTE, inputSqlQuery);
            targetServer = hostFinder.findServerSql(inputSqlQuery);
            LOG.debug("Forwarding SQL query to broker [%s]", targetServer.getHost());
        } catch (IOException e) {
            handleQueryParseException(request, response, objectMapper, e, false);
            return;
        } catch (Exception e) {
            handleException(response, objectMapper, e);
            return;
        }
    } else {
        targetServer = hostFinder.pickDefaultServer();
        LOG.debug("Forwarding query to broker [%s]", targetServer.getHost());
    }
    request.setAttribute(HOST_ATTRIBUTE, targetServer.getHost());
    request.setAttribute(SCHEME_ATTRIBUTE, targetServer.getScheme());
    doService(request, response);
}

Example 4 with SqlQuery

use of org.apache.druid.sql.http.SqlQuery in project druid by druid-io.

the class ITSqlCancelTest method testCancelValidQuery.

@Test
public void testCancelValidQuery() throws Exception {
    final String queryId = "sql-cancel-test";
    final List<Future<StatusResponseHolder>> queryResponseFutures = new ArrayList<>();
    for (int i = 0; i < NUM_QUERIES; i++) {
        queryResponseFutures.add(sqlClient.queryAsync(sqlHelper.getQueryURL(config.getRouterUrl()), new SqlQuery(QUERY, null, false, false, false, ImmutableMap.of(BaseQuery.SQL_QUERY_ID, queryId), null)));
    }
    // Wait until the sqlLifecycle is authorized and registered
    Thread.sleep(1000);
    final HttpResponseStatus responseStatus = sqlClient.cancelQuery(sqlHelper.getCancelUrl(config.getRouterUrl(), queryId), 1000);
    if (!responseStatus.equals(HttpResponseStatus.ACCEPTED)) {
        throw new RE("Failed to cancel query [%s]. Response code was [%s]", queryId, responseStatus);
    }
    for (Future<StatusResponseHolder> queryResponseFuture : queryResponseFutures) {
        final StatusResponseHolder queryResponse = queryResponseFuture.get(1, TimeUnit.SECONDS);
        if (!queryResponse.getStatus().equals(HttpResponseStatus.INTERNAL_SERVER_ERROR)) {
            throw new ISE("Query is not canceled after cancel request");
        }
        QueryException queryException = jsonMapper.readValue(queryResponse.getContent(), QueryException.class);
        if (!QueryInterruptedException.QUERY_CANCELLED.equals(queryException.getErrorCode())) {
            throw new ISE("Expected error code [%s], actual [%s]", QueryInterruptedException.QUERY_CANCELLED, queryException.getErrorCode());
        }
    }
}

Example 5 with SqlQuery

use of org.apache.druid.sql.http.SqlQuery in project druid by druid-io.

the class ITSqlCancelTest method testCancelInvalidQuery.

@Test
public void testCancelInvalidQuery() throws Exception {
    final Future<StatusResponseHolder> queryResponseFuture = sqlClient.queryAsync(sqlHelper.getQueryURL(config.getRouterUrl()), new SqlQuery(QUERY, null, false, false, false, ImmutableMap.of(BaseQuery.SQL_QUERY_ID, "validId"), null));
    // Wait until the sqlLifecycle is authorized and registered
    Thread.sleep(1000);
    final HttpResponseStatus responseStatus = sqlClient.cancelQuery(sqlHelper.getCancelUrl(config.getRouterUrl(), "invalidId"), 1000);
    if (!responseStatus.equals(HttpResponseStatus.NOT_FOUND)) {
        throw new RE("Expected http response [%s], actual response [%s]", HttpResponseStatus.NOT_FOUND, responseStatus);
    }
    final StatusResponseHolder queryResponse = queryResponseFuture.get(30, TimeUnit.SECONDS);
    if (!queryResponse.getStatus().equals(HttpResponseStatus.OK)) {
        throw new ISE("Cancel request failed with status[%s] and content[%s]", queryResponse.getStatus(), queryResponse.getContent());
    }
}