SNOW-2900795: Make MAX_BUFFER_SIZE in SnowflakeFileTranferAgent Configurable

src/main/java/net/snowflake/client/api/connection/UploadStreamConfig.java

@@ -15,6 +15,7 @@
  *   UploadStreamConfig config = UploadStreamConfig.builder()
  *       .setDestPrefix("data/2024")
  *       .setCompressData(true)
+ *       .setFileBackedBufferThreshold(1024 * 1024 * 1024) // 1 GiB
  *       .build();
  *
  *   connection.uploadStream("@my_stage", "uploaded_data.csv", dataStream, config);
@@ -26,6 +27,7 @@
 public class UploadStreamConfig {
   private final String destPrefix;
   private final boolean compressData;
+  private final int fileBackedBufferThreshold;
 
   /**
    * Private constructor. Use {@link Builder} to create instances.
@@ -35,6 +37,7 @@ public class UploadStreamConfig {
   private UploadStreamConfig(Builder builder) {
     this.destPrefix = builder.destPrefix;
     this.compressData = builder.compressData;
+    this.fileBackedBufferThreshold = builder.fileBackedBufferThreshold;
   }
 
   /**
@@ -55,6 +58,16 @@ public boolean isCompressData() {
     return compressData;
   }
 
+  /**
+   * Gets the threshold in bytes at which the internal buffer switches from heap memory to a temp
+   * file on disk during compression and digest computation.
+   *
+   * @return the threshold in bytes (always positive; default is 128 MiB)
+   */
+  public int getFileBackedBufferThreshold() {
+    return fileBackedBufferThreshold;
+  }
+
   /**
    * Creates a new builder instance.
    *
@@ -76,12 +89,14 @@ public static Builder builder() {
    * UploadStreamConfig config = UploadStreamConfig.builder()
    *     .setDestPrefix("data/2024")
    *     .setCompressData(true)
+   *     .setFileBackedBufferThreshold(1024 * 1024 * 1024) // 1 GiB
    *     .build();
    * }</pre>
    */
   public static class Builder {
     private String destPrefix;
     private boolean compressData = true;
+    private int fileBackedBufferThreshold = 1 << 27; // 128 MiB
 
     /** Private constructor. Use {@link UploadStreamConfig#builder()} instead. */
     private Builder() {}
@@ -126,6 +141,29 @@ public Builder setCompressData(boolean compressData) {
       return this;
     }
 
+    /**
+     * Sets the threshold in bytes at which the internal buffer switches from heap memory to a temp
+     * file on disk during compression and digest computation.
+     *
+     * <p>Below this threshold, stream data is held entirely in JVM heap memory. Above it, the data
+     * spills to a temporary file on local disk, which reduces heap pressure at the cost of disk
+     * I/O.
+     *
+     * <p>The default is 128 MiB ({@code 1 << 27}). Lower values reduce heap usage at the cost of
+     * disk I/O; higher values keep more data in memory and avoid temp file creation.
+     *
+     * @param fileBackedBufferThreshold the threshold in bytes; must be positive
+     * @return this builder instance
+     * @throws IllegalArgumentException if the value is not positive
+     */
+    public Builder setFileBackedBufferThreshold(int fileBackedBufferThreshold) {
+      if (fileBackedBufferThreshold <= 0) {
+        throw new IllegalArgumentException("fileBackedBufferThreshold must be positive");
+      }
+      this.fileBackedBufferThreshold = fileBackedBufferThreshold;
+      return this;
+    }
+
     /**
      * Builds the {@link UploadStreamConfig} instance.
      *
@@ -144,6 +182,8 @@ public String toString() {
         + '\''
         + ", compressData="
         + compressData
+        + ", fileBackedBufferThreshold="
+        + fileBackedBufferThreshold
         + '}';
   }
 }

src/main/java/net/snowflake/client/internal/api/implementation/connection/SnowflakeConnectionImpl.java

@@ -891,7 +891,12 @@ public void uploadStream(
       throw new IllegalArgumentException("UploadStreamConfig cannot be null");
     }
     uploadStreamInternal(
-        stageName, config.getDestPrefix(), inputStream, destFileName, config.isCompressData());
+        stageName,
+        config.getDestPrefix(),
+        inputStream,
+        destFileName,
+        config.isCompressData(),
+        config.getFileBackedBufferThreshold());
   }
 
   /**
@@ -912,7 +917,7 @@ public void uploadStream(
   public void compressAndUploadStream(
       String stageName, String destPrefix, InputStream inputStream, String destFileName)
       throws SQLException {
-    uploadStreamInternal(stageName, destPrefix, inputStream, destFileName, true);
+    uploadStreamInternal(stageName, destPrefix, inputStream, destFileName, true, null);
   }
 
   /**
@@ -930,14 +935,18 @@ public void compressAndUploadStream(
    * @param inputStream input stream from which the data will be uploaded
    * @param destFileName destination file name to use
    * @param compressData whether compression is requested fore uploading data
+   * @param fileBackedBufferThreshold threshold (in bytes) at which the internal buffer used for
+   *     compression and digest computation switches from heap memory to a temp file on disk; pass
+   *     {@code null} to use the driver default
    * @throws SQLException raises if any error occurs
    */
   private void uploadStreamInternal(
       String stageName,
       String destPrefix,
       InputStream inputStream,
       String destFileName,
-      boolean compressData)
+      boolean compressData,
+      Integer fileBackedBufferThreshold)
       throws SQLException {
     logger.debug(
         "Upload data from stream: stageName={}" + ", destPrefix={}, destFileName={}",
@@ -991,6 +1000,9 @@ private void uploadStreamInternal(
     transferAgent.setSourceStream(inputStream);
     transferAgent.setDestFileNameForStreamSource(destFileName);
     transferAgent.setCompressSourceFromStream(compressData);
+    if (fileBackedBufferThreshold != null) {
+      transferAgent.setFileBackedBufferThreshold(fileBackedBufferThreshold);
+    }
     transferAgent.execute();
 
     stmt.close();

src/main/java/net/snowflake/client/internal/jdbc/SFBaseFileTransferAgent.java

@@ -164,6 +164,8 @@ public DownloadCommandEncryptionFacade(
   protected boolean showEncryptionParameter;
   protected List<Object> statusRows = new ArrayList<>();
   protected CommandType commandType = CommandType.UPLOAD;
+  protected int fileBackedBufferThreshold =
+      SnowflakeFileTransferAgent.DEFAULT_FILE_BACKED_BUFFER_THRESHOLD;
   private int currentRowIndex;
 
   /**
@@ -256,6 +258,20 @@ public void setCompressSourceFromStream(boolean compressSourceFromStream) {
     this.compressSourceFromStream = compressSourceFromStream;
   }
 
+  /**
+   * Sets the threshold in bytes at which the internal buffer switches from heap memory to a temp
+   * file on disk during compression and digest computation.
+   *
+   * @param fileBackedBufferThreshold the threshold in bytes; must be positive
+   * @throws IllegalArgumentException if the value is not positive
+   */
+  public void setFileBackedBufferThreshold(int fileBackedBufferThreshold) {
+    if (fileBackedBufferThreshold <= 0) {
+      throw new IllegalArgumentException("fileBackedBufferThreshold must be positive");
+    }
+    this.fileBackedBufferThreshold = fileBackedBufferThreshold;
+  }
+
   /**
    * Run the PUT/GET command, if a command has been set.
    *