Merge branch 'master' into oauth-code-flow

snowflakedb · Jan 10, 2025 · fa3761b · fa3761b
2 parents cdff04e + dcac111
commit fa3761b
Show file tree

Hide file tree

Showing 16 changed files with 195 additions and 23 deletions.
diff --git a/parent-pom.xml b/parent-pom.xml
@@ -486,6 +486,17 @@
         <artifactId>arrow-memory-netty-buffer-patch</artifactId>
         <version>${arrow.version}</version>
         <scope>runtime</scope>
+        <exclusions>
+          <!-- We explicitly add netty dependencies in higher version, here only excluding for removing snyk issue since snyk does not understand our dependency tree -->
+          <exclusion>
+            <groupId>io.netty</groupId>
+            <artifactId>netty-common</artifactId>
+          </exclusion>
+          <exclusion>
+            <groupId>io.netty</groupId>
+            <artifactId>netty-buffer</artifactId>
+          </exclusion>
+        </exclusions>
       </dependency>
       <dependency>
         <groupId>org.apache.arrow</groupId>

diff --git a/src/main/java/net/snowflake/client/core/SFBaseSession.java b/src/main/java/net/snowflake/client/core/SFBaseSession.java
@@ -146,6 +146,8 @@ public abstract class SFBaseSession {
 
   private boolean isJdbcArrowTreatDecimalAsInt = true;
 
+  private boolean supportImplicitAsyncQueryTimeout = false;
+
   protected SFBaseSession(SFConnectionHandler sfConnectionHandler) {
     this.sfConnectionHandler = sfConnectionHandler;
   }
@@ -1314,4 +1316,22 @@ public SFConnectionHandler getSfConnectionHandler() {
   public boolean getEnableReturnTimestampWithTimeZone() {
     return enableReturnTimestampWithTimeZone;
   }
+
+  /**
+   * @return True if query timeout should be set on the server side for async queries. False by
+   *     default.
+   */
+  @SnowflakeJdbcInternalApi
+  public boolean getSupportImplicitAsyncQueryTimeout() {
+    return supportImplicitAsyncQueryTimeout;
+  }
+
+  /**
+   * @param supportImplicitAsyncQueryTimeout Setting supportImplicitAsyncQueryTimeout to true allows
+   *     for query timeout to be set on the server side.
+   */
+  @SnowflakeJdbcInternalApi
+  public void setSupportImplicitAsyncQueryTimeout(boolean supportImplicitAsyncQueryTimeout) {
+    this.supportImplicitAsyncQueryTimeout = supportImplicitAsyncQueryTimeout;
+  }
 }
diff --git a/src/main/java/net/snowflake/client/core/SFBaseStatement.java b/src/main/java/net/snowflake/client/core/SFBaseStatement.java
@@ -39,9 +39,13 @@ public abstract class SFBaseStatement {
   public void addProperty(String propertyName, Object propertyValue) throws SFException {
     statementParametersMap.put(propertyName, propertyValue);
 
-    // for query timeout, we implement it on client side for now
     if ("query_timeout".equalsIgnoreCase(propertyName)) {
+      // Client side implementation
       queryTimeout = (Integer) propertyValue;
+      if (this.getSFBaseSession().getSupportImplicitAsyncQueryTimeout()) {
+        // Set server parameter for supporting query timeout on async queries
+        statementParametersMap.put("STATEMENT_TIMEOUT_IN_SECONDS", (Integer) propertyValue);
+      }
     }
 
     // check if the number of session properties exceed limit

diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/src/main/java/net/snowflake/client/core/SFSession.java
@@ -553,6 +553,12 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
           }
           break;
 
+        case SUPPORT_IMPLICIT_ASYNC_QUERY_TIMEOUT:
+          if (propertyValue != null) {
+            setSupportImplicitAsyncQueryTimeout(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
@@ -121,8 +121,12 @@ public enum SFSessionProperty {
   HTTP_CLIENT_SOCKET_TIMEOUT("HTTP_CLIENT_SOCKET_TIMEOUT", false, Integer.class),
 
   JAVA_LOGGING_CONSOLE_STD_OUT("JAVA_LOGGING_CONSOLE_STD_OUT", false, Boolean.class),
+
   JAVA_LOGGING_CONSOLE_STD_OUT_THRESHOLD(
-      "JAVA_LOGGING_CONSOLE_STD_OUT_THRESHOLD", false, String.class);
+      "JAVA_LOGGING_CONSOLE_STD_OUT_THRESHOLD", false, String.class),
+
+  SUPPORT_IMPLICIT_ASYNC_QUERY_TIMEOUT(
+      "SUPPORT_IMPLICIT_ASYNC_QUERY_TIMEOUT", false, Boolean.class);
 
   // property key in string
   private String propertyKey;

diff --git a/src/main/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverter.java b/src/main/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverter.java
@@ -45,7 +45,10 @@ public ThreeFieldStructToTimestampTZConverter(
 
   @Override
   public boolean isNull(int index) {
-    return epochs.isNull(index);
+    return structVector.isNull(index)
+        || epochs.isNull(index)
+        || fractions.isNull(index)
+        || timeZoneIndices.isNull(index);
   }
 
   @Override
@@ -54,7 +57,7 @@ public String toString(int index) throws SFException {
       throw new SFException(ErrorCode.INTERNAL_ERROR, "missing timestamp TZ formatter");
     }
     try {
-      Timestamp ts = epochs.isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true);
+      Timestamp ts = isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true);
       return ts == null
           ? null
           : context.getTimestampTZFormatter().format(ts, timeZone, context.getScale(columnIndex));
@@ -65,7 +68,7 @@ public String toString(int index) throws SFException {
 
   @Override
   public byte[] toBytes(int index) throws SFException {
-    if (epochs.isNull(index)) {
+    if (isNull(index)) {
       return null;
     }
     throw new SFException(
@@ -79,7 +82,7 @@ public Object toObject(int index) throws SFException {
 
   @Override
   public Timestamp toTimestamp(int index, TimeZone tz) throws SFException {
-    return epochs.isNull(index) ? null : getTimestamp(index, tz, false);
+    return isNull(index) ? null : getTimestamp(index, tz, false);
   }
 
   private Timestamp getTimestamp(int index, TimeZone tz, boolean fromToString) throws SFException {
@@ -98,7 +101,7 @@ private Timestamp getTimestamp(int index, TimeZone tz, boolean fromToString) thr
 
   @Override
   public Date toDate(int index, TimeZone tz, boolean dateFormat) throws SFException {
-    if (epochs.isNull(index)) {
+    if (isNull(index)) {
       return null;
     }
     Timestamp ts = getTimestamp(index, TimeZone.getDefault(), false);
@@ -116,7 +119,7 @@ public Time toTime(int index) throws SFException {
 
   @Override
   public boolean toBoolean(int index) throws SFException {
-    if (epochs.isNull(index)) {
+    if (isNull(index)) {
       return false;
     }
     Timestamp val = toTimestamp(index, TimeZone.getDefault());
@@ -127,7 +130,7 @@ public boolean toBoolean(int index) throws SFException {
 
   @Override
   public short toShort(int rowIndex) throws SFException {
-    if (epochs.isNull(rowIndex)) {
+    if (isNull(rowIndex)) {
       return 0;
     }
     throw new SFException(

diff --git a/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverter.java b/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverter.java
@@ -41,7 +41,7 @@ public TwoFieldStructToTimestampLTZConverter(
 
   @Override
   public boolean isNull(int index) {
-    return epochs.isNull(index);
+    return structVector.isNull(index) || epochs.isNull(index) || fractions.isNull(index);
   }
 
   @Override
@@ -51,7 +51,7 @@ public String toString(int index) throws SFException {
     }
 
     try {
-      Timestamp ts = epochs.isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true);
+      Timestamp ts = isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true);
 
       return ts == null
           ? null
@@ -82,7 +82,7 @@ private Timestamp getTimestamp(int index, TimeZone tz, boolean fromToString) thr
 
   @Override
   public byte[] toBytes(int index) throws SFException {
-    if (epochs.isNull(index)) {
+    if (isNull(index)) {
       return null;
     }
     throw new SFException(
@@ -111,7 +111,7 @@ public Time toTime(int index) throws SFException {
 
   @Override
   public boolean toBoolean(int index) throws SFException {
-    if (epochs.isNull(index)) {
+    if (isNull(index)) {
       return false;
     }
     Timestamp val = toTimestamp(index, TimeZone.getDefault());

diff --git a/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverter.java b/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverter.java
@@ -42,7 +42,7 @@ public TwoFieldStructToTimestampNTZConverter(
 
   @Override
   public boolean isNull(int index) {
-    return epochs.isNull(index);
+    return structVector.isNull(index) || epochs.isNull(index) || fractions.isNull(index);
   }
 
   @Override
@@ -51,7 +51,7 @@ public String toString(int index) throws SFException {
       throw new SFException(ErrorCode.INTERNAL_ERROR, "missing timestamp NTZ formatter");
     }
     try {
-      Timestamp ts = epochs.isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true);
+      Timestamp ts = isNull(index) ? null : getTimestamp(index, TimeZone.getDefault(), true);
 
       return ts == null
           ? null
@@ -92,7 +92,7 @@ private Timestamp getTimestamp(int index, TimeZone tz, boolean fromToString) thr
 
   @Override
   public byte[] toBytes(int index) throws SFException {
-    if (epochs.isNull(index)) {
+    if (isNull(index)) {
       return null;
     }
     throw new SFException(
@@ -119,7 +119,7 @@ public Time toTime(int index) throws SFException {
 
   @Override
   public boolean toBoolean(int index) throws SFException {
-    if (epochs.isNull(index)) {
+    if (isNull(index)) {
       return false;
     }
     Timestamp val = toTimestamp(index, TimeZone.getDefault());

diff --git a/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverter.java b/src/main/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverter.java
@@ -37,7 +37,7 @@ public TwoFieldStructToTimestampTZConverter(
 
   @Override
   public boolean isNull(int index) {
-    return epochs.isNull(index);
+    return structVector.isNull(index) || epochs.isNull(index) || timeZoneIndices.isNull(index);
   }
 
   @Override
@@ -59,7 +59,7 @@ public Object toObject(int index) throws SFException {
 
   @Override
   public Timestamp toTimestamp(int index, TimeZone tz) throws SFException {
-    return epochs.isNull(index) ? null : getTimestamp(index, tz);
+    return isNull(index) ? null : getTimestamp(index, tz);
   }
 
   private Timestamp getTimestamp(int index, TimeZone tz) throws SFException {
@@ -76,7 +76,7 @@ private Timestamp getTimestamp(int index, TimeZone tz) throws SFException {
 
   @Override
   public Date toDate(int index, TimeZone tz, boolean dateFormat) throws SFException {
-    if (epochs.isNull(index)) {
+    if (isNull(index)) {
       return null;
     }
     Timestamp ts = getTimestamp(index, TimeZone.getDefault());
@@ -94,7 +94,7 @@ public Time toTime(int index) throws SFException {
 
   @Override
   public boolean toBoolean(int index) throws SFException {
-    if (epochs.isNull(index)) {
+    if (isNull(index)) {
       return false;
     }
     Timestamp val = toTimestamp(index, TimeZone.getDefault());
@@ -105,7 +105,7 @@ public boolean toBoolean(int index) throws SFException {
 
   @Override
   public byte[] toBytes(int index) throws SFException {
-    if (epochs.isNull(index)) {
+    if (isNull(index)) {
       return null;
     }
     throw new SFException(
@@ -114,7 +114,7 @@ public byte[] toBytes(int index) throws SFException {
 
   @Override
   public short toShort(int rowIndex) throws SFException {
-    if (epochs.isNull(rowIndex)) {
+    if (isNull(rowIndex)) {
       return 0;
     }
     throw new SFException(

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeStatement.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeStatement.java
@@ -58,4 +58,12 @@ public interface SnowflakeStatement {
    * @throws SQLException if an error is encountered
    */
   void resultSetMetadataHandler(SFBaseResultSet resultSet) throws SQLException;
+
+  /**
+   * Sets the query timeout when running an async query.
+   *
+   * @param seconds The number of seconds until timeout.
+   * @throws SQLException if an error is encountered
+   */
+  void setAsyncQueryTimeout(int seconds) throws SQLException;
 }
diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeStatementV1.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeStatementV1.java
@@ -883,6 +883,21 @@ public void setQueryTimeout(int seconds) throws SQLException {
     }
   }
 
+  @Override
+  public void setAsyncQueryTimeout(int seconds) throws SQLException {
+    logger.trace("setAsyncQueryTimeout(int seconds)", false);
+    raiseSQLExceptionIfStatementIsClosed();
+
+    try {
+      if (this.sfBaseStatement != null) {
+        this.sfBaseStatement.addProperty("STATEMENT_TIMEOUT_IN_SECONDS", seconds);
+      }
+    } catch (SFException ex) {
+      throw new SnowflakeSQLException(
+          ex.getCause(), ex.getSqlState(), ex.getVendorCode(), ex.getParams());
+    }
+  }
+
   @Override
   public boolean isWrapperFor(Class<?> iface) throws SQLException {
     logger.trace("isWrapperFor(Class<?> iface)", false);
@@ -1260,6 +1275,9 @@ public void setParameter(String name, Object value) throws SQLException {}
     @Override
     public void setQueryTimeout(int seconds) throws SQLException {}
 
+    @Override
+    public void setAsyncQueryTimeout(int seconds) throws SQLException {}
+
     @Override
     public boolean isWrapperFor(Class<?> iface) throws SQLException {
       logger.trace("isWrapperFor(Class<?> iface)", false);

diff --git a/...test/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverterTest.java b/...test/java/net/snowflake/client/core/arrow/ThreeFieldStructToTimestampTZConverterTest.java
@@ -169,9 +169,11 @@ public void testTimestampTZ(
         seconds.setSafe(j, testSecondsInt64[i]);
         nanos.setSafe(j, testNanos[i]);
         timeZoneIdx.setSafe(j, testTimeZoneIndices[i++]);
+        structVector.setIndexDefined(j);
       }
       j++;
     }
+    structVector.setValueCount(j);
 
     ArrowVectorConverter converter =
         new ThreeFieldStructToTimestampTZConverter(structVector, 0, this);

diff --git a/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverterTest.java b/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampLTZConverterTest.java
@@ -152,9 +152,11 @@ public void testTimestampLTZ(
       } else {
         epochs.setSafe(j, testSecondsInt64[i]);
         fractions.setSafe(j, testNanoSecs[i++]);
+        structVector.setIndexDefined(j);
       }
       j++;
     }
+    structVector.setValueCount(j);
 
     ArrowVectorConverter converter =
         new TwoFieldStructToTimestampLTZConverter(structVector, 0, this);

diff --git a/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverterTest.java b/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampNTZConverterTest.java
@@ -175,9 +175,11 @@ public void testTimestampNTZ(
       } else {
         epochs.setSafe(j, testSecondsInt64[i]);
         fractions.setSafe(j, testNanoSecs[i++]);
+        structVector.setIndexDefined(j);
       }
       j++;
     }
+    structVector.setValueCount(j);
 
     ArrowVectorConverter converter =
         new TwoFieldStructToTimestampNTZConverter(structVector, 0, this);

diff --git a/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverterTest.java b/src/test/java/net/snowflake/client/core/arrow/TwoFieldStructToTimestampTZConverterTest.java
@@ -114,9 +114,11 @@ public void testTimestampTZ(String tz) throws SFException {
       } else {
         epoch.setSafe(j, testEpochesInt64[i]);
         timeZoneIdx.setSafe(j, testTimeZoneIndices[i++]);
+        structVector.setIndexDefined(j);
       }
       j++;
     }
+    structVector.setValueCount(j);
 
     ArrowVectorConverter converter =
         new TwoFieldStructToTimestampTZConverter(structVector, 0, this);