Merge branch 'main' into handle-es-missing-index-edge-case

.github/CODEOWNERS

@@ -13,9 +13,9 @@
 
 /amundsen-kube-helm/ @feng-tao @jornh @javamonkey79 @mgorsk1 @verdan @sewardgw @youngyjd
 
-common/*.py        @feng-tao @jinhyukchang @allisonsuarez @verdan @bolkedebruin @mgorsk1 @dorianj @youngyjd @dechoma @sewardgw @dkunitsk @kristenarmes @ozandogrultan @B-T-D
-databuilder/*.py   @feng-tao @jinhyukchang @allisonsuarez @dorianj @verdan @mgorsk1 @youngyjd @dechoma @sewardgw @dkunitsk @kristenarmes @ozandogrultan @B-T-D
-frontend/amundsen_application/static/    @Golodhros @ttannis @allisonsuarez @feng-tao @dorianj @verdan @mgorsk1 @sewardgw @youngyjd @dechoma @kristenarmes @ozandogrultan @B-T-D
-frontend/*.py    @ttannis @feng-tao @allisonsuarez @dorianj @verdan @mgorsk1 @sewardgw @youngyjd @dkunitsk @dechoma @kristenarmes @ozandogrultan @B-T-D
-metadata/*.py    @feng-tao @jinhyukchang @allisonsuarez @verdan @bolkedebruin @mgorsk1 @dorianj @youngyjd @dechoma @sewardgw @dkunitsk @kristenarmes @ozandogrultan @B-T-D
-search/*.py      @feng-tao @jinhyukchang @allisonsuarez @verdan @bolkedebruin @mgorsk1 @dorianj @youngyjd @dechoma @sewardgw @dkunitsk @kristenarmes @ozandogrultan @B-T-D
+common/*.py        @feng-tao @jinhyukchang @allisonsuarez @verdan @bolkedebruin @mgorsk1 @youngyjd @dechoma @sewardgw @dkunitsk @kristenarmes @ozandogrultan @B-T-D
+databuilder/*.py   @feng-tao @jinhyukchang @allisonsuarez @verdan @mgorsk1 @youngyjd @dechoma @sewardgw @dkunitsk @kristenarmes @ozandogrultan @B-T-D
+frontend/amundsen_application/static/    @Golodhros @ttannis @allisonsuarez @feng-tao @verdan @mgorsk1 @sewardgw @youngyjd @dechoma @kristenarmes @ozandogrultan @B-T-D
+frontend/*.py    @ttannis @feng-tao @allisonsuarez @verdan @mgorsk1 @sewardgw @youngyjd @dkunitsk @dechoma @kristenarmes @ozandogrultan @B-T-D
+metadata/*.py    @feng-tao @jinhyukchang @allisonsuarez @verdan @bolkedebruin @mgorsk1 @youngyjd @dechoma @sewardgw @dkunitsk @kristenarmes @ozandogrultan @B-T-D
+search/*.py      @feng-tao @jinhyukchang @allisonsuarez @verdan @bolkedebruin @mgorsk1 @youngyjd @dechoma @sewardgw @dkunitsk @kristenarmes @ozandogrultan @B-T-D

OWNERS.md

@@ -11,7 +11,6 @@ routing PRs, questions, etc. to the right place.
 - Diksha Thakur (https://github.com/dikshathakur3119)
 - Dmitriy Kunitskiy (https://github.com/dkunitsk)
 - Dominik Choma (https://github.com/dechoma)
-- Dorian Johnson (https://github.com/dorianj)
 - Grant Seward (https://github.com/sewardgw)
 - Jin Hyuk Chang (https://github.com/jinhyukchang)
 - Junda Yang (https://github.com/youngyjd)
@@ -22,3 +21,6 @@ routing PRs, questions, etc. to the right place.
 - Tamika Tannis (https://github.com/ttannis)
 - Tao Feng (https://github.com/feng-tao)
 - Verdan Mahmood (https://github.com/verdan)
+
+# Amundsen Emeritus Contributors
+- Dorian Johnson (https://github.com/dorianj)

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),

databuilder/setup.py

@@ -5,7 +5,7 @@
 
 from setuptools import find_packages, setup
 
-__version__ = '7.4.4'
+__version__ = '7.4.5'
 
 requirements_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
                                  'requirements.txt')

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)
             }
         ]
 

metadata/metadata_service/proxy/neo4j_proxy.py

@@ -159,8 +159,9 @@ def get_table(self, *, table_uri: str) -> Table:
         cols, last_neo4j_record = self._exec_col_query(table_uri)
 
         readers = self._exec_usage_query(table_uri)
+        owners = self._exec_owners_query(table_uri)
 
-        wmk_results, table_writer, table_apps, timestamp_value, owners, tags, source, \
+        wmk_results, table_writer, table_apps, timestamp_value, tags, source, \
             badges, prog_descs, resource_reports = self._exec_table_query(table_uri)
 
         joins, filters = self._exec_table_query_query(table_uri)
@@ -340,22 +341,41 @@ def _exec_usage_query(self, table_uri: str) -> List[Reader]:
 
         return readers
 
+    @timer_with_counter
+    def _exec_owners_query(self, table_uri: str) -> List[User]:
+        # Return Value: List[User]
+        owners_query = textwrap.dedent("""
+            MATCH (owner:User)<-[:OWNER]-(tbl:Table {key: $tbl_key})
+            RETURN collect(distinct owner) as owner_records
+        """)
+        owners_neo4j_records = self._execute_cypher_query(statement=owners_query,
+                                                          param_dict={'tbl_key': table_uri})
+
+        owners_neo4j_records = get_single_record(owners_neo4j_records)
+
+        owners = []  # type: List[User]
+        for owner_neo4j_record in owners_neo4j_records.get('owner_records', []):
+            owner_data = self._get_user_details(user_id=owner_neo4j_record['email'])
+            owner = self._build_user_from_record(record=owner_data)
+            owners.append(owner)
+
+        return owners
+
     @timer_with_counter
     def _exec_table_query(self, table_uri: str) -> Tuple:
         """
         Queries one Cypher record with watermark list, Application,
-        ,timestamp, owner records and tag records.
+        ,timestamp, and tag records.
         """
 
-        # Return Value: (Watermark Results, Table Writer, Last Updated Timestamp, owner records, tag records)
+        # Return Value: (Watermark Results, Table Writer, Last Updated Timestamp, tag records)
 
         table_level_query = textwrap.dedent("""\
         MATCH (tbl:Table {key: $tbl_key})
         OPTIONAL MATCH (wmk:Watermark)-[:BELONG_TO_TABLE]->(tbl)
         OPTIONAL MATCH (app_producer:Application)-[:GENERATES]->(tbl)
         OPTIONAL MATCH (app_consumer:Application)-[:CONSUMES]->(tbl)
         OPTIONAL MATCH (tbl)-[:LAST_UPDATED_AT]->(t:Timestamp)
-        OPTIONAL MATCH (owner:User)<-[:OWNER]-(tbl)
         OPTIONAL MATCH (tbl)-[:TAGGED_BY]->(tag:Tag{tag_type: $tag_normal_type})
         OPTIONAL MATCH (tbl)-[:HAS_BADGE]->(badge:Badge)
         OPTIONAL MATCH (tbl)-[:SOURCE]->(src:Source)
@@ -365,7 +385,6 @@ def _exec_table_query(self, table_uri: str) -> Tuple:
         collect(distinct app_producer) as producing_apps,
         collect(distinct app_consumer) as consuming_apps,
         t.last_updated_timestamp as last_updated_timestamp,
-        collect(distinct owner) as owner_records,
         collect(distinct tag) as tag_records,
         collect(distinct badge) as badge_records,
         src,
@@ -405,12 +424,6 @@ def _exec_table_query(self, table_uri: str) -> Tuple:
 
         timestamp_value = table_records['last_updated_timestamp']
 
-        owner_record = []
-
-        for owner in table_records.get('owner_records', []):
-            owner_data = self._get_user_details(user_id=owner['email'])
-            owner_record.append(self._build_user_from_record(record=owner_data))
-
         src = None
 
         if table_records['src']:
@@ -423,7 +436,7 @@ def _exec_table_query(self, table_uri: str) -> Tuple:
 
         resource_reports = self._extract_resource_reports_from_query(table_records.get('resource_reports', []))
 
-        return wmk_results, table_writer, table_apps, timestamp_value, owner_record,\
+        return wmk_results, table_writer, table_apps, timestamp_value,\
             tags, src, badges, prog_descriptions, resource_reports
 
     @timer_with_counter
@@ -434,7 +447,7 @@ def _exec_table_query_query(self, table_uri: str) -> Tuple:
         on the table.
         """
 
-        # Return Value: (Watermark Results, Table Writer, Last Updated Timestamp, owner records, tag records)
+        # Return Value: (Watermark Results, Table Writer, Last Updated Timestamp, tag records)
         table_query_level_query = textwrap.dedent("""
         MATCH (tbl:Table {key: $tbl_key})
         OPTIONAL MATCH (tbl)-[:COLUMN]->(col:Column)-[COLUMN_JOINS_WITH]->(j:Join)

metadata/setup.py

@@ -5,7 +5,7 @@
 
 from setuptools import find_packages, setup
 
-__version__ = '3.12.1'
+__version__ = '3.12.2'
 
 requirements_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'requirements.txt')
 with open(requirements_path) as requirements_file:

metadata/tests/unit/proxy/test_neo4j_proxy.py

@@ -171,13 +171,6 @@ def setUp(self) -> None:
                 }
             ],
             'last_updated_timestamp': 1,
-            'owner_records': [
-                {
-                    'key': '[email protected]',
-                    'email': '[email protected]',
-                    'updated_at': 0,
-                }
-            ],
             'tag_records': [
                 {
                     'key': 'test',
@@ -236,6 +229,14 @@ def setUp(self) -> None:
             ]
         }]
 
+        owners_results = [{'owner_records': [
+            {
+                'key': '[email protected]',
+                'email': '[email protected]',
+                'updated_at': 0,
+            }
+        ], }]
+
         last_updated_timestamp = '01'
 
         self.col_usage_return_value = [
@@ -250,6 +251,8 @@ def setUp(self) -> None:
 
         self.table_common_usage = table_common_usage
 
+        self.owners_return_value = owners_results
+
         self.col_bar_id_1_expected_type_metadata = self._get_col_bar_id_1_expected_type_metadata()
         self.col_bar_id_2_expected_type_metadata = self._get_col_bar_id_2_expected_type_metadata()
 
@@ -355,9 +358,11 @@ def test_health_neo4j(self) -> None:
 
     def test_get_table(self) -> None:
         with patch.object(GraphDatabase, 'driver'), patch.object(Neo4jProxy, '_execute_cypher_query') as mock_execute:
+            # mock database return values such that we match ordering of queries executed in Neo4jProxy.get_table
             mock_execute.side_effect = [
                 self.col_usage_return_value,
                 [],
+                self.owners_return_value,
                 self.table_level_return_value,
                 self.table_common_usage,
                 []
@@ -445,6 +450,7 @@ def test_get_table_view_only(self) -> None:
             mock_execute.side_effect = [
                 col_usage_return_value,
                 [],
+                self.owners_return_value,
                 self.table_level_return_value,
                 self.table_common_usage,
                 []