From 4ea45d3be3082d15838e7e43e69ae195fd0e23c9 Mon Sep 17 00:00:00 2001
From: Sriharsha Chintalapani <harshach@users.noreply.github.com>
Date: Tue, 30 Nov 2021 20:08:30 -0800
Subject: [PATCH] Fix #1496: Follow-up: fix for sample-data after user classes
 refactored (#1497)

* Fix #1496: Follow-up: fix for sample-data after user classes refactored

* Fix #1496: Follow-up: fix for sample-data after user classes refactored
---
 .../src/metadata/ingestion/sink/metadata_rest.py    | 13 +++++++------
 .../src/metadata/ingestion/source/sample_data.py    |  8 ++++++--
 2 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/ingestion/src/metadata/ingestion/sink/metadata_rest.py b/ingestion/src/metadata/ingestion/sink/metadata_rest.py
index e606a408bab..158b94bcac3 100644
--- a/ingestion/src/metadata/ingestion/sink/metadata_rest.py
+++ b/ingestion/src/metadata/ingestion/sink/metadata_rest.py
@@ -350,7 +350,6 @@ class MetadataRestSink(Sink):
 
     def write_ml_model(self, model: MlModel):
         try:
-            logger.info(model)
             model_request = CreateMlModelEntityRequest(
                 name=model.name,
                 displayName=model.displayName,
@@ -377,7 +376,7 @@ class MetadataRestSink(Sink):
         )
         try:
             r = self.metadata.create_or_update(metadata_team)
-            instance_id = r.id.__root__
+            instance_id = str(r.id.__root__)
             self.team_entities[team.name] = instance_id
         except Exception as err:
             logger.error(traceback.format_exc())
@@ -386,9 +385,11 @@ class MetadataRestSink(Sink):
 
     def write_users(self, record: User):
         teams = []
-        for team in record.teams:
-            self._create_team(team)
-            teams.append([self.team_entities[team.name]])
+        for team in record.teams.__root__:
+            if team.name not in self.team_entities:
+                self._create_team(team)
+            teams.append(self.team_entities[team.name])
+
         metadata_user = CreateUserEntityRequest(
             name=record.name.__root__,
             displayName=record.name.__root__,
@@ -398,7 +399,7 @@ class MetadataRestSink(Sink):
         try:
             self.metadata.create_or_update(metadata_user)
             self.status.records_written(record.name.__root__)
-            logger.info("Sink: {}".format(record.name))
+            logger.info("Sink: {}".format(record.name.__root__))
         except Exception as err:
             logger.error(traceback.format_exc())
             logger.error(traceback.print_exc())
diff --git a/ingestion/src/metadata/ingestion/source/sample_data.py b/ingestion/src/metadata/ingestion/source/sample_data.py
index 3432f2a82b9..1e307ea3c85 100644
--- a/ingestion/src/metadata/ingestion/source/sample_data.py
+++ b/ingestion/src/metadata/ingestion/source/sample_data.py
@@ -35,6 +35,7 @@ from metadata.generated.schema.entity.data.mlmodel import MlModel
 from metadata.generated.schema.entity.data.pipeline import Pipeline
 from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.entity.teams.user import User
+from metadata.generated.schema.type.basic import Href
 from metadata.generated.schema.type.entityLineage import EntitiesEdge
 from metadata.generated.schema.type.entityReference import EntityReference
 from metadata.ingestion.api.common import Record
@@ -419,16 +420,19 @@ class SampleDataSource(Source):
             )
             yield model_ev
 
-    def ingest_users(self) -> Iterable[CreateUserEntityRequest]:
+    def ingest_users(self) -> Iterable[User]:
         try:
             for user in self.users["users"]:
-                teams = [EntityReference(id=uuid.uuid4(), name=user.teams)]
+                teams = [
+                    EntityReference(id=uuid.uuid4(), name=user["teams"], type="team")
+                ]
                 user_metadata = User(
                     id=uuid.uuid4(),
                     name=user["email"],
                     displayName=user["displayName"],
                     email=user["email"],
                     teams=teams,
+                    href=Href(__root__="http://localhost"),
                 )
                 yield user_metadata
         except Exception as err: