SNOW-1890076: Ensure correct file is downloaded as stream

Author: sfc-gh-dprzybysz

src/main/java/net/snowflake/client/jdbc/ErrorCode.java

@@ -84,7 +84,8 @@ public enum ErrorCode {
   GCP_SERVICE_ERROR(200061, SqlState.SYSTEM_ERROR),
   AUTHENTICATOR_REQUEST_TIMEOUT(200062, SqlState.CONNECTION_EXCEPTION),
   INVALID_STRUCT_DATA(200063, SqlState.DATA_EXCEPTION),
-  DISABLEOCSP_INSECUREMODE_VALUE_MISMATCH(200064, SqlState.INVALID_PARAMETER_VALUE);
+  DISABLEOCSP_INSECUREMODE_VALUE_MISMATCH(200064, SqlState.INVALID_PARAMETER_VALUE),
+  TOO_MANY_FILES_TO_DOWNLOAD_AS_STREAM(200065, SqlState.DATA_EXCEPTION);
 
   public static final String errorMessageResource = "net.snowflake.client.jdbc.jdbc_error_messages";
 

src/main/java/net/snowflake/client/jdbc/SnowflakeConnectionV1.java

@@ -61,12 +61,14 @@ public class SnowflakeConnectionV1 implements Connection, SnowflakeConnection {
 
   /** Refer to all created and open statements from this connection */
   private final Set<Statement> openStatements = ConcurrentHashMap.newKeySet();
+
   // Injected delay for the purpose of connection timeout testing
   // Any statement execution will sleep for the specified number of milliseconds
   private final AtomicInteger _injectedDelay = new AtomicInteger(0);
   private boolean isClosed;
   private SQLWarning sqlWarnings = null;
   private List<DriverPropertyInfo> missingProperties = null;
+
   /**
    * Amount of milliseconds a user is willing to tolerate for network related issues (e.g. HTTP
    * 503/504) or database transient issues (e.g. GS not responding)
@@ -76,12 +78,14 @@ public class SnowflakeConnectionV1 implements Connection, SnowflakeConnection {
    * <p>Default: 300 seconds
    */
   private int networkTimeoutInMilli = 0; // in milliseconds
+
   /* this should be set to Connection.TRANSACTION_READ_COMMITTED
    * There may not be many implications here since the call to
    * setTransactionIsolation doesn't do anything.
    */
   private int transactionIsolation = Connection.TRANSACTION_NONE;
   private SFBaseSession sfSession;
+
   /** The SnowflakeConnectionImpl that provides the underlying physical-layer implementation */
   private SFConnectionHandler sfConnectionHandler;
 
@@ -1037,6 +1041,11 @@ public InputStream downloadStream(String stageName, String sourceFileName, boole
     // no file will be downloaded to this location
     getCommand.append(" file:///tmp/ /*jdbc download stream*/");
 
+    // We cannot match whole sourceFileName since it may be different e.g. for git repositories so
+    // we match only raw filename
+    String[] split = sourceFileName.split("/");
+    getCommand.append(" PATTERN=\".*").append(split[split.length - 1]).append("\"");
+
     SFBaseFileTransferAgent transferAgent =
         sfConnectionHandler.getFileTransferAgent(getCommand.toString(), stmt.getSFBaseStatement());
 

src/main/java/net/snowflake/client/jdbc/SnowflakeFileTransferAgent.java

@@ -1696,6 +1696,18 @@ public InputStream downloadStream(String fileName) throws SnowflakeSQLException
     remoteLocation remoteLocation = extractLocationAndPath(stageInfo.getLocation());
 
     // when downloading files as stream there should be only one file in source files
+    // let's fail fast when more than one file matches instead of fetching random one
+    if (sourceFiles.size() > 1) {
+      throw new SnowflakeSQLException(
+          queryID,
+          SqlState.NO_DATA,
+          ErrorCode.TOO_MANY_FILES_TO_DOWNLOAD_AS_STREAM.getMessageCode(),
+          session,
+          "There are more than one file matching "
+              + fileName
+              + ": "
+              + String.join(",", sourceFiles));
+    }
     String sourceLocation =
         sourceFiles.stream()
             .findFirst()

src/test/java/net/snowflake/client/jdbc/StreamLatestIT.java

@@ -7,18 +7,22 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.junit.jupiter.api.Assertions.fail;
 
+import java.io.BufferedReader;
 import java.io.BufferedWriter;
 import java.io.File;
 import java.io.FileWriter;
 import java.io.IOException;
 import java.io.InputStream;
+import java.io.InputStreamReader;
 import java.io.StringWriter;
 import java.nio.charset.StandardCharsets;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.Properties;
+import java.util.UUID;
+import java.util.stream.Collectors;
 import net.snowflake.client.annotations.DontRunOnGithubActions;
 import net.snowflake.client.category.TestTags;
 import org.apache.commons.io.IOUtils;
@@ -237,4 +241,72 @@ public void testSpecialCharactersInFileName() throws SQLException, IOException {
       }
     }
   }
+
+  @Test
+  public void shouldDownloadStreamInDeterministicWay() throws Exception {
+    try (Connection conn = getConnection();
+        Statement stat = conn.createStatement()) {
+      String randomStage = "test" + UUID.randomUUID().toString().replaceAll("-", "");
+      try {
+        stat.execute("CREATE OR REPLACE STAGE " + randomStage);
+        String randomDir = UUID.randomUUID().toString();
+        String sourceFilePathWithoutExtension = getClass().getResource("/test_file").getPath();
+        String sourceFilePathWithExtension = getClass().getResource("/test_file.csv").getPath();
+        String stageDest = String.format("@%s/%s", randomStage, randomDir);
+        putFile(stat, sourceFilePathWithExtension, stageDest, false);
+        putFile(stat, sourceFilePathWithoutExtension, stageDest, false);
+        putFile(stat, sourceFilePathWithExtension, stageDest, true);
+        putFile(stat, sourceFilePathWithoutExtension, stageDest, true);
+        expectsFilesOnStage(stat, stageDest, 4);
+        String stageName = "@" + randomStage;
+        downloadStreamExpectingContent(
+            conn, stageName, randomDir + "/test_file.gz", true, "I am a file without extension");
+        downloadStreamExpectingContent(
+            conn, stageName, randomDir + "/test_file.csv.gz", true, "I am a file with extension");
+        downloadStreamExpectingContent(
+            conn, stageName, randomDir + "/test_file", false, "I am a file without extension");
+        downloadStreamExpectingContent(
+            conn, stageName, randomDir + "/test_file.csv", false, "I am a file with extension");
+      } finally {
+        stat.execute("DROP STAGE IF EXISTS " + randomStage);
+      }
+    }
+  }
+
+  private static void downloadStreamExpectingContent(
+      Connection conn,
+      String stageName,
+      String fileName,
+      boolean decompress,
+      String expectedFileContent)
+      throws IOException, SQLException {
+    try (InputStream inputStream =
+            conn.unwrap(SnowflakeConnectionV1.class)
+                .downloadStream(stageName, fileName, decompress);
+        InputStreamReader isr = new InputStreamReader(inputStream);
+        BufferedReader br = new BufferedReader(isr)) {
+      String content = br.lines().collect(Collectors.joining("\n"));
+      assertEquals(expectedFileContent, content);
+    }
+  }
+
+  private static void expectsFilesOnStage(Statement stat, String stageDest, int expectCount)
+      throws SQLException {
+    int filesInStageDir = 0;
+    try (ResultSet rs = stat.executeQuery("LIST " + stageDest)) {
+      while (rs.next()) {
+        ++filesInStageDir;
+      }
+    }
+    assertEquals(expectCount, filesInStageDir);
+  }
+
+  private static boolean putFile(
+      Statement stat, String localFileName, String stageDest, boolean autoCompress)
+      throws SQLException {
+    return stat.execute(
+        String.format(
+            "PUT file://%s %s AUTO_COMPRESS=%s",
+            localFileName, stageDest, String.valueOf(autoCompress).toUpperCase()));
+  }
 }

src/test/resources/test_file

@@ -0,0 +1 @@
+I am a file without extension

src/test/resources/test_file.csv

@@ -0,0 +1 @@
+I am a file with extension