feat: Column badges in Atlas Proxy (#263)

feat: Column badges in Atlas Proxy (#263) Signed-off-by: mgorsk1 <gorskimariusz13@gmail.com>
amundsen-io · Mar 3, 2021 · 8c4b5a6 · 8c4b5a6
1 parent b2ed7b2
commit 8c4b5a6
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 6 deletions.
diff --git 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
@@ -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
@@ -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,