fix response close for getRowsUpdated(ClickHouse#1538)

wooEnrico · Jan 30, 2024 · f9851e7 · f9851e7
1 parent bbd6808
commit f9851e7
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/clickhouse-r2dbc/src/main/java/com/clickhouse/r2dbc/ClickHouseResult.java b/clickhouse-r2dbc/src/main/java/com/clickhouse/r2dbc/ClickHouseResult.java
@@ -31,10 +31,18 @@ public class ClickHouseResult implements Result {
                                 .map(rec -> ClickHousePair.of(resp.getColumns(), rec))))
                     .map(pair -> new ClickHouseRow(pair.getRight(), pair.getLeft()))
                 .map(RowSegment::new);
-        this.updatedCount =  Mono.just(response).map(ClickHouseResponse::getSummary)
-                .map(ClickHouseResponseSummary::getProgress)
-                .map(ClickHouseResponseSummary.Progress::getWrittenRows)
-                .map(UpdateCount::new);
+
+        this.updatedCount = Mono.using(() -> response,
+                resp -> Mono.just(response).map(ClickHouseResponse::getSummary)
+                        .map(ClickHouseResponseSummary::getProgress)
+                        .map(ClickHouseResponseSummary.Progress::getWrittenRows)
+                        .map(UpdateCount::new),
+                resp -> {
+                    if (!resp.isClosed()) {
+                        resp.close();
+                    }
+                });
+
         this.segments = Flux.concat(this.updatedCount, this.rowSegments);
     }
 

diff --git a/clickhouse-r2dbc/src/main/java/com/clickhouse/r2dbc/ClickHouseResult091.java b/clickhouse-r2dbc/src/main/java/com/clickhouse/r2dbc/ClickHouseResult091.java
@@ -31,10 +31,16 @@ class ClickHouseResult implements Result {
                                 .map(rec -> ClickHousePair.of(resp.getColumns(), rec))))
                     .map(pair -> new ClickHouseRow(pair.getRight(), pair.getLeft()))
                 .map(RowSegment::new);
-        this.updatedCount =  Mono.just(response).map(ClickHouseResponse::getSummary)
-                .map(ClickHouseResponseSummary::getProgress)
-                .map(ClickHouseResponseSummary.Progress::getWrittenRows)
-                .map(UpdateCount::new);
+        this.updatedCount = Mono.using(() -> response,
+                resp -> Mono.just(response).map(ClickHouseResponse::getSummary)
+                        .map(ClickHouseResponseSummary::getProgress)
+                        .map(ClickHouseResponseSummary.Progress::getWrittenRows)
+                        .map(UpdateCount::new),
+                resp -> {
+                    if (!resp.isClosed()) {
+                        resp.close();
+                    }
+                });
         this.segments = Flux.concat(this.updatedCount, this.rowSegments);
     }