diff --git a/databuilder/models/badge.py b/databuilder/models/badge.py index e18e28da4..9322b82ba 100644 --- 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 index 20c8afc51..94b285622 100644 --- 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 index ef02a7add..90284e934 100644 --- 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 index b8acf4398..5b43e1ebc 100644 --- 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,11 +51,8 @@ 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: @@ -66,11 +60,8 @@ def test_bad_entity_label(self) -> None: 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: