From 2a13e1926a20d6f7bc52a4b101910a7e6bb406ed Mon Sep 17 00:00:00 2001 From: Ilesh garish Date: Wed, 4 Oct 2023 04:29:29 +0000 Subject: [PATCH 1/3] Added connection property to disable PUT/GET command --- .../snowflake/client/core/SFBaseSession.java | 18 ++++++++++++++++++ .../net/snowflake/client/core/SFSession.java | 6 ++++++ .../client/core/SFSessionProperty.java | 2 ++ .../net/snowflake/client/core/SFStatement.java | 8 ++++++++ .../net/snowflake/client/core/SessionUtil.java | 5 +++++ 5 files changed, 39 insertions(+) diff --git a/src/main/java/net/snowflake/client/core/SFBaseSession.java b/src/main/java/net/snowflake/client/core/SFBaseSession.java index fa32ac19a..173ecb215 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 + protected boolean enablePutGet = true; + private Map commonParameters; protected SFBaseSession(SFConnectionHandler sfConnectionHandler) { @@ -682,6 +688,18 @@ 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 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..6300acc5e 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) { + enablePutGet = 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..ddc11c6d0 100644 --- a/src/main/java/net/snowflake/client/core/SFStatement.java +++ b/src/main/java/net/snowflake/client/core/SFStatement.java @@ -123,6 +123,14 @@ private SFBaseResultSet executeQuery( // snowflake specific client side commands if (isFileTransfer(trimmedSql)) { + if ((session != null && !session.getJdbcEnablePutGet()) // Server side value + || (session != null && !session.getEnablePutGet()) // Connection string value + ) { + // 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()); From bfc59276089065b28a900372f20eef7ccb286de9 Mon Sep 17 00:00:00 2001 From: Ilesh garish Date: Thu, 5 Oct 2023 01:54:14 +0000 Subject: [PATCH 2/3] Refactor the code as per review comment --- src/main/java/net/snowflake/client/core/SFBaseSession.java | 6 +++++- src/main/java/net/snowflake/client/core/SFSession.java | 2 +- src/main/java/net/snowflake/client/core/SFStatement.java | 5 ++--- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/src/main/java/net/snowflake/client/core/SFBaseSession.java b/src/main/java/net/snowflake/client/core/SFBaseSession.java index 173ecb215..22b4e553f 100644 --- a/src/main/java/net/snowflake/client/core/SFBaseSession.java +++ b/src/main/java/net/snowflake/client/core/SFBaseSession.java @@ -121,7 +121,7 @@ public abstract class SFBaseSession { private boolean jdbcEnablePutGet = true; // Connection string setting - protected boolean enablePutGet = true; + private boolean enablePutGet = true; private Map commonParameters; @@ -700,6 +700,10 @@ 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 6300acc5e..5537cde72 100644 --- a/src/main/java/net/snowflake/client/core/SFSession.java +++ b/src/main/java/net/snowflake/client/core/SFSession.java @@ -365,7 +365,7 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro case ENABLE_PUT_GET: if (propertyValue != null) { - enablePutGet = getBooleanValue(propertyValue); + putEnablePutGet(getBooleanValue(propertyValue)); } break; diff --git a/src/main/java/net/snowflake/client/core/SFStatement.java b/src/main/java/net/snowflake/client/core/SFStatement.java index ddc11c6d0..b251d2f25 100644 --- a/src/main/java/net/snowflake/client/core/SFStatement.java +++ b/src/main/java/net/snowflake/client/core/SFStatement.java @@ -123,9 +123,8 @@ private SFBaseResultSet executeQuery( // snowflake specific client side commands if (isFileTransfer(trimmedSql)) { - if ((session != null && !session.getJdbcEnablePutGet()) // Server side value - || (session != null && !session.getEnablePutGet()) // Connection string value - ) { + // 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."); From c612916dbcef9c4a05dfd31773b629d5be564ef8 Mon Sep 17 00:00:00 2001 From: Ilesh garish Date: Thu, 5 Oct 2023 03:55:01 +0000 Subject: [PATCH 3/3] Added test for PUT and GET for disable state --- .../client/jdbc/SnowflakeDriverLatestIT.java | 79 +++++++++++++++++++ 1 file changed, 79 insertions(+) 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