Merge pull request #341 from splunk/csc-multibyte

Custom search command support for multibyte characters in Python 3

CHANGELOG.md

@@ -1,12 +1,17 @@
 # Splunk SDK for Python Changelog
 
+## Version 1.6.14
+
+### Bug fix
+* `SearchCommand` now correctly supports multibyte characters in Python 3.
+
 ## Version 1.6.13
 
 ### Bug fix
-* Fixed regression in mod inputs which resulted in error ’file' object has no attribute 'readable’, by not forcing to text/bytes in mod inputs event writer any longer.  
+* Fixed regression in mod inputs which resulted in error ’file' object has no attribute 'readable’, by not forcing to text/bytes in mod inputs event writer any longer.
 
 ### Minor changes
-* Minor updates to the splunklib search commands to support Python3 
+* Minor updates to the splunklib search commands to support Python3
 
 ## Version 1.6.12
 
@@ -22,25 +27,25 @@
 
 ### Bug Fix
 
-* Fix custom search command V2 failures on Windows for Python3 
+* Fix custom search command V2 failures on Windows for Python3
 
 ## Version 1.6.10
 
 ### Bug Fix
 
-* Fix long type gets wrong values on windows for python 2 
+* Fix long type gets wrong values on windows for python 2
 
 ## Version 1.6.9
 
 ### Bug Fix
 
-* Fix buffered input in python 3 
+* Fix buffered input in python 3
 
 ## Version 1.6.8
 
 ### Bug Fix
 
-* Fix custom search command on python 3 on windows 
+* Fix custom search command on python 3 on windows
 
 ## Version 1.6.7
 

README.md

@@ -3,7 +3,7 @@
 
 # The Splunk Software Development Kit for Python
 
-#### Version 1.6.13
+#### Version 1.6.14
 
 The Splunk Software Development Kit (SDK) for Python contains library code and
 examples designed to enable developers to build applications using Splunk.

examples/searchcommands_app/setup.py

@@ -111,8 +111,8 @@ def splunk_restart(uri, auth):
 
 
 class AnalyzeCommand(Command):
-    """ 
-    setup.py command to run code coverage of the test suite. 
+    """
+    setup.py command to run code coverage of the test suite.
 
     """
     description = 'Create an HTML coverage report from running the full test suite.'
@@ -367,8 +367,8 @@ def _link_debug_client(self):
 
 
 class TestCommand(Command):
-    """ 
-    setup.py command to run the whole test suite. 
+    """
+    setup.py command to run the whole test suite.
 
     """
     description = 'Run full test suite.'
