Fix bug: count sql statement leak (#1640)

z7658329 · web-flow · commit 299449b1ac9b · 2021-11-10T19:47:14.000+08:00
fix resultset leak for count query fixed #1626
diff --git a/hugegraph-core/src/main/java/com/baidu/hugegraph/backend/store/BackendEntry.java b/hugegraph-core/src/main/java/com/baidu/hugegraph/backend/store/BackendEntry.java
@@ -102,7 +102,7 @@ public default boolean olap() {
         return false;
     }
 
-    public interface BackendIterator<T> extends Iterator<T> {
+    public interface BackendIterator<T> extends Iterator<T>, AutoCloseable {
 
         public void close();
 
diff --git a/hugegraph-mysql/src/main/java/com/baidu/hugegraph/backend/store/mysql/MysqlTable.java b/hugegraph-mysql/src/main/java/com/baidu/hugegraph/backend/store/mysql/MysqlTable.java
@@ -340,6 +340,8 @@ public Number queryNumber(Session session, Query query) {
                 return IteratorUtils.of(rs.resultSet().getLong(1));
             } catch (SQLException e) {
                 throw new BackendException(e);
+            } finally {
+                rs.close();
             }
         });
         return aggregate.reduce(results);
diff --git a/hugegraph-rocksdb/src/main/java/com/baidu/hugegraph/backend/store/rocksdb/RocksDBTable.java b/hugegraph-rocksdb/src/main/java/com/baidu/hugegraph/backend/store/rocksdb/RocksDBTable.java
@@ -129,11 +129,12 @@ public Number queryNumber(Session session, Query query) {
 
         assert aggregate.func() == AggregateFunc.COUNT;
         assert query.noLimit();
-        Iterator<BackendColumn> results = this.queryBy(session, query);
-        if (results instanceof Countable) {
-            return ((Countable) results).count();
+        try (BackendColumnIterator results = this.queryBy(session, query)) {
+            if (results instanceof Countable) {
+                return ((Countable) results).count();
+            }
+            return IteratorUtils.count(results);
         }
-        return IteratorUtils.count(results);
     }
 
     @Override

Original file line number	Diff line number	Diff line change
`@@ -102,7 +102,7 @@ public default boolean olap() {`
`102`	`102`	`return false;`
`103`	`103`	`}`
`104`	`104`
`105`		`- public interface BackendIterator<T> extends Iterator<T> {`
	`105`	`+ public interface BackendIterator<T> extends Iterator<T>, AutoCloseable {`
`106`	`106`
`107`	`107`	`public void close();`
`108`	`108`
Original file line number	Diff line number	Diff line change
`@@ -340,6 +340,8 @@ public Number queryNumber(Session session, Query query) {`
`340`	`340`	`return IteratorUtils.of(rs.resultSet().getLong(1));`
`341`	`341`	`} catch (SQLException e) {`
`342`	`342`	`throw new BackendException(e);`
	`343`	`+ } finally {`
	`344`	`+ rs.close();`
`343`	`345`	`}`
`344`	`346`	`});`
`345`	`347`	`return aggregate.reduce(results);`