Limit the max number of intervals we remove in one stream_remove call

nilmdb/server/nilmdb.py

@@ -83,8 +83,11 @@ _sql_schema_updates = {
 class NilmDB(object):
     verbose = 0
 
-    def __init__(self, basepath, max_results=None,
-                 max_removals=None, bulkdata_args=None):
+    def __init__(self, basepath,
+                 max_results=None,
+                 max_removals=None,
+                 max_int_removals=None,
+                 bulkdata_args=None):
         """Initialize NilmDB at the given basepath.
         Other arguments are for debugging / testing:
 
@@ -92,7 +95,10 @@ class NilmDB(object):
         stream_intervals or stream_extract response.
 
         'max_removals' is the max rows to delete at once
-        in stream_move.
+        in stream_remove.
+
+        'max_int_removals' is the max intervals to delete
+        at once in stream_remove.
 
         'bulkdata_args' is kwargs for the bulkdata module.
         """
@@ -134,6 +140,9 @@ class NilmDB(object):
         # Remove up to this many rows per call to stream_remove.
         self.max_removals = max_removals or 1048576
 
+        # Remove up to this many intervals per call to stream_remove.
+        self.max_int_removals = max_int_removals or 4096
+
     def get_basepath(self):
         return self.basepath
 
@@ -643,6 +652,7 @@ class NilmDB(object):
         to_remove = Interval(start, end)
         removed = 0
         remaining = self.max_removals
+        int_remaining = self.max_int_removals
         restart = None
 
         # Can't remove intervals from within the iterator, so we need to
@@ -653,6 +663,11 @@ class NilmDB(object):
         remove_end = None
 
         for (dbint, orig) in all_candidates:
+            # Stop if we've hit the max number of interval removals
+            if int_remaining <= 0:
+                restart = dbint.start
+                break
+
             # Find row start and end
             row_start = self._find_start(table, dbint)
             row_end = self._find_end(table, dbint)
@@ -689,6 +704,7 @@ class NilmDB(object):
             # Count how many were removed
             removed += row_end - row_start
             remaining -= row_end - row_start
+            int_remaining -= 1
 
             if restart is not None:
                 break