Merge pull request #1188 from zhicwu/master

Ensure inner output stream will be closed
ClickHouse · Jan 10, 2023 · 5cf13b2 · 5cf13b2
2 parents 9c478dc + d1fc0cf
commit 5cf13b2
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 17 deletions.
diff --git a/clickhouse-client/src/main/java/com/clickhouse/client/stream/Lz4OutputStream.java b/clickhouse-client/src/main/java/com/clickhouse/client/stream/Lz4OutputStream.java
@@ -82,4 +82,16 @@ public void flush() throws IOException {
         }
         output.flush();
     }
+
+    @Override
+    public void close() throws IOException {
+        if (!closed) {
+            try {
+                // flush before closing the inner output stream
+                super.close();
+            } finally {
+                output.close();
+            }
+        }
+    }
 }
diff --git a/clickhouse-client/src/main/java/com/clickhouse/client/stream/WrappedOutputStream.java b/clickhouse-client/src/main/java/com/clickhouse/client/stream/WrappedOutputStream.java
@@ -36,4 +36,16 @@ public void flush() throws IOException {
         }
         output.flush();
     }
+
+    @Override
+    public void close() throws IOException {
+        if (!closed) {
+            try {
+                // flush before closing the inner output stream
+                super.close();
+            } finally {
+                output.close();
+            }
+        }
+    }
 }
diff --git a/clickhouse-client/src/test/java/com/clickhouse/client/ClientIntegrationTest.java b/clickhouse-client/src/test/java/com/clickhouse/client/ClientIntegrationTest.java
@@ -1761,11 +1761,11 @@ public void testErrorDuringInsert() throws ClickHouseException {
             Assert.fail("Insert should be aborted");
         } catch (UncheckedIOException e) {
             ClickHouseException ex = ClickHouseException.of(e, server);
-            Assert.assertEquals(ex.getErrorCode(), 395);
+            Assert.assertEquals(ex.getErrorCode(), 395, "Expected error code 395 but we got: " + ex.getMessage());
             Assert.assertTrue(ex.getCause() instanceof IOException, "Should end up with IOException");
             success = false;
         } catch (ClickHouseException e) {
-            Assert.assertEquals(e.getErrorCode(), 395);
+            Assert.assertEquals(e.getErrorCode(), 395, "Expected error code 395 but we got: " + e.getMessage());
             Assert.assertTrue(e.getCause() instanceof IOException, "Should end up with IOException");
             success = false;
         }
@@ -1838,7 +1838,8 @@ public void testSessionLock() throws ClickHouseException {
             try (ClickHouseResponse resp = req2.executeAndWait()) {
                 Assert.fail("Should fail due to session is locked by previous query");
             } catch (ClickHouseException e) {
-                Assert.assertEquals(e.getErrorCode(), ClickHouseException.ERROR_SESSION_IS_LOCKED);
+                Assert.assertEquals(e.getErrorCode(), ClickHouseException.ERROR_SESSION_IS_LOCKED,
+                        "Expected error code 373 but we got: " + e.getMessage());
             }
             new Thread(() -> {
                 try {
@@ -1880,7 +1881,8 @@ public void testAbortTransaction() throws ClickHouseException {
                 checkRowCount(txRequest, tableName, 0);
                 Assert.fail("Should fail as the transaction is invalid");
             } catch (ClickHouseException e) {
-                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION);
+                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION,
+                        "Expected error code 649 but we got: " + e.getMessage());
             }
         }
     }
@@ -2098,7 +2100,8 @@ public void testTransactionSnapshot() throws ClickHouseException {
             try {
                 req2.getTransaction().snapshot(5);
             } catch (ClickHouseTransactionException e) {
-                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION);
+                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION,
+                        "Expected error code 649 but we got: " + e.getMessage());
             }
 
             req1.getTransaction().commit();
@@ -2108,7 +2111,8 @@ public void testTransactionSnapshot() throws ClickHouseException {
             try {
                 req1.getTransaction().snapshot(5);
             } catch (ClickHouseTransactionException e) {
-                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION);
+                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION,
+                        "Expected error code 649 but we got: " + e.getMessage());
             }
         }
     }
