Review comments

snowflakedb · Aug 3, 2023 · f05e91f · f05e91f
1 parent 617821b
commit f05e91f
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 21 deletions.
diff --git a/src/main/java/net/snowflake/client/core/HttpUtil.java b/src/main/java/net/snowflake/client/core/HttpUtil.java
@@ -899,4 +899,19 @@ static int convertSystemPropertyToIntValue(String systemProperty, int defaultVal
     }
     return returnVal;
   }
+
+  /**
+   * Helper function to attach additional headers to a request if present. This takes a (nullable)
+   * map of headers in <name,value> format and adds them to the incoming request using addHeader.
+   *
+   * @param request The request to add headers to. Must not be null.
+   * @param additionalHeaders The headers to add. May be null.
+   */
+  static void applyAdditionalHeaders(
+      HttpRequestBase request, Map<String, String> additionalHeaders) {
+    Map<String, String> additionalHttpHeaders = additionalHeaders;
+    if (additionalHttpHeaders != null) {
+      additionalHttpHeaders.forEach(request::addHeader);
+    }
+  }
 }
diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java
@@ -588,10 +588,7 @@ private static SFLoginOutput newSession(
       postRequest = new HttpPost(loginURI);
 
       // Add custom headers before adding common headers
-      Map<String, String> additionalHttpHeaders = loginInput.getAdditionalHttpHeaders();
-      if (additionalHttpHeaders != null) {
-        additionalHttpHeaders.forEach(postRequest::addHeader);
-      }
+      HttpUtil.applyAdditionalHeaders(postRequest, loginInput.getAdditionalHttpHeaders());
 
       // attach the login info json body to the post request
       StringEntity input = new StringEntity(json, StandardCharsets.UTF_8);
@@ -904,10 +901,7 @@ private static SFLoginOutput tokenRequest(SFLoginInput loginInput, TokenRequestT
       postRequest = new HttpPost(uriBuilder.build());
 
       // Add custom headers before adding common headers
-      Map<String, String> additionalHttpHeaders = loginInput.getAdditionalHttpHeaders();
-      if (additionalHttpHeaders != null) {
-        additionalHttpHeaders.forEach(postRequest::addHeader);
-      }
+      HttpUtil.applyAdditionalHeaders(postRequest, loginInput.getAdditionalHttpHeaders());
     } catch (URISyntaxException ex) {
       logger.error("Exception when creating http request", ex);
 
@@ -1021,10 +1015,7 @@ static void closeSession(SFLoginInput loginInput) throws SFException, SnowflakeS
       postRequest = new HttpPost(uriBuilder.build());
 
       // Add custom headers before adding common headers
-      Map<String, String> additionalHttpHeaders = loginInput.getAdditionalHttpHeaders();
-      if (additionalHttpHeaders != null) {
-        additionalHttpHeaders.forEach(postRequest::addHeader);
-      }
+      HttpUtil.applyAdditionalHeaders(postRequest, loginInput.getAdditionalHttpHeaders());
 
       postRequest.setHeader(
           SF_HEADER_AUTHORIZATION,

diff --git a/src/main/java/net/snowflake/client/core/StmtUtil.java b/src/main/java/net/snowflake/client/core/StmtUtil.java
@@ -307,9 +307,9 @@ public static StmtOutput execute(StmtInput stmtInput, ExecTimeTelemetryData exec
         }
 
         httpRequest = new HttpPost(uriBuilder.build());
-        if (stmtInput.additionalHttpHeaders != null) {
-          stmtInput.additionalHttpHeaders.forEach(httpRequest::addHeader);
-        }
+
+        // Add custom headers before adding common headers
+        HttpUtil.applyAdditionalHeaders(httpRequest, stmtInput.additionalHttpHeaders);
 
         /*
          * sequence id is only needed for old query API, when old query API
@@ -600,9 +600,8 @@ protected static String getQueryResult(String getResultPath, StmtInput stmtInput
       uriBuilder.addParameter(SF_QUERY_REQUEST_ID, UUIDUtils.getUUID().toString());
 
       httpRequest = new HttpGet(uriBuilder.build());
-      if (stmtInput.additionalHttpHeaders != null) {
-        stmtInput.additionalHttpHeaders.forEach(httpRequest::addHeader);
-      }
+      // Add custom headers before adding common headers
+      HttpUtil.applyAdditionalHeaders(httpRequest, stmtInput.additionalHttpHeaders);
 
       httpRequest.addHeader("accept", stmtInput.mediaType);
 
@@ -704,9 +703,8 @@ public static void cancel(StmtInput stmtInput) throws SFException, SnowflakeSQLE
       uriBuilder.addParameter(SF_QUERY_REQUEST_ID, UUIDUtils.getUUID().toString());
 
       httpRequest = new HttpPost(uriBuilder.build());
-      if (stmtInput.additionalHttpHeaders != null) {
-        stmtInput.additionalHttpHeaders.forEach(httpRequest::addHeader);
-      }
+      // Add custom headers before adding common headers
+      HttpUtil.applyAdditionalHeaders(httpRequest, stmtInput.additionalHttpHeaders);
 
       /*
        * The JSON input has two fields: sqlText and requestId