Skip to content

Commit

Permalink
feat(cli): add sibling helper CLI command (#9744)
Browse files Browse the repository at this point in the history
  • Loading branch information
anshbansal authored Jan 30, 2024
1 parent 8559eba commit aa98f48
Showing 1 changed file with 39 additions and 1 deletion.
40 changes: 39 additions & 1 deletion metadata-ingestion/src/datahub/cli/specific/dataset_cli.py
Original file line number Diff line number Diff line change
@@ -1,12 +1,15 @@
import json
import logging
from pathlib import Path
from typing import Set, Tuple

import click
from click_default_group import DefaultGroup

from datahub.api.entities.dataset.dataset import Dataset
from datahub.ingestion.graph.client import get_default_graph
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
from datahub.metadata.com.linkedin.pegasus2avro.common import Siblings
from datahub.telemetry import telemetry
from datahub.upgrade import upgrade

Expand Down Expand Up @@ -65,3 +68,38 @@ def get(urn: str, to_file: str) -> None:
click.secho(f"Dataset yaml written to {to_file}", fg="green")
else:
click.secho(f"Dataset {urn} does not exist")


@dataset.command()
@click.option("--urn", required=True, type=str, help="URN of primary sibling")
@click.option(
"--sibling-urns",
required=True,
type=str,
help="URN of secondary sibling(s)",
multiple=True,
)
@telemetry.with_telemetry()
def add_sibling(urn: str, sibling_urns: Tuple[str]) -> None:
all_urns = set()
all_urns.add(urn)
for sibling_urn in sibling_urns:
all_urns.add(sibling_urn)
with get_default_graph() as graph:
for _urn in all_urns:
_emit_sibling(graph, urn, _urn, all_urns)


def _emit_sibling(
graph: DataHubGraph, primary_urn: str, urn: str, all_urns: Set[str]
) -> None:
siblings = []
for sibling_urn in all_urns:
if sibling_urn != urn:
siblings.append(sibling_urn)
graph.emit(
MetadataChangeProposalWrapper(
entityUrn=urn,
aspect=Siblings(primary=primary_urn == urn, siblings=sorted(siblings)),
)
)

0 comments on commit aa98f48

Please sign in to comment.