mirror of
https://github.com/datahub-project/datahub.git
synced 2025-07-03 23:28:11 +00:00
1042 lines
39 KiB
Python
1042 lines
39 KiB
Python
import glob
|
|
import json
|
|
import logging
|
|
import os
|
|
import re
|
|
import shutil
|
|
import unittest.mock
|
|
from dataclasses import Field, dataclass, field
|
|
from datetime import datetime, timezone
|
|
from enum import auto
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterable, List, Optional, Tuple
|
|
|
|
import avro.schema
|
|
import click
|
|
from utils import should_write_json_file
|
|
|
|
from datahub.configuration.common import ConfigEnum, PermissiveConfigModel
|
|
from datahub.emitter.mce_builder import (
|
|
make_data_platform_urn,
|
|
make_dataset_urn,
|
|
make_schema_field_urn,
|
|
)
|
|
from datahub.emitter.mcp import MetadataChangeProposalWrapper
|
|
from datahub.emitter.rest_emitter import DatahubRestEmitter
|
|
from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
|
|
from datahub.ingestion.api.sink import NoopWriteCallback
|
|
from datahub.ingestion.extractor.schema_util import avro_schema_to_mce_fields
|
|
from datahub.ingestion.sink.file import FileSink, FileSinkConfig
|
|
from datahub.metadata.schema_classes import (
|
|
BrowsePathEntryClass,
|
|
BrowsePathsClass,
|
|
BrowsePathsV2Class,
|
|
DatasetPropertiesClass,
|
|
ForeignKeyConstraintClass,
|
|
GlobalTagsClass,
|
|
OtherSchemaClass,
|
|
SchemaFieldClass as SchemaField,
|
|
SchemaFieldDataTypeClass,
|
|
SchemaMetadataClass,
|
|
StringTypeClass,
|
|
SubTypesClass,
|
|
TagAssociationClass,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# TODO: Support generating docs for each event type in entity registry.
|
|
|
|
|
|
def capitalize_first(something: str) -> str:
|
|
return something[0:1].upper() + something[1:]
|
|
|
|
|
|
class EntityCategory(ConfigEnum):
|
|
CORE = auto()
|
|
INTERNAL = auto()
|
|
|
|
|
|
@dataclass
|
|
class EntityDefinition:
|
|
name: str
|
|
keyAspect: str
|
|
aspects: List[str] = field(default_factory=list)
|
|
aspect_map: Optional[Dict[str, Any]] = None
|
|
relationship_map: Optional[Dict[str, str]] = None
|
|
doc: Optional[str] = None
|
|
doc_file_contents: Optional[str] = None
|
|
# entities are by default in the CORE category unless specified otherwise
|
|
category: EntityCategory = EntityCategory.CORE
|
|
priority: Optional[int] = None
|
|
# schema: Optional[avro.schema.Schema] = None
|
|
# logical_schema: Optional[avro.schema.Schema] = None
|
|
|
|
# @validator("name")
|
|
# def lower_everything(cls, v: str) -> str:
|
|
# return v.lower()
|
|
|
|
@property
|
|
def display_name(self):
|
|
return capitalize_first(self.name)
|
|
|
|
|
|
@dataclass
|
|
class AspectDefinition:
|
|
name: str
|
|
EntityUrns: Optional[List[str]] = None
|
|
schema: Optional[avro.schema.Schema] = None
|
|
type: Optional[str] = None
|
|
|
|
|
|
@dataclass
|
|
class EventDefinition:
|
|
name: str
|
|
|
|
|
|
# New dataclasses for lineage representation
|
|
@dataclass
|
|
class LineageRelationship:
|
|
name: str
|
|
entityTypes: List[str]
|
|
isLineage: bool = True
|
|
|
|
|
|
@dataclass
|
|
class LineageField:
|
|
name: str
|
|
path: str
|
|
isLineage: bool = True
|
|
relationship: Optional[LineageRelationship] = None
|
|
|
|
|
|
@dataclass
|
|
class LineageAspect:
|
|
aspect: str
|
|
fields: List[LineageField]
|
|
|
|
|
|
@dataclass
|
|
class LineageEntity:
|
|
aspects: Dict[str, LineageAspect]
|
|
|
|
|
|
@dataclass
|
|
class LineageData:
|
|
entities: Dict[str, LineageEntity]
|
|
|
|
|
|
entity_registry: Dict[str, EntityDefinition] = {}
|
|
|
|
|
|
def get_aspects_from_snapshot(
|
|
snapshot_schema: avro.schema.RecordSchema,
|
|
) -> Dict[str, AspectDefinition]:
|
|
union_schema: avro.schema.UnionSchema = snapshot_schema.fields[1].type.items
|
|
aspect_map = {}
|
|
for aspect_schema in union_schema.schemas:
|
|
if "Aspect" in aspect_schema.props:
|
|
aspectDef = AspectDefinition(
|
|
schema=aspect_schema,
|
|
name=aspect_schema.props["Aspect"].get("name"),
|
|
)
|
|
aspect_map[aspectDef.name] = aspectDef
|
|
|
|
return aspect_map
|
|
|
|
|
|
aspect_registry: Dict[str, AspectDefinition] = {}
|
|
|
|
# A holder for generated docs
|
|
generated_documentation: Dict[str, str] = {}
|
|
|
|
|
|
# Patch add_name method to NOT complain about duplicate names
|
|
def add_name(self, name_attr, space_attr, new_schema):
|
|
to_add = avro.schema.Name(name_attr, space_attr, self.default_namespace)
|
|
|
|
if self.names:
|
|
self.names[to_add.fullname] = new_schema
|
|
return to_add
|
|
|
|
|
|
def load_schema_file(schema_file: str) -> None:
|
|
logger.debug(f"Loading schema file: {schema_file}")
|
|
|
|
with open(schema_file) as f:
|
|
raw_schema_text = f.read()
|
|
|
|
avro_schema = avro.schema.parse(raw_schema_text)
|
|
|
|
if (
|
|
isinstance(avro_schema, avro.schema.RecordSchema)
|
|
and "Aspect" in avro_schema.other_props
|
|
):
|
|
# probably an aspect schema
|
|
record_schema: avro.schema.RecordSchema = avro_schema
|
|
aspect_def = record_schema.get_prop("Aspect")
|
|
aspect_definition = AspectDefinition(**aspect_def)
|
|
|
|
aspect_definition.schema = record_schema
|
|
aspect_registry[aspect_definition.name] = aspect_definition
|
|
logger.debug(f"Loaded aspect schema: {aspect_definition.name}")
|
|
elif avro_schema.name == "MetadataChangeEvent":
|
|
# probably an MCE schema
|
|
field: Field = avro_schema.fields[1]
|
|
assert isinstance(field.type, avro.schema.UnionSchema)
|
|
for member_schema in field.type.schemas:
|
|
if "Entity" in member_schema.props:
|
|
entity_def = member_schema.props["Entity"]
|
|
entity_name = entity_def["name"]
|
|
entity_definition = entity_registry.get(
|
|
entity_name, EntityDefinition(**entity_def)
|
|
)
|
|
entity_definition.aspect_map = get_aspects_from_snapshot(member_schema)
|
|
all_aspects = [a for a in entity_definition.aspect_map]
|
|
# in terms of order, we prefer the aspects from snapshot over the aspects from the config registry
|
|
# so we flip the aspect list here
|
|
for aspect_name in entity_definition.aspects:
|
|
if aspect_name not in all_aspects:
|
|
all_aspects.append(aspect_name)
|
|
entity_definition.aspects = all_aspects
|
|
entity_registry[entity_name] = entity_definition
|
|
logger.debug(f"Loaded entity schema: {entity_name} with aspects: {all_aspects}")
|
|
else:
|
|
logger.debug(f"Ignoring schema {schema_file}")
|
|
|
|
|
|
def extract_lineage_fields_from_schema(
|
|
schema: avro.schema.Schema,
|
|
current_path: str = ""
|
|
) -> List[LineageField]:
|
|
"""
|
|
Recursively extract lineage fields from an Avro schema.
|
|
|
|
Args:
|
|
schema: The Avro schema to analyze
|
|
current_path: The current field path (for nested fields)
|
|
|
|
Returns:
|
|
List of LineageField objects found in the schema
|
|
"""
|
|
lineage_fields = []
|
|
|
|
if isinstance(schema, avro.schema.RecordSchema):
|
|
logger.debug(f"Analyzing record schema at path: {current_path}")
|
|
for field in schema.fields:
|
|
field_path = f"{current_path}.{field.name}" if current_path else field.name
|
|
logger.debug(f"Analyzing field: {field.name} at path: {field_path}")
|
|
|
|
# Check if this field has lineage properties
|
|
is_lineage = False
|
|
relationship_info = None
|
|
|
|
# Check for isLineage property
|
|
if hasattr(field, 'other_props') and field.other_props:
|
|
is_lineage = field.other_props.get('isLineage', False)
|
|
if is_lineage:
|
|
logger.debug(f"Found isLineage=true for field: {field_path}")
|
|
|
|
# Check for Relationship property
|
|
if 'Relationship' in field.other_props:
|
|
relationship_data = field.other_props['Relationship']
|
|
logger.debug(f"Found Relationship property for field: {field_path}: {relationship_data}")
|
|
|
|
# Handle both direct relationship and path-based relationship
|
|
if 'entityTypes' in relationship_data:
|
|
# Direct relationship
|
|
relationship_is_lineage = relationship_data.get('isLineage', False)
|
|
relationship_info = LineageRelationship(
|
|
name=relationship_data.get('name', ''),
|
|
entityTypes=relationship_data.get('entityTypes', []),
|
|
isLineage=relationship_is_lineage
|
|
)
|
|
is_lineage = is_lineage or relationship_is_lineage
|
|
else:
|
|
# Path-based relationship - find the actual relationship data
|
|
for _, value in relationship_data.items():
|
|
if isinstance(value, dict) and 'entityTypes' in value:
|
|
relationship_is_lineage = value.get('isLineage', False)
|
|
relationship_info = LineageRelationship(
|
|
name=value.get('name', ''),
|
|
entityTypes=value.get('entityTypes', []),
|
|
isLineage=relationship_is_lineage
|
|
)
|
|
is_lineage = is_lineage or relationship_is_lineage
|
|
break
|
|
|
|
# If this field is lineage, add it to the results
|
|
if is_lineage:
|
|
lineage_field = LineageField(
|
|
name=field.name,
|
|
path=field_path,
|
|
isLineage=True,
|
|
relationship=relationship_info
|
|
)
|
|
lineage_fields.append(lineage_field)
|
|
logger.debug(f"Added lineage field: {field_path}")
|
|
|
|
# Recursively check nested fields
|
|
nested_fields = extract_lineage_fields_from_schema(field.type, field_path)
|
|
lineage_fields.extend(nested_fields)
|
|
|
|
elif isinstance(schema, avro.schema.ArraySchema):
|
|
logger.debug(f"Analyzing array schema at path: {current_path}")
|
|
# For arrays, check the items schema
|
|
nested_fields = extract_lineage_fields_from_schema(schema.items, current_path)
|
|
lineage_fields.extend(nested_fields)
|
|
|
|
elif isinstance(schema, avro.schema.UnionSchema):
|
|
logger.debug(f"Analyzing union schema at path: {current_path}")
|
|
# For unions, check all possible schemas
|
|
for union_schema in schema.schemas:
|
|
nested_fields = extract_lineage_fields_from_schema(union_schema, current_path)
|
|
lineage_fields.extend(nested_fields)
|
|
|
|
return lineage_fields
|
|
|
|
|
|
def extract_lineage_fields() -> LineageData:
|
|
"""
|
|
Extract lineage fields from all aspects in the aspect registry.
|
|
|
|
Returns:
|
|
LineageData containing all lineage information organized by entity and aspect
|
|
"""
|
|
logger.info("Starting lineage field extraction")
|
|
|
|
lineage_data = LineageData(entities={})
|
|
|
|
# Group aspects by entity
|
|
entity_aspects: Dict[str, List[str]] = {}
|
|
for entity_name, entity_def in entity_registry.items():
|
|
entity_aspects[entity_name] = entity_def.aspects
|
|
logger.debug(f"Entity {entity_name} has aspects: {entity_def.aspects}")
|
|
|
|
# Process each aspect
|
|
for aspect_name, aspect_def in aspect_registry.items():
|
|
logger.debug(f"Processing aspect: {aspect_name}")
|
|
|
|
if not aspect_def.schema:
|
|
logger.warning(f"Aspect {aspect_name} has no schema, skipping")
|
|
continue
|
|
|
|
# Extract lineage fields from this aspect
|
|
lineage_fields = extract_lineage_fields_from_schema(aspect_def.schema)
|
|
|
|
if lineage_fields:
|
|
logger.info(f"Found {len(lineage_fields)} lineage fields in aspect {aspect_name}")
|
|
|
|
# Find which entities use this aspect
|
|
for entity_name, entity_aspect_list in entity_aspects.items():
|
|
if aspect_name in entity_aspect_list:
|
|
logger.debug(f"Aspect {aspect_name} is used by entity {entity_name}")
|
|
|
|
# Initialize entity if not exists
|
|
if entity_name not in lineage_data.entities:
|
|
lineage_data.entities[entity_name] = LineageEntity(aspects={})
|
|
|
|
# Add aspect with lineage fields
|
|
lineage_aspect = LineageAspect(
|
|
aspect=aspect_name,
|
|
fields=lineage_fields
|
|
)
|
|
lineage_data.entities[entity_name].aspects[aspect_name] = lineage_aspect
|
|
else:
|
|
logger.debug(f"No lineage fields found in aspect {aspect_name}")
|
|
|
|
# Log summary
|
|
total_entities_with_lineage = len(lineage_data.entities)
|
|
total_aspects_with_lineage = sum(len(entity.aspects) for entity in lineage_data.entities.values())
|
|
total_lineage_fields = sum(
|
|
len(aspect.fields)
|
|
for entity in lineage_data.entities.values()
|
|
for aspect in entity.aspects.values()
|
|
)
|
|
|
|
logger.info("Lineage extraction complete:")
|
|
logger.info(f" - Entities with lineage: {total_entities_with_lineage}")
|
|
logger.info(f" - Aspects with lineage: {total_aspects_with_lineage}")
|
|
logger.info(f" - Total lineage fields: {total_lineage_fields}")
|
|
|
|
return lineage_data
|
|
|
|
|
|
def generate_lineage_json(lineage_data: LineageData) -> str:
|
|
"""
|
|
Generate JSON representation of lineage data.
|
|
|
|
Args:
|
|
lineage_data: The lineage data to convert to JSON
|
|
|
|
Returns:
|
|
JSON string representation
|
|
"""
|
|
logger.info("Generating lineage JSON")
|
|
|
|
# Convert dataclasses to dictionaries
|
|
def lineage_field_to_dict(field: LineageField) -> Dict[str, Any]:
|
|
result = {
|
|
"name": field.name,
|
|
"path": field.path,
|
|
"isLineage": field.isLineage
|
|
}
|
|
if field.relationship:
|
|
result["relationship"] = {
|
|
"name": field.relationship.name,
|
|
"entityTypes": field.relationship.entityTypes,
|
|
"isLineage": field.relationship.isLineage
|
|
}
|
|
return result
|
|
|
|
def lineage_aspect_to_dict(aspect: LineageAspect) -> Dict[str, Any]:
|
|
return {
|
|
"aspect": aspect.aspect,
|
|
"fields": [lineage_field_to_dict(field) for field in aspect.fields]
|
|
}
|
|
|
|
def lineage_entity_to_dict(entity: LineageEntity) -> Dict[str, Any]:
|
|
return {
|
|
aspect_name: lineage_aspect_to_dict(aspect)
|
|
for aspect_name, aspect in entity.aspects.items()
|
|
}
|
|
|
|
# Build the final JSON structure
|
|
json_data = {
|
|
"entities": {
|
|
entity_name: lineage_entity_to_dict(entity)
|
|
for entity_name, entity in lineage_data.entities.items()
|
|
}
|
|
}
|
|
|
|
json_data["generated_by"] = "metadata-ingestion/scripts/modeldocgen.py"
|
|
json_data["generated_at"] = datetime.now(timezone.utc).isoformat()
|
|
|
|
json_string = json.dumps(json_data, indent=2)
|
|
logger.info(f"Generated lineage JSON with {len(json_string)} characters")
|
|
|
|
return json_string
|
|
|
|
|
|
@dataclass
|
|
class Relationship:
|
|
name: str
|
|
src: str
|
|
dst: str
|
|
doc: Optional[str] = None
|
|
id: Optional[str] = None
|
|
|
|
|
|
@dataclass
|
|
class RelationshipAdjacency:
|
|
self_loop: List[Relationship] = field(default_factory=list)
|
|
incoming: List[Relationship] = field(default_factory=list)
|
|
outgoing: List[Relationship] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class RelationshipGraph:
|
|
map: Dict[str, RelationshipAdjacency] = field(default_factory=dict)
|
|
|
|
def add_edge(
|
|
self, src: str, dst: str, label: str, reason: str, edge_id: Optional[str] = None
|
|
) -> None:
|
|
relnship = Relationship(
|
|
label, src, dst, reason, id=edge_id or f"{src}:{label}:{dst}:{reason}"
|
|
)
|
|
|
|
if src == dst:
|
|
adjacency = self.map.get(src, RelationshipAdjacency())
|
|
for reln in adjacency.self_loop:
|
|
if relnship.id == reln.id:
|
|
print(f"Skipping adding edge since ids match {reln.id}")
|
|
return
|
|
adjacency.self_loop.append(relnship)
|
|
self.map[src] = adjacency
|
|
else:
|
|
adjacency = self.map.get(src, RelationshipAdjacency())
|
|
for reln in adjacency.outgoing:
|
|
if relnship.id == reln.id:
|
|
logger.info(f"Skipping adding edge since ids match {reln.id}")
|
|
return
|
|
|
|
adjacency.outgoing.append(relnship)
|
|
self.map[src] = adjacency
|
|
|
|
adjacency = self.map.get(dst, RelationshipAdjacency())
|
|
for reln in adjacency.incoming:
|
|
if relnship.id == reln.id:
|
|
logger.info(f"Skipping adding edge since ids match {reln.id}")
|
|
return
|
|
|
|
adjacency.incoming.append(relnship)
|
|
self.map[dst] = adjacency
|
|
|
|
def get_adjacency(self, node: str) -> RelationshipAdjacency:
|
|
return self.map.get(node, RelationshipAdjacency())
|
|
|
|
|
|
def make_relnship_docs(relationships: List[Relationship], direction: str) -> str:
|
|
doc = ""
|
|
map: Dict[str, List[Relationship]] = {}
|
|
for relnship in relationships:
|
|
map[relnship.name] = map.get(relnship.name, [])
|
|
map[relnship.name].append(relnship)
|
|
for rel_name, relnships in map.items():
|
|
doc += f"\n- {rel_name}\n"
|
|
for relnship in relnships:
|
|
doc += f"\n - {relnship.dst if direction == 'outgoing' else relnship.src}{relnship.doc or ''}"
|
|
return doc
|
|
|
|
|
|
def make_entity_docs(entity_display_name: str, graph: RelationshipGraph) -> str:
|
|
entity_name = entity_display_name[0:1].lower() + entity_display_name[1:]
|
|
entity_def: Optional[EntityDefinition] = entity_registry.get(entity_name)
|
|
if entity_def:
|
|
doc = entity_def.doc_file_contents or (
|
|
f"# {entity_def.display_name}\n{entity_def.doc}"
|
|
if entity_def.doc
|
|
else f"# {entity_def.display_name}"
|
|
)
|
|
# create aspects section
|
|
aspects_section = "\n## Aspects\n" if entity_def.aspects else ""
|
|
|
|
deprecated_aspects_section = ""
|
|
timeseries_aspects_section = ""
|
|
|
|
for aspect in entity_def.aspects or []:
|
|
aspect_definition: AspectDefinition = aspect_registry[aspect]
|
|
assert aspect_definition
|
|
assert aspect_definition.schema
|
|
deprecated_message = (
|
|
" (Deprecated)"
|
|
if aspect_definition.schema.get_prop("Deprecated")
|
|
else ""
|
|
)
|
|
timeseries_qualifier = (
|
|
" (Timeseries)" if aspect_definition.type == "timeseries" else ""
|
|
)
|
|
this_aspect_doc = ""
|
|
this_aspect_doc += (
|
|
f"\n### {aspect}{deprecated_message}{timeseries_qualifier}\n"
|
|
)
|
|
this_aspect_doc += f"{aspect_definition.schema.get_prop('doc')}\n"
|
|
this_aspect_doc += "<details>\n<summary>Schema</summary>\n\n"
|
|
# breakpoint()
|
|
this_aspect_doc += f"```javascript\n{json.dumps(aspect_definition.schema.to_json(), indent=2)}\n```\n</details>\n"
|
|
|
|
if deprecated_message:
|
|
deprecated_aspects_section += this_aspect_doc
|
|
elif timeseries_qualifier:
|
|
timeseries_aspects_section += this_aspect_doc
|
|
else:
|
|
aspects_section += this_aspect_doc
|
|
|
|
aspects_section += timeseries_aspects_section + deprecated_aspects_section
|
|
|
|
# create relationships section
|
|
relationships_section = "\n## Relationships\n"
|
|
adjacency = graph.get_adjacency(entity_def.display_name)
|
|
if adjacency.self_loop:
|
|
relationships_section += "\n### Self\nThese are the relationships to itself, stored in this entity's aspects"
|
|
for relnship in adjacency.self_loop:
|
|
relationships_section += (
|
|
f"\n- {relnship.name} ({relnship.doc[1:] if relnship.doc else ''})"
|
|
)
|
|
|
|
if adjacency.outgoing:
|
|
relationships_section += "\n### Outgoing\nThese are the relationships stored in this entity's aspects"
|
|
relationships_section += make_relnship_docs(
|
|
adjacency.outgoing, direction="outgoing"
|
|
)
|
|
|
|
if adjacency.incoming:
|
|
relationships_section += "\n### Incoming\nThese are the relationships stored in other entity's aspects"
|
|
relationships_section += make_relnship_docs(
|
|
adjacency.incoming, direction="incoming"
|
|
)
|
|
|
|
# create global metadata graph
|
|
global_graph_url = "https://github.com/datahub-project/static-assets/raw/main/imgs/datahub-metadata-model.png"
|
|
global_graph_section = (
|
|
f"\n## [Global Metadata Model]({global_graph_url})"
|
|
+ f"\n"
|
|
)
|
|
final_doc = doc + aspects_section + relationships_section + global_graph_section
|
|
generated_documentation[entity_name] = final_doc
|
|
return final_doc
|
|
else:
|
|
raise Exception(f"Failed to find information for entity: {entity_name}")
|
|
|
|
|
|
def generate_stitched_record(
|
|
relnships_graph: RelationshipGraph,
|
|
) -> Iterable[MetadataChangeProposalWrapper]:
|
|
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
|
|
|
|
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(
|
|
f["type"],
|
|
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="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 aspect_info.get("type", "") == "timeseries":
|
|
# f_field.globalTags = f_field.globalTags or GlobalTagsClass(
|
|
# tags=[]
|
|
# )
|
|
f_field.globalTags.tags.append(
|
|
TagAssociationClass(tag="urn:li:tag:Temporal")
|
|
)
|
|
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 _, 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="datahub",
|
|
name=destination_entity_name,
|
|
)
|
|
fkey = ForeignKeyConstraintClass(
|
|
name=relnship_name,
|
|
foreignDataset=foreign_dataset_urn,
|
|
foreignFields=[
|
|
make_schema_field_urn(foreign_dataset_urn, "urn")
|
|
],
|
|
sourceFields=[
|
|
make_schema_field_urn(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)}",
|
|
)
|
|
|
|
dataset_urn = make_dataset_urn(
|
|
platform="datahub",
|
|
name=entity_display_name,
|
|
)
|
|
|
|
yield from MetadataChangeProposalWrapper.construct_many(
|
|
entityUrn=dataset_urn,
|
|
aspects=[
|
|
SchemaMetadataClass(
|
|
schemaName=str(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,
|
|
),
|
|
GlobalTagsClass(
|
|
tags=[TagAssociationClass(tag="urn:li:tag:Entity")]
|
|
),
|
|
BrowsePathsClass([f"/prod/datahub/entities/{entity_display_name}"]),
|
|
BrowsePathsV2Class(
|
|
[
|
|
BrowsePathEntryClass(id="entities"),
|
|
BrowsePathEntryClass(id=entity_display_name),
|
|
]
|
|
),
|
|
DatasetPropertiesClass(
|
|
description=make_entity_docs(
|
|
dataset_urn.split(":")[-1].split(",")[1], relnships_graph
|
|
)
|
|
),
|
|
SubTypesClass(typeNames=["entity"]),
|
|
],
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class EntityAspectName:
|
|
entityName: str
|
|
aspectName: str
|
|
|
|
|
|
class AspectPluginConfig(PermissiveConfigModel):
|
|
className: str
|
|
enabled: bool
|
|
supportedEntityAspectNames: List[EntityAspectName] = []
|
|
packageScan: Optional[List[str]] = None
|
|
supportedOperations: Optional[List[str]] = None
|
|
|
|
|
|
class PluginConfiguration(PermissiveConfigModel):
|
|
aspectPayloadValidators: Optional[List[AspectPluginConfig]] = None
|
|
mutationHooks: Optional[List[AspectPluginConfig]] = None
|
|
mclSideEffects: Optional[List[AspectPluginConfig]] = None
|
|
mcpSideEffects: Optional[List[AspectPluginConfig]] = None
|
|
|
|
|
|
class EntityRegistry(PermissiveConfigModel):
|
|
entities: List[EntityDefinition]
|
|
events: Optional[List[EventDefinition]]
|
|
plugins: Optional[PluginConfiguration] = None
|
|
|
|
|
|
def load_registry_file(registry_file: str) -> Dict[str, EntityDefinition]:
|
|
import yaml
|
|
|
|
with open(registry_file, "r") as f:
|
|
registry = EntityRegistry.parse_obj(yaml.safe_load(f))
|
|
index: int = 0
|
|
for entity_def in registry.entities:
|
|
index += 1
|
|
entity_def.priority = index
|
|
entity_registry[entity_def.name] = entity_def
|
|
return entity_registry
|
|
|
|
|
|
def get_sorted_entity_names(
|
|
entity_names: List[Tuple[str, EntityDefinition]]
|
|
) -> List[Tuple[str, List[str]]]:
|
|
"""
|
|
Sort entity names by category and priority for documentation generation.
|
|
|
|
This function organizes entities into a structured order for generating
|
|
documentation. Entities are grouped by category (CORE vs INTERNAL) and
|
|
within each category, sorted by priority and then alphabetically.
|
|
|
|
Business Logic:
|
|
- CORE entities are displayed first, followed by INTERNAL entities
|
|
- Within each category, entities with priority values are sorted first
|
|
- Priority entities are sorted by their priority value (lower numbers = higher priority)
|
|
- Non-priority entities are sorted alphabetically after priority entities
|
|
- Zero and negative priority values are treated as valid priorities
|
|
|
|
Args:
|
|
entity_names: List of tuples containing (entity_name, EntityDefinition)
|
|
|
|
Returns:
|
|
List of tuples containing (EntityCategory, List[str]) where:
|
|
- First tuple: (EntityCategory.CORE, sorted_core_entity_names)
|
|
- Second tuple: (EntityCategory.INTERNAL, sorted_internal_entity_names)
|
|
|
|
Example:
|
|
Input: [
|
|
("dataset", EntityDefinition(priority=2, category=CORE)),
|
|
("table", EntityDefinition(priority=None, category=CORE)),
|
|
("internal", EntityDefinition(priority=1, category=INTERNAL))
|
|
]
|
|
Output: [
|
|
(EntityCategory.CORE, ["dataset", "table"]),
|
|
(EntityCategory.INTERNAL, ["internal"])
|
|
]
|
|
"""
|
|
core_entities = [
|
|
(x, y) for (x, y) in entity_names if y.category == EntityCategory.CORE
|
|
]
|
|
priority_bearing_core_entities = [(x, y) for (x, y) in core_entities if y.priority]
|
|
priority_bearing_core_entities.sort(key=lambda t: t[1].priority)
|
|
priority_bearing_core_entities = [x for (x, y) in priority_bearing_core_entities]
|
|
|
|
non_priority_core_entities = [x for (x, y) in core_entities if not y.priority]
|
|
non_priority_core_entities.sort()
|
|
|
|
internal_entities = [
|
|
(x, y) for (x, y) in entity_names if y.category == EntityCategory.INTERNAL
|
|
]
|
|
priority_bearing_internal_entities = [
|
|
x for (x, y) in internal_entities if y.priority
|
|
]
|
|
|
|
non_priority_internal_entities = [
|
|
x for (x, y) in internal_entities if not y.priority
|
|
]
|
|
|
|
sorted_entities = [
|
|
(
|
|
EntityCategory.CORE,
|
|
priority_bearing_core_entities + non_priority_core_entities,
|
|
),
|
|
(
|
|
EntityCategory.INTERNAL,
|
|
priority_bearing_internal_entities + non_priority_internal_entities,
|
|
),
|
|
]
|
|
|
|
return sorted_entities
|
|
|
|
|
|
@click.command()
|
|
@click.argument("schemas_root", type=click.Path(exists=True), required=True)
|
|
@click.option("--registry", type=click.Path(exists=True), required=True)
|
|
@click.option("--generated-docs-dir", type=click.Path(exists=True), required=True)
|
|
@click.option("--server", type=str, required=False)
|
|
@click.option("--file", type=str, required=False)
|
|
@click.option(
|
|
"--dot", type=str, required=False, help="generate a dot file representing the graph"
|
|
)
|
|
@click.option("--png", type=str, required=False)
|
|
@click.option("--extra-docs", type=str, required=False)
|
|
@click.option("--lineage-output", type=str, required=False, help="generate lineage JSON file")
|
|
def generate(
|
|
schemas_root: str,
|
|
registry: str,
|
|
generated_docs_dir: str,
|
|
server: Optional[str],
|
|
file: Optional[str],
|
|
dot: Optional[str],
|
|
png: Optional[str],
|
|
extra_docs: Optional[str],
|
|
lineage_output: Optional[str],
|
|
) -> None:
|
|
logger.info(f"server = {server}")
|
|
logger.info(f"file = {file}")
|
|
logger.info(f"dot = {dot}")
|
|
logger.info(f"png = {png}")
|
|
logger.info(f"lineage_output = {lineage_output}")
|
|
|
|
entity_extra_docs = {}
|
|
if extra_docs:
|
|
for path in glob.glob(f"{extra_docs}/**/*.md", recursive=True):
|
|
m = re.search("/docs/entities/(.*)/*.md", path)
|
|
if m:
|
|
entity_name = m.group(1)
|
|
with open(path, "r") as doc_file:
|
|
file_contents = doc_file.read()
|
|
entity_extra_docs[entity_name] = file_contents
|
|
|
|
# registry file
|
|
load_registry_file(registry)
|
|
|
|
# schema files
|
|
for schema_file in Path(schemas_root).glob("**/*.avsc"):
|
|
if (
|
|
schema_file.name in {"MetadataChangeEvent.avsc"}
|
|
or json.loads(schema_file.read_text()).get("Aspect") is not None
|
|
):
|
|
load_schema_file(str(schema_file))
|
|
|
|
if entity_extra_docs:
|
|
for entity_name in entity_extra_docs:
|
|
|
|
entity_registry[entity_name].doc_file_contents = entity_extra_docs[
|
|
entity_name
|
|
]
|
|
|
|
if lineage_output:
|
|
logger.info(f"Generating lineage JSON to {lineage_output}")
|
|
try:
|
|
lineage_data = extract_lineage_fields()
|
|
lineage_json = generate_lineage_json(lineage_data)
|
|
|
|
output_path = Path(lineage_output)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
new_json_data = json.loads(lineage_json)
|
|
|
|
write_file = should_write_json_file(output_path, new_json_data, "lineage file")
|
|
|
|
if write_file:
|
|
with open(output_path, 'w') as f:
|
|
f.write(lineage_json)
|
|
logger.info(f"Successfully wrote lineage JSON to {lineage_output}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to generate lineage JSON: {e}")
|
|
raise
|
|
|
|
relationship_graph = RelationshipGraph()
|
|
mcps = list(generate_stitched_record(relationship_graph))
|
|
|
|
shutil.rmtree(f"{generated_docs_dir}/entities", ignore_errors=True)
|
|
entity_names = [(x, entity_registry[x]) for x in generated_documentation]
|
|
|
|
sorted_entity_names = get_sorted_entity_names(entity_names)
|
|
|
|
index = 0
|
|
for _, sorted_entities in sorted_entity_names:
|
|
for entity_name in sorted_entities:
|
|
|
|
entity_dir = f"{generated_docs_dir}/entities/"
|
|
|
|
os.makedirs(entity_dir, exist_ok=True)
|
|
|
|
with open(f"{entity_dir}/{entity_name}.md", "w") as fp:
|
|
fp.write("---\n")
|
|
fp.write(f"sidebar_position: {index}\n")
|
|
fp.write("---\n")
|
|
fp.write(generated_documentation[entity_name])
|
|
index += 1
|
|
|
|
if file:
|
|
logger.info(f"Will write events to {file}")
|
|
Path(file).parent.mkdir(parents=True, exist_ok=True)
|
|
fileSink = FileSink(
|
|
PipelineContext(run_id="generated-metaModel"),
|
|
FileSinkConfig(filename=file),
|
|
)
|
|
for e in mcps:
|
|
fileSink.write_record_async(
|
|
RecordEnvelope(e, metadata={}), write_callback=NoopWriteCallback()
|
|
)
|
|
fileSink.close()
|
|
pipeline_config = {
|
|
"source": {
|
|
"type": "file",
|
|
"config": {"filename": file},
|
|
},
|
|
"sink": {
|
|
"type": "datahub-rest",
|
|
"config": {
|
|
"server": "${DATAHUB_SERVER:-http://localhost:8080}",
|
|
"token": "${DATAHUB_TOKEN:-}",
|
|
},
|
|
},
|
|
"run_id": "modeldoc-generated",
|
|
}
|
|
pipeline_file = Path(file).parent.absolute() / "pipeline.yml"
|
|
with open(pipeline_file, "w") as f:
|
|
json.dump(pipeline_config, f, indent=2)
|
|
logger.info(f"Wrote pipeline to {pipeline_file}")
|
|
|
|
if server:
|
|
logger.info(f"Will send events to {server}")
|
|
assert server.startswith("http://"), "server address must start with http://"
|
|
emitter = DatahubRestEmitter(gms_server=server)
|
|
emitter.test_connection()
|
|
for e in mcps:
|
|
emitter.emit(e)
|
|
|
|
if dot:
|
|
logger.info(f"Will write dot file to {dot}")
|
|
|
|
import pydot
|
|
|
|
graph = pydot.Dot("my_graph", graph_type="graph")
|
|
for node, adjacency in relationship_graph.map.items():
|
|
my_node = pydot.Node(
|
|
node,
|
|
label=node,
|
|
shape="box",
|
|
)
|
|
graph.add_node(my_node)
|
|
if adjacency.self_loop:
|
|
for relnship in adjacency.self_loop:
|
|
graph.add_edge(
|
|
pydot.Edge(
|
|
src=relnship.src, dst=relnship.dst, label=relnship.name
|
|
)
|
|
)
|
|
if adjacency.outgoing:
|
|
for relnship in adjacency.outgoing:
|
|
graph.add_edge(
|
|
pydot.Edge(
|
|
src=relnship.src, dst=relnship.dst, label=relnship.name
|
|
)
|
|
)
|
|
Path(dot).parent.mkdir(parents=True, exist_ok=True)
|
|
graph.write_raw(dot)
|
|
if png:
|
|
try:
|
|
graph.write_png(png)
|
|
except Exception as e:
|
|
logger.error(
|
|
"Failed to create png file. Do you have graphviz installed?"
|
|
)
|
|
raise e
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logger.setLevel("INFO")
|
|
generate()
|