Add BrokerClient implementation

server/src/main/java/org/apache/druid/discovery/BrokerClient.java

@@ -29,6 +29,7 @@
 import org.apache.druid.java.util.http.client.Request;
 import org.apache.druid.java.util.http.client.response.StringFullResponseHandler;
 import org.apache.druid.java.util.http.client.response.StringFullResponseHolder;
+import org.apache.druid.rpc.ServiceClient;
 import org.jboss.netty.channel.ChannelException;
 import org.jboss.netty.handler.codec.http.HttpMethod;
 import org.jboss.netty.handler.codec.http.HttpResponseStatus;
@@ -41,7 +42,10 @@
 
 /**
  * This class facilitates interaction with Broker.
+ * Note that this should be removed and reconciled with org.apache.druid.sql.client.BrokerClient, which has the
+ * built-in functionality of {@link ServiceClient}, and proper Guice and service discovery wired in.
  */
+@Deprecated
 public class BrokerClient
 {
   private static final int MAX_RETRIES = 5;

server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java

@@ -48,7 +48,7 @@
 public class ServiceClientModule implements DruidModule
 {
   private static final int CONNECT_EXEC_THREADS = 4;
-  private static final int CLIENT_MAX_ATTEMPTS = 6;
+  protected static final int CLIENT_MAX_ATTEMPTS = 6;
 
   @Override
   public void configure(Binder binder)

sql/src/main/java/org/apache/druid/sql/calcite/planner/ExplainAttributes.java

@@ -19,12 +19,14 @@
 
 package org.apache.druid.sql.calcite.planner;
 
+import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import org.apache.druid.java.util.common.granularity.Granularity;
 
 import javax.annotation.Nullable;
 import java.util.List;
+import java.util.Objects;
 
 /**
  * ExplainAttributes holds the attributes of a SQL statement that is used in the EXPLAIN PLAN result.
@@ -45,6 +47,7 @@ public final class ExplainAttributes
   @Nullable
   private final String replaceTimeChunks;
 
+  @JsonCreator
   public ExplainAttributes(
       @JsonProperty("statementType") final String statementType,
       @JsonProperty("targetDataSource") @Nullable final String targetDataSource,
@@ -117,6 +120,31 @@ public String getReplaceTimeChunks()
     return replaceTimeChunks;
   }
 
+  @Override
+  public boolean equals(Object o)
+  {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    ExplainAttributes that = (ExplainAttributes) o;
+    return Objects.equals(statementType, that.statementType) && Objects.equals(
+        targetDataSource,
+        that.targetDataSource
+    ) && Objects.equals(partitionedBy, that.partitionedBy) && Objects.equals(
+        clusteredBy,
+        that.clusteredBy
+    ) && Objects.equals(replaceTimeChunks, that.replaceTimeChunks);
+  }
+
+  @Override
+  public int hashCode()
+  {
+    return Objects.hash(statementType, targetDataSource, partitionedBy, clusteredBy, replaceTimeChunks);
+  }
+
   @Override
   public String toString()
   {

sql/src/main/java/org/apache/druid/sql/client/Broker.java

@@ -0,0 +1,34 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql.client;
+
+import com.google.inject.BindingAnnotation;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+@BindingAnnotation
+@Target({ElementType.FIELD, ElementType.PARAMETER, ElementType.METHOD})
+@Retention(RetentionPolicy.RUNTIME)
+public @interface Broker
+{
+}

sql/src/main/java/org/apache/druid/sql/client/BrokerClient.java

@@ -0,0 +1,51 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql.client;
+
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.sql.http.ExplainPlanInformation;
+import org.apache.druid.sql.http.SqlQuery;
+import org.apache.druid.sql.http.SqlTaskStatus;
+
+import java.util.List;
+
+/**
+ * High-level Broker client.
+ * <p>
+ * All methods return futures, enabling asynchronous logic. If you want a synchronous response, use
+ * {@code FutureUtils.get} or {@code FutureUtils.getUnchecked}.
+ * Futures resolve to exceptions in the manner described by {@link org.apache.druid.rpc.ServiceClient#asyncRequest}.
+ * </p>
+ * Typically acquired via Guice, where it is registered using {@link org.apache.druid.rpc.guice.ServiceClientModule}.
+ */
+public interface BrokerClient
+{
+  /**
+   * Submit the given {@code sqlQuery} to the Broker's SQL task endpoint.
+   */
+  ListenableFuture<SqlTaskStatus> submitSqlTask(SqlQuery sqlQuery);
+
+  /**
+   * Fetches the explain plan information for the given {@code sqlQuery}.
+   *
+   * @param sqlQuery the SQL query for which the {@code EXPLAIN PLAN FOR} information is to be fetched
+   */
+  ListenableFuture<List<ExplainPlanInformation>> fetchExplainPlanInformation(SqlQuery sqlQuery);
+}

sql/src/main/java/org/apache/druid/sql/client/BrokerClientImpl.java

@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql.client;
+
+import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.util.concurrent.ListenableFuture;
+import org.apache.druid.common.guava.FutureUtils;
+import org.apache.druid.java.util.common.StringUtils;
+import org.apache.druid.java.util.common.jackson.JacksonUtils;
+import org.apache.druid.java.util.http.client.response.BytesFullResponseHandler;
+import org.apache.druid.rpc.RequestBuilder;
+import org.apache.druid.rpc.ServiceClient;
+import org.apache.druid.sql.http.ExplainPlanInformation;
+import org.apache.druid.sql.http.SqlQuery;
+import org.apache.druid.sql.http.SqlTaskStatus;
+import org.jboss.netty.handler.codec.http.HttpMethod;
+
+import java.util.List;
+
+public class BrokerClientImpl implements BrokerClient
+{
+  private final ServiceClient client;
+  private final ObjectMapper jsonMapper;
+
+  public BrokerClientImpl(final ServiceClient client, final ObjectMapper jsonMapper)
+  {
+    this.client = client;
+    this.jsonMapper = jsonMapper;
+  }
+
+  @Override
+  public ListenableFuture<SqlTaskStatus> submitSqlTask(final SqlQuery sqlQuery)
+  {
+    return FutureUtils.transform(
+        client.asyncRequest(
+            new RequestBuilder(HttpMethod.POST, "/druid/v2/sql/task/")
+                .jsonContent(jsonMapper, sqlQuery),
+            new BytesFullResponseHandler()
+        ),
+        holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), SqlTaskStatus.class)
+    );
+  }
+
+  @Override
+  public ListenableFuture<List<ExplainPlanInformation>> fetchExplainPlanInformation(final SqlQuery sqlQuery)
+  {
+    final SqlQuery explainSqlQuery = new SqlQuery(
+        StringUtils.format("EXPLAIN PLAN FOR %s", sqlQuery.getQuery()),
+        null,
+        false,
+        false,
+        false,
+        null,
+        null
+    );
+    return FutureUtils.transform(
+        client.asyncRequest(
+            new RequestBuilder(HttpMethod.POST, "/druid/v2/sql/task/")
+                .jsonContent(jsonMapper, explainSqlQuery),
+            new BytesFullResponseHandler()
+        ),
+        holder -> JacksonUtils.readValue(jsonMapper, holder.getContent(), new TypeReference<List<ExplainPlanInformation>>() {})
+    );
+  }
+}
+

