SNOW-1196041: Ignore default credentials in GCS client (#1656)

Author: sfc-gh-dprzybysz

parent-pom.xml

@@ -38,6 +38,7 @@
     <commons.pool.version>1.5.4</commons.pool.version>
     <c3p0.version>0.9.5.4</c3p0.version>
     <google.api.grpc.version>2.22.0</google.api.grpc.version>
+    <google.auth.library.oauth2.http.version>1.19.0</google.auth.library.oauth2.http.version>
     <google.cloud.core.version>2.21.0</google.cloud.core.version>
     <google.cloud.storage.version>2.22.6</google.cloud.storage.version>
     <google.code.gson.version>2.10.1</google.code.gson.version>
@@ -161,6 +162,11 @@
         <artifactId>proto-google-common-protos</artifactId>
         <version>${google.api.grpc.version}</version>
       </dependency>
+      <dependency>
+        <groupId>com.google.auth</groupId>
+        <artifactId>google-auth-library-oauth2-http</artifactId>
+        <version>${google.auth.library.oauth2.http.version}</version>
+      </dependency>
       <dependency>
         <groupId>com.google.cloud</groupId>
         <artifactId>google-cloud-core</artifactId>
@@ -514,6 +520,10 @@
       <groupId>com.google.api</groupId>
       <artifactId>gax</artifactId>
     </dependency>
+    <dependency>
+      <groupId>com.google.auth</groupId>
+      <artifactId>google-auth-library-oauth2-http</artifactId>
+    </dependency>
     <dependency>
       <groupId>com.google.cloud</groupId>
       <artifactId>google-cloud-core</artifactId>

src/main/java/net/snowflake/client/core/SFBaseSession.java

@@ -137,6 +137,9 @@ public abstract class SFBaseSession {
   // we need to allow for it to maintain backwards compatibility.
   private boolean enablePatternSearch = true;
 
+  /** Disable lookup for default credentials by GCS library */
+  private boolean disableGcsDefaultCredentials = false;
+
   private Map<String, Object> commonParameters;
 
   protected SFBaseSession(SFConnectionHandler sfConnectionHandler) {
@@ -735,6 +738,14 @@ public void setEnablePatternSearch(boolean enablePatternSearch) {
     this.enablePatternSearch = enablePatternSearch;
   }
 
+  public boolean getDisableGcsDefaultCredentials() {
+    return disableGcsDefaultCredentials;
+  }
+
+  public void setDisableGcsDefaultCredentials(boolean disableGcsDefaultCredentials) {
+    this.disableGcsDefaultCredentials = disableGcsDefaultCredentials;
+  }
+
   public int getClientResultChunkSize() {
     return clientResultChunkSize;
   }

src/main/java/net/snowflake/client/core/SFSession.java

@@ -469,6 +469,11 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
             setEnablePatternSearch(getBooleanValue(propertyValue));
           }
           break;
+        case DISABLE_GCS_DEFAULT_CREDENTIALS:
+          if (propertyValue != null) {
+            setDisableGcsDefaultCredentials(getBooleanValue(propertyValue));
+          }
+          break;
 
         default:
           break;

src/main/java/net/snowflake/client/core/SFSessionProperty.java

@@ -78,7 +78,9 @@ public enum SFSessionProperty {
 
   RETRY_TIMEOUT("retryTimeout", false, Integer.class),
 
-  ENABLE_PATTERN_SEARCH("enablePatternSearch", false, Boolean.class);
+  ENABLE_PATTERN_SEARCH("enablePatternSearch", false, Boolean.class),
+
+  DISABLE_GCS_DEFAULT_CREDENTIALS("disableGcsDefaultCredentials", false, Boolean.class);
 
   // property key in string
   private String propertyKey;

src/main/java/net/snowflake/client/jdbc/SnowflakeUtil.java

@@ -36,6 +36,7 @@
 import net.snowflake.client.core.OCSPMode;
 import net.snowflake.client.core.SFBaseSession;
 import net.snowflake.client.core.SFSessionProperty;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
 import net.snowflake.common.core.SqlState;
@@ -50,7 +51,7 @@
  */
 public class SnowflakeUtil {
 
-  static final SFLogger logger = SFLoggerFactory.getLogger(RestRequest.class);
+  private static final SFLogger logger = SFLoggerFactory.getLogger(SnowflakeUtil.class);
 
   /** Additional data types not covered by standard JDBC */
   public static final int EXTRA_TYPES_TIMESTAMP_LTZ = 50000;
@@ -754,4 +755,21 @@ public static Time getTimeInSessionTimezone(Long time, int nanos) {
     ts.setTime(c.getTimeInMillis());
     return ts;
   }
+
+  /**
+   * Helper function to convert system properties to boolean
+   *
+   * @param systemProperty name of the system property
+   * @param defaultValue default value used
+   * @return the value of the system property as boolean, else the default value
+   */
+  @SnowflakeJdbcInternalApi
+  public static boolean convertSystemPropertyToBooleanValue(
+      String systemProperty, boolean defaultValue) {
+    String systemPropertyValue = systemGetProperty(systemProperty);
+    if (systemPropertyValue != null) {
+      return Boolean.parseBoolean(systemPropertyValue);
+    }
+    return defaultValue;
+  }
 }

src/main/java/net/snowflake/client/jdbc/cloud/storage/SnowflakeGCSClient.java

@@ -12,6 +12,8 @@
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.api.gax.paging.Page;
 import com.google.api.gax.rpc.FixedHeaderProvider;
+import com.google.auth.oauth2.AccessToken;
+import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.storage.Blob;
 import com.google.cloud.storage.BlobId;
 import com.google.cloud.storage.BlobInfo;
@@ -43,6 +45,7 @@
 import net.snowflake.client.core.ObjectMapperFactory;
 import net.snowflake.client.core.SFSession;
 import net.snowflake.client.core.SFSessionProperty;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
 import net.snowflake.client.jdbc.ErrorCode;
 import net.snowflake.client.jdbc.FileBackedOutputStream;
 import net.snowflake.client.jdbc.MatDesc;
@@ -74,6 +77,10 @@
  * @author ppaulus
  */
 public class SnowflakeGCSClient implements SnowflakeStorageClient {
+  @SnowflakeJdbcInternalApi
+  public static final String DISABLE_GCS_DEFAULT_CREDENTIALS_PROPERTY_NAME =
+      "net.snowflake.jdbc.disableGcsDefaultCredentials";
+
   private static final String GCS_ENCRYPTIONDATAPROP = "encryptiondata";
   private static final String localFileSep = systemGetProperty("file.separator");
   private static final String GCS_METADATA_PREFIX = "x-goog-meta-";
@@ -1200,13 +1207,19 @@ private void setupGCSClient(
     try {
       String accessToken = (String) stage.getCredentials().get("GCS_ACCESS_TOKEN");
       if (accessToken != null) {
+        // We are authenticated with an oauth access token.
+        StorageOptions.Builder builder = StorageOptions.newBuilder();
+        if (areDisabledGcsDefaultCredentials(session)) {
+          logger.debug(
+              "Adding explicit credentials to avoid default credential lookup by the GCS client");
+          builder.setCredentials(GoogleCredentials.create(new AccessToken(accessToken, null)));
+        }
+
         // Using GoogleCredential with access token will cause IllegalStateException when the token
         // is expired and trying to refresh, which cause error cannot be caught. Instead, set a
         // header so we can caught the error code.
-
-        // We are authenticated with an oauth access token.
         this.gcsClient =
-            StorageOptions.newBuilder()
+            builder
                 .setHeaderProvider(
                     FixedHeaderProvider.create("Authorization", "Bearer " + accessToken))
                 .build()
@@ -1234,6 +1247,12 @@ private void setupGCSClient(
     }
   }
 
+  private static boolean areDisabledGcsDefaultCredentials(SFSession session) {
+    return session != null && session.getDisableGcsDefaultCredentials()
+        || SnowflakeUtil.convertSystemPropertyToBooleanValue(
+            DISABLE_GCS_DEFAULT_CREDENTIALS_PROPERTY_NAME, false);
+  }
+
   private static boolean isSuccessStatusCode(int code) {
     return code < 300 && code >= 200;
   }

src/main/java/net/snowflake/client/jdbc/cloud/storage/StorageClientFactory.java

@@ -22,7 +22,7 @@
  */
 public class StorageClientFactory {
 
-  private static final SFLogger logger = SFLoggerFactory.getLogger(SnowflakeS3Client.class);
+  private static final SFLogger logger = SFLoggerFactory.getLogger(StorageClientFactory.class);
 
   private static StorageClientFactory factory;
 

src/test/java/net/snowflake/client/jdbc/ConnectionLatestIT.java

@@ -403,10 +403,10 @@ public void testQueryStatusErrorMessageAndErrorCodeChangeOnAsyncQuery() throws S
       SnowflakeResultSet sfResultSet = rs1.unwrap(SnowflakeResultSet.class);
       // status should change state to RUNNING and then to SUCCESS
       await()
-          .atMost(Duration.ofSeconds(5))
+          .atMost(Duration.ofSeconds(10))
           .until(() -> sfResultSet.getStatusV2().getStatus(), equalTo(QueryStatus.RUNNING));
       await()
-          .atMost(Duration.ofSeconds(5))
+          .atMost(Duration.ofSeconds(50))
           .until(() -> sfResultSet.getStatusV2().getStatus(), equalTo(QueryStatus.SUCCESS));
     }
   }