Search in sources :

Example 6 with Server

use of org.apache.druid.client.selector.Server 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);
}
Also used : SqlQuery(org.apache.druid.sql.http.SqlQuery) Server(org.apache.druid.client.selector.Server) Query(org.apache.druid.query.Query) SqlQuery(org.apache.druid.sql.http.SqlQuery) Service(org.apache.calcite.avatica.remote.Service) IOException(java.io.IOException) Map(java.util.Map) ImmutableMap(com.google.common.collect.ImmutableMap) ObjectMapper(com.fasterxml.jackson.databind.ObjectMapper) ServletException(javax.servlet.ServletException) QueryInterruptedException(org.apache.druid.query.QueryInterruptedException) JsonProcessingException(com.fasterxml.jackson.core.JsonProcessingException) IOException(java.io.IOException)

Example 7 with Server

use of org.apache.druid.client.selector.Server in project druid by druid-io.

the class ServerDiscoverySelectorTest method testPickIPv6Bracket.

@Test
public void testPickIPv6Bracket() throws Exception {
    final String address = "[2001:0db8:0000:0000:0000:ff00:0042:8329]";
    EasyMock.expect(serviceProvider.getInstance()).andReturn(instance).anyTimes();
    EasyMock.expect(instance.getAddress()).andReturn(address).anyTimes();
    EasyMock.expect(instance.getPort()).andReturn(PORT).anyTimes();
    EasyMock.expect(instance.getSslPort()).andReturn(-1).anyTimes();
    EasyMock.replay(instance, serviceProvider);
    Server server = serverDiscoverySelector.pick();
    Assert.assertEquals(PORT, server.getPort());
    Assert.assertEquals(address, server.getAddress());
    Assert.assertTrue(server.getHost().contains(Integer.toString(PORT)));
    Assert.assertTrue(server.getHost().contains(address));
    Assert.assertEquals("http", server.getScheme());
    EasyMock.verify(instance, serviceProvider);
    final URI uri = new URI(server.getScheme(), null, server.getAddress(), server.getPort(), "/druid/indexer/v1/action", null, null);
    Assert.assertEquals(PORT, uri.getPort());
    Assert.assertEquals(address, uri.getHost());
    Assert.assertEquals("http", uri.getScheme());
}
Also used : Server(org.apache.druid.client.selector.Server) URI(java.net.URI) Test(org.junit.Test)

Example 8 with Server

use of org.apache.druid.client.selector.Server in project druid by druid-io.

the class ServerDiscoverySelectorTest method testPickWithNullInstance.

@Test
public void testPickWithNullInstance() throws Exception {
    EasyMock.expect(serviceProvider.getInstance()).andReturn(null).anyTimes();
    EasyMock.replay(serviceProvider);
    Server server = serverDiscoverySelector.pick();
    Assert.assertNull(server);
    EasyMock.verify(serviceProvider);
}
Also used : Server(org.apache.druid.client.selector.Server) Test(org.junit.Test)

Example 9 with Server

use of org.apache.druid.client.selector.Server in project druid by druid-io.

the class ServerDiscoverySelectorTest method testPickWithSslPort.

@Test
public void testPickWithSslPort() throws Exception {
    EasyMock.expect(serviceProvider.getInstance()).andReturn(instance).anyTimes();
    EasyMock.expect(instance.getAddress()).andReturn(ADDRESS).anyTimes();
    EasyMock.expect(instance.getPort()).andReturn(PORT).anyTimes();
    EasyMock.expect(instance.getSslPort()).andReturn(SSL_PORT).anyTimes();
    EasyMock.replay(instance, serviceProvider);
    Server server = serverDiscoverySelector.pick();
    Assert.assertEquals(SSL_PORT, server.getPort());
    Assert.assertEquals(ADDRESS, server.getAddress());
    Assert.assertTrue(server.getHost().contains(Integer.toString(SSL_PORT)));
    Assert.assertTrue(server.getHost().contains(ADDRESS));
    Assert.assertEquals("https", server.getScheme());
    EasyMock.verify(instance, serviceProvider);
    final URI uri = new URI(server.getScheme(), null, server.getAddress(), server.getPort(), "/druid/indexer/v1/action", null, null);
    Assert.assertEquals(SSL_PORT, uri.getPort());
    Assert.assertEquals(ADDRESS, uri.getHost());
    Assert.assertEquals("https", uri.getScheme());
}
Also used : Server(org.apache.druid.client.selector.Server) URI(java.net.URI) Test(org.junit.Test)

Example 10 with Server

use of org.apache.druid.client.selector.Server in project druid by druid-io.

the class ServerDiscoverySelectorTest method testPickIPv6.

@Test
public void testPickIPv6() throws Exception {
    final String address = "2001:0db8:0000:0000:0000:ff00:0042:8329";
    EasyMock.expect(serviceProvider.getInstance()).andReturn(instance).anyTimes();
    EasyMock.expect(instance.getAddress()).andReturn(address).anyTimes();
    EasyMock.expect(instance.getPort()).andReturn(PORT).anyTimes();
    EasyMock.expect(instance.getSslPort()).andReturn(-1).anyTimes();
    EasyMock.replay(instance, serviceProvider);
    Server server = serverDiscoverySelector.pick();
    Assert.assertEquals(PORT, server.getPort());
    Assert.assertEquals(address, server.getAddress());
    Assert.assertTrue(server.getHost().contains(Integer.toString(PORT)));
    Assert.assertTrue(server.getHost().contains(address));
    Assert.assertEquals("http", server.getScheme());
    EasyMock.verify(instance, serviceProvider);
    final URI uri = new URI(server.getScheme(), null, server.getAddress(), server.getPort(), "/druid/indexer/v1/action", null, null);
    Assert.assertEquals(PORT, uri.getPort());
    Assert.assertEquals(StringUtils.format("[%s]", address), uri.getHost());
    Assert.assertEquals("http", uri.getScheme());
}
Also used : Server(org.apache.druid.client.selector.Server) URI(java.net.URI) Test(org.junit.Test)

Aggregations

Server (org.apache.druid.client.selector.Server)20 Test (org.junit.Test)10 URI (java.net.URI)5 IOException (java.io.IOException)3 Query (org.apache.druid.query.Query)3 JsonProcessingException (com.fasterxml.jackson.core.JsonProcessingException)2 ObjectMapper (com.fasterxml.jackson.databind.ObjectMapper)2 ImmutableMap (com.google.common.collect.ImmutableMap)2 HashMap (java.util.HashMap)2 Map (java.util.Map)2 ServletException (javax.servlet.ServletException)2 Service (org.apache.calcite.avatica.remote.Service)2 DruidServer (org.apache.druid.client.DruidServer)2 QueryInterruptedException (org.apache.druid.query.QueryInterruptedException)2 CountAggregatorFactory (org.apache.druid.query.aggregation.CountAggregatorFactory)2 SqlQuery (org.apache.druid.sql.http.SqlQuery)2 SmileMediaTypes (com.fasterxml.jackson.jaxrs.smile.SmileMediaTypes)1 VisibleForTesting (com.google.common.annotations.VisibleForTesting)1 Inject (com.google.inject.Inject)1 Provider (com.google.inject.Provider)1