DIP-151 Make list of attributes configurable via Env var

janrain · Mar 18, 2016 · aac924a · aac924a
1 parent 675e593
commit aac924a
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 21 deletions.
diff --git a/bluekai/bluekai_tsv.py b/bluekai/bluekai_tsv.py
@@ -1,15 +1,17 @@
-def fromRecords(records):
-    return str.join('', fromRecordsIterator(records))
+def fromRecords(records, keys=[]):
+    return str.join('', fromRecordsIterator(records, keys))
 
-def fromRecordsIterator(records):
+def fromRecordsIterator(records, keys=[]):
     for record in records:
-        yield fromRecord(record)
+        yield fromRecord(record, keys)
 
-def fromRecord(record):
+def fromRecord(record, keys=[]):
 
     uuid = record.get('uuid')
 
-    items = record.items()
+    sub_record = { key: record[key] for key in keys if key in record }
+
+    items = sub_record.items()
     items = filter(lambda pair: pair[0] != "uuid", items)
     items = map(lambda pair: "{}={}".format(pair[0], pair[1]), items)
     items = str.join('|', items)

diff --git a/bluekai/config.py b/bluekai/config.py
@@ -15,6 +15,7 @@
     'JANRAIN_CLIENT_SECRET': '',
     'JANRAIN_SCHEMA_NAME': 'user',
     'JANRAIN_BATCH_SIZE': 1000,
+    'JANRAIN_ATTRIBUTE_KEYS': "",
     'REMOTE_HOST': '',
     'REMOTE_PORT': 22,
     'REMOTE_USERNAME': '',

diff --git a/bluekai/job.py b/bluekai/job.py
@@ -36,6 +36,8 @@ def do_job(job, writter, config, logger, datalib, converter):
 
     records_iterator = recordsNewerThan(capture_app, config, last_updated)
 
+    attribute_keys = config['JANRAIN_ATTRIBUTE_KEYS'].split(',')
+
     with writter as fp:
 
         record_num = 0
@@ -46,7 +48,8 @@ def do_job(job, writter, config, logger, datalib, converter):
                 last_updated = fromRecordDateTime(record['lastUpdated'])
                 job.lastupdated = last_updated
 
-            row = converter(record)
+
+            row = converter(record, attribute_keys)
 
             fp.write(row)
 

diff --git a/tests/test_bluekai_tsv.py b/tests/test_bluekai_tsv.py
@@ -7,26 +7,54 @@ class bluekai_test(TestCase):
   def test_fromRecord(self):
     record = { "uuid": "a-b-c", "key1": "value1", "key2":"value2" }
     record = OrderedDict(sorted(record.items(), key=lambda t: t[0]))
-    actual = fromRecord(record)
-    expected = "a-b-c\tkey1=value1|key2=value2\n"
+    actual = fromRecord(record,["key2","key1"])
+    expected1 = "a-b-c\tkey1=value1|key2=value2\n"
+    expected2 = "a-b-c\tkey2=value2|key1=value1\n"
+    self.assertIn(actual, [expected1, expected2])
+
+  def test_fromRecord_default_keys(self):
+    record = { "uuid": "a-b-c", "key1": "value1", "key2":"value2" }
+    record = OrderedDict(sorted(record.items(), key=lambda t: t[0]))
+    actual = fromRecord(record,)
+    expected = "a-b-c\t\n"
     self.assertEqual(actual, expected)
 
   def test_fromRecordsIterator(self):
-    record1 = { "uuid": "a-b-c", "key1": "value1" }
-    record2 = { "uuid": "x-y-z", "key2": "value2" }
+    record1 = { "uuid": "a-b-c", "key": "value1" }
+    record2 = { "uuid": "x-y-z", "key": "value2" }
+    records = [record1, record2]
+    records_iterator = fromRecordsIterator(records, ["key"])
+    expected1 = "a-b-c\tkey=value1\n"
+    expected2 = "x-y-z\tkey=value2\n"
+    self.assertEqual(records_iterator.__next__(), expected1)
+    self.assertEqual(records_iterator.__next__(), expected2)
+    with self.assertRaises(StopIteration):
+        records_iterator.__next__()
+
+  def test_fromRecordsIterator_default_keys(self):
+    record1 = { "uuid": "a-b-c", "key": "value1" }
+    record2 = { "uuid": "x-y-z", "key": "value2" }
     records = [record1, record2]
     records_iterator = fromRecordsIterator(records)
-    expected1 = "a-b-c\tkey1=value1\n"
-    expected2 = "x-y-z\tkey2=value2\n"
+    expected1 = "a-b-c\t\n"
+    expected2 = "x-y-z\t\n"
     self.assertEqual(records_iterator.__next__(), expected1)
     self.assertEqual(records_iterator.__next__(), expected2)
     with self.assertRaises(StopIteration):
         records_iterator.__next__()
 
   def test_fromRecords(self):
-    record1 = { "uuid": "a-b-c", "key1": "value1" }
-    record2 = { "uuid": "x-y-z", "key2": "value2" }
+    record1 = { "uuid": "a-b-c", "key": "value1" }
+    record2 = { "uuid": "x-y-z", "key": "value2" }
+    records = [record1, record2]
+    actual = fromRecords(records, ["key"])
+    expected = "a-b-c\tkey=value1\nx-y-z\tkey=value2\n"
+    self.assertEqual(actual, expected)
+
+  def test_fromRecords_default_keys(self):
+    record1 = { "uuid": "a-b-c", "key": "value1" }
+    record2 = { "uuid": "x-y-z", "key": "value2" }
     records = [record1, record2]
-    actual = fromRecords(records)
-    expected = "a-b-c\tkey1=value1\nx-y-z\tkey2=value2\n"
+    actual = fromRecords(records, [])
+    expected = "a-b-c\t\nx-y-z\t\n"
     self.assertEqual(actual, expected)
diff --git a/tests/test_job.py b/tests/test_job.py
@@ -50,6 +50,7 @@ def setUp(self):
             'JANRAIN_CLIENT_ID': 'test_janrain_client_id',
             'JANRAIN_CLIENT_SECRET': 'test_janrain_client_secret',
             'JANRAIN_SCHEMA_NAME': 'test_janrain_schema_name',
+            'JANRAIN_ATTRIBUTE_KEYS': '',
             'JANRAIN_BATCH_SIZE': 2,
         }
         self.jobModel_mock = Mock()
@@ -79,10 +80,10 @@ def test_run(self):
         self.writter_mock.__enter__.assert_called_once_with()
         self.writter_mock.__exit__.assert_called_once_with(None, None, None)
         self.converter_mock.assert_has_calls([
-            call(self.records_iterator[0]),
-            call(self.records_iterator[1]),
-            call(self.records_iterator[2]),
-            call(self.records_iterator[3]),
+            call(self.records_iterator[0], ['']),
+            call(self.records_iterator[1], ['']),
+            call(self.records_iterator[2], ['']),
+            call(self.records_iterator[3], ['']),
         ])
         self.logger_mock.debug.assert_has_calls([
             call("wrote record 2"),