Fix bug: count sql statement leak (#1640)

fix resultset leak for count query fixed #1626
apache · Nov 10, 2021 · 299449b · 299449b
1 parent d9b70d7
commit 299449b
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 5 deletions.
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