Avoid hard-coded http scheme for Pinot controller url

trinodb · Aug 10, 2021 · 9d1e6ac · 9d1e6ac
1 parent a46e7cd
commit 9d1e6ac
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotConfig.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/PinotConfig.java
@@ -21,18 +21,20 @@
 
 import javax.validation.constraints.NotNull;
 
+import java.net.URI;
 import java.util.List;
 import java.util.concurrent.TimeUnit;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableList.toImmutableList;
 
 public class PinotConfig
 {
     private static final Splitter LIST_SPLITTER = Splitter.on(",").trimResults().omitEmptyStrings();
 
     private int maxConnectionsPerServer = 30;
 
-    private List<String> controllerUrls = ImmutableList.of();
+    private List<URI> controllerUrls = ImmutableList.of();
 
     private Duration idleTimeout = new Duration(5, TimeUnit.MINUTES);
     private Duration connectionTimeout = new Duration(1, TimeUnit.MINUTES);
@@ -52,15 +54,17 @@ public class PinotConfig
     private int maxRowsPerSplitForSegmentQueries = 50_000;
 
     @NotNull
-    public List<String> getControllerUrls()
+    public List<URI> getControllerUrls()
     {
         return controllerUrls;
     }
 
     @Config("pinot.controller-urls")
     public PinotConfig setControllerUrls(String controllerUrl)
     {
-        this.controllerUrls = LIST_SPLITTER.splitToList(controllerUrl);
+        this.controllerUrls = LIST_SPLITTER.splitToList(controllerUrl).stream()
+                .map(PinotConfig::stringToUri)
+                .collect(toImmutableList());
         return this;
     }
 
@@ -257,4 +261,12 @@ public PinotConfig setMaxRowsPerSplitForSegmentQueries(int maxRowsPerSplitForSeg
         this.maxRowsPerSplitForSegmentQueries = maxRowsPerSplitForSegmentQueries;
         return this;
     }
+
+    private static URI stringToUri(String server)
+    {
+        if (server.startsWith("http://") || server.startsWith("https://")) {
+            return URI.create(server);
+        }
+        return URI.create("http://" + server);
+    }
 }
diff --git a/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/client/PinotClient.java b/plugin/trino-pinot/src/main/java/io/trino/plugin/pinot/client/PinotClient.java
@@ -68,6 +68,7 @@
 
 import static com.google.common.base.Preconditions.checkState;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom;
 import static io.airlift.http.client.JsonResponseHandler.createJsonResponseHandler;
 import static io.airlift.json.JsonCodec.listJsonCodec;
 import static io.airlift.json.JsonCodec.mapJsonCodec;
@@ -96,7 +97,7 @@ public class PinotClient
     private static final String REQUEST_PAYLOAD_TEMPLATE = "{\"sql\" : \"%s\" }";
     private static final String QUERY_URL_TEMPLATE = "http://%s/query/sql";
 
-    private final List<String> controllerUrls;
+    private final List<URI> controllerUrls;
     private final HttpClient httpClient;
     private final PinotHostMapper pinotHostMapper;
 
@@ -178,7 +179,7 @@ protected <T> T doHttpActionWithHeadersJson(Request.Builder requestBuilder, Opti
     private <T> T sendHttpGetToControllerJson(String path, JsonCodec<T> codec)
     {
         return doHttpActionWithHeadersJson(
-                Request.Builder.prepareGet().setUri(URI.create(format("http://%s/%s", getControllerUrl(), path))),
+                Request.Builder.prepareGet().setUri(uriBuilderFrom(getControllerUrl()).appendPath(path).build()),
                 Optional.empty(),
                 codec);
     }
@@ -191,7 +192,7 @@ private <T> T sendHttpGetToBrokerJson(String table, String path, JsonCodec<T> co
                 codec);
     }
 
-    private String getControllerUrl()
+    private URI getControllerUrl()
     {
         return controllerUrls.get(ThreadLocalRandom.current().nextInt(controllerUrls.size()));
     }