Add message representation (singer-io#6)

mdelaurentis · web-flow · commit e0e5fc7d965b · 2017-02-24T14:25:25.000-05:00
* Working on message classes

* Use namedtuple for message classes

* Add unit tests

* Pep8

* Add circle.yml

* Use vars

* Don't use namedtuple

* Don't use deprecated method

* Pep8
diff --git a/circle.yml b/circle.yml
@@ -0,0 +1,5 @@
+machine:
+  python:
+    version: 3.4.4
+    
+  
diff --git a/singer/__init__.py b/singer/__init__.py
@@ -4,11 +4,61 @@
 import logging
 import logging.config
 
-from collections import namedtuple
 
+class Message(object):
+    def __init__(self, **kwargs):
+        for k in self.attr_list:
+            if k not in kwargs:
+                raise ValueError("missing {}".format(k))
+            setattr(self, k, kwargs[k])
 
-def _writeline(s):
-    sys.stdout.write(s + '\n')
+    def asdict(self):
+        res = {k: getattr(self, k) for k in self.attr_list}
+        res['type'] = self._type
+        return res
+
+    def __eq__(self, other):
+        return self.asdict() == other.asdict()
+
+    def __repr__(self):
+        attrstr = ", ".join(
+            "{}={}".format(k, getattr(self, k)) for k in self.attr_list)
+        return "{}({})".format(self.__class__.__name__, attrstr)
+
+    def tojson(self):
+        return json.dumps(self.asdict())
+
+
+class RecordMessage(Message):
+    _type = 'RECORD'
+    attr_list = ['stream', 'record']
+
+
+class SchemaMessage(Message):
+    _type = 'SCHEMA'
+    attr_list = ['stream', 'schema', 'key_properties']
+
+
+class StateMessage(Message):
+    _type = 'STATE'
+    attr_list = ['value']
+
+
+def to_json(message):
+    m = vars(message)
+    if isinstance(message, RecordMessage):
+        m['type'] = 'RECORD'
+    elif isinstance(message, SchemaMessage):
+        m['type'] = 'SCHEMA'
+    elif isinstance(message, StateMessage):
+        m['type'] = 'STATE'
+    else:
+        raise Exception('Unrecognized message {}'.format(message))
+    return json.dumps(m)
+
+
+def _write_message(message):
+    sys.stdout.write(to_json(message) + '\n')
     sys.stdout.flush()
 
 
@@ -17,9 +67,7 @@ def write_record(stream_name, record):
 
     >>> write_record("users", {"id": 2, "email": "mike@stitchdata.com"})
     """
-    _writeline(json.dumps({'type': 'RECORD',
-                           'stream': stream_name,
-                           'record': record}))
+    _write_message(RecordMessage(stream=stream_name, record=record))
 
 
 def write_records(stream_name, records):
@@ -45,19 +93,44 @@ def write_schema(stream_name, schema, key_properties):
         key_properties = [key_properties]
     if not isinstance(key_properties, list):
         raise Exception("key_properties must be a string or list of strings")
-    _writeline(json.dumps({'type': 'SCHEMA',
-                           'stream': stream_name,
-                           'key_properties': key_properties,
-                           'schema': schema}))
+    _write_message(
+        SchemaMessage(
+            stream=stream_name,
+            schema=schema,
+            key_properties=key_properties))
 
 
 def write_state(value):
     """Write a state message.
 
     >>> write_state({'last_updated_at': '2017-02-14T09:21:00'})
     """
-    _writeline(json.dumps({'type': 'STATE',
-                           'value': value}))
+    _write_message(StateMessage(value))
+
+
+def _required_key(msg, k):
+    if k not in msg:
+        raise Exception("Message is missing required key '{}': {}".format(
+            k, msg))
+    return msg[k]
+
+
+def parse_message(s):
+    """Parse a message string into a Message object."""
+    o = json.loads(s)
+    t = _required_key(o, 'type')
+
+    if t == 'RECORD':
+        return RecordMessage(stream=_required_key(o, 'stream'),
+                             record=_required_key(o, 'record'))
+
+    elif t == 'SCHEMA':
+        return SchemaMessage(stream=_required_key(o, 'stream'),
+                             schema=_required_key(o, 'schema'),
+                             key_properties=_required_key(o, 'key_properties'))
+
+    elif t == 'STATE':
+        return StateMessage(value=_required_key(o, 'value'))
 
 
 def get_logger():
diff --git a/tests/__init__.py b/tests/__init__.py
diff --git a/tests/test_singer.py b/tests/test_singer.py
@@ -0,0 +1,79 @@
+import singer
+import unittest
+
+
+class TestSinger(unittest.TestCase):
+    def test_parse_message_record_good(self):
+        message = singer.parse_message(
+            '{"type": "RECORD", "record": {"name": "foo"}, "stream": "users"}')
+        self.assertEqual(
+            message,
+            singer.RecordMessage(record={'name': 'foo'}, stream='users'))
+
+    def test_parse_message_record_missing_record(self):
+        with self.assertRaises(Exception):
+            singer.parse_message('{"type": "RECORD", "stream": "users"}')
+
+    def test_parse_message_record_missing_stream(self):
+        with self.assertRaises(Exception):
+            singer.parse_message(
+                '{"type": "RECORD", "record": {"name": "foo"}}')
+
+    def test_parse_message_schema_good(self):
+        message = singer.parse_message('{"type": "SCHEMA", "stream": "users", "schema": {"type": "object", "properties": {"name": {"type": "string"}}}, "key_properties": ["name"]}')  # nopep8
+        self.assertEqual(
+            message,
+            singer.SchemaMessage(
+                stream='users',
+                key_properties=['name'],
+                schema={'type': 'object',
+                        'properties': {
+                            'name': {'type': 'string'}}}))
+
+    def test_parse_message_schema_missing_stream(self):
+        with self.assertRaises(Exception):
+            message = singer.parse_message('{"type": "SCHEMA", "schema": {"type": "object", "properties": {"name": {"type": "string"}}}, "key_properties": ["name"]}')  # nopep8
+
+    def test_parse_message_schema_missing_schema(self):
+        with self.assertRaises(Exception):
+            message = singer.parse_message(
+                '{"type": "SCHEMA", "stream": "users", "key_properties": ["name"]}')  # nopep8
+
+    def test_parse_message_schema_missing_key_properties(self):
+        with self.assertRaises(Exception):
+            message = singer.parse_message('{"type": "SCHEMA", "stream": "users", "schema": {"type": "object", "properties": {"name": {"type": "string"}}}}')  # nopep8
+
+    def test_parse_message_state_good(self):
+        message = singer.parse_message(
+            '{"type": "STATE", "value": {"seq": 1}}')
+        self.assertEqual(message, singer.StateMessage(value={'seq': 1}))
+
+    def test_parse_message_state_missing_value(self):
+        with self.assertRaises(Exception):
+            singer.parse_message('{"type": "STATE"}')
+
+    def test_round_trip(self):
+
+        record_message = singer.RecordMessage(
+            record={'name': 'foo'},
+            stream='users')
+
+        schema_message = singer.SchemaMessage(
+            stream='users',
+            key_properties=['name'],
+            schema={'type': 'object',
+                    'properties': {
+                        'name': {'type': 'string'}}})
+
+        state_message = singer.StateMessage(value={'seq': 1})
+
+        self.assertEqual(record_message,
+                         singer.parse_message(singer.to_json(record_message)))
+        self.assertEqual(schema_message,
+                         singer.parse_message(singer.to_json(schema_message)))
+        self.assertEqual(state_message,
+                         singer.parse_message(singer.to_json(state_message)))
+
+
+if __name__ == '__main__':
+    unittest.main()

-Original file line number
+Diff line change
@@ @@ -0,0 +1,5 @@ @@
 +machine:
 +  python:
 +    version: 3.4.4
++
++