Extend user model to accept any KV attribute (amundsen-io#95)

databuilder/databuilder/models/user.py

@@ -1,3 +1,4 @@
+import copy
 from typing import Union, Dict, Any  # noqa: F401
 
 from databuilder.models.neo4j_csv_serde import Neo4jCsvSerializable, NODE_KEY, \
@@ -40,6 +41,7 @@ def __init__(self,
                  slack_id='',  # type: str
                  is_active=True,  # type: bool
                  updated_at=0,  # type: int
+                 **kwargs  # type: Dict
                  ):
         # type: (...) -> None
         """
@@ -57,6 +59,7 @@ def __init__(self,
         :param updated_at: everytime we update the node, we will push the timestamp.
                            then we will have a cron job to update the ex-employee nodes based on
                            the case if this timestamp hasn't been updated for two weeks.
+        :param kwargs: Any K/V attributes we want to update the
         """
         self.first_name = first_name
         self.last_name = last_name
@@ -72,6 +75,9 @@ def __init__(self,
         self.slack_id = slack_id
         self.is_active = is_active
         self.updated_at = updated_at
+        self.attrs = None
+        if kwargs:
+            self.attrs = copy.deepcopy(kwargs)
 
         self._node_iter = iter(self.create_nodes())
         self._rel_iter = iter(self.create_relation())
@@ -124,6 +130,12 @@ def create_nodes(self):
         result_node[User.USER_NODE_SLACK_ID] = self.slack_id if self.slack_id else ''
         result_node[User.USER_NODE_UPDATED_AT] = self.updated_at if self.updated_at else 0
 
+        if self.attrs:
+            print (self.attrs)
+            for k, v in self.attrs.items():
+                if k not in result_node:
+                    result_node[k] = v
+
         return [result_node]
 
     def create_relation(self):

databuilder/setup.py

@@ -1,7 +1,7 @@
 from setuptools import setup, find_packages
 
 
-__version__ = '1.3.3'
+__version__ = '1.3.4'
 
 
 setup(

databuilder/tests/unit/models/test_user.py

@@ -33,6 +33,25 @@ def test_create_nodes(self):
         nodes = self.user.create_nodes()
         self.assertEquals(len(nodes), 1)
 
+    def test_create_node_additional_attr(self):
+        test_user = User(first_name='test_first',
+                         last_name='test_last',
+                         name='test_first test_last',
+                         email='test@email.com',
+                         github_username='github_test',
+                         team_name='test_team',
+                         employee_type='FTE',
+                         manager_email='test_manager@email.com',
+                         slack_id='slack',
+                         is_active=True,
+                         updated_at=1,
+                         role='SWE',
+                         enable_notify=True)
+        nodes = test_user.create_nodes()
+        self.assertEqual(nodes[0]['email'], 'test@email.com')
+        self.assertEqual(nodes[0]['role'], 'SWE')
+        self.assertTrue(nodes[0]['enable_notify'])
+
     def test_create_relation(self):
         # type: () -> None
         relations = self.user.create_relation()