diff --git a/metadata/metadata_service/proxy/atlas_proxy.py b/metadata/metadata_service/proxy/atlas_proxy.py index 6ca791c7f0..f9ad7f0cc4 100644 --- a/metadata/metadata_service/proxy/atlas_proxy.py +++ b/metadata/metadata_service/proxy/atlas_proxy.py @@ -11,9 +11,10 @@ from amundsen_common.models.dashboard import DashboardSummary from amundsen_common.models.lineage import Lineage from amundsen_common.models.popular_table import PopularTable -from amundsen_common.models.table import (Column, ProgrammaticDescription, - Reader, ResourceReport, Stat, Table, - Tag, User, Watermark) +from amundsen_common.models.table import (Badge, Column, + ProgrammaticDescription, Reader, + ResourceReport, Stat, Table, Tag, + User, Watermark) from amundsen_common.models.user import User as UserEntity from atlasclient.client import Atlas from atlasclient.exceptions import BadRequest, Conflict, NotFound @@ -331,6 +332,14 @@ def _serialize_columns(self, *, entity: EntityUniqueAttribute) -> \ col_attrs = col_entity[self.ATTRS_KEY] statistics = list() + badges = list() + + for column_classification in col_entity.get('classifications') or list(): + if column_classification.get('entityStatus') == Status.ACTIVE: + name = column_classification.get('typeName') + + badges.append(Badge(badge_name=name, category='default')) + for stats in col_attrs.get('statistics') or list(): stats_attrs = stats['attributes'] @@ -369,6 +378,7 @@ def _serialize_columns(self, *, entity: EntityUniqueAttribute) -> \ col_type=col_attrs.get('type') or col_attrs.get('dataType') or col_attrs.get('data_type'), sort_order=col_attrs.get('position') or 9999, stats=statistics, + badges=badges ) ) return sorted(columns, key=lambda item: item.sort_order) diff --git a/metadata/tests/unit/proxy/fixtures/atlas_test_data.py b/metadata/tests/unit/proxy/fixtures/atlas_test_data.py index 0e03a3c791..6edac41e33 100644 --- a/metadata/tests/unit/proxy/fixtures/atlas_test_data.py +++ b/metadata/tests/unit/proxy/fixtures/atlas_test_data.py @@ -54,6 +54,16 @@ class Data: }}, ] }, + "classifications": [ + { + "typeName": "active_col_badge", + "entityStatus": "ACTIVE" + }, + { + "typeName": "inactive_col_badge", + "entityStatus": "INACTIVE" + } + ] } test_column_inactive = copy.deepcopy(test_column) diff --git a/metadata/tests/unit/proxy/test_atlas_proxy.py b/metadata/tests/unit/proxy/test_atlas_proxy.py index 0214e7139f..6fc08783b7 100644 --- a/metadata/tests/unit/proxy/test_atlas_proxy.py +++ b/metadata/tests/unit/proxy/test_atlas_proxy.py @@ -7,8 +7,9 @@ from unittest.mock import MagicMock, patch from amundsen_common.models.popular_table import PopularTable -from amundsen_common.models.table import (Column, ProgrammaticDescription, - Reader, Stat, Table, Tag, User) +from amundsen_common.models.table import (Badge, Column, + ProgrammaticDescription, Reader, + Stat, Table, Tag, User) from atlasclient.exceptions import BadRequest from metadata_service import create_app @@ -136,7 +137,8 @@ def _get_table(self, custom_stats_format: bool = False) -> None: description='column description', col_type='Managed', sort_order=col_attrs['position'], - stats=exp_col_stats) + stats=exp_col_stats, + badges=[Badge(category='default', badge_name='active_col_badge')]) expected = Table(database=self.entity_type, cluster=self.cluster,