Add support for new-style gzip-compressed objects (#88)

Author: whd

moztelemetry/heka/message_parser.py

@@ -8,6 +8,7 @@
 import boto
 import snappy
 import ujson as json
+import json as standard_json
 from cStringIO import StringIO
 from google.protobuf.message import DecodeError
 
@@ -24,24 +25,34 @@ def parse_heka_message(message):
 
 
 def _parse_heka_record(record):
-    if record.message.payload:
-        result = json.loads(record.message.payload)
-    else:
-        result = {}
-    result["meta"] = {
+    result = {"meta": {
         # TODO: uuid, logger, severity, env_version, pid
         "Timestamp": record.message.timestamp,
         "Type":      record.message.type,
         "Hostname":  record.message.hostname,
-    }
+    }}
+
+    if record.message.payload:
+        payload = _parse_json(record.message.payload)
+    else:
+        payload = {}
 
     for field in record.message.fields:
         name = field.name.split('.')
         value = field.value_string
         if field.value_type == 1:
-            # TODO: handle bytes in a way that doesn't cause problems with JSON
-            # value = field.value_bytes
-            continue
+            # Non-UTF8 bytes fields are currently not supported
+            try:
+                string = field.value_bytes[0].decode('utf-8')
+            except UnicodeDecodeError:
+                continue
+            # Special case: the submission field (bytes) replaces the top level
+            # Payload in the hindsight-based infra
+            if name[0] == 'submission':
+                payload = _parse_json(string)
+                continue
+            # handle bytes in a way that doesn't cause problems with JSON
+            value = [string]
         elif field.value_type == 2:
             value = field.value_integer
         elif field.value_type == 3:
@@ -54,7 +65,12 @@ def _parse_heka_record(record):
         else:
             _add_field(result, name, value)
 
-    return result
+    # merge results back into the payload/submission
+    # payload may contain NULLed out structures in the new infra and must
+    # therefore be the receiver of the update
+    payload.update(result)
+
+    return payload
 
 
 def _add_field(container, keys, value):
@@ -68,6 +84,15 @@ def _add_field(container, keys, value):
     _add_field(container[key], keys, value)
 
 
+def _parse_json(string):
+    try:
+        result = json.loads(string)
+    except:
+        # Fall back to the standard parser if ujson fails
+        result = standard_json.loads(string)
+    return result
+
+
 def _lazyjson(content):
     if not isinstance(content, basestring):
         raise ValueError("Argument must be a string.")
@@ -88,7 +113,7 @@ class WrapperType(type(default)):
     def wrap(method_name):
         def _wrap(*args, **kwargs):
             if not hasattr(WrapperType, '__cache__'):
-                setattr(WrapperType, '__cache__', json.loads(content))
+                setattr(WrapperType, '__cache__', _parse_json(content))
 
             cached = WrapperType.__cache__
             method = getattr(cached, method_name)

moztelemetry/spark.py

@@ -10,6 +10,8 @@
 from .dataset import Dataset
 from .histogram import Histogram
 
+from .util.streaming_gzip import streaming_gzip_wrapper
+
 logger = logging.getLogger(__name__)
 
 if not boto.config.has_section('Boto'):

moztelemetry/store.py

@@ -5,6 +5,8 @@
 
 import boto3
 
+from .util.streaming_gzip import streaming_gzip_wrapper
+
 
 class S3Store:
 
@@ -36,7 +38,11 @@ def get_key(self, key):
         try:
             # get_key must return a file-like object because that's what's
             # required by parse_heka_message
-            return bucket.Object(key).get()['Body']
+            s3object = bucket.Object(key).get()
+            if s3object.get('ContentEncoding') == "gzip":
+                return streaming_gzip_wrapper(s3object['Body'])
+            else:
+                return s3object['Body']
         except:
             raise Exception('Error retrieving key "{}" from S3'.format(key))
 

moztelemetry/util/__init__.py

@@ -0,0 +1,3 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, you can obtain one at http://mozilla.org/MPL/2.0/.

moztelemetry/util/streaming_gzip.py

@@ -0,0 +1,62 @@
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, you can obtain one at http://mozilla.org/MPL/2.0/.
+
+from gzip import GzipFile
+from StringIO import StringIO
+from os import SEEK_SET, SEEK_END
+
+# boto's s3.Key and botocore's StreamingBody do not support the seek()
+# operation, which we need in order to pass the stream along to a
+# GzipFile for gzip decompression. However, the implementation of
+# GzipFile only uses seek() in combination with tell() on its fileobj
+# to determine whether EOF has been reached when parsing a new
+# member. Since in our case we know:
+# - there is only one member per fileobj
+# - it does not begin at EOF
+# - the only remaining data after decompression will be the final 8 bytes of
+#   gzip member verification metadata
+# we can create:
+# - a modified GzipFile that does not need to use seek() to perform the
+#   final CRC and size verification checks
+# - a wrapper around StreamingBody that ensures the first member passes the
+#   not-EOF check without needing to actually perform a seek()
+# Combining these two shims gives us a way to transparently perform streaming
+# gzip decompression of s3 objects.
+
+# Refer to the implementation of GzipFile for more details.
+
+
+class StreamingGzipFile(GzipFile):
+    def __init__(self, **args):
+        super(StreamingGzipFile, self).__init__(**args)
+
+    def _read_eof(self):
+        remainder = self.decompress.unused_data + self.fileobj.read()
+        assert(len(remainder) == 8)
+        self.fileobj.close()
+        self.fileobj = StringIO(remainder)
+        super(StreamingGzipFile, self)._read_eof()
+
+
+class GzipStreamingBody:
+    def __init__(self, sb):
+        self.sb = sb
+        self.eof = False
+
+    def tell(self):
+        return int(self.eof)
+
+    def close(self):
+        return self.sb.close()
+
+    def read(self, n=-1):
+        return self.sb.read(n)
+
+    def seek(self, offset, whence=SEEK_SET):
+        if whence == SEEK_END:
+            self.eof = True
+
+
+def streaming_gzip_wrapper(body):
+    return StreamingGzipFile(fileobj=GzipStreamingBody(body))

tests/data/test_gzip.heka

@@ -6,12 +6,15 @@
 # file, You can obtain one at http://mozilla.org/MPL/2.0/.
 import json
 from moztelemetry.heka import message_parser
+from moztelemetry.util.streaming_gzip import streaming_gzip_wrapper
 
 
 def test_unpack(data_dir):
-    for t in ["plain", "snappy", "mixed"]:
+    for t in ["plain", "snappy", "mixed", "gzip", "gzip_mixed"]:
         filename = "{}/test_{}.heka".format(data_dir, t)
         with open(filename, "rb") as o:
+            if "gzip" in t:
+                o = streaming_gzip_wrapper(o)
             msg = 0
             for r, b in message_parser.unpack(o, try_snappy=True):
                 j = json.loads(r.message.payload)
@@ -21,11 +24,14 @@ def test_unpack(data_dir):
 
 
 def test_unpack_nosnappy(data_dir):
-    expected_counts = {"plain": 10, "snappy": 0, "mixed": 5}
+    expected_counts = {"plain": 10, "snappy": 0, "mixed": 5,
+                       "gzip": 10, "gzip_mixed": 5}
     for t in expected_counts.keys():
         count = 0
         filename = "{}/test_{}.heka".format(data_dir, t)
         with open(filename, "rb") as o:
+            if "gzip" in t:
+                o = streaming_gzip_wrapper(o)
             try:
                 for r, b in message_parser.unpack(o, try_snappy=False):
                     count += 1
@@ -35,13 +41,16 @@ def test_unpack_nosnappy(data_dir):
 
 
 def test_unpack_strict(data_dir):
-    expected_exceptions = {"plain": False, "snappy": True, "mixed": True}
+    expected_exceptions = {"plain": False, "snappy": True, "mixed": True,
+                           "gzip": False, "gzip_mixed": True}
     for t in expected_exceptions.keys():
         count = 0
         filename = "{}/test_{}.heka".format(data_dir, t)
         threw = False
         got_err = False
         with open(filename, "rb") as o:
+            if "gzip" in t:
+                o = streaming_gzip_wrapper(o)
             try:
                 for r, b in message_parser.unpack(o, strict=True, try_snappy=False):
                     if r.error is not None: