Add BrokerClient implementation (#17382)

This patch is extracted from PR 17353. Changes: - Added BrokerClient and BrokerClientImpl to the sql package that leverages the ServiceClient functionality; similar to OverlordClient and CoordinatorClient implementations in the server module. - For now, only two broker API stubs are added: submitSqlTask() and fetchExplainPlan(). - Added a new POJO class ExplainPlan that encapsulates explain plan info. - Deprecated org.apache.druid.discovery.BrokerClient in favor of the new BrokerClient in this patch. - Clean up ExplainAttributesTest a bit and added serde verification.
apache · Oct 21, 2024 · 187e21a · 187e21a
1 parent 5da9949
commit 187e21a
Show file tree

Hide file tree

Showing 13 changed files with 953 additions and 33 deletions.
diff --git a/server/src/main/java/org/apache/druid/discovery/BrokerClient.java b/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;

diff --git a/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java b/server/src/main/java/org/apache/druid/rpc/guice/ServiceClientModule.java
@@ -47,8 +47,8 @@
 
 public class ServiceClientModule implements DruidModule
 {
+  public static final int CLIENT_MAX_ATTEMPTS = 6;
   private static final int CONNECT_EXEC_THREADS = 4;
-  private static final int CLIENT_MAX_ATTEMPTS = 6;
 
   @Override
   public void configure(Binder binder)
@@ -59,11 +59,9 @@ public void configure(Binder binder)
   @Provides
   @LazySingleton
   @EscalatedGlobal
-  public ServiceClientFactory makeServiceClientFactory(@EscalatedGlobal final HttpClient httpClient)
+  public ServiceClientFactory getServiceClientFactory(@EscalatedGlobal final HttpClient httpClient)
   {
-    final ScheduledExecutorService connectExec =
-        ScheduledExecutors.fixed(CONNECT_EXEC_THREADS, "ServiceClientFactory-%d");
-    return new ServiceClientFactoryImpl(httpClient, connectExec);
+    return makeServiceClientFactory(httpClient);
   }
 
   @Provides
@@ -117,4 +115,11 @@ public CoordinatorClient makeCoordinatorClient(
         jsonMapper
     );
   }
+
+  public static ServiceClientFactory makeServiceClientFactory(@EscalatedGlobal final HttpClient httpClient)
+  {
+    final ScheduledExecutorService connectExec =
+        ScheduledExecutors.fixed(CONNECT_EXEC_THREADS, "ServiceClientFactory-%d");
+    return new ServiceClientFactoryImpl(httpClient, connectExec);
+  }
 }
diff --git a/server/src/test/java/org/apache/druid/rpc/guice/ServiceClientModuleTest.java b/server/src/test/java/org/apache/druid/rpc/guice/ServiceClientModuleTest.java
@@ -0,0 +1,99 @@
+/*
+ * 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.rpc.guice;
+
+import com.google.common.collect.ImmutableList;
+import com.google.inject.Guice;
+import com.google.inject.Injector;
+import org.apache.druid.client.coordinator.CoordinatorClient;
+import org.apache.druid.discovery.DruidNodeDiscoveryProvider;
+import org.apache.druid.guice.DruidGuiceExtensions;
+import org.apache.druid.guice.LifecycleModule;
+import org.apache.druid.guice.annotations.EscalatedGlobal;
+import org.apache.druid.jackson.JacksonModule;
+import org.apache.druid.java.util.http.client.HttpClient;
+import org.apache.druid.rpc.ServiceClientFactory;
+import org.apache.druid.rpc.ServiceLocator;
+import org.apache.druid.rpc.indexing.OverlordClient;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+import static org.junit.Assert.assertNotNull;
+
+public class ServiceClientModuleTest
+{
+  private Injector injector;
+
+  @Rule
+  public MockitoRule mockitoRule = MockitoJUnit.rule();
+
+  @Mock
+  private HttpClient httpClient;
+
+  @Mock
+  private DruidNodeDiscoveryProvider discoveryProvider;
+
+  @Mock
+  private ServiceLocator serviceLocator;
+
+  @Mock
+  private ServiceClientFactory serviceClientFactory;
+
+  @Before
+  public void setUp()
+  {
+    injector = Guice.createInjector(
+        ImmutableList.of(
+            new DruidGuiceExtensions(),
+            new LifecycleModule(),
+          new JacksonModule(),
+          new ServiceClientModule(),
+          binder -> {
+            binder.bind(HttpClient.class).annotatedWith(EscalatedGlobal.class).toInstance(httpClient);
+            binder.bind(ServiceLocator.class).toInstance(serviceLocator);
+            binder.bind(DruidNodeDiscoveryProvider.class).toInstance(discoveryProvider);
+            binder.bind(ServiceClientFactory.class).toInstance(serviceClientFactory);
+          }
+          )
+    );
+  }
+
+  @Test
+  public void testGetServiceClientFactory()
+  {
+    assertNotNull(injector.getInstance(ServiceClientFactory.class));
+  }
+
+  @Test
+  public void testGetOverlordClient()
+  {
+    assertNotNull(injector.getInstance(OverlordClient.class));
+  }
+
+  @Test
+  public void testGetCoordinatorClient()
+  {
+    assertNotNull(injector.getInstance(CoordinatorClient.class));
+  }
+}
diff --git a/sql/src/main/java/org/apache/druid/sql/calcite/planner/ExplainAttributes.java b/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,29 @@ 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()
   {

diff --git a/sql/src/main/java/org/apache/druid/sql/client/Broker.java b/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
+{
+}
diff --git a/sql/src/main/java/org/apache/druid/sql/client/BrokerClient.java b/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.ExplainPlan;
+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 for the given {@code sqlQuery} from the Broker's SQL task endpoint.
+   *
+   * @param sqlQuery the SQL query for which the {@code EXPLAIN PLAN FOR} information is to be fetched
+   */
+  ListenableFuture<List<ExplainPlan>> fetchExplainPlan(SqlQuery sqlQuery);
+}
diff --git a/sql/src/main/java/org/apache/druid/sql/client/BrokerClientImpl.java b/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.ExplainPlan;
+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<ExplainPlan>> fetchExplainPlan(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<ExplainPlan>>() {})
+    );
+  }
+}
+