Making bookmark (follow) to utilize resource type (#121)

* Make user follow resource generic * Unit test * Update * Update
amundsen-io · Mar 26, 2020 · 3dbaa47 · 3dbaa47
1 parent c360660
commit 3dbaa47
Show file tree

Hide file tree

Showing 13 changed files with 230 additions and 93 deletions.
diff --git a/metadata_service/__init__.py b/metadata_service/__init__.py
@@ -103,7 +103,7 @@ def create_app(*, config_module_class: str) -> Flask:
     api.add_resource(UserFollowsAPI,
                      '/user/<path:user_id>/follow/')
     api.add_resource(UserFollowAPI,
-                     '/user/<path:user_id>/follow/<resource_type>/<path:table_uri>')
+                     '/user/<path:user_id>/follow/<resource_type>/<path:resource_id>')
     api.add_resource(UserOwnsAPI,
                      '/user/<path:user_id>/own/')
     api.add_resource(UserOwnAPI,

diff --git a/metadata_service/api/swagger_doc/user/follow_delete.yml b/metadata_service/api/swagger_doc/user/follow_delete.yml
@@ -18,7 +18,7 @@ parameters:
     schema:
       type: string
     required: true
-  - name: table_uri
+  - name: resource_id
     in: path
     example: 'hive://gold.test_schema/test_table1'
     type: string

diff --git a/metadata_service/api/swagger_doc/user/follow_put.yml b/metadata_service/api/swagger_doc/user/follow_put.yml
@@ -18,7 +18,7 @@ parameters:
     schema:
       type: string
     required: true
-  - name: table_uri
+  - name: resource_id
     in: path
     example: 'hive://gold.test_schema/test_table1'
     type: string

diff --git a/metadata_service/api/user.py b/metadata_service/api/user.py
@@ -12,6 +12,7 @@
 from metadata_service.exception import NotFoundException
 from metadata_service.proxy import get_proxy_client
 from metadata_service.util import UserResourceRel
+from metadata_service.entity.resource_type import to_resource_type
 
 LOGGER = logging.getLogger(__name__)
 
@@ -78,51 +79,56 @@ def __init__(self) -> None:
         self.client = get_proxy_client()
 
     @swag_from('swagger_doc/user/follow_put.yml')
-    def put(self, user_id: str, resource_type: str, table_uri: str) -> Iterable[Union[Mapping, int, None]]:
+    def put(self, user_id: str, resource_type: str, resource_id: str) -> Iterable[Union[Mapping, int, None]]:
         """
         Create the follow relationship between user and resources.
-        todo: It will need to refactor all neo4j proxy api to take a type argument.
 
         :param user_id:
         :param table_uri:
         :return:
         """
         try:
-            self.client.add_table_relation_by_user(table_uri=table_uri,
-                                                   user_email=user_id,
-                                                   relation_type=UserResourceRel.follow)
-            return {'message': 'The user {} for table_uri {} '
+            self.client.add_resource_relation_by_user(id=resource_id,
+                                                      user_id=user_id,
+                                                      relation_type=UserResourceRel.follow,
+                                                      resource_type=to_resource_type(label=resource_type))
+
+            return {'message': 'The user {} for id {} resource type {} '
                                'is added successfully'.format(user_id,
-                                                              table_uri)}, HTTPStatus.OK
+                                                              resource_id,
+                                                              resource_type)}, HTTPStatus.OK
         except Exception as e:
             LOGGER.exception('UserFollowAPI PUT Failed')
-            return {'message': 'The user {} for table_uri {} '
+            return {'message': 'The user {} for id {} resource type {}'
                                'is not added successfully'.format(user_id,
-                                                                  table_uri)}, \
+                                                                  resource_id,
+                                                                  resource_type)}, \
                 HTTPStatus.INTERNAL_SERVER_ERROR
 
     @swag_from('swagger_doc/user/follow_delete.yml')
-    def delete(self, user_id: str, resource_type: str, table_uri: str) -> Iterable[Union[Mapping, int, None]]:
+    def delete(self, user_id: str, resource_type: str, resource_id: str) -> Iterable[Union[Mapping, int, None]]:
         """
         Delete the follow relationship between user and resources.
-        todo: It will need to refactor all neo4j proxy api to take a type argument.
 
         :param user_id:
         :param table_uri:
         :return:
         """
         try:
-            self.client.delete_table_relation_by_user(table_uri=table_uri,
-                                                      user_email=user_id,
-                                                      relation_type=UserResourceRel.follow)
-            return {'message': 'The user following {} for table_uri {} '
+            self.client.delete_resource_relation_by_user(id=resource_id,
+                                                         user_id=user_id,
+                                                         relation_type=UserResourceRel.follow,
+                                                         resource_type=to_resource_type(label=resource_type))
+            return {'message': 'The user following {} for id {} resource type {} '
                                'is deleted successfully'.format(user_id,
-                                                                table_uri)}, HTTPStatus.OK
+                                                                resource_id,
+                                                                resource_type)}, HTTPStatus.OK
         except Exception as e:
             LOGGER.exception('UserFollowAPI DELETE Failed')
-            return {'message': 'The user {} for table_uri {} '
+            return {'message': 'The user {} for id {} resource type {} '
                                'is not deleted successfully'.format(user_id,
-                                                                    table_uri)}, \
+                                                                    resource_id,
+                                                                    resource_type)}, \
                 HTTPStatus.INTERNAL_SERVER_ERROR
 
 

diff --git a/metadata_service/entity/resource_type.py b/metadata_service/entity/resource_type.py
@@ -5,3 +5,7 @@ class ResourceType(Enum):
     Table = auto()
     Dashboard = auto()
     User = auto()
+
+
+def to_resource_type(*, label: str) -> ResourceType:
+    return ResourceType[label.title()]
diff --git a/metadata_service/proxy/atlas_proxy.py b/metadata_service/proxy/atlas_proxy.py
@@ -566,19 +566,44 @@ def get_table_by_user_relation(self, *, user_email: str, relation_type: UserReso
     def get_frequently_used_tables(self, *, user_email: str) -> Dict[str, Any]:
         pass
 
-    def add_table_relation_by_user(self, *,
-                                   table_uri: str,
-                                   user_email: str,
-                                   relation_type: UserResourceRel) -> None:
+    def add_resource_relation_by_user(self, *,
+                                      id: str,
+                                      user_id: str,
+                                      relation_type: UserResourceRel,
+                                      resource_type: ResourceType) -> None:
+
+        if resource_type is not ResourceType.Table:
+            raise NotImplemented('resource type {} is not supported'.format(resource_type))
+
+        self._add_table_relation_by_user(table_uri=id,
+                                         user_email=user_id,
+                                         relation_type=relation_type)
+
+    def _add_table_relation_by_user(self, *,
+                                    table_uri: str,
+                                    user_email: str,
+                                    relation_type: UserResourceRel) -> None:
 
         entity = self._get_reader_entity(table_uri=table_uri, user_id=user_email)
         entity.entity[self.ATTRS_KEY][self.BKMARKS_KEY] = True
         entity.update()
 
-    def delete_table_relation_by_user(self, *,
-                                      table_uri: str,
-                                      user_email: str,
-                                      relation_type: UserResourceRel) -> None:
+    def delete_resource_relation_by_user(self, *,
+                                         id: str,
+                                         user_id: str,
+                                         relation_type: UserResourceRel,
+                                         resource_type: ResourceType) -> None:
+        if resource_type is not ResourceType.Table:
+            raise NotImplemented('resource type {} is not supported'.format(resource_type))
+
+        self._delete_table_relation_by_user(table_uri=id,
+                                            user_email=user_id,
+                                            relation_type=relation_type)
+
+    def _delete_table_relation_by_user(self, *,
+                                       table_uri: str,
+                                       user_email: str,
+                                       relation_type: UserResourceRel) -> None:
         entity = self._get_reader_entity(table_uri=table_uri, user_id=user_email)
         entity.entity[self.ATTRS_KEY][self.BKMARKS_KEY] = False
         entity.update()

diff --git a/metadata_service/proxy/base_proxy.py b/metadata_service/proxy/base_proxy.py
@@ -91,17 +91,19 @@ def get_frequently_used_tables(self, *, user_email: str) -> Dict[str, Any]:
         pass
 
     @abstractmethod
-    def add_table_relation_by_user(self, *,
-                                   table_uri: str,
-                                   user_email: str,
-                                   relation_type: UserResourceRel) -> None:
+    def add_resource_relation_by_user(self, *,
+                                      id: str,
+                                      user_id: str,
+                                      relation_type: UserResourceRel,
+                                      resource_type: ResourceType) -> None:
         pass
 
     @abstractmethod
-    def delete_table_relation_by_user(self, *,
-                                      table_uri: str,
-                                      user_email: str,
-                                      relation_type: UserResourceRel) -> None:
+    def delete_resource_relation_by_user(self, *,
+                                         id: str,
+                                         user_id: str,
+                                         relation_type: UserResourceRel,
+                                         resource_type: ResourceType) -> None:
         pass
 
     @abstractmethod

diff --git a/metadata_service/proxy/gremlin_proxy.py b/metadata_service/proxy/gremlin_proxy.py
@@ -147,16 +147,18 @@ def get_table_by_user_relation(self, *, user_email: str,
     def get_frequently_used_tables(self, *, user_email: str) -> Dict[str, Any]:
         pass
 
-    def add_table_relation_by_user(self, *,
-                                   table_uri: str,
-                                   user_email: str,
-                                   relation_type: UserResourceRel) -> None:
-        pass
-
-    def delete_table_relation_by_user(self, *,
-                                      table_uri: str,
-                                      user_email: str,
-                                      relation_type: UserResourceRel) -> None:
+    def add_resource_relation_by_user(self, *,
+                                      id: str,
+                                      user_id: str,
+                                      relation_type: UserResourceRel,
+                                      resource_type: ResourceType) -> None:
+        pass
+
+    def delete_resource_relation_by_user(self, *,
+                                         id: str,
+                                         user_id: str,
+                                         relation_type: UserResourceRel,
+                                         resource_type: ResourceType) -> None:
         pass
 
     def get_dashboard(self,

diff --git a/metadata_service/proxy/neo4j_proxy.py b/metadata_service/proxy/neo4j_proxy.py
@@ -821,6 +821,38 @@ def _build_user_from_record(record: dict, manager_name: str = '') -> UserEntity:
                           role_name=record.get('role_name'),
                           manager_fullname=manager_name)
 
+    @staticmethod
+    def _get_user_resource_relationship_clause(relation_type: UserResourceRel, id: str = None,
+                                               user_key: str = None,
+                                               resource_type: ResourceType = ResourceType.Table) -> str:
+        """
+        Returns the relationship clause of a cypher query between users and tables
+        The User node is 'usr', the table node is 'tbl', and the relationship is 'rel'
+        e.g. (usr:User)-[rel:READ]->(tbl:Table), (usr)-[rel:READ]->(tbl)
+        """
+        resource_matcher: str = ''
+        user_matcher: str = ''
+
+        if id is not None:
+            resource_matcher += ':{}'.format(resource_type.name)
+            if id != '':
+                resource_matcher += ' {key: $resource_key}'
+
+        if user_key is not None:
+            user_matcher += ':User'
+            if user_key != '':
+                user_matcher += ' {key: $user_key}'
+
+        if relation_type == UserResourceRel.follow:
+            relation = f'(usr{user_matcher})-[rel:FOLLOW]->(resource{resource_matcher})'
+        elif relation_type == UserResourceRel.own:
+            relation = f'(usr{user_matcher})<-[rel:OWNER]-(resource{resource_matcher})'
+        elif relation_type == UserResourceRel.read:
+            relation = f'(usr{user_matcher})-[rel:READ]->(resource{resource_matcher})'
+        else:
+            raise NotImplementedError(f'The relation type {relation_type} is not defined!')
+        return relation
+
     @staticmethod
     def _get_user_table_relationship_clause(relation_type: UserResourceRel, tbl_key: str = None,
                                             user_key: str = None) -> str:
@@ -920,17 +952,18 @@ def get_frequently_used_tables(self, *, user_email: str) -> Dict[str, Any]:
         return {'table': results}
 
     @timer_with_counter
-    def add_table_relation_by_user(self, *,
-                                   table_uri: str,
-                                   user_email: str,
-                                   relation_type: UserResourceRel) -> None:
+    def add_resource_relation_by_user(self, *,
+                                      id: str,
+                                      user_id: str,
+                                      relation_type: UserResourceRel,
+                                      resource_type: ResourceType) -> None:
         """
         Update table user informations.
         1. Do a upsert of the user node.
         2. Do a upsert of the relation/reverse-relation edge.
 
         :param table_uri:
-        :param user_email:
+        :param user_id:
         :param relation_type:
         :return:
         """
@@ -940,26 +973,26 @@ def add_table_relation_by_user(self, *,
         on CREATE SET u={email: $user_email, key: $user_email}
         """)
 
-        user_email_clause = f'key: "{user_email}"'
-        tbl_key = f'key: "{table_uri}"'
+        rel_clause: str = self._get_user_resource_relationship_clause(relation_type=relation_type,
+                                                                      resource_type=resource_type)
 
-        rel_clause: str = self._get_user_table_relationship_clause(relation_type=relation_type)
-        upsert_user_relation_query = textwrap.dedent(f"""
-        MATCH (usr:User {{{user_email_clause}}}), (tbl:Table {{{tbl_key}}})
+        upsert_user_relation_query = textwrap.dedent("""
+        MATCH (usr:User {{key: $user_key}}), (resource:{resource_type} {{key: $resource_key}})
         MERGE {rel_clause}
-        RETURN usr.key, tbl.key
-        """)
+        RETURN usr.key, resource.key
+        """.format(resource_type=resource_type.name,
+                   rel_clause=rel_clause))
 
         try:
             tx = self._driver.session().begin_transaction()
             # upsert the node
-            tx.run(upsert_user_query, {'user_email': user_email})
-            result = tx.run(upsert_user_relation_query, {})
+            tx.run(upsert_user_query, {'user_email': user_id})
+            result = tx.run(upsert_user_relation_query, {'user_key': user_id, 'resource_key': id})
 
             if not result.single():
                 raise RuntimeError('Failed to create relation between '
-                                   'user {user} and table {tbl}'.format(user=user_email,
-                                                                        tbl=table_uri))
+                                   'user {user} and resource {id}'.format(user=user_id,
+                                                                          id=id))
             tx.commit()
         except Exception as e:
             if not tx.closed():
@@ -970,30 +1003,33 @@ def add_table_relation_by_user(self, *,
             tx.close()
 
     @timer_with_counter
-    def delete_table_relation_by_user(self, *,
-                                      table_uri: str,
-                                      user_email: str,
-                                      relation_type: UserResourceRel) -> None:
+    def delete_resource_relation_by_user(self, *,
+                                         id: str,
+                                         user_id: str,
+                                         relation_type: UserResourceRel,
+                                         resource_type: ResourceType) -> None:
         """
         Delete the relationship between user and resources.
 
         :param table_uri:
-        :param user_email:
+        :param user_id:
         :param relation_type:
         :return:
         """
-        rel_clause: str = self._get_user_table_relationship_clause(relation_type=relation_type,
-                                                                   user_key=user_email,
-                                                                   tbl_key=table_uri)
+        rel_clause: str = self._get_user_resource_relationship_clause(relation_type=relation_type,
+                                                                      resource_type=resource_type,
+                                                                      user_key=user_id,
+                                                                      id=id
+                                                                      )
 
-        delete_query = textwrap.dedent(f"""
+        delete_query = textwrap.dedent("""
         MATCH {rel_clause}
         DELETE rel
-        """)
+        """.format(rel_clause=rel_clause))
 
         try:
             tx = self._driver.session().begin_transaction()
-            tx.run(delete_query, {})
+            tx.run(delete_query, {'user_key': user_id, 'resource_key': id})
             tx.commit()
         except Exception as e:
             # propagate the exception back to api

diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 from setuptools import setup, find_packages
 
-__version__ = '2.3.1'
+__version__ = '2.4.0'
 
 
 requirements_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'requirements.txt')