Skip to content

Commit

Permalink
feat(sdk): structured properties - add support for listing
Browse files Browse the repository at this point in the history
  • Loading branch information
shirshanka committed Jan 7, 2025
1 parent 03e3f46 commit 261f987
Show file tree
Hide file tree
Showing 3 changed files with 108 additions and 1 deletion.
Original file line number Diff line number Diff line change
@@ -0,0 +1,12 @@
# Usage: python3 list_structured_properties.py
# Expected Output: List of structured properties
# This script lists all structured properties in DataHub
from datahub.api.entities.structuredproperties.structuredproperties import (
StructuredProperties,
)
from datahub.ingestion.graph.client import get_default_graph

with get_default_graph() as graph:
structuredproperties = StructuredProperties.list(graph)
for structuredproperty in structuredproperties:
print(structuredproperty.dict())
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import logging
from enum import Enum
from pathlib import Path
from typing import List, Optional
from typing import Iterable, List, Optional

import yaml
from pydantic import validator
Expand Down Expand Up @@ -226,3 +226,14 @@ def to_yaml(
yaml.indent(mapping=2, sequence=4, offset=2)
yaml.default_flow_style = False
yaml.dump(self.dict(), fp)

@staticmethod
def list_urns(graph: DataHubGraph) -> Iterable[str]:
return graph.get_urns_by_filter(

Check warning on line 232 in metadata-ingestion/src/datahub/api/entities/structuredproperties/structuredproperties.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/api/entities/structuredproperties/structuredproperties.py#L232

Added line #L232 was not covered by tests
entity_types=["structuredProperty"],
)

@staticmethod
def list(graph: DataHubGraph) -> Iterable["StructuredProperties"]:
for urn in StructuredProperties.list_urns(graph):
yield StructuredProperties.from_datahub(graph, urn)

Check warning on line 239 in metadata-ingestion/src/datahub/api/entities/structuredproperties/structuredproperties.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/api/entities/structuredproperties/structuredproperties.py#L238-L239

Added lines #L238 - L239 were not covered by tests
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
import json
import logging
from pathlib import Path
from typing import Iterable

import click
from click_default_group import DefaultGroup
from ruamel.yaml import YAML

from datahub.api.entities.structuredproperties.structuredproperties import (
StructuredProperties,
Expand Down Expand Up @@ -61,3 +63,85 @@ def get(urn: str, to_file: str) -> None:
)
else:
click.secho(f"Structured property {urn} does not exist")


@properties.command(
name="list",
)
@click.option("--details/--no-details", is_flag=True, default=True)
@click.option("--to-file", required=False, type=str)
@telemetry.with_telemetry()
def list(details: bool, to_file: str) -> None:
"""List structured properties in DataHub"""

def to_yaml_list(

Check warning on line 77 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L77

Added line #L77 was not covered by tests
objects: Iterable[StructuredProperties], # iterable of objects to dump
file: Path,
) -> None:
# if file exists, first we read it
yaml = YAML(typ="rt") # default, if not specfied, is 'rt' (round-trip)
yaml.indent(mapping=2, sequence=4, offset=2)
yaml.default_flow_style = False
serialized_objects = []
if file.exists():
with open(file, "r") as fp:
existing_objects = yaml.load(fp) # this is a list of dicts
existing_objects = [

Check warning on line 89 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L82-L89

Added lines #L82 - L89 were not covered by tests
StructuredProperties.parse_obj(obj) for obj in existing_objects
]
objects = [obj for obj in objects]

Check warning on line 92 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L92

Added line #L92 was not covered by tests
# do a positional update of the existing objects
existing_urns = {obj.urn for obj in existing_objects}

Check warning on line 94 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L94

Added line #L94 was not covered by tests
# existing_urns = {obj["urn"] if "urn" in obj else f"urn:li:structuredProperty:{obj['id']}" for obj in existing_objects}
for i, obj in enumerate(existing_objects):

Check warning on line 96 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L96

Added line #L96 was not covered by tests
# existing_urn = obj["urn"] if "urn" in obj else f"urn:li:structuredProperty:{obj['id']}"
existing_urn = obj.urn

Check warning on line 98 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L98

Added line #L98 was not covered by tests
# breakpoint()
if existing_urn in {obj.urn for obj in objects}:
existing_objects[i] = next(

Check warning on line 101 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L100-L101

Added lines #L100 - L101 were not covered by tests
obj.dict(exclude_unset=True, exclude_none=True)
for obj in objects
if obj.urn == existing_urn
)
new_objects = [

Check warning on line 106 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L106

Added line #L106 was not covered by tests
obj.dict(exclude_unset=True, exclude_none=True)
for obj in objects
if obj.urn not in existing_urns
]
serialized_objects = existing_objects + new_objects

Check warning on line 111 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L111

Added line #L111 was not covered by tests
else:
serialized_objects = [

Check warning on line 113 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L113

Added line #L113 was not covered by tests
obj.dict(exclude_unset=True, exclude_none=True) for obj in objects
]

with open(file, "w") as fp:
yaml.dump(serialized_objects, fp)

Check warning on line 118 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L117-L118

Added lines #L117 - L118 were not covered by tests

with get_default_graph() as graph:
if details:
logger.info(

Check warning on line 122 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L120-L122

Added lines #L120 - L122 were not covered by tests
"Listing structured properties with details. Use --no-details for urns only"
)
structuredproperties = StructuredProperties.list(graph)
if to_file:
to_yaml_list(structuredproperties, Path(to_file))

Check warning on line 127 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L125-L127

Added lines #L125 - L127 were not covered by tests
else:
for structuredproperty in structuredproperties:
click.secho(

Check warning on line 130 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L129-L130

Added lines #L129 - L130 were not covered by tests
f"{json.dumps(structuredproperty.dict(exclude_unset=True, exclude_none=True), indent=2)}"
)
else:
logger.info(

Check warning on line 134 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L134

Added line #L134 was not covered by tests
"Listing structured property urns only, use --details for more information"
)
structured_property_urns = StructuredProperties.list_urns(graph)
if to_file:
with open(to_file, "w") as f:
for urn in structured_property_urns:
f.write(f"{urn}\n")
click.secho(

Check warning on line 142 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L137-L142

Added lines #L137 - L142 were not covered by tests
f"Structured property urns written to {to_file}", fg="green"
)
else:
for urn in structured_property_urns:
click.secho(f"{urn}")

Check warning on line 147 in metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py

View check run for this annotation

Codecov / codecov/patch

metadata-ingestion/src/datahub/cli/specific/structuredproperties_cli.py#L146-L147

Added lines #L146 - L147 were not covered by tests

0 comments on commit 261f987

Please sign in to comment.