refine compress, add test

Author: crimson-gao

aliyun/log/compress.py

@@ -0,0 +1,74 @@
+
+import zlib
+from enum import Enum
+from .logexception import LogException
+from .util import Util
+
+import lz4.block
+
+
+def lz_decompress(raw_size, data):
+    return lz4.block.decompress(data, uncompressed_size=raw_size)
+
+
+def lz_compresss(data):
+    return lz4.block.compress(data)[4:]
+
+
+class CompressType(Enum):
+    UNCOMPRESSED = 1
+    LZ4 = 2
+
+    @staticmethod
+    def default_compress_type():
+        return CompressType.LZ4
+
+    def __str__(self):
+        if self == CompressType.LZ4:
+            return 'lz4'
+        return ''
+
+
+class Compressor():
+    @staticmethod
+    def compress(data, compress_type):
+        # type: (bytes, CompressType) -> bytes
+        if compress_type == CompressType.LZ4:
+            return lz_compresss(data)
+
+        if compress_type == CompressType.UNCOMPRESSED:
+            return data
+
+        raise LogException('UnknownCompressType',
+                           'Unknown compress type: ' + compress_type)
+
+    @staticmethod
+    def decompress(data, raw_size, compress_type):
+        # type: (bytes, int, CompressType) -> bytes
+        if compress_type == CompressType.LZ4:
+            return lz_decompress(raw_size, data)
+
+        if compress_type == CompressType.UNCOMPRESSED:
+            return data
+
+        raise LogException('UnknownCompressType',
+                           'Unknown compress type: ' + compress_type)
+
+    @staticmethod
+    def decompress_response(header, response):
+        # type: (dict, bytes) -> bytes
+        compress_type_str = Util.h_v_td(
+            header, 'x-log-compresstype', '').lower()
+
+        if compress_type_str == 'lz4':
+            raw_size = int(Util.h_v_t(header, 'x-log-bodyrawsize'))
+            return Compressor.decompress(response, raw_size, CompressType.LZ4)
+
+        if compress_type_str == '':
+            return response
+
+        if compress_type_str in ('gzip', 'deflate'):
+            return zlib.decompress(response)
+
+        raise LogException('UnknownCompressType', 'Unknown compress type: ' +
+                           compress_type_str, resp_header=header, resp_body=response)

aliyun/log/logclient.py

@@ -62,15 +62,13 @@
 from .export_response import *
 from .common_response import *
 from .auth import *
+from .compress import CompressType, Compressor
 
 logger = logging.getLogger(__name__)
 
 if six.PY3:
     xrange = range
 
-lz4_available = Util.is_lz4_available()
-if lz4_available:
-    from .util import lz_decompress, lz_compresss
 
 CONNECTION_TIME_OUT = 120
 MAX_LIST_PAGING_SIZE = 500
@@ -384,13 +382,11 @@ def put_log_raw(self, project, logstore, log_group, compress=None):
         raw_body_size = len(body)
         headers = {'x-log-bodyrawsize': str(raw_body_size), 'Content-Type': 'application/x-protobuf'}
 
-        if compress is None or compress:
-            if lz4_available:
-                headers['x-log-compresstype'] = 'lz4'
-                body = lz_compresss(body)
-            else:
-                headers['x-log-compresstype'] = 'deflate'
-                body = zlib.compress(body)
+        need_compress = compress is None or compress
+        if need_compress:
+            compress_type = CompressType.default_compress_type()
+            headers['x-log-compresstype'] = str(compress_type)
+            body = Compressor.compress(body, compress_type)
 
         params = {}
         resource = '/logstores/' + logstore + "/shards/lb"
@@ -445,16 +441,12 @@ def put_logs(self, request):
                                    len(body) / 1024.0 / 1024))
 
         headers = {'x-log-bodyrawsize': str(len(body)), 'Content-Type': 'application/x-protobuf'}
-        is_compress = request.get_compress()
-
-        compress_data = None
-        if is_compress:
-            if lz4_available:
-                headers['x-log-compresstype'] = 'lz4'
-                compress_data = lz_compresss(body)
-            else:
-                headers['x-log-compresstype'] = 'deflate'
-                compress_data = zlib.compress(body)
+        
+        need_compress = request.get_compress() is None or request.get_compress()
+        if need_compress:
+            compress_type = CompressType.default_compress_type()
+            headers['x-log-compresstype'] = str(compress_type)
+            body = Compressor.compress(body, compress_type)
 
         params = {}
         logstore = request.get_logstore()
