Fix generating commands

aij · Shakeel · commit 9304b0a02f11 · 2020-02-03T11:06:00.000-08:00
diff --git a/splunklib/searchcommands/generating_command.py b/splunklib/searchcommands/generating_command.py
@@ -194,14 +194,6 @@ def generate(self):
         raise NotImplementedError('GeneratingCommand.generate(self)')
 
 
-    def __generate_chunk(self, unused_input):
-        count = 0
-        for row in generate():
-            yield row
-            count += 1
-            if count == self._record_writer._maxresultrows:
-                # count = 0
-                return
 
     def _execute(self, ifile, process):
         """ Execution loop
@@ -213,12 +205,22 @@ def _execute(self, ifile, process):
 
         """
         if self._protocol_version == 2:
-            self._execute_v2(ifile, self.__generate_chunk)
+            self._execute_v2(ifile, self.generate())
         else:
             assert self._protocol_version == 1
             self._record_writer.write_records(self.generate())
         self.finish()
 
+    def _execute_chunk_v2(self, process, chunk):
+        count = 0
+        for row in process:
+            self._record_writer.write_record(row)
+            count += 1
+            if count == self._record_writer._maxresultrows:
+                self._finished = False
+                return
+        self._finished = True
+
     # endregion
 
     # region Types
diff --git a/splunklib/searchcommands/search_command.py b/splunklib/searchcommands/search_command.py
@@ -945,7 +945,7 @@ def _execute_v2(self, ifile, process):
 
             self._execute_chunk_v2(process, result)
 
-            self._record_writer.write_chunk()
+            self._record_writer.write_chunk(finished=self._finished)
 
     def _execute_chunk_v2(self, process, chunk):
             metadata, body = chunk
