diff --git a/src/main/java/net/snowflake/client/core/SFBaseSession.java b/src/main/java/net/snowflake/client/core/SFBaseSession.java index fa32ac19a..22b4e553f 100644 --- 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 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 index c85aa30a8..5537cde72 100644 --- 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 index 48f80db72..8609cbdc5 100644 --- 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 index 56b9d2362..b251d2f25 100644 --- 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 index c4a3ead7a..f50586b42 100644 --- 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 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 index 3de23b2bb..2cea0e5fb 100644 --- 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 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 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