sql/src/main/java/org/apache/druid/sql/guice/BrokerServiceModule.java

@@ -0,0 +1,74 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.druid.sql.guice;
+
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.inject.Provides;
+import org.apache.druid.discovery.DruidNodeDiscoveryProvider;
+import org.apache.druid.discovery.NodeRole;
+import org.apache.druid.guice.LazySingleton;
+import org.apache.druid.guice.ManageLifecycle;
+import org.apache.druid.guice.annotations.EscalatedGlobal;
+import org.apache.druid.guice.annotations.Json;
+import org.apache.druid.rpc.DiscoveryServiceLocator;
+import org.apache.druid.rpc.ServiceClientFactory;
+import org.apache.druid.rpc.ServiceLocator;
+import org.apache.druid.rpc.StandardRetryPolicy;
+import org.apache.druid.rpc.guice.ServiceClientModule;
+import org.apache.druid.sql.client.Broker;
+import org.apache.druid.sql.client.BrokerClient;
+import org.apache.druid.sql.client.BrokerClientImpl;
+
+/**
+ * Module that processes can bind to if they require a {@link BrokerClient}.
+ * <p>
+ * This extend {@link ServiceClientModule} because the {@link BrokerClient} requires
+ * classes present in the sql module.
+ * </p>
+ */
+public class BrokerServiceModule extends ServiceClientModule
+{
+  @Provides
+  @ManageLifecycle
+  @Broker
+  public ServiceLocator makeBrokerServiceLocator(final DruidNodeDiscoveryProvider discoveryProvider)
+  {
+    return new DiscoveryServiceLocator(discoveryProvider, NodeRole.BROKER);
+  }
+
+  @Provides
+  @LazySingleton
+  public BrokerClient makeBrokerClient(
+      @Json final ObjectMapper jsonMapper,
+      @EscalatedGlobal final ServiceClientFactory clientFactory,
+      @Broker final ServiceLocator serviceLocator
+  )
+  {
+    return new BrokerClientImpl(
+        clientFactory.makeClient(
+            NodeRole.BROKER.getJsonName(),
+            serviceLocator,
+            StandardRetryPolicy.builder().maxAttempts(CLIENT_MAX_ATTEMPTS).build()
+        ),
+        jsonMapper
+    );
+  }
+}
+