fix

Author: hantmac

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

@@ -16,6 +16,8 @@
 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;
@@ -442,60 +444,39 @@ public void TestStageFileRemovedAfterBatchInsert() throws SQLException {
             s.execute("use " + dbName);
             s.execute("create or replace table t_stage_cleanup(a int, b string)");
 
-            try (TrackingPreparedStatement ps = new TrackingPreparedStatement((DatabendConnection) c, "insert into t_stage_cleanup values")) {
+            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});
+            }
 
-                String location = ps.getLastAttachmentLocation();
-                Assert.assertNotNull(location);
-                System.out.println("[DEBUG] uploaded stage file: " + location);
-                String dir = location.substring(0, location.lastIndexOf('/') + 1);
-                System.out.println("location dir is:"+dir);
-                try (ResultSet rs = s.executeQuery("LIST " + dir)) {
-                    if (rs.next()) {
-                        Assert.fail("Stage directory not empty after batch insert, unexpected entry: " + rs.getString(1));
-                    }
-                } catch (SQLException e) {
-                    if (e.getErrorCode() != 1003) {
-                        throw e;
-                    }
-                } finally {
-                    try {
-                        System.out.println("[DEBUG] drop stage path: " + location);
-                        s.execute("REMOVE " + location);
-                    } catch (SQLException ignore) {
-                        // best-effort cleanup
-                    }
-                }
+            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());
+            }
 
-                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));
                 }
             }
-        }
-    }
-
-    private static final class TrackingPreparedStatement extends DatabendPreparedStatement {
-        private StageAttachment lastAttachment;
-
-        TrackingPreparedStatement(DatabendConnection connection, String sql) throws SQLException {
-            super(connection, stmt -> {}, sql);
-        }
-
-        @Override
-        boolean dropStageAttachment(StageAttachment attachment) {
-            lastAttachment = attachment;
-            return super.dropStageAttachment(attachment);
-        }
-
-        String getLastAttachmentLocation() {
-            return lastAttachment == null ? null : lastAttachment.getLocation();
+            Set<String> diff = new HashSet<>(after);
+            diff.removeAll(before);
+            if (!diff.isEmpty()) {
+                Assert.fail("Stage has unexpected leftover entries: " + diff);
+            }
         }
     }