@@ -2142,31 +2146,35 @@ public void testTransactionTimeout() throws ClickHouseException {
                 Assert.fail("Query should fail due to session timed out");
             } catch (ClickHouseException e) {
                 // session not found(since it's timed out)
-                Assert.assertEquals(e.getErrorCode(), ClickHouseException.ERROR_SESSION_NOT_FOUND);
+                Assert.assertEquals(e.getErrorCode(), ClickHouseException.ERROR_SESSION_NOT_FOUND,
+                        "Expected error code 372 but we got: " + e.getMessage());
             }
             Assert.assertEquals(tx.getState(), ClickHouseTransaction.ACTIVE);
 
             try {
                 tx.commit();
                 Assert.fail("Should fail to commit due to session timed out");
             } catch (ClickHouseTransactionException e) {
-                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION);
+                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION,
+                        "Expected error code 649 but we got: " + e.getMessage());
             }
             Assert.assertEquals(tx.getState(), ClickHouseTransaction.FAILED);
 
             try {
                 tx.rollback();
                 Assert.fail("Should fail to roll back due to session timed out");
             } catch (ClickHouseTransactionException e) {
-                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION);
+                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION,
+                        "Expected error code 649 but we got: " + e.getMessage());
             }
             Assert.assertEquals(tx.getState(), ClickHouseTransaction.FAILED);
 
             try {
                 tx.begin();
                 Assert.fail("Should fail to restart due to session timed out");
             } catch (ClickHouseTransactionException e) {
-                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION);
+                Assert.assertEquals(e.getErrorCode(), ClickHouseTransactionException.ERROR_INVALID_TRANSACTION,
+                        "Expected error code 649 but we got: " + e.getMessage());
             }
             Assert.assertEquals(tx.getState(), ClickHouseTransaction.FAILED);
 
@@ -2191,7 +2199,8 @@ public void testTransactionTimeout() throws ClickHouseException {
                 checkRowCount(request, tableName, 3);
                 Assert.fail("Should fail to query due to session timed out");
             } catch (ClickHouseException e) {
-                Assert.assertEquals(e.getErrorCode(), ClickHouseException.ERROR_SESSION_NOT_FOUND);
+                Assert.assertEquals(e.getErrorCode(), ClickHouseException.ERROR_SESSION_NOT_FOUND,
+                        "Expected error code 372 but we got: " + e.getMessage());
             }
             Assert.assertEquals(request.getTransaction().getState(), ClickHouseTransaction.ACTIVE);
         }

diff --git a/clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHouseConnectionTest.java b/clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/ClickHouseConnectionTest.java
@@ -80,7 +80,7 @@ public void testAutoCommitMode() throws SQLException {
                         + " from numbers(100000)");
             } catch (SQLException e) {
                 if (i % 3 == 0 || i % 5 == 0) {
-                    Assert.assertEquals(e.getErrorCode(), 395);
+                    Assert.assertEquals(e.getErrorCode(), 395, "Expected error code 395 but we got: " + e.getMessage());
                 } else {
                     Assert.fail("Should not have exception");
                 }
@@ -157,7 +157,7 @@ public void testReadOnly() throws SQLException {
                     exp = e;
                 }
                 Assert.assertNotNull(exp, "Should fail with SQL exception");
-                Assert.assertEquals(exp.getErrorCode(), 164);
+                Assert.assertEquals(exp.getErrorCode(), 164, "Expected error code 164 but we got: " + exp.getMessage());
             }
 
             conn.setReadOnly(false);
@@ -181,7 +181,7 @@ public void testReadOnly() throws SQLException {
                 exp = e;
             }
             Assert.assertNotNull(exp, "Should fail with SQL exception");
-            Assert.assertEquals(exp.getErrorCode(), 164);
+            Assert.assertEquals(exp.getErrorCode(), 164, "Expected error code 164 but we got: " + exp.getMessage());
 
             exp = null;
             try {
@@ -191,7 +191,7 @@ public void testReadOnly() throws SQLException {
                 exp = e;
             }
             Assert.assertNotNull(exp, "Should fail with SQL exception");
-            Assert.assertEquals(exp.getErrorCode(), 164);
+            Assert.assertEquals(exp.getErrorCode(), 164, "Expected error code 164 but we got: " + exp.getMessage());
         }
 
         props.setProperty("user", "readonly2");
@@ -209,7 +209,7 @@ public void testReadOnly() throws SQLException {
                 exp = e;
             }
             Assert.assertNotNull(exp, "Should fail with SQL exception");
-            Assert.assertEquals(exp.getErrorCode(), 164);
+            Assert.assertEquals(exp.getErrorCode(), 164, "Expected error code 164 but we got: " + exp.getMessage());
 
             conn.setReadOnly(true);
             Assert.assertTrue(conn.isReadOnly(), "Connection should be readonly");
@@ -233,7 +233,7 @@ public void testReadOnly() throws SQLException {
                 exp = e;
             }
             Assert.assertNotNull(exp, "Should fail with SQL exception");
-            Assert.assertEquals(exp.getErrorCode(), 164);
+            Assert.assertEquals(exp.getErrorCode(), 164, "Expected error code 164 but we got: " + exp.getMessage());
         }
     }