Merge pull request #370 from databendlabs/feat/drop-stage-test

feat: add drop stage test

Author: hantmac

databend-jdbc/src/main/java/com/databend/jdbc/cloud/DatabendPresignClientV1.java

@@ -124,12 +124,12 @@ private void uploadFromStream(InputStream inputStream, String stageName, String
     private void uploadFromStream(InputStream inputStream, Headers headers, String presignedUrl, long fileSize)
             throws IOException
     {
-        logger.info("Starting upload: size=" + fileSize + " bytes, url=" + presignedUrl);
+        logger.fine("Starting upload: size=" + fileSize + " bytes, url=" + presignedUrl);
         long startTime = System.currentTimeMillis();
         try {
             Request r = putRequest(headers, presignedUrl, inputStream, fileSize);
             executeInternal(r, true);
-            logger.info("Upload completed in " + (System.currentTimeMillis() - startTime) + "ms");
+            logger.fine("Upload completed in " + (System.currentTimeMillis() - startTime) + "ms");
         }
         catch (IOException e) {
             logger.severe("Upload failed after " + (System.currentTimeMillis() - startTime) + "ms: " + e.getMessage());

databend-jdbc/src/test/java/com/databend/jdbc/TestPrepareStatement.java

@@ -16,6 +16,9 @@
 import java.sql.Statement;
 import java.sql.Timestamp;
 import java.sql.Types;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.UUID;
 import java.util.Arrays;
 import java.util.Properties;
 import java.util.TimeZone;
@@ -431,6 +434,52 @@ public void testAllPreparedStatement() throws SQLException {
         }
     }
 
+    @Test(groups = "IT")
+    public void TestStageFileRemovedAfterBatchInsert() throws SQLException {
+        String dbName = ("stage_cleanup_" + UUID.randomUUID()).replace("-", "");
+        try (Connection c = Utils.createConnection();
+             Statement s = c.createStatement()) {
+            c.setAutoCommit(false);
+            s.execute("create or replace database " + dbName);
+            s.execute("use " + dbName);
+            s.execute("create or replace table t_stage_cleanup(a int, b string)");
+
+            Set<String> before = new HashSet<>();
+            try (ResultSet rs = s.executeQuery("LIST @~/")) {
+                while (rs.next()) {
+                    before.add(rs.getString(1));
+                }
+            }
+
+            try (PreparedStatement ps = c.prepareStatement("insert into t_stage_cleanup values")) {
+                ps.setInt(1, 1);
+                ps.setString(2, "hello");
+                ps.addBatch();
+                int[] counts = ps.executeBatch();
+                Assert.assertEquals(counts, new int[] {1});
+            }
+
+            try (ResultSet rs = s.executeQuery("SELECT a, b FROM t_stage_cleanup")) {
+                Assert.assertTrue(rs.next());
+                Assert.assertEquals(rs.getInt(1), 1);
+                Assert.assertEquals(rs.getString(2), "hello");
+                Assert.assertFalse(rs.next());
+            }
+
+            Set<String> after = new HashSet<>();
+            try (ResultSet rs = s.executeQuery("LIST @~/")) {
+                while (rs.next()) {
+                    after.add(rs.getString(1));
+                }
+            }
+            Set<String> diff = new HashSet<>(after);
+            diff.removeAll(before);
+            if (!diff.isEmpty()) {
+                Assert.fail("Stage has unexpected leftover entries: " + diff);
+            }
+        }
+    }
+
     @Test(groups = "IT")
     public void shouldBuildStageAttachmentWithFileFormatOptions() throws SQLException {
         Connection conn = Utils.createConnection();