def set_metadata(dataset_name: str,
fields: List,
platform: str = "api") -> SchemaMetadata:
canonical_schema: List[SchemaField] = []
for column in fields:
field = SchemaField(
fieldPath=column,
nativeDataType="str",
type=SchemaFieldDataTypeClass(type=StringTypeClass()),
description="",
recursive=False,
)
canonical_schema.append(field)
actor = "urn:li:corpuser:etl"
sys_time = int(time.time() * 1000)
schema_metadata = SchemaMetadata(
schemaName=dataset_name,
platform=f"urn:li:dataPlatform:{platform}",
version=0,
hash="",
platformSchema=OtherSchemaClass(rawSchema=""),
created=AuditStamp(time=sys_time, actor=actor),
lastModified=AuditStamp(time=sys_time, actor=actor),
fields=canonical_schema,
)
return schema_metadata
env="PROD"),
aspectName="schemaMetadata",
aspect=SchemaMetadataClass(
schemaName="customer", # not used
platform=make_data_platform_urn(
"hive"), # important <- platform must be an urn
version=
0, # when the source system has a notion of versioning of schemas, insert this in, otherwise leave as 0
hash=
"", # when the source system has a notion of unique schemas identified via hash, include a hash, else leave it as empty string
platformSchema=OtherSchemaClass(
rawSchema="__insert raw schema here__"),
fields=[
SchemaFieldClass(
fieldPath="address.zipcode",
type=SchemaFieldDataTypeClass(type=StringTypeClass()),
nativeDataType=
"VARCHAR(50)", # use this to provide the type of the field in the source system's vernacular
description=
"This is the zipcode of the address. Specified using extended form and limited to addresses in the United States",
),
SchemaFieldClass(
fieldPath="address.street",
type=SchemaFieldDataTypeClass(type=StringTypeClass()),
nativeDataType="VARCHAR(100)",
description="Street corresponding to the address",
),
SchemaFieldClass(
fieldPath="last_sold_date",
type=SchemaFieldDataTypeClass(type=DateTypeClass()),
nativeDataType="Date",
def generate_stitched_record(relnships_graph: RelationshipGraph) -> List[Any]:
def strip_types(field_path: str) -> str:
final_path = field_path
final_path = re.sub(r"(\[type=[a-zA-Z]+\]\.)", "", final_path)
final_path = re.sub(r"^\[version=2.0\]\.", "", final_path)
return final_path
datasets: List[DatasetSnapshotClass] = []
for entity_name, entity_def in entity_registry.items():
entity_display_name = entity_def.display_name
entity_fields = []
for aspect_name in entity_def.aspects:
if aspect_name not in aspect_registry:
print(
f"Did not find aspect name: {aspect_name} in aspect_registry"
)
continue
# all aspects should have a schema
aspect_schema = aspect_registry[aspect_name].schema
assert aspect_schema
entity_fields.append({
"type": aspect_schema.to_json(),
"name": aspect_name,
})
if entity_fields:
names = avro.schema.Names()
field_objects = []
for f in entity_fields:
field = avro.schema.Field(
type=f["type"],
name=f["name"],
has_default=False,
)
field_objects.append(field)
with unittest.mock.patch("avro.schema.Names.add_name", add_name):
entity_avro_schema = avro.schema.RecordSchema(
name=entity_name,
namespace="datahub.metadata.model",
names=names,
fields=[],
)
entity_avro_schema.set_prop("fields", field_objects)
rawSchema = json.dumps(entity_avro_schema.to_json())
# always add the URN which is the primary key
urn_field = SchemaField(
fieldPath="urn",
type=SchemaFieldDataTypeClass(type=StringTypeClass()),
nativeDataType="string",
nullable=False,
isPartOfKey=True,
description=
f"The primary identifier for the {entity_name} entity. See the {entity_def.keyAspect} field to understand the structure of this urn.",
)
schema_fields: List[SchemaField] = [
urn_field
] + avro_schema_to_mce_fields(rawSchema)
foreign_keys: List[ForeignKeyConstraintClass] = []
source_dataset_urn = make_dataset_urn(
platform=make_data_platform_urn("datahub"),
name=f"{entity_display_name}",
)
for f_field in schema_fields:
if f_field.jsonProps:
json_dict = json.loads(f_field.jsonProps)
if "Aspect" in json_dict:
aspect_info = json_dict["Aspect"]
f_field.globalTags = f_field.globalTags or GlobalTagsClass(
tags=[])
f_field.globalTags.tags.append(
TagAssociationClass(tag="urn:li:tag:Aspect"))
# if this is the key aspect, also add primary-key
if entity_def.keyAspect == aspect_info.get("name"):
f_field.isPartOfKey = True
if "timeseries" == aspect_info.get("type", ""):
# f_field.globalTags = f_field.globalTags or GlobalTagsClass(
# tags=[]
# )
f_field.globalTags.tags.append(
TagAssociationClass(tag="urn:li:tag:Temporal"))
import pdb
# breakpoint()
if "Searchable" in json_dict:
f_field.globalTags = f_field.globalTags or GlobalTagsClass(
tags=[])
f_field.globalTags.tags.append(
TagAssociationClass(tag="urn:li:tag:Searchable"))
if "Relationship" in json_dict:
relationship_info = json_dict["Relationship"]
# detect if we have relationship specified at leaf level or thru path specs
if "entityTypes" not in relationship_info:
# path spec
assert (
len(relationship_info.keys()) == 1
), "We should never have more than one path spec assigned to a relationship annotation"
final_info = None
for k, v in relationship_info.items():
final_info = v
relationship_info = final_info
assert "entityTypes" in relationship_info
entity_types: List[str] = relationship_info.get(
"entityTypes", [])
relnship_name = relationship_info.get("name", None)
for entity_type in entity_types:
destination_entity_name = capitalize_first(
entity_type)
foreign_dataset_urn = make_dataset_urn(
platform=make_data_platform_urn("datahub"),
name=destination_entity_name,
)
fkey = ForeignKeyConstraintClass(
name=relnship_name,
foreignDataset=foreign_dataset_urn,
foreignFields=[
f"urn:li:schemaField:({foreign_dataset_urn}, urn)"
],
sourceFields=[
f"urn:li:schemaField:({source_dataset_urn},{f_field.fieldPath})"
],
)
foreign_keys.append(fkey)
relnships_graph.add_edge(
entity_display_name,
destination_entity_name,
fkey.name,
f" via `{strip_types(f_field.fieldPath)}`",
edge_id=
f"{entity_display_name}:{fkey.name}:{destination_entity_name}:{strip_types(f_field.fieldPath)}",
)
schemaMetadata = SchemaMetadataClass(
schemaName=f"{entity_name}",
platform=make_data_platform_urn("datahub"),
platformSchema=OtherSchemaClass(rawSchema=rawSchema),
fields=schema_fields,
version=0,
hash="",
foreignKeys=foreign_keys if foreign_keys else None,
)
dataset = DatasetSnapshotClass(
urn=make_dataset_urn(
platform=make_data_platform_urn("datahub"),
name=f"{entity_display_name}",
),
aspects=[
schemaMetadata,
GlobalTagsClass(
tags=[TagAssociationClass(tag="urn:li:tag:Entity")]),
BrowsePathsClass(
[f"/prod/datahub/entities/{entity_display_name}"]),
],
)
datasets.append(dataset)
events: List[Union[MetadataChangeEventClass,
MetadataChangeProposalWrapper]] = []
for d in datasets:
entity_name = d.urn.split(":")[-1].split(",")[1]
d.aspects.append(
DatasetPropertiesClass(
description=make_entity_docs(entity_name, relnships_graph)))
mce = MetadataChangeEventClass(proposedSnapshot=d, )
events.append(mce)
mcp = MetadataChangeProposalWrapper(
entityType="dataset",
changeType=ChangeTypeClass.UPSERT,
entityUrn=d.urn,
aspectName="subTypes",
aspect=SubTypesClass(typeNames=["entity"]),
)
events.append(mcp)
return events