@@ -465,11 +457,7 @@ def put_logs(self, request):
         else:
             resource = '/logstores/' + logstore + "/shards/lb"
 
-        if is_compress:
-            (resp, header) = self._send('POST', project, compress_data, resource, params, headers)
-        else:
-            (resp, header) = self._send('POST', project, body, resource, params, headers)
-
+        (resp, header) = self._send('POST', project, body, resource, params, headers)
         return PutLogsResponse(header, resp)
 
     def list_logstores(self, request):
@@ -654,12 +642,12 @@ def get_log(self, project, logstore, from_time, to_time, topic=None,
                 params['forward'] = forward
                 body_str = six.b(json.dumps(params))
                 headers["x-log-bodyrawsize"] = str(len(body_str))
-                accept_encoding = "lz4" if lz4_available else "deflate"
+                accept_encoding = str(CompressType.default_compress_type())
                 headers['Accept-Encoding'] = accept_encoding
 
                 (resp, header) = self._send("POST", project, body_str, resource, None, headers, respons_body_type=accept_encoding)
 
-                raw_data = Util.uncompress_response(header, resp)
+                raw_data = Compressor.decompress_response(header, resp)
                 exJson = self._loadJson(200, header, raw_data, requestId=Util.h_v_td(header, 'x-log-requestid', ''))
                 exJson = Util.convert_unicode_to_str(exJson)
                 ret = GetLogsResponse(exJson, header)
@@ -1128,11 +1116,10 @@ def pull_logs(self, project_name, logstore_name, shard_id, cursor, count=None, e
         """
 
         headers = {}
-        if compress is None or compress:
-            if lz4_available:
-                headers['Accept-Encoding'] = 'lz4'
-            else:
-                headers['Accept-Encoding'] = 'gzip'
+
+        need_compress = compress is None or compress
+        if need_compress:
+            headers['Accept-Encoding'] = str(CompressType.default_compress_type())
         else:
             headers['Accept-Encoding'] = ''
 
@@ -1155,18 +1142,9 @@ def pull_logs(self, project_name, logstore_name, shard_id, cursor, count=None, e
         raw_size = int(Util.h_v_t(header, 'x-log-bodyrawsize'))
         if raw_size <= 0:
             return PullLogResponse(None, header)
-        compress_type = Util.h_v_td(header, 'x-log-compresstype', '').lower()
-        if compress_type == 'lz4':
-            if lz4_available:
-                raw_data = lz_decompress(raw_size, resp)
-                return PullLogResponse(raw_data, header)
-            else:
-                raise LogException("ClientHasNoLz4", "There's no Lz4 lib available to decompress the response", resp_header=header, resp_body=resp)
-        elif compress_type in ('gzip', 'deflate'):
-            raw_data = zlib.decompress(resp)
-            return PullLogResponse(raw_data, header)
-        else:
-            return PullLogResponse(resp, header)
+
+        raw_data = Compressor.decompress_response(header, resp)
+        return PullLogResponse(raw_data, header)
 
     def pull_log(self, project_name, logstore_name, shard_id, from_time, to_time, batch_size=None, compress=None, query=None):
         """ batch pull log data from log service using time-range

aliyun/log/putlogsrequest.py

@@ -29,7 +29,7 @@ class PutLogsRequest(LogRequest):
     :param hashKey: put data with set hash, the data will be send to shard whose range contains the hashKey
 
     :type compress: bool
-    :param compress: if need to compress the logs
+    :param compress: if need to compress the logs, default is True
 
     :type logtags: list
     :param logtags: list of key:value tag pair , [(tag_key_1,tag_value_1) , (tag_key_2,tag_value_2)]

aliyun/log/util.py

@@ -17,27 +17,7 @@
 import re
 import logging
 import json
-import struct
-import zlib
 
-try:
-    import lz4
-    if hasattr(lz4, 'loads') and hasattr(lz4, 'dumps'):
-        def lz_decompress(raw_size, data):
-            return lz4.loads(struct.pack('<I', raw_size) + data)
-
-        def lz_compresss(data):
-            return lz4.dumps(data)[4:]
-    else:
-        import lz4.block
-        def lz_decompress(raw_size, data):
-            return lz4.block.decompress(data, uncompressed_size=raw_size)
-
-        def lz_compresss(data):
-            return lz4.block.compress(data)[4:]
-    lz4_available = True
-except ImportError:
-    lz4_available = False
 
 logger = logging.getLogger(__name__)
 
@@ -91,11 +71,6 @@ def get_host_ip(logHost):
             if s:
                 s.close()
 
-    @staticmethod
-    def compress_data(data):
-        import zlib
-        return zlib.compress(data, 6)
-
     @staticmethod
     def cal_md5(content):
         return hashlib.md5(content).hexdigest().upper()
@@ -164,9 +139,6 @@ def convert_unicode_to_str(data):
 
         return data
 
-    @staticmethod
-    def is_lz4_available():
-        return lz4_available
 
     @staticmethod
     def h_v_t(header, key):
@@ -208,23 +180,7 @@ def v_or_d(v, default=None):
             else return v
         """
         return v if v is not None else default
-
-    @staticmethod
-    def uncompress_response(header, response):
-        """
-        Supported compress type [lz4, gzip, deflate]
-        """
-        compress_type = Util.h_v_td(header, 'x-log-compresstype', '').lower()
-        if compress_type == 'lz4':
-            raw_size = int(Util.h_v_t(header, 'x-log-bodyrawsize'))
-            if Util.is_lz4_available():
-                return lz_decompress(raw_size, response)
-            else:
-                raise LogException("ClientHasNoLz4", "There's no Lz4 lib available to decompress the response", resp_header=header, resp_body=response)
-        elif compress_type in ('gzip', 'deflate'):
-            return zlib.decompress(response)
-        else:
-            raise LogException('UnknownCompressType', 'Unknown compress type: ' + compress_type, resp_header=header, resp_body=response)
+    
 
 ZERO = timedelta(0)
 

tests/compress_test.py

@@ -0,0 +1,66 @@
+import time
+from aliyun.log.compress import Compressor, CompressType
+from aliyun.log import LogClient, PutLogsRequest, LogItem, LogGroup
+
+
+def test_lz4():
+    text = b'sadsadsa189634o2??ASBKHD'
+    compressed = Compressor.compress(text, CompressType.LZ4)
+    raw_size = len(text)
+    uncompressed = Compressor.decompress(
+        compressed, raw_size, CompressType.LZ4)
+
+    assert text == uncompressed, "The decompressed data does not match the original"
+
+
+# setup
+project = ''
+logstore = ''
+access_key_id = ''
+access_key_secret = ''
+endpoint = ''
+client = LogClient(endpoint, access_key_id, access_key_secret)
+
+
+def test_pull_logs():
+    shards = client.list_shards(project, logstore).get_shards_info()
+    shard = shards[0]['shardID']
+    cursor = client.get_cursor(project, logstore, shard, 'begin').get_cursor()
+    resp = client.pull_logs(project, logstore, shard, cursor, compress=None)
+    resp.log_print()
+    client.pull_logs(project, logstore, shard, cursor, compress=True)
+    client.pull_logs(project, logstore, shard, cursor, compress=False)
+
+
+def test_put_log_raw():
+    log_group = LogGroup()
+    log = log_group.Logs.add()
+    log.Time = int(time.time())
+    for i in range(10):
+        content = log.Contents.add()
+        content.Key = client._get_unicode('test' + str(i))
+        content.Value = client._get_binary('test' + str(i))
+    client.put_log_raw(project, logstore, log_group, True)
+    client.put_log_raw(project, logstore, log_group, False)
+    client.put_log_raw(project, logstore, log_group, None)
+
+
+def test_put_logs():
+    client.put_logs(PutLogsRequest(project, logstore, 'test',
+                    '', [LogItem(contents=[('ghello', 'test')])], compress=True))
+    client.put_logs(PutLogsRequest(project, logstore, 'test',
+                    '', [LogItem(contents=[('ghello', 'test')])], compress=None))
+    client.put_logs(PutLogsRequest(project, logstore, 'test',
+                    '', [LogItem(contents=[('ghello', 'test')])], compress=False))
+
+
+def test_get_log():
+    client._get_logs_v2_enabled = True
+    from_time = int(time.time()) - 100
+    to_time = int(time.time())
+    resp = client.get_log(project, logstore, from_time, to_time, query='*')
+    resp.log_print()
+
+    client._get_logs_v2_enabled = False
+    resp = client.get_log(project, logstore, from_time, to_time, query='*')
+    resp.log_print()

tests/logclient/test_get_logs.py

@@ -18,7 +18,6 @@ def setUp(self):
             self.endpoint, self.access_key_id, self.access_key_secret)
         self.to_time = int(time.time())
         self.from_time = self.to_time - 3600
-        print('lz4 is available:', Util.is_lz4_available())
 
     def tearDown(self):
         pass