1 files changed, 18 insertions, 9 deletions

diff --git a/ml/Query.h b/ml/Query.h
index 24c5fa384..78d117003 100644
--- a/ml/Query.h
+++ b/ml/Query.h
@@ -7,8 +7,8 @@ namespace ml {
 
 class Query {
 public:
-    Query(RRDDIM *RD) : RD(RD) {
-        Ops = &RD->tiers[0]->query_ops;
+    Query(RRDDIM *RD) : RD(RD), Initialized(false) {
+        Ops = RD->tiers[0]->query_ops;
     }
 
     time_t latestTime() {
@@ -20,27 +20,36 @@ public:
     }
 
     void init(time_t AfterT, time_t BeforeT) {
-        Ops->init(RD->tiers[0]->db_metric_handle, &Handle, AfterT, BeforeT, TIER_QUERY_FETCH_SUM);
+        Ops->init(RD->tiers[0]->db_metric_handle, &Handle, AfterT, BeforeT);
+        Initialized = true;
+        points_read = 0;
     }
 
     bool isFinished() {
         return Ops->is_finished(&Handle);
     }
 
+    ~Query() {
+        if (Initialized) {
+            Ops->finalize(&Handle);
+            global_statistics_ml_query_completed(points_read);
+            points_read = 0;
+        }
+    }
+
     std::pair<time_t, CalculatedNumber> nextMetric() {
+        points_read++;
         STORAGE_POINT sp = Ops->next_metric(&Handle);
         return { sp.start_time, sp.sum / sp.count };
     }
 
-    ~Query() {
-        Ops->finalize(&Handle);
-    }
-
 private:
     RRDDIM *RD;
+    bool Initialized;
+    size_t points_read;
 
-    struct rrddim_query_ops *Ops;
-    struct rrddim_query_handle Handle;
+    struct storage_engine_query_ops *Ops;
+    struct storage_engine_query_handle Handle;
 };
 
 } // namespace ml