Set Presto User in header of HTTP request of Presto Node Info API (#451)

qubole · Sep 11, 2020 · e6e4ad9 · e6e4ad9
1 parent 7e797a5
commit e6e4ad9
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 12 deletions.
diff --git a/rubix-presto/src/main/java/com/qubole/rubix/presto/PrestoClusterManager.java b/rubix-presto/src/main/java/com/qubole/rubix/presto/PrestoClusterManager.java
@@ -38,11 +38,14 @@
 import java.util.Objects;
 import java.util.Set;
 
+import static java.util.Objects.requireNonNull;
+
 /**
  * Created by stagra on 14/1/16.
  */
 public class PrestoClusterManager extends ClusterManager
 {
+  private static final String DEFAULT_USER = "rubix";
   private int serverPort = 8081;
   private String serverAddress = "localhost";
 
@@ -69,9 +72,7 @@ public List<String> getNodesInternal()
       URL allNodesRequest = getNodeUrl();
       URL failedNodesRequest = getFailedNodeUrl();
 
-      HttpURLConnection allHttpCon = (HttpURLConnection) allNodesRequest.openConnection();
-      allHttpCon.setConnectTimeout(500); //ms
-      allHttpCon.setRequestMethod("GET");
+      HttpURLConnection allHttpCon = getHttpURLConnection(allNodesRequest);
 
       int allNodesResponseCode = allHttpCon.getResponseCode();
 
@@ -105,9 +106,7 @@ public List<String> getNodesInternal()
         allHttpCon.disconnect();
       }
 
-      HttpURLConnection failHttpConn = (HttpURLConnection) failedNodesRequest.openConnection();
-      failHttpConn.setConnectTimeout(500);    //ms
-      failHttpConn.setRequestMethod("GET");
+      HttpURLConnection failHttpConn = getHttpURLConnection(failedNodesRequest);
       int failedNodesResponseCode = failHttpConn.getResponseCode();
       // check on failed nodes
       try {
@@ -164,6 +163,17 @@ public List<String> getNodesInternal()
     }
   }
 
+  private HttpURLConnection getHttpURLConnection(URL urlRequest)
+          throws IOException
+  {
+    requireNonNull(urlRequest, "urlRequest is null");
+    HttpURLConnection allHttpCon = (HttpURLConnection) urlRequest.openConnection();
+    allHttpCon.setConnectTimeout(500); //ms
+    allHttpCon.setRequestMethod("GET");
+    allHttpCon.setRequestProperty("X-Presto-User", DEFAULT_USER);
+    return allHttpCon;
+  }
+
   @Override
   public ClusterType getClusterType()
   {

diff --git a/rubix-prestosql/src/main/java/com/qubole/rubix/prestosql/PrestoClusterManager.java b/rubix-prestosql/src/main/java/com/qubole/rubix/prestosql/PrestoClusterManager.java
@@ -50,6 +50,7 @@
  */
 public class PrestoClusterManager extends ClusterManager
 {
+  private static final String DEFAULT_USER = "rubix";
   private int serverPort = 8081;
   private String serverAddress = "localhost";
 
@@ -80,9 +81,7 @@ public List<String> getNodesInternal()
       URL allNodesRequest = getNodeUrl();
       URL failedNodesRequest = getFailedNodeUrl();
 
-      HttpURLConnection allHttpCon = (HttpURLConnection) allNodesRequest.openConnection();
-      allHttpCon.setConnectTimeout(500); //ms
-      allHttpCon.setRequestMethod("GET");
+      HttpURLConnection allHttpCon = getHttpURLConnection(allNodesRequest);
 
       int allNodesResponseCode = allHttpCon.getResponseCode();
 
@@ -116,9 +115,7 @@ public List<String> getNodesInternal()
         allHttpCon.disconnect();
       }
 
-      HttpURLConnection failHttpConn = (HttpURLConnection) failedNodesRequest.openConnection();
-      failHttpConn.setConnectTimeout(500);    //ms
-      failHttpConn.setRequestMethod("GET");
+      HttpURLConnection failHttpConn = getHttpURLConnection(failedNodesRequest);
       int failedNodesResponseCode = failHttpConn.getResponseCode();
       // check on failed nodes
       try {
@@ -172,6 +169,17 @@ public List<String> getNodesInternal()
     }
   }
 
+  private HttpURLConnection getHttpURLConnection(URL urlRequest)
+          throws IOException
+  {
+    requireNonNull(urlRequest, "urlRequest is null");
+    HttpURLConnection allHttpCon = (HttpURLConnection) urlRequest.openConnection();
+    allHttpCon.setConnectTimeout(500); //ms
+    allHttpCon.setRequestMethod("GET");
+    allHttpCon.setRequestProperty("X-Presto-User", DEFAULT_USER);
+    return allHttpCon;
+  }
+
   private List<String> getNodesFromNodeManager()
   {
     requireNonNull(nodeManager, "nodeManager is null");