Extend TableMetadata model to accept KV pairs attributes (#72)

databuilder/models/table_metadata.py

+import copy
 from collections import namedtuple
 
 from typing import Iterable, Any, Union, Iterator, Dict, Set  # noqa: F401
@@ -106,6 +107,7 @@ class TableMetadata(Neo4jCsvSerializable):
                  description,  # type: Union[str, None]
                  columns=None,  # type: Iterable[ColumnMetadata]
                  is_view=False,  # type: bool
+                 **kwargs  # type: Dict
                  ):
         # type: (...) -> None
         """
@@ -116,6 +118,8 @@ class TableMetadata(Neo4jCsvSerializable):
         :param name:
         :param description:
         :param columns:
+        :param is_view: Indicate whether the table is a view or not
+        :param kwargs: Put additional attributes to the table model if there is any.
         """
         self.database = database
         self.cluster = cluster
@@ -124,6 +128,10 @@ class TableMetadata(Neo4jCsvSerializable):
         self.description = description
         self.columns = columns if columns else []
         self.is_view = is_view
+        self.attrs = None
+        if kwargs:
+            self.attrs = copy.deepcopy(kwargs)
+
         self._node_iterator = self._create_next_node()
         self._relation_iterator = self._create_next_relation()
 
@@ -191,10 +199,16 @@ class TableMetadata(Neo4jCsvSerializable):
 
     def _create_next_node(self):
         # type: () -> Iterator[Any]
-        yield {NODE_LABEL: TableMetadata.TABLE_NODE_LABEL,
+
+        table_node = {NODE_LABEL: TableMetadata.TABLE_NODE_LABEL,
                       NODE_KEY: self._get_table_key(),
                       TableMetadata.TABLE_NAME: self.name,
                       TableMetadata.IS_VIEW: self.is_view}
+        if self.attrs:
+            for k, v in self.attrs.items():
+                if k not in table_node:
+                    table_node[k] = v
+        yield table_node
 
         if self.description:
             yield {NODE_LABEL: DESCRIPTION_NODE_LABEL,

requirements.txt

@@ -49,7 +49,7 @@ six==1.11.0
 sqlalchemy==1.3.0
 wheel==0.31.1
 neo4j-driver==1.7.2
-neotime==1.0.0
+neotime==1.7.1
 pytz==2018.4
 antlr4-python2-runtime==4.7.1
 statsd==3.2.1

setup.py

 from setuptools import setup, find_packages
 
 
-__version__ = '1.2.3'
+__version__ = '1.2.4'
 
 
 setup(

tests/unit/models/test_table_metadata.py

@@ -23,6 +23,14 @@ class TestTableMetadata(unittest.TestCase):
             ColumnMetadata('etl_created_at', 'description of etl_created_at', 'timestamp', 4),
             ColumnMetadata('ds', None, 'varchar', 5)])
 
+        self.table_metadata3 = TableMetadata('hive', 'gold', 'test_schema3', 'test_table3', 'test_table3', [
+            ColumnMetadata('test_id1', 'description of test_table1', 'bigint', 0),
+            ColumnMetadata('test_id2', 'description of test_id2', 'bigint', 1),
+            ColumnMetadata('is_active', None, 'boolean', 2),
+            ColumnMetadata('source', 'description of source', 'varchar', 3),
+            ColumnMetadata('etl_created_at', 'description of etl_created_at', 'timestamp', 4),
+            ColumnMetadata('ds', None, 'varchar', 5)], is_view=False, attr1='uri', attr2='attr2')
+
         self.expected_nodes_deduped = [
             {'name': 'test_table1', 'KEY': 'hive://gold.test_schema1/test_table1', 'LABEL': 'Table',
              'is_view:UNQUOTED': False},
@@ -132,6 +140,15 @@ class TestTableMetadata(unittest.TestCase):
 
         self.assertEqual(self.expected_rels_deduped, actual)
 
+        node_row = self.table_metadata3.next_node()
+        t2_actual = []
+        while node_row:
+            t2_actual.append(node_row)
+            node_row = self.table_metadata3.next_node()
+
+        self.assertEqual(t2_actual[0].get('attr1'), 'uri')
+        self.assertEqual(t2_actual[0].get('attr2'), 'attr2')
+
 
 if __name__ == '__main__':
     unittest.main()