@@ -439,7 +439,7 @@ def run(self):
     setup(
         description='Custom Search Command examples',
         name=os.path.basename(project_dir),
-        version='1.6.13',
+        version='1.6.14',
         author='Splunk, Inc.',
         author_email='[email protected]',
         url='http://github.com/splunk/splunk-sdk-python',

splunklib/__init__.py

@@ -16,5 +16,5 @@
 
 from __future__ import absolute_import
 from splunklib.six.moves import map
-__version_info__ = (1, 6, 13)
+__version_info__ = (1, 6, 14)
 __version__ = ".".join(map(str, __version_info__))

splunklib/binding.py

@@ -1378,7 +1378,7 @@ def request(url, message, **kwargs):
         head = {
             "Content-Length": str(len(body)),
             "Host": host,
-            "User-Agent": "splunk-sdk-python/1.6.13",
+            "User-Agent": "splunk-sdk-python/1.6.14",
             "Accept": "*/*",
             "Connection": "Close",
         } # defaults

splunklib/searchcommands/generating_command.py

@@ -204,7 +204,7 @@ def _execute(self, ifile, process):
 
         """
         if self._protocol_version == 2:
-            result = self._read_chunk(ifile)
+            result = self._read_chunk(self._as_binary_stream(ifile))
 
             if not result:
                 return

splunklib/searchcommands/search_command.py

@@ -656,7 +656,7 @@ def _process_protocol_v2(self, argv, ifile, ofile):
         # noinspection PyBroadException
         try:
             debug('Reading metadata')
-            metadata, body = self._read_chunk(ifile)
+            metadata, body = self._read_chunk(self._as_binary_stream(ifile))
 
             action = getattr(metadata, 'action', None)
 
@@ -850,17 +850,29 @@ def _execute(self, ifile, process):
         self.finish()
 
     @staticmethod
-    def _read_chunk(ifile):
+    def _as_binary_stream(ifile):
+        if six.PY2:
+            return ifile
+
+        try:
+            return ifile.buffer
+        except AttributeError as error:
+            raise RuntimeError('Failed to get underlying buffer: {}'.format(error))
+
+    @staticmethod
+    def _read_chunk(istream):
         # noinspection PyBroadException
+        assert isinstance(istream.read(0), six.binary_type), 'Stream must be binary'
+
         try:
-            header = ifile.readline()
+            header = istream.readline()
         except Exception as error:
             raise RuntimeError('Failed to read transport header: {}'.format(error))
 
         if not header:
             return None
 
-        match = SearchCommand._header.match(header)
+        match = SearchCommand._header.match(six.ensure_str(header))
 
         if match is None:
             raise RuntimeError('Failed to parse transport header: {}'.format(header))
@@ -870,14 +882,14 @@ def _read_chunk(ifile):
         body_length = int(body_length)
 
         try:
-            metadata = ifile.read(metadata_length)
+            metadata = istream.read(metadata_length)
         except Exception as error:
             raise RuntimeError('Failed to read metadata of length {}: {}'.format(metadata_length, error))
 
         decoder = MetadataDecoder()
 
         try:
-            metadata = decoder.decode(metadata)
+            metadata = decoder.decode(six.ensure_str(metadata))
         except Exception as error:
             raise RuntimeError('Failed to parse metadata of length {}: {}'.format(metadata_length, error))
 
@@ -887,11 +899,11 @@ def _read_chunk(ifile):
         body = ""
         try:
             if body_length > 0:
-                body = ifile.read(body_length)
+                body = istream.read(body_length)
         except Exception as error:
             raise RuntimeError('Failed to read body of length {}: {}'.format(body_length, error))
 
-        return metadata, body
+        return metadata, six.ensure_str(body)
 
     _header = re.compile(r'chunked\s+1.0\s*,\s*(\d+)\s*,\s*(\d+)\s*\n')
 
@@ -922,9 +934,10 @@ def _records_protocol_v1(self, ifile):
             yield record
 
     def _records_protocol_v2(self, ifile):
+        istream = self._as_binary_stream(ifile)
 
         while True:
-            result = self._read_chunk(ifile)
+            result = self._read_chunk(istream)
 
             if not result:
                 return

tests/searchcommands/test_multibyte_processing.py

@@ -0,0 +1,39 @@
+import io
+import gzip
+import sys
+
+from os import path
+
+from splunklib import six
+from splunklib.searchcommands import StreamingCommand, Configuration
+
+
+def build_test_command():
+    @Configuration()
+    class TestSearchCommand(StreamingCommand):
+        def stream(self, records):
+            for record in records:
+                yield record
+
+    return TestSearchCommand()
+
+
+def get_input_file(name):
+    return path.join(
+        path.dirname(path.dirname(__file__)), 'data', 'custom_search', name + '.gz')
+
+
+def test_multibyte_chunked():
+    data = gzip.open(get_input_file("multibyte_input"))
+    if not six.PY2:
+        data = io.TextIOWrapper(data)
+    cmd = build_test_command()
+    cmd._process_protocol_v2(sys.argv, data, sys.stdout)
+
+
+def test_v1_searchcommand():
+    data = gzip.open(get_input_file("v1_search_input"))
+    if not six.PY2:
+        data = io.TextIOWrapper(data)
+    cmd = build_test_command()
+    cmd._process_protocol_v1(["test_script.py", "__EXECUTE__"], data, sys.stdout)

tests/searchcommands/test_search_command.py

@@ -33,8 +33,21 @@
 import os
 import re
 
+from io import TextIOWrapper
+
 import pytest
 
+def build_command_input(getinfo_metadata, execute_metadata, execute_body):
+    input = ('chunked 1.0,{},0\n{}'.format(len(six.ensure_binary(getinfo_metadata)), getinfo_metadata) +
+             'chunked 1.0,{},{}\n{}{}'.format(len(six.ensure_binary(execute_metadata)), len(six.ensure_binary(execute_body)), execute_metadata, execute_body))
+
+    ifile = BytesIO(six.ensure_binary(input))
+
+    if not six.PY2:
+        ifile = TextIOWrapper(ifile)
+
+    return ifile
+
 @Configuration()
 class TestCommand(SearchCommand):
 
@@ -428,11 +441,9 @@ def test_process_scpv2(self):
             show_configuration=('true' if show_configuration is True else 'false'))
 
         execute_metadata = '{"action":"execute","finished":true}'
-        execute_body = 'test\r\ndata\r\n'
+        execute_body = 'test\r\ndata\r\n测试\r\n'
 
-        ifile = StringIO(
-            'chunked 1.0,{},0\n{}'.format(len(getinfo_metadata), getinfo_metadata) +
-            'chunked 1.0,{},{}\n{}{}'.format(len(execute_metadata), len(execute_body), execute_metadata, execute_body))
+        ifile = build_command_input(getinfo_metadata, execute_metadata, execute_body)
 
         command = TestCommand()
         result = BytesIO()
@@ -455,12 +466,17 @@ def test_process_scpv2(self):
         self.assertEqual(command.required_option_1, 'value_1')
         self.assertEqual(command.required_option_2, 'value_2')
 
-        self.assertEqual(
+        expected = (
             'chunked 1.0,68,0\n'
             '{"inspector":{"messages":[["INFO","test command configuration: "]]}}\n'
-            'chunked 1.0,17,23\n'
+            'chunked 1.0,17,32\n'
             '{"finished":true}test,__mv_test\r\n'
-            'data,\r\n',
+            'data,\r\n'
+            '测试,\r\n'
+        )
+
+        self.assertEqual(
+            expected,
             result.getvalue().decode('utf-8'))
 
         self.assertEqual(command.protocol_version, 2)
@@ -620,11 +636,9 @@ def test_process_scpv2(self):
             show_configuration=show_configuration)
 
         execute_metadata = '{"action":"execute","finished":true}'
-        execute_body = 'test\r\ndata\r\n'
+        execute_body = 'test\r\ndata\r\n测试\r\n'
 
-        ifile = StringIO(
-            'chunked 1.0,{},0\n{}'.format(len(getinfo_metadata), getinfo_metadata) +
-            'chunked 1.0,{},{}\n{}{}'.format(len(execute_metadata), len(execute_body), execute_metadata, execute_body))
+        ifile = build_command_input(getinfo_metadata, execute_metadata, execute_body)
 
         command = TestCommand()
         result = BytesIO()
@@ -666,11 +680,9 @@ def test_process_scpv2(self):
             show_configuration=('true' if show_configuration is True else 'false'))
 
         execute_metadata = '{"action":"execute","finished":true}'
-        execute_body = 'action\r\nraise_exception\r\n'
+        execute_body = 'action\r\nraise_exception\r\n测试\r\n'
 
-        ifile = StringIO(
-            'chunked 1.0,{},0\n{}'.format(len(getinfo_metadata), getinfo_metadata) +
-            'chunked 1.0,{},{}\n{}{}'.format(len(execute_metadata), len(execute_body), execute_metadata, execute_body))
+        ifile = build_command_input(getinfo_metadata, execute_metadata, execute_body)
 
         command = TestCommand()
         result = BytesIO()