Keep chunk requests and responses synchronized

Author: aij

splunklib/searchcommands/generating_command.py

@@ -194,6 +194,16 @@ 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
 
@@ -204,18 +214,10 @@ def _execute(self, ifile, process):
 
         """
         if self._protocol_version == 2:
-            result = self._read_chunk(ifile)
-
-            if not result:
-                return
-
-            metadata, body = result
-            action = getattr(metadata, 'action', None)
-
-            if action != 'execute':
-                raise RuntimeError('Expected execute action, not {}'.format(action))
-
-        self._record_writer.write_records(self.generate())
+            self._execute_v2(ifile, self.__generate_chunk)
+        else:
+            assert self._protocol_version == 1
+            self._record_writer.write_records(self.generate())
         self.finish()
 
     # endregion

splunklib/searchcommands/internals.py

@@ -654,6 +654,7 @@ def _write_record(self, record):
         self._record_count += 1
 
         if self._record_count >= self._maxresultrows:
+
             self.flush(partial=True)
 
     try:
@@ -748,42 +749,39 @@ class RecordWriterV2(RecordWriter):
     def flush(self, finished=None, partial=None):
 
         RecordWriter.flush(self, finished, partial)  # validates arguments and the state of this instance
-        inspector = self._inspector
 
-        if partial:
+        if partial or not finished:
             # Don't flush partial chunks, since the SCP v2 protocol does not
             # provide a way to send partial chunks yet.
             return
 
-        if self._flushed is False:
-
-            self._total_record_count += self._record_count
-            self._chunk_count += 1
-
-            # TODO: DVPL-6448: splunklib.searchcommands | Add support for partial: true when it is implemented in
-            # ChunkedExternProcessor (See SPL-103525)
-            #
-            # We will need to replace the following block of code with this block:
-            #
-            # metadata = [
-            #     ('inspector', self._inspector if len(self._inspector) else None),
-            #     ('finished', finished),
-            #     ('partial', partial)]
-
-            if len(inspector) == 0:
-                inspector = None
+        #if finished is True:
+        self.write_chunk(finished=True)
 
-            if partial is True:
-                finished = False
-
-            metadata = [item for item in (('inspector', inspector), ('finished', finished))]
-            self._write_chunk(metadata, self._buffer.getvalue())
-            self._clear()
-
-        elif finished is True:
-            self._write_chunk((('finished', True),), '')
-
-        self._finished = finished is True
+    def write_chunk(self, finished=None):
+        inspector = self._inspector
+        self._total_record_count += self._record_count
+        self._chunk_count += 1
+
+        # TODO: DVPL-6448: splunklib.searchcommands | Add support for partial: true when it is implemented in
+        # ChunkedExternProcessor (See SPL-103525)
+        #
+        # We will need to replace the following block of code with this block:
+        #
+        # metadata = [
+        #     ('inspector', self._inspector if len(self._inspector) else None),
+        #     ('finished', finished),
+        #     ('partial', partial)]
+
+        if len(inspector) == 0:
+            inspector = None
+
+        #if partial is True:
+        #    finished = False
+
+        metadata = [item for item in (('inspector', inspector), ('finished', finished))]
+        self._write_chunk(metadata, self._buffer.getvalue())
+        self._clear()
 
     def write_metadata(self, configuration):
         self._ensure_validity()

splunklib/searchcommands/search_command.py

@@ -776,7 +776,7 @@ def _process_protocol_v2(self, argv, ifile, ofile):
         # noinspection PyBroadException
         try:
             debug('Executing under protocol_version=2')
-            self._records = self._records_protocol_v2
+            #self._records = self._records_protocol_v2
             self._metadata.action = 'execute'
             self._execute(ifile, None)
         except SystemExit:
@@ -833,6 +833,8 @@ def _decode_list(mv):
 
     _encoded_value = re.compile(r'\$(?P<item>(?:\$\$|[^$])*)\$(?:;|$)')  # matches a single value in an encoded list
 
+    # Note: Subclasses must override this method so that it can be called
+    # called as self._execute(ifile, None)
     def _execute(self, ifile, process):
         """ Default processing loop
 
@@ -846,8 +848,12 @@ def _execute(self, ifile, process):
         :rtype: NoneType
 
         """
-        self._record_writer.write_records(process(self._records(ifile)))
-        self.finish()
+        if self.protocol_version == 1:
+            self._record_writer.write_records(process(self._records(ifile)))
+            self.finish()
+        else:
+            assert self._protocol_version == 2
+            self._execute_v2(ifile, process)
 
     @staticmethod
     def _read_chunk(ifile):
@@ -896,7 +902,9 @@ def _read_chunk(ifile):
     _header = re.compile(r'chunked\s+1.0\s*,\s*(\d+)\s*,\s*(\d+)\s*\n')
 
     def _records_protocol_v1(self, ifile):
+        return self._read_csv_records(ifile)
 
+    def _read_csv_records(self, ifile):
         reader = csv.reader(ifile, dialect=CsvDialect)
 
         try:
@@ -921,7 +929,7 @@ def _records_protocol_v1(self, ifile):
                     record[fieldname] = value
             yield record
 
-    def _records_protocol_v2(self, ifile):
+    def _execute_v2(self, ifile, process):
 
         while True:
             result = self._read_chunk(ifile)
@@ -931,41 +939,20 @@ def _records_protocol_v2(self, ifile):
 
             metadata, body = result
             action = getattr(metadata, 'action', None)
-
             if action != 'execute':
                 raise RuntimeError('Expected execute action, not {}'.format(action))
-
-            finished = getattr(metadata, 'finished', False)
             self._record_writer.is_flushed = False
 
-            if len(body) > 0:
-                reader = csv.reader(StringIO(body), dialect=CsvDialect)
+            self._execute_chunk_v2(process, result)
 
-                try:
-                    fieldnames = next(reader)
-                except StopIteration:
-                    return
+            self._record_writer.write_chunk()
 
-                mv_fieldnames = dict([(name, name[len('__mv_'):]) for name in fieldnames if name.startswith('__mv_')])
-
-                if len(mv_fieldnames) == 0:
-                    for values in reader:
-                        yield OrderedDict(izip(fieldnames, values))
-                else:
-                    for values in reader:
-                        record = OrderedDict()
-                        for fieldname, value in izip(fieldnames, values):
-                            if fieldname.startswith('__mv_'):
-                                if len(value) > 0:
-                                    record[mv_fieldnames[fieldname]] = self._decode_list(value)
-                            elif fieldname not in record:
-                                record[fieldname] = value
-                        yield record
-
-            if finished:
-                return
+    def _execute_chunk_v2(self, process, chunk):
+            metadata, body = chunk
+            if len(body) > 0:
+                records = self._read_csv_records(StringIO(body))
+                self._record_writer.write_records(process(records))
 
-            self.flush()
 
     def _report_unexpected_error(self):