def add_entity_to_container(container_key: KeyType, entity_type: str,
entity_urn: str) -> Iterable[MetadataWorkUnit]:
container_urn = make_container_urn(guid=container_key.guid(), )
mcp = MetadataChangeProposalWrapper(
entityType=entity_type,
changeType=ChangeTypeClass.UPSERT,
entityUrn=entity_urn,
aspectName="container",
aspect=ContainerClass(container=f"{container_urn}"),
)
wu = MetadataWorkUnit(id=f"container-{container_urn}-to-{entity_urn}",
mcp=mcp)
yield wu
def add_dataset_to_container(
container_key: KeyType,
dataset_urn: str) -> Iterable[Union[MetadataWorkUnit]]:
container_urn = make_container_urn(guid=container_key.guid(), )
mcp = MetadataChangeProposalWrapper(
entityType="dataset",
changeType=ChangeTypeClass.UPSERT,
entityUrn=f"{dataset_urn}",
aspectName="container",
aspect=ContainerClass(container=f"{container_urn}"),
# aspect=ContainerKeyClass(guid=schema_container_key.guid())
)
wu = MetadataWorkUnit(id=f"container-{container_urn}-to-{dataset_urn}",
mcp=mcp)
yield wu
def test_sql_common_state() -> None:
state1 = BaseSQLAlchemyCheckpointState()
test_table_urn = make_dataset_urn("test_platform", "db1.test_table1", "test")
state1.add_table_urn(test_table_urn)
test_view_urn = make_dataset_urn("test_platform", "db1.test_view1", "test")
state1.add_view_urn(test_view_urn)
test_container_urn = make_container_urn("test_container")
state1.add_container_guid(test_container_urn)
state2 = BaseSQLAlchemyCheckpointState()
table_urns_diff = list(state1.get_table_urns_not_in(state2))
assert len(table_urns_diff) == 1 and table_urns_diff[0] == test_table_urn
view_urns_diff = list(state1.get_view_urns_not_in(state2))
assert len(view_urns_diff) == 1 and view_urns_diff[0] == test_view_urn
container_urns_diff = list(state1.get_container_urns_not_in(state2))
assert (
len(container_urns_diff) == 1 and container_urns_diff[0] == test_container_urn
)
def gen_containers(
container_key: KeyType,
name: str,
sub_types: List[str],
parent_container_key: Optional[PlatformKey] = None,
domain_urn: Optional[str] = None,
description: Optional[str] = None,
owner_urn: Optional[str] = None,
external_url: Optional[str] = None,
tags: Optional[List[str]] = None,
) -> Iterable[MetadataWorkUnit]:
container_urn = make_container_urn(guid=container_key.guid(), )
mcp = MetadataChangeProposalWrapper(
entityType="container",
changeType=ChangeTypeClass.UPSERT,
entityUrn=f"{container_urn}",
# entityKeyAspect=ContainerKeyClass(guid=schema_container_key.guid()),
aspectName="containerProperties",
aspect=ContainerProperties(
name=name,
description=description,
customProperties=container_key.dict(exclude_none=True,
by_alias=True),
externalUrl=external_url,
),
)
wu = MetadataWorkUnit(id=f"container-info-{name}-{container_urn}", mcp=mcp)
yield wu
mcp = MetadataChangeProposalWrapper(
entityType="container",
changeType=ChangeTypeClass.UPSERT,
entityUrn=f"{container_urn}",
# entityKeyAspect=ContainerKeyClass(guid=schema_container_key.guid()),
aspectName="dataPlatformInstance",
aspect=DataPlatformInstance(
platform=f"{make_data_platform_urn(container_key.platform)}", ),
)
wu = MetadataWorkUnit(
id=f"container-platforminstance-{name}-{container_urn}", mcp=mcp)
yield wu
# Set subtype
subtype_mcp = MetadataChangeProposalWrapper(
entityType="container",
changeType=ChangeTypeClass.UPSERT,
entityUrn=f"{container_urn}",
# entityKeyAspect=ContainerKeyClass(guid=schema_container_key.guid()),
aspectName="subTypes",
aspect=SubTypesClass(typeNames=sub_types),
)
wu = MetadataWorkUnit(id=f"container-subtypes-{name}-{container_urn}",
mcp=subtype_mcp)
yield wu
if domain_urn:
yield from add_domain_to_entity_wu(
entity_type="container",
entity_urn=container_urn,
domain_urn=domain_urn,
)
if owner_urn:
yield from add_owner_to_entity_wu(
entity_type="container",
entity_urn=container_urn,
owner_urn=owner_urn,
)
if tags:
yield from add_tags_to_entity_wu(
entity_type="container",
entity_urn=container_urn,
tags=tags,
)
if parent_container_key:
parent_container_urn = make_container_urn(
guid=parent_container_key.guid(), )
# Set database container
parent_container_mcp = MetadataChangeProposalWrapper(
entityType="container",
changeType=ChangeTypeClass.UPSERT,
entityUrn=f"{container_urn}",
# entityKeyAspect=ContainerKeyClass(guid=schema_container_key.guid()),
aspectName="container",
aspect=ContainerClass(container=parent_container_urn),
# aspect=ContainerKeyClass(guid=database_container_key.guid())
)
wu = MetadataWorkUnit(
id=
f"container-parent-container-{name}-{container_urn}-{parent_container_urn}",
mcp=parent_container_mcp,
)
yield wu
def _get_container_urns_not_in(encoded_urns_1: List[str],
encoded_urns_2: List[str]) -> Iterable[str]:
difference = set(encoded_urns_1) - set(encoded_urns_2)
for guid in difference:
yield make_container_urn(guid)