Fix: Add rds badge association records (#2179)

Signed-off-by: Sahithi Velma <[email protected]> Co-authored-by: Sahithi Velma <[email protected]>
amundsen-io · Aug 11, 2023 · c130ab3 · c130ab3
1 parent 8c3ceb7
commit c130ab3
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/databuilder/databuilder/models/badge.py b/databuilder/databuilder/models/badge.py
@@ -8,6 +8,9 @@
 from amundsen_common.utils.atlas import AtlasCommonParams, AtlasCommonTypes
 from amundsen_rds.models import RDSModel
 from amundsen_rds.models.badge import Badge as RDSBadge
+from amundsen_rds.models.column import ColumnBadge as RDSColumnBadge
+from amundsen_rds.models.dashboard import DashboardBadge as RDSDashboardBadge
+from amundsen_rds.models.table import TableBadge as RDSTableBadge
 
 from databuilder.models.atlas_entity import AtlasEntity
 from databuilder.models.atlas_relationship import AtlasRelationship
@@ -156,6 +159,16 @@ def _create_record_iterator(self) -> Iterator[RDSModel]:
         for record in records:
             yield record
 
+            if self.start_label == 'Table':
+                table_badge_record = RDSTableBadge(table_rk=self.start_key, badge_rk=record.rk)
+                yield table_badge_record
+            elif self.start_label == 'Column':
+                column_badge_record = RDSColumnBadge(column_rk=self.start_key, badge_rk=record.rk)
+                yield column_badge_record
+            elif self.start_label == 'Dashboard':
+                dashboard_badge_record = RDSDashboardBadge(dashboard_rk=self.start_key, badge_rk=record.rk)
+                yield dashboard_badge_record
+
     def _create_atlas_classification_entity(self, badge: Badge) -> AtlasEntity:
         attrs_mapping = [
             (AtlasCommonParams.qualified_name, badge.name),

diff --git a/databuilder/tests/unit/models/test_badge.py b/databuilder/tests/unit/models/test_badge.py
@@ -228,9 +228,17 @@ def test_create_records(self) -> None:
                 'rk': BadgeMetadata.BADGE_KEY_FORMAT.format(badge=badge1.name),
                 'category': badge1.category
             },
+            {
+                'column_rk': 'hive://default.base/test/ds',
+                'badge_rk': BadgeMetadata.BADGE_KEY_FORMAT.format(badge=badge1.name)
+            },
             {
                 'rk': BadgeMetadata.BADGE_KEY_FORMAT.format(badge=badge2.name),
                 'category': badge2.category
+            },
+            {
+                'column_rk': 'hive://default.base/test/ds',
+                'badge_rk': BadgeMetadata.BADGE_KEY_FORMAT.format(badge=badge2.name)
             }
         ]