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
This commit is contained in:
Sriharsha Chintalapani 2021-11-30 20:08:30 -08:00 committed by GitHub
parent 3ae3eec3d0
commit 4ea45d3be3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 13 additions and 8 deletions

View File

@ -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:
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]])
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())

View File

@ -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: