Add ! suffix and enhance error handling

Author: zhicwu

clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/internal/ClickHouseStatementImpl.java

@@ -184,36 +184,48 @@ protected ClickHouseResponse processSqlStatement(ClickHouseSqlStatement stmt) th
             }
             return ClickHouseResponse.EMPTY;
         } else if (connection.getJdbcConfig().useLocalFile() && stmt.hasFile()) {
-            final String file = ClickHouseUtils.unescape(stmt.getFile());
+            String file = ClickHouseUtils.unescape(stmt.getFile());
+            boolean suffix = file.lastIndexOf('!') == file.length() - 1;
+            if (suffix) {
+                file = file.substring(0, file.length() - 1);
+            }
             if (stmt.getStatementType() == StatementType.SELECT) {
-                ClickHouseFile f = ClickHouseFile.of(file);
-                if (f.getFile().exists()) {
+                ClickHouseFile f = ClickHouseFile.of(ClickHouseUtils.getFile(file));
+                if (!suffix && f.getFile().exists()) {
                     throw SqlExceptionUtils
                             .clientError(ClickHouseUtils.format("Output file [%s] already exists!", f.getFile()));
                 }
 
+                f = getFile(f, stmt);
                 final ClickHouseResponseSummary summary = new ClickHouseResponseSummary(null, null);
-                try (ClickHouseResponse response = request.query(stmt.getSQL()).output(getFile(f, stmt))
-                        .executeAndWait()) {
+                try (ClickHouseResponse response = request.query(stmt.getSQL()).output(f).executeAndWait()) {
                     summary.add(response.getSummary());
                 } catch (ClickHouseException e) {
                     throw SqlExceptionUtils.handle(e);
                 }
                 return ClickHouseSimpleResponse.of(getConfig(),
                         Arrays.asList(ClickHouseColumn.of("file", ClickHouseDataType.String, false),
-                                ClickHouseColumn.of("rows", ClickHouseDataType.UInt64, false),
+                                ClickHouseColumn.of("format", ClickHouseDataType.String, false),
+                                ClickHouseColumn.of("compression", ClickHouseDataType.String, false),
+                                ClickHouseColumn.of("level", ClickHouseDataType.Int32, false),
                                 ClickHouseColumn.of("bytes", ClickHouseDataType.UInt64, false)),
-                        new Object[][] {
-                                { file, summary.getReadRows(), summary.getReadBytes() }
-                        }, summary);
+                        new Object[][] { { file, f.getFormat().name(),
+                                f.hasCompression() ? f.getCompressionAlgorithm().encoding() : "none",
+                                f.getCompressionLevel(), f.getFile().length() } },
+                        summary);
             } else if (stmt.getStatementType() == StatementType.INSERT) {
                 final Mutation m = request.write().query(stmt.getSQL());
                 final ClickHouseResponseSummary summary = new ClickHouseResponseSummary(null, null);
                 try {
                     for (Path p : ClickHouseUtils.findFiles(file)) {
                         ClickHouseFile f = ClickHouseFile.of(p.toFile());
                         if (!f.getFile().exists()) {
-                            log.warn("Skip [%s] as it does not exist", f);
+                            if (suffix) {
+                                throw SqlExceptionUtils
+                                        .clientError(ClickHouseUtils.format("File [%s] does not exist!", f.getFile()));
+                            } else {
+                                log.warn("Skip [%s] as it does not exist - perhaps it was just deleted somehow?", f);
+                            }
                         } else {
                             try (ClickHouseResponse response = m.data(getFile(f, stmt)).executeAndWait()) {
                                 summary.add(response.getSummary());
@@ -224,7 +236,7 @@ protected ClickHouseResponse processSqlStatement(ClickHouseSqlStatement stmt) th
                             }
                         }
                     }
-                    if (summary.getUpdateCount() == 0) {
+                    if (suffix && summary.getUpdateCount() == 0) {
                         throw SqlExceptionUtils.clientError("No file imported: " + file);
                     }
                     summary.seal();

clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHousePreparedStatementTest.java

@@ -1260,6 +1260,15 @@ public void testOutFileAndInFile() throws SQLException {
                     Assert.assertFalse(rs.next());
                 }
                 Assert.assertTrue(f.exists());
+
+                stmt.setString(1, f.getName());
+                Assert.assertThrows(SQLException.class, () -> stmt.executeQuery());
+                stmt.setString(1, f.getName() + "!");
+                try (ResultSet rs = stmt.executeQuery()) {
+                    Assert.assertTrue(rs.next());
+                    Assert.assertEquals(rs.getString(1), f.getName());
+                    Assert.assertFalse(rs.next());
+                }
             }
 
             try (PreparedStatement stmt = conn
@@ -1268,12 +1277,14 @@ public void testOutFileAndInFile() throws SQLException {
                 stmt.addBatch();
                 stmt.setString(1, f.getName());
                 stmt.addBatch();
+                stmt.setString(1, f.getName() + "!");
+                stmt.addBatch();
                 stmt.executeBatch();
             }
 
             try (ResultSet rs = s.executeQuery("select count(1), uniqExact(n) from test_load_infile_with_params")) {
                 Assert.assertTrue(rs.next(), "Should have at least one row");
-                Assert.assertEquals(rs.getInt(1), 999 * 2);
+                Assert.assertEquals(rs.getInt(1), 999 * 3);
                 Assert.assertEquals(rs.getInt(2), 999);
                 Assert.assertFalse(rs.next(), "Should have only one row");
             }

clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHouseStatementTest.java

@@ -159,15 +159,27 @@ public void testOutFileAndInFile() throws SQLException {
         }
 
         try (ClickHouseConnection conn = newConnection(props)) {
+            String sql1 = "select number n, toString(n) from numbers(1234) into outfile '" + f1.getName() + "'";
+            try (ClickHouseStatement stmt = conn.createStatement(); ResultSet rs = stmt.executeQuery(sql1)) {
+                Assert.assertTrue(rs.next());
+                Assert.assertFalse(rs.next());
+                Assert.assertTrue(f1.exists());
+                // end up with exception because the file already exists
+                Assert.assertThrows(SQLException.class, () -> stmt.executeQuery(sql1));
+            }
+
+            // try again with ! suffix to override the existing file
+            String sql = "select number n, toString(n) from numbers(1234) into outfile '" + f1.getName() + "!'";
             try (ClickHouseStatement stmt = conn.createStatement();
-                    ResultSet rs = stmt.executeQuery(
-                            "select number n, toString(n) from numbers(1234) into outfile '" + f1.getName() + "'")) {
+                    ResultSet rs = stmt.executeQuery(sql)) {
                 Assert.assertTrue(rs.next());
                 Assert.assertFalse(rs.next());
+                Assert.assertTrue(f1.exists());
             }
+
+            sql = "select number n, toString(n) from numbers(4321) into outfile '" + f2.getName() + "!'";
             try (ClickHouseStatement stmt = conn.createStatement();
-                    ResultSet rs = stmt.executeQuery(
-                            "select number n, toString(n) from numbers(4321) into outfile '" + f2.getName() + "'")) {
+                    ResultSet rs = stmt.executeQuery(sql)) {
                 Assert.assertTrue(rs.next());
                 Assert.assertFalse(rs.next());
             }
@@ -182,6 +194,19 @@ public void testOutFileAndInFile() throws SQLException {
                     Assert.assertFalse(rs.next());
                 }
             }
+
+            try (ClickHouseStatement stmt = conn.createStatement()) {
+                // it's fine when no record was inserted
+                stmt.executeUpdate("insert into test_load_infile from infile 'non-existent.csv'");
+                // unless suffix ! was added...
+                Assert.assertThrows(SQLException.class,
+                        () -> stmt.executeUpdate("insert into test_load_infile from infile 'non-existent.csv!'"));
+                try (ResultSet rs = stmt.executeQuery("select count(1) from test_load_infile")) {
+                    Assert.assertTrue(rs.next());
+                    Assert.assertEquals(rs.getInt(1), 5555);
+                    Assert.assertFalse(rs.next());
+                }
+            }
         }
     }