chore: Add repr function to BadgeMetadata (#397)

* Add repr function for BadgeMetadata Signed-off-by: dikshathakur3119 <[email protected]> * remove badge instance from repr function Signed-off-by: dikshathakur3119 <[email protected]> * bump version Signed-off-by: dikshathakur3119 <[email protected]> * Delete db,, schema from BadgeMetadata Signed-off-by: dikshathakur3119 <[email protected]> * Remove cluster Signed-off-by: dikshathakur3119 <[email protected]> * fix lint Signed-off-by: dikshathakur3119 <[email protected]>
amundsen-io · Nov 2, 2020 · 052ce5d · 052ce5d
1 parent 8a559f5
commit 052ce5d
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 28 deletions.
diff --git a/databuilder/models/badge.py b/databuilder/models/badge.py
@@ -32,19 +32,12 @@ class BadgeMetadata(Neo4jCsvSerializable):
     INVERSE_BADGE_RELATION_TYPE = 'BADGE_FOR'
 
     def __init__(self,
-                 db_name: str,
-                 schema: str,
                  start_label: str,  # Table, Dashboard, Column
                  start_key: str,
                  badges: List[Badge],
-                 cluster: str = 'gold',  # is this what we want as default for badges..?
                  ):
         self.badges = badges
 
-        self.db = db_name.lower()
-        self.schema = schema.lower()
-        self.cluster = cluster.lower()
-
         table_key_pattern = re.compile('[a-z]+://[a-zA-Z0-9_.-]+.[a-zA-Z0-9_.-]+/[a-zA-Z0-9_.-]+')
         dashboard_key_pattern = re.compile('[a-z]+_dashboard://[a-zA-Z0-9_.-]+.[a-zA-Z0-9_.-]+/[a-zA-Z0-9_.-]+')
         column_key_pattern = re.compile('[a-z]+://[a-zA-Z0-9_.-]+.[a-zA-Z0-9_.-]+/[a-zA-Z0-9_.-]+/[a-zA-Z0-9_.-]+')
@@ -65,6 +58,10 @@ def __init__(self,
         self._node_iter = iter(self.create_nodes())
         self._relation_iter = iter(self.create_relation())
 
+    def __repr__(self) -> str:
+        return 'BadgeMetadata({!r}, {!r})'.format(self.start_label,
+                                                  self.start_key)
+
     def create_next_node(self) -> Optional[Dict[str, Any]]:
         # return the string representation of the data
         try:

diff --git a/databuilder/models/table_metadata.py b/databuilder/models/table_metadata.py
@@ -378,12 +378,9 @@ def _create_next_node(self) -> Iterator[Any]:  # noqa: C901
                 yield col.description.get_node_dict(node_key)
 
             if col.badges:
-                badge_metadata = BadgeMetadata(db_name=self._get_database_key(),
-                                               schema=self._get_schema_key(),
-                                               start_label=ColumnMetadata.COLUMN_NODE_LABEL,
+                badge_metadata = BadgeMetadata(start_label=ColumnMetadata.COLUMN_NODE_LABEL,
                                                start_key=self._get_col_key(col),
-                                               badges=col.badges,
-                                               cluster=self._get_cluster_key())
+                                               badges=col.badges)
                 badge_nodes = badge_metadata.create_nodes()
                 for node in badge_nodes:
                     yield node
@@ -457,12 +454,9 @@ def _create_next_relation(self) -> Iterator[Any]:
                                                    self._get_col_key(col),
                                                    self._get_col_description_key(col, col.description))
             if col.badges:
-                badge_metadata = BadgeMetadata(db_name=self._get_database_key(),
-                                               schema=self._get_schema_key(),
-                                               start_label=ColumnMetadata.COLUMN_NODE_LABEL,
+                badge_metadata = BadgeMetadata(start_label=ColumnMetadata.COLUMN_NODE_LABEL,
                                                start_key=self._get_col_key(col),
-                                               badges=col.badges,
-                                               cluster=self._get_cluster_key())
+                                               badges=col.badges)
                 badge_relations = badge_metadata.create_relation()
                 for relation in badge_relations:
                     yield relation

diff --git a/setup.py b/setup.py
@@ -4,7 +4,7 @@
 from setuptools import setup, find_packages
 
 
-__version__ = '4.0.1'
+__version__ = '4.0.2'
 
 
 requirements = [

diff --git a/tests/unit/models/test_badge.py b/tests/unit/models/test_badge.py
@@ -19,11 +19,8 @@
 class TestBadge(unittest.TestCase):
     def setUp(self) -> None:
         super(TestBadge, self).setUp()
-        self.badge_metada = BadgeMetadata(db_name='hive',
-                                          schema=SCHEMA,
-                                          start_label='Column',
+        self.badge_metada = BadgeMetadata(start_label='Column',
                                           start_key='hive://default.base/test/ds',
-                                          cluster=CLUSTER,
                                           badges=[badge1, badge2])
 
     def test_get_badge_key(self) -> None:
@@ -54,23 +51,17 @@ def test_bad_key_entity_match(self) -> None:
 
         self.assertRaises(Exception,
                           BadgeMetadata,
-                          db_name='hive',
-                          schema=SCHEMA,
                           start_label=column_label,
                           start_key=table_key,
-                          cluster=CLUSTER,
                           badges=[badge1, badge2])
 
     def test_bad_entity_label(self) -> None:
         user_label = 'User'
         table_key = 'hive://default.base/test'
         self.assertRaises(Exception,
                           BadgeMetadata,
-                          db_name='hive',
-                          schema=SCHEMA,
                           start_label=user_label,
                           start_key=table_key,
-                          cluster=CLUSTER,
                           badges=[badge1, badge2])
 
     def test_create_relation(self) -> None: