#1051: migrated all usages of models.MetadataUser & models.User classes to g… (#1473)

* migrated all usages of models.MetadataUser & models.User classes to generated.schema.teams.user.User classes instead.

* formatted file with black

Co-authored-by: omri.alon <omri.alon@myheritage.com>

ingestion/src/metadata/ingestion/models/user.py

@@ -22,115 +22,6 @@ from metadata.ingestion.models.json_serializable import JsonSerializable
 UNQUOTED_SUFFIX = ":UNQUOTED"
 
 
-class User(JsonSerializable):
-    """
-    User model. This model doesn't define any relationship.
-    """
-
-    USER_NODE_LABEL = "User"
-    USER_NODE_KEY_FORMAT = "{email}"
-    USER_NODE_EMAIL = "email"
-    USER_NODE_FIRST_NAME = "first_name"
-    USER_NODE_LAST_NAME = "last_name"
-    USER_NODE_FULL_NAME = "full_name"
-    USER_NODE_GITHUB_NAME = "github_username"
-    USER_NODE_TEAM = "team_name"
-    USER_NODE_EMPLOYEE_TYPE = "employee_type"
-    USER_NODE_MANAGER_EMAIL = "manager_email"
-    USER_NODE_SLACK_ID = "slack_id"
-    USER_NODE_IS_ACTIVE = "is_active{}".format(
-        UNQUOTED_SUFFIX
-    )  # bool value needs to be unquoted when publish to neo4j
-    USER_NODE_UPDATED_AT = "updated_at"
-    USER_NODE_ROLE_NAME = "role_name"
-
-    USER_MANAGER_RELATION_TYPE = "MANAGE_BY"
-    MANAGER_USER_RELATION_TYPE = "MANAGE"
-
-    def __init__(
-        self,
-        email: str,
-        first_name: str = "",
-        last_name: str = "",
-        name: str = "",
-        github_username: str = "",
-        team_name: str = "",
-        employee_type: str = "",
-        manager_email: str = "",
-        slack_id: str = "",
-        is_active: bool = True,
-        updated_at: int = 0,
-        role_name: str = "",
-        do_not_update_empty_attribute: bool = False,
-        **kwargs: Any
-    ) -> None:
-        """
-        This class models user node for Amundsen people.
-
-        :param first_name:
-        :param last_name:
-        :param name:
-        :param email:
-        :param github_username:
-        :param team_name:
-        :param employee_type:
-        :param manager_email:
-        :param is_active:
-        :param updated_at: everytime we update the node, we will push the timestamp.
-                           then we will have a cron job to update the ex-employee nodes based on
-                           the case if this timestamp hasn't been updated for two weeks.
-        :param role_name: the role_name of the user (e.g swe)
-        :param do_not_update_empty_attribute: If False, all empty or not defined params will be overwritten with
-        empty string.
-        :param kwargs: Any K/V attributes we want to update the
-        """
-        self.first_name = first_name
-        self.last_name = last_name
-        self.name = name
-
-        self.email = email
-        self.github_username = github_username
-        # todo: team will be a separate node once Amundsen People supports team
-        self.team_name = team_name
-        self.manager_email = manager_email
-        self.employee_type = employee_type
-        # this attr not available in team service, either update team service, update with FE
-        self.slack_id = slack_id
-        self.is_active = is_active
-        self.updated_at = updated_at
-        self.role_name = role_name
-        self.do_not_update_empty_attribute = do_not_update_empty_attribute
-        self.attrs = None
-        if kwargs:
-            self.attrs = copy.deepcopy(kwargs)
-
-
-class MetadataUser(JsonSerializable):
-    """
-    Catalog User model. This model doesn't define any relationship.
-    """
-
-    def __init__(
-        self,
-        name: str,
-        display_name: str,
-        email: str,
-        timezone: str = "PST",
-        is_bot: bool = False,
-        teams: [] = None,
-        **kwargs: Any
-    ) -> None:
-        """ """
-        self.name = name
-        self.display_name = display_name
-        self.email = email
-        self.timezone = timezone
-        self.is_bot = is_bot
-        self.teams = teams
-        if kwargs:
-            self.attrs = copy.deepcopy(kwargs)
-
-
 class MetadataOrg(JsonSerializable):
     """
     Catalog Org Model

ingestion/src/metadata/ingestion/sink/ldap_rest_users.py

@@ -16,9 +16,10 @@
 import logging
 
 from metadata.config.common import ConfigModel
+from metadata.generated.schema.api.teams.createUser import CreateUserEntityRequest
+from metadata.generated.schema.entity.teams.user import User
 from metadata.ingestion.api.common import Record, WorkflowContext
 from metadata.ingestion.api.sink import Sink, SinkStatus
-from metadata.ingestion.models.user import MetadataUser
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
 
@@ -57,13 +58,11 @@ class LdapRestUsersSink(Sink):
     def write_record(self, record: Record) -> None:
         self._create_user(record)
 
-    def _create_user(self, record: MetadataUser) -> None:
+    def _create_user(self, record: User) -> None:
-        metadata_user = MetadataUser(
+        metadata_user = CreateUserEntityRequest(
-            name=record.github_username[0],
+            name=record.name, displayName=record.displayName, email=record.email
-            display_name=record.name[0],
-            email=record.email[0],
         )
-        self.rest.post(self.api_users, data=metadata_user.to_json())
+        self.rest.post(self.api_users, data=metadata_user.json())
         self.status.records_written(record.name[0])
 
     def get_status(self):

ingestion/src/metadata/ingestion/source/ldap_users.py

@@ -19,9 +19,9 @@ from typing import Iterable
 from ldap3 import ALL, LEVEL, Connection, Server
 
 from metadata.config.common import ConfigModel
+from metadata.generated.schema.api.teams.createUser import CreateUserEntityRequest
 from metadata.ingestion.api.common import WorkflowContext
 from metadata.ingestion.api.source import Source, SourceStatus
-from metadata.ingestion.models.user import MetadataUser, User
 from metadata.ingestion.ometa.openmetadata_rest import MetadataServerConfig
 
 logger = logging.getLogger(__name__)
@@ -84,19 +84,12 @@ class LdapUsersSource(Source):
         metadata_config = MetadataServerConfig.parse_obj(metadata_config_dict)
         return cls(ctx, config, metadata_config)
 
-    def next_record(self) -> Iterable[MetadataUser]:
+    def next_record(self) -> Iterable[CreateUserEntityRequest]:
         for user in self.users:
-            user_metadata = User(
+            user_metadata = CreateUserEntityRequest(
-                user["attributes"]["mail"],
+                email=user["attributes"]["mail"],
-                user["attributes"]["givenName"],
+                displayName=user["attributes"]["cn"],
-                user["attributes"]["sn"],
+                name=user["attributes"]["givenName"],
-                user["attributes"]["cn"],
-                user["attributes"]["uid"],
-                "",
-                "",
-                "",
-                True,
-                0,
             )
             self.status.scanned(user_metadata.name)
             yield user_metadata

ingestion/tests/integration/ldap/test_ldap_crud.py

@@ -17,16 +17,16 @@ import logging
 import time
 from datetime import datetime
 
-
 import pytest
 import requests
 from ldap3 import ALL, Connection, Server
 
-from metadata.ingestion.models.user import MetadataUser, User
+from metadata.generated.schema.api.teams.createUser import CreateUserEntityRequest
 
 headers = {"Content-type": "application/json"}
 url = "http://localhost:8585/api/v1/users"
 
+
 def sleep(timeout_s):
     print(f"sleeping for {timeout_s} seconds")
     n = len(str(timeout_s))
@@ -35,6 +35,7 @@ def sleep(timeout_s):
         time.sleep(1)
     print(f"{'':>{n}}", end="\n", flush=True)
 
+
 def read_user_by_name(name: str):
     r = requests.get(url + "/name/" + name)
     r.raise_for_status()
@@ -48,6 +49,7 @@ def status(r):
     else:
         return 0
 
+
 def ldap_connection():
     s = Server("ldap://localhost:389", get_info=ALL)
     c = Connection(s, user="cn=admin,dc=example,dc=com", password="ldappassword")
@@ -57,11 +59,13 @@ def ldap_connection():
         return False
     return [True, c]
 
+
 def is_ldap_listening(openldap_service):
     c = openldap_service
     if "listening" in str(c):
         return True
 
+
 @pytest.fixture(scope="session")
 def openldap_service(docker_ip, docker_services):
     """Ensure that Docker service is up and responsive."""
@@ -75,16 +79,22 @@ def openldap_service(docker_ip, docker_services):
     )
     return conn
 
+
 @pytest.fixture(scope="session")
 def ldap_user_entry(openldap_service):
     c = openldap_service
-    c.search('cn=John Doe,ou=users,dc=example,dc=com', '(objectclass=person)', attributes=['*'])
+    c.search(
+        "cn=John Doe,ou=users,dc=example,dc=com",
+        "(objectclass=person)",
+        attributes=["*"],
+    )
     if c.entries:
         return c.entries[0]
     else:
         logging.error("OpenLDAP not running")
         assert 0
 
+
 @pytest.fixture(scope="session")
 def datetime_suffix():
     # Openmetadata doesn't delete users; it deactivates them.
@@ -96,22 +106,12 @@ def datetime_suffix():
 
 
 def test_insert_user(ldap_user_entry, datetime_suffix):
-    user = User(
+    metadata_user = CreateUserEntityRequest(
-        str(ldap_user_entry['mail']),
+        name=str(ldap_user_entry["uid"]) + datetime_suffix,
-        str(ldap_user_entry['givenName']),
+        displayName=str(ldap_user_entry["cn"]),
-        str(ldap_user_entry['sn']),
+        email=str(ldap_user_entry["mail"]),
-        str(ldap_user_entry['cn']),
-        str(ldap_user_entry['uid']) + datetime_suffix,
-        "",
-        "",
-        "",
-        True,
-        0,
     )
-    metadata_user = MetadataUser(
+    r = requests.post(url, data=metadata_user.json(), headers=headers)
-        name=user.github_username, display_name=user.name, email=user.email
-    )
-    r = requests.post(url, data=metadata_user.to_json(), headers=headers)
     r.raise_for_status()
     if r.status_code == 200 or r.status_code == 201:
         assert 1
@@ -120,22 +120,22 @@ def test_insert_user(ldap_user_entry, datetime_suffix):
 
 
 def test_read_user(ldap_user_entry, datetime_suffix):
-    assert read_user_by_name(str(ldap_user_entry['uid']) + datetime_suffix)[0]
+    assert read_user_by_name(str(ldap_user_entry["uid"]) + datetime_suffix)[0]
 
 
 def test_update_user(ldap_user_entry, datetime_suffix):
-    user = read_user_by_name(str(ldap_user_entry['uid']) + datetime_suffix)
+    user = read_user_by_name(str(ldap_user_entry["uid"]) + datetime_suffix)
     user[1]["displayName"] = "Jane Doe"
-    metadata_user = MetadataUser(
+    metadata_user = CreateUserEntityRequest(
         name=user[1]["name"],
-        display_name=user[1]["displayName"],
+        displayName=user[1]["displayName"],
-        email=user[1]["name"],
+        email=user[1]["email"],
     )
-    r = requests.patch(url, data=metadata_user.to_json(), headers=headers)
+    r = requests.patch(url, data=metadata_user.json(), headers=headers)
 
 
 def test_delete_user(ldap_user_entry, datetime_suffix):
-    r = read_user_by_name(str(ldap_user_entry['uid']) + datetime_suffix)
+    r = read_user_by_name(str(ldap_user_entry["uid"]) + datetime_suffix)
     r = requests.delete(url + "/{}".format(r[1]["id"]))
     r.raise_for_status()
     assert 1