SNOW-887071 : Added connection property to disable PUT/GET command (#…

…1523) * Added connection property to disable PUT/GET command * Refactor the code as per review comment * Added test for PUT and GET for disable state
snowflakedb · Oct 5, 2023 · c997f9a · c997f9a
1 parent a1ba3d1
commit c997f9a
Show file tree

Hide file tree

Showing 6 changed files with 121 additions and 0 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SFBaseSession.java b/src/main/java/net/snowflake/client/core/SFBaseSession.java
@@ -117,6 +117,12 @@ public abstract class SFBaseSession {
   // Whether enable returning timestamp with timezone as data type
   private boolean enableReturnTimestampWithTimeZone = true;
 
+  // Server side value
+  private boolean jdbcEnablePutGet = true;
+
+  // Connection string setting
+  private boolean enablePutGet = true;
+
   private Map<String, Object> commonParameters;
 
   protected SFBaseSession(SFConnectionHandler sfConnectionHandler) {
@@ -682,6 +688,22 @@ public void setQueryContextCacheSize(int queryContextCacheSize) {
     this.queryContextCacheSize = queryContextCacheSize;
   }
 
+  public boolean getJdbcEnablePutGet() {
+    return jdbcEnablePutGet;
+  }
+
+  public void setJdbcEnablePutGet(boolean jdbcEnablePutGet) {
+    this.jdbcEnablePutGet = jdbcEnablePutGet;
+  }
+
+  public boolean getEnablePutGet() {
+    return enablePutGet;
+  }
+
+  public boolean putEnablePutGet(boolean enablePutGet) {
+    return this.enablePutGet = enablePutGet;
+  }
+
   public int getClientResultChunkSize() {
     return clientResultChunkSize;
   }

diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java
@@ -363,6 +363,12 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
           }
           break;
 
+        case ENABLE_PUT_GET:
+          if (propertyValue != null) {
+            putEnablePutGet(getBooleanValue(propertyValue));
+          }
+          break;
+
         default:
           break;
       }

diff --git a/src/main/java/net/snowflake/client/core/SFSessionProperty.java b/src/main/java/net/snowflake/client/core/SFSessionProperty.java
@@ -71,6 +71,8 @@ public enum SFSessionProperty {
 
   MAX_HTTP_RETRIES("maxHttpRetries", false, Integer.class),
 
+  ENABLE_PUT_GET("enablePutGet", false, Boolean.class),
+
   PUT_GET_MAX_RETRIES("putGetMaxRetries", false, Integer.class);
 
   // property key in string

diff --git a/src/main/java/net/snowflake/client/core/SFStatement.java b/src/main/java/net/snowflake/client/core/SFStatement.java
@@ -123,6 +123,13 @@ private SFBaseResultSet executeQuery(
 
     // snowflake specific client side commands
     if (isFileTransfer(trimmedSql)) {
+      // Server side value or Connection string value is false then disable the PUT/GET command
+      if ((session != null && !(session.getJdbcEnablePutGet() && session.getEnablePutGet()))) {
+        // PUT/GET command disabled either on server side or in the client connection string
+        logger.debug("Executing file transfer locally is disabled: {}", sql);
+        throw new SnowflakeSQLException("File transfers have been disabled.");
+      }
+
       // PUT/GET command
       logger.debug("Executing file transfer locally: {}", sql);
 

diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/src/main/java/net/snowflake/client/core/SessionUtil.java
@@ -84,6 +84,7 @@ public class SessionUtil {
   public static final String CLIENT_MEMORY_LIMIT_JVM = "net.snowflake.jdbc.clientMemoryLimit";
   public static final String CLIENT_MEMORY_LIMIT = "CLIENT_MEMORY_LIMIT";
   public static final String QUERY_CONTEXT_CACHE_SIZE = "QUERY_CONTEXT_CACHE_SIZE";
+  public static final String JDBC_ENABLE_PUT_GET = "JDBC_ENABLE_PUT_GET";
   public static final String CLIENT_PREFETCH_THREADS_JVM =
       "net.snowflake.jdbc.clientPrefetchThreads";
   public static final String CLIENT_PREFETCH_THREADS = "CLIENT_PREFETCH_THREADS";
@@ -1547,6 +1548,10 @@ static void updateSfDriverParamValues(Map<String, Object> parameters, SFBaseSess
         if (session != null) {
           session.setQueryContextCacheSize((int) entry.getValue());
         }
+      } else if (JDBC_ENABLE_PUT_GET.equalsIgnoreCase(entry.getKey())) {
+        if (session != null) {
+          session.setJdbcEnablePutGet(SFLoginInput.getBooleanValue(entry.getValue()));
+        }
       } else {
         if (session != null) {
           session.setOtherParameter(entry.getKey(), entry.getValue());

diff --git a/src/test/java/net/snowflake/client/jdbc/SnowflakeDriverLatestIT.java b/src/test/java/net/snowflake/client/jdbc/SnowflakeDriverLatestIT.java
@@ -536,6 +536,85 @@ public void testPutOverwriteFalseNoDigest() throws Throwable {
     }
   }
 
+  /**
+   * Tests PUT disable test
+   *
+   * @throws Throwable
+   */
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class)
+  public void testPutDisable() throws Throwable {
+    Connection connection = null;
+    Statement statement = null;
+
+    // create a file
+    File file = tmpFolder.newFile("testfile99.csv");
+    BufferedWriter bw = new BufferedWriter(new FileWriter(file));
+    bw.write("This content won't be uploaded as PUT is disabled.");
+    bw.close();
+
+    String sourceFilePathOriginal = file.getCanonicalPath();
+
+    Properties paramProperties = new Properties();
+    paramProperties.put("enablePutGet", false);
+
+    List<String> accounts = Arrays.asList(null, "s3testaccount", "azureaccount", "gcpaccount");
+    for (int i = 0; i < accounts.size(); i++) {
+      try {
+        connection = getConnection(accounts.get(i), paramProperties);
+
+        statement = connection.createStatement();
+
+        statement.execute("PUT file://" + sourceFilePathOriginal + " @testPutGet_disable_stage");
+
+        assertTrue("Shouldn't come here", false);
+      } catch (Exception ex) {
+        // Expected
+        assertTrue(ex.getMessage().equalsIgnoreCase("File transfers have been disabled."));
+      } finally {
+        statement.close();
+      }
+    }
+  }
+
+  /**
+   * Tests GET disable test
+   *
+   * @throws Throwable
+   */
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubAction.class)
+  public void testGetDisable() throws Throwable {
+    Connection connection = null;
+    Statement statement = null;
+
+    // create a folder
+    File destFolder = tmpFolder.newFolder();
+    String destFolderCanonicalPath = destFolder.getCanonicalPath();
+
+    Properties paramProperties = new Properties();
+    paramProperties.put("enablePutGet", false);
+
+    List<String> accounts = Arrays.asList(null, "s3testaccount", "azureaccount", "gcpaccount");
+    for (int i = 0; i < accounts.size(); i++) {
+      try {
+        connection = getConnection(accounts.get(i), paramProperties);
+
+        statement = connection.createStatement();
+
+        statement.execute(
+            "GET @testPutGet_disable_stage 'file://" + destFolderCanonicalPath + "' parallel=8");
+
+        assertTrue("Shouldn't come here", false);
+      } catch (Exception ex) {
+        // Expected
+        assertTrue(ex.getMessage().equalsIgnoreCase("File transfers have been disabled."));
+      } finally {
+        statement.close();
+      }
+    }
+  }
+
   /**
    * Test NULL in LIMIT and OFFSET with Snow-76376 enabled this should be handled as without LIMIT
    * and OFFSET