from typing import Dict, List, Literal, Union, Optional, Any
from airflow.utils.context import Context
from config import Config
from util import BaseUtil
DATAHUB_SECRET_MANAGER_CONNECTION_ID = (
f"airflow/connections/{Config.DATAHUB_CONNECTION_ID}"
)
DATAHUB_DEFAULT_DOMAIN_NAME = "Data Engineering"
TILESERVER_PLATFORM = "tileservergl"
INGESTION_USER = "ingestion"
PlatformType = Literal[
"athena", "postgres", "airflow", "tileservergl", "bossplatform", "grapi", "game"
]
DataHubEnvType = Literal[
"DEV",
"STG",
"PROD",
]
[docs]
class DataHubUtil:
"""A set of the utility functions to interact with DataHub."""
@staticmethod
def _get_datahub_token() -> str:
"""
Private method to get the DataHub access token.
:return: An access token
:rtype: _type_
"""
import json
secret_value = BaseUtil.get_aws_secret(
secret_id=DATAHUB_SECRET_MANAGER_CONNECTION_ID
)
return json.loads(secret_value)["password"]
[docs]
@staticmethod
def get_emitter():
"""
Returns a DataHubRestEmitter object.
This function retrieves an access token using the `_get_datahub_token` method from the `DataHubUtil` class.
It then creates and returns a `DataHubRestEmitter` object using the retrieved access token and the `DATAHUB_GMS_ENDPOINT`
configuration value.
:return: A `DataHubRestEmitter` object.
:rtype: DataHubRestEmitter
"""
from datahub.emitter.rest_emitter import DataHubRestEmitter
access_token = DataHubUtil._get_datahub_token()
return DataHubRestEmitter(
gms_server=Config.DATAHUB_GMS_ENDPOINT, token=access_token
)
[docs]
@staticmethod
def get_graph():
"""
Returns a DataHubGraph object.
:return: DataHubGraph object.
"""
from datahub.ingestion.graph.client import DatahubClientConfig, DataHubGraph
access_token = DataHubUtil._get_datahub_token()
return DataHubGraph(
DatahubClientConfig(server=Config.DATAHUB_GMS_ENDPOINT, token=access_token)
)
[docs]
@staticmethod
def dataset_urn(
platform: PlatformType,
name: str,
is_production: bool = Config.PRODUCTION,
) -> str:
"""
returns generated dataset urn considering Ready Data platform settings
:param platform: the platform where the dataset is located
:type platform: Literal["athena","postgres","airflow", "tileservergl", "bossplatform", "grapi"]
:param name: name of the table or dataset or dag, task
:type name: str
:param is_production: the running environment., defaults to Config.PRODUCTION
:type is_production: bool, optional
:return: a dataset urn matching with datahub schema
:rtype: str
"""
from datahub.metadata.schema_classes import FabricTypeClass
import datahub.emitter.mce_builder as builder
if is_production:
env = FabricTypeClass.PROD
else:
env = FabricTypeClass.DEV
if platform == "athena":
name = f"{Config.DATA_LAKE_DB}.{name}"
urn = builder.make_dataset_urn(platform, name, env)
elif platform == "bossplatform" or platform == "grapi":
name = f"bossplatform.{name}"
urn = builder.make_dataset_urn(platform, name, env)
elif platform == "game":
name = f"game.{name}"
urn = builder.make_dataset_urn(platform, name, env)
elif platform == "airflow":
urn = builder.make_data_flow_urn(
orchestrator="airflow", flow_id=name, cluster=env
)
else:
urn = builder.make_dataset_urn(platform, name, env)
return urn
[docs]
@staticmethod
def lineage_urn(
platform: PlatformType,
name: str,
is_production: bool = Config.PRODUCTION,
):
"""
_summary_
:param platform: the platform where the inlet or outlet is located
:type platform: Literal["athena", "postgres", "airflow", "tileservergl", "bossplatform", "grapi"]
:param name: the name of the inlet or outlet
:type name: str
:param is_production: defaults to Config.PRODUCTION
:type is_production: bool, optional
:return: _description_
:rtype: _type_
"""
from datahub_airflow_plugin.entities import Urn
dataset_urn = DataHubUtil.dataset_urn(platform, name, is_production)
return Urn(dataset_urn)
[docs]
@staticmethod
def domain_urn(
domain: str = DATAHUB_DEFAULT_DOMAIN_NAME,
) -> str:
"""
Returns the domain URN for the given domain.
:param domain: The domain name., defaults to DATAHUB_DEFAULT_DOMAIN_NAME
:type domain: str, optional
:return: The domain URN.
:rtype: str
"""
import datahub.emitter.mce_builder as builder
return builder.make_domain_urn(domain.lower())
[docs]
@staticmethod
def create_domain(
name: str,
description: str,
parent_domain_urn: str = None,
) -> str:
"""
create a domain in datahub
:param name: name of the domain
:type name: str
:param description: description of the domain
:type description: str
:param parent_domain_urn: the parent domain urn. It should be a urn, defaults to None
:type parent_domain_urn: str, optional
:return: the new domain urn
:rtype: str
"""
from datahub.emitter.mce_builder import make_domain_urn
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.metadata.schema_classes import (
ChangeTypeClass,
DomainPropertiesClass,
)
domain_urn = make_domain_urn(name.lower())
domain_properties_aspect = DomainPropertiesClass(
name=name, description=description, parentDomain=parent_domain_urn
)
event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
entityType="domain",
changeType=ChangeTypeClass.UPSERT,
entityUrn=domain_urn,
aspect=domain_properties_aspect,
)
rest_emitter = DataHubUtil.get_emitter()
rest_emitter.emit(event)
print(f"Domain {name} created")
return domain_urn
[docs]
@staticmethod
def tag_urn(
name: str,
) -> str:
"""
generate tag urn
:param name: name of the tag
:type name: str
:return: tag urn
:rtype: str
"""
from datahub.emitter.mce_builder import make_tag_urn
return make_tag_urn(name.lower())
[docs]
@staticmethod
def add_tag_to_dataset(
dataset_urn: str,
tag_urn: str,
):
"""
add a tag to a dataset
:param dataset_urn: the dataset urn. It should be a urn generated using the `dataset_urn` method of the `DataHubUtil` class. or a valid urn
:type dataset_urn: str
:param tag_urn: the tag urn. It should be a urn generated using the `tag_urn` method of the `DataHubUtil` class. or a valid urn
:type tag_urn: str
"""
from datahub.metadata.schema_classes import GlobalTagsClass, TagAssociationClass
from datahub.emitter.mcp import MetadataChangeProposalWrapper
graph = DataHubUtil.get_graph()
current_tags = graph.get_aspect(
entity_urn=dataset_urn,
aspect_type=GlobalTagsClass,
)
tag_association_to_add = TagAssociationClass(tag=tag_urn)
need_write = False
if current_tags:
if tag_urn not in [x.tag for x in current_tags.tags]:
# tags exist, but this tag is not present in the current tags
current_tags.tags.append(TagAssociationClass(tag_urn))
need_write = True
else:
# create a brand new tags aspect
current_tags = GlobalTagsClass(tags=[tag_association_to_add])
need_write = True
if need_write:
event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
entityUrn=dataset_urn,
aspect=current_tags,
)
graph.emit(event)
print(f"Tag {tag_urn} added to dataset {dataset_urn}")
else:
print(f"Tag {tag_urn} already exists, omitting write")
[docs]
@staticmethod
def create_tag(
name: str,
description: str,
) -> str:
"""
generates a new tag urn
:param name: name of the tag
:type name: str
:param description: description of the tag
:type description: str
:return: the newly created tag urn
:rtype: str
"""
from datahub.emitter.mcp import MetadataChangeProposalWrapper
# Imports for metadata model classes
from datahub.metadata.schema_classes import TagPropertiesClass
tag_urn = DataHubUtil.tag_urn(name)
tag_properties_aspect = TagPropertiesClass(
name=name,
description=description,
)
event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
entityUrn=tag_urn,
aspect=tag_properties_aspect,
)
# Create rest emitter
rest_emitter = DataHubUtil.get_emitter()
rest_emitter.emit(event)
print(f"Tag {name} created")
return tag_urn
[docs]
@staticmethod
def user_urn(
username: str,
) -> str:
"""
generates and returns user urn
:param username: name of the user. If you pass dag owner it should include the full name like `majid@ready.net`
:type username: str
:return: user urn
:rtype: str
"""
from datahub.emitter.mce_builder import make_user_urn
return make_user_urn(username.lower())
[docs]
@staticmethod
def field_urn(
field_name: str,
dataset_urn: str,
) -> str:
"""
Gets a field urn for a given field name and dataset urn.
:param field_name: field name
:type field_name: str
:param dataset_urn: dataset urn. It should be a urn generated using the `dataset_urn` method of the `DataHubUtil` class. or a valid urn
:type dataset_urn: str
:return: a field urn
:rtype: str
"""
import datahub.emitter.mce_builder as builder
return builder.make_schema_field_urn(dataset_urn, field_name)
[docs]
@staticmethod
def add_link_to_dataset(
dataset_urn: str,
link_to_add: str,
link_description: str,
user_id: str = INGESTION_USER,
):
"""
Adds a link to the institutional memory of a dataset.
:param dataset_urn: It should be a urn generated using the `dataset_urn` method of the `DataHubUtil` class. or a valid urn
:type dataset_urn: str
:param link_to_add: A link url formatted as a string.
:type link_to_add: str
:param link_description: link description
:type link_description: str
:param user_id: the user who generated the link. If you pass dag owner it should include the full name like `majid@ready.net`, defaults to "ingestion"
:type user_id: str, optional
"""
import time
from datahub.emitter.mcp import MetadataChangeProposalWrapper
# Imports for metadata model classes
from datahub.metadata.schema_classes import (
AuditStampClass,
InstitutionalMemoryClass,
InstitutionalMemoryMetadataClass,
)
current_timestamp = DataHubUtil.get_time_stamp(user=user_id)
institutional_memory_element = InstitutionalMemoryMetadataClass(
url=link_to_add,
description=link_description,
createStamp=current_timestamp,
)
graph = DataHubUtil.get_graph()
current_institutional_memory = graph.get_aspect(
entity_urn=dataset_urn,
aspect_type=InstitutionalMemoryClass,
)
need_write = False
if current_institutional_memory:
if link_description not in [
x.description for x in current_institutional_memory.elements
]:
current_institutional_memory.elements.append(
institutional_memory_element
)
need_write = True
else:
# create a brand new institutional memory aspect
current_institutional_memory = InstitutionalMemoryClass(
elements=[institutional_memory_element]
)
need_write = True
if need_write:
event = MetadataChangeProposalWrapper(
entityUrn=dataset_urn,
aspect=current_institutional_memory,
)
graph.emit(event)
print(f"Link {link_to_add} added to dataset {dataset_urn}")
else:
print(f"Link {link_to_add} already exists and is identical, omitting write")
[docs]
@staticmethod
def add_documentation_to_dataset(
dataset_urn: str,
documentation_to_add: str,
user_id=INGESTION_USER,
):
"""Adds a documentation to the institutional memory of a dataset.
:param dataset_urn: the item urn. should be a urn generated using the `dataset_urn` method of the `DataHubUtil` class. or a valid urn
:type dataset_urn: str
:param documentation_to_add: the text. it can be md formatted text
:type documentation_to_add: str
:param user_id: the user who created it. If you pass dag owner it should include the full name like `majid@ready.net`, defaults to "ingestion"
:type user_id: str, optional
"""
import time
from datahub.emitter.mcp import MetadataChangeProposalWrapper
# Imports for metadata model classes
from datahub.metadata.schema_classes import (
AuditStampClass,
EditableDatasetPropertiesClass,
)
actor = DataHubUtil.user_urn(user_id)
# Some helpful variables to fill out objects later
now = int(time.time() * 1000) # milliseconds since epoch
current_timestamp = AuditStampClass(time=now, actor=actor)
graph = DataHubUtil.get_graph()
current_editable_properties = graph.get_aspect(
entity_urn=dataset_urn, aspect_type=EditableDatasetPropertiesClass
)
need_write = False
if current_editable_properties:
if documentation_to_add != current_editable_properties.description:
current_editable_properties.description = documentation_to_add
need_write = True
else:
# create a brand new editable dataset properties aspect
current_editable_properties = EditableDatasetPropertiesClass(
created=current_timestamp, description=documentation_to_add
)
need_write = True
if need_write:
event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
entityUrn=dataset_urn,
aspect=current_editable_properties,
)
graph.emit(event)
print(f"Documentation added to dataset {dataset_urn}")
else:
print("Documentation already exists and is identical, omitting write")
[docs]
@staticmethod
def add_owner_to_dataset(
dataset_urn: str,
owner_urn: str,
ownership_type="TECHNICAL_OWNER",
):
"""Adds an owner to a dataset.
:param dataset_urn: dataset urn. It should be a urn generated using the `dataset_urn` method of the `DataHubUtil` class. or a valid urn
:type dataset_urn: str
:param owner_urn: owner urn. It should be a urn generated using the `user_urn` method of the `DataHubUtil` class. or a valid urn
:type owner_urn: str
:param user_id: the user who created it. If you pass dag owner it should include the full name like
:type user_id: str, optional
:param ownership_type: the type of ownership. It should be a valid ownership type, defaults to OwnershipTypeClass.TECHNICAL_OWNER
:type ownership_type: OwnershipTypeClass, optional
"""
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from typing import Optional
from datahub.metadata.schema_classes import (
OwnerClass,
OwnershipClass,
)
print(f"Adding owner {owner_urn} to dataset {dataset_urn}")
# Some objects to help with conditional pathways later
owner_class_to_add = OwnerClass(
owner=owner_urn,
type=ownership_type,
)
ownership_to_add = OwnershipClass(owners=[owner_class_to_add])
# First we get the current owners
graph = DataHubUtil.get_graph()
current_owners: Optional[OwnershipClass] = graph.get_aspect(
entity_urn=dataset_urn,
aspect_type=OwnershipClass,
)
need_write = False
if current_owners:
if (owner_urn, ownership_type) not in [
(x.owner, x.type) for x in current_owners.owners
]:
# owners exist, but this owner is not present in the current owners
current_owners.owners.append(owner_class_to_add)
need_write = True
else:
# create a brand new ownership aspect
current_owners = ownership_to_add
need_write = True
if need_write:
event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
entityUrn=dataset_urn,
aspect=current_owners,
)
graph.emit(event)
print(
f"Owner {owner_urn}, type {ownership_type} added to dataset {dataset_urn}"
)
else:
print(f"Owner {owner_urn} already exists, omitting write")
[docs]
@staticmethod
def add_properties_to_dataset(
dataset_urn: str,
properties: dict[str, str],
):
"""
Add properties to a dataset
:param properties: the properties to add
:type properties: dict[str,str]
:param dataset_urn: the dataset urn. It should be a urn generated using the `dataset_urn` method of the `DataHubUtil` class. or a valid urn
:type dataset_urn: str
"""
from datahub.specific.dataset import DatasetPatchBuilder
dataset_patch_builder = DatasetPatchBuilder(dataset_urn)
for keys, values in properties.items():
dataset_patch_builder.add_custom_property(keys, values)
emitter = DataHubUtil.get_emitter()
for patch_mcp in dataset_patch_builder.build():
emitter.emit(patch_mcp)
print(f"Added properties to dataset {dataset_urn}")
[docs]
@staticmethod
def add_domain_to_dataset(
dataset_urn: str,
domain_urn: str,
):
"""
Adds a domain to a dataset.
:param dataset_urn: The dataset urn. It should be a valid urn or generated using the `dataset_urn` method of the `DataHubUtil` class.
:type dataset_urn: str
:param domain_urn: Domain urn. It should be a valid urn or generated using the `domain_urn` method of the `DataHubUtil` class.
:type domain_urn: str
"""
from datahub.metadata.schema_classes import DomainsClass, ChangeTypeClass
from datahub.emitter.mcp import MetadataChangeProposalWrapper
# Create the Domains aspect
domains_aspect = DomainsClass(domains=[domain_urn])
# Create the MetadataChangeProposalWrapper event
event = MetadataChangeProposalWrapper(
entityUrn=dataset_urn,
aspect=domains_aspect,
changeType=ChangeTypeClass.UPSERT,
)
DataHubUtil.get_emitter().emit(event)
print(f"Domain {domain_urn} added to dataset {dataset_urn}")
[docs]
@staticmethod
def get_data_product(
data_product_id: str,
domain_urn: str,
):
"""
Get a data product from DataHub.
:param data_product_id: a data product id
:type data_product_id: str
:param domain_urn: the domain urn. It should be a urn generated using the `domain_urn` method of the `DataHubUtil` class. or a valid urn
:type domain_urn: urn
:return: a data product entity
:rtype: _type_
"""
from datahub.api.entities.dataproduct.dataproduct import DataProduct
if not domain_urn:
print("No domain urn provided")
return None
graph = DataHubUtil.get_graph()
data_product_to_query = DataProduct(
id=data_product_id,
domain=domain_urn,
)
return graph.get_entity_raw(data_product_to_query.urn)
[docs]
@staticmethod
def merge_data_product_assets(
data_product_id,
assets,
domain_urn,
):
"""
Merge assets of a data product with existing assets.
:param data_product_id: the id of the data product used to create it. does not need to be a urn
:type data_product_id: str
:param assets: list of the assets. It should be generated using dataset_urn or be a valid urn
:type assets: list[urn]
:param domain_urn: the domain urn. It should be a urn generated using the `domain_urn` method of the `DataHubUtil` class. or a valid urn
:type domain_urn: urn
:return: a list of assets
:rtype: _type_
"""
from datahub.api.entities.dataproduct.dataproduct import DataProduct
if not domain_urn:
print("No domain urn provided")
return assets
existing_data_product = DataHubUtil.get_data_product(
data_product_id, domain_urn
)
if existing_data_product:
existing_assets = (
existing_data_product.get("aspects", {})
.get("dataProductProperties", {})
.get("value", {})
.get("assets", [])
)
for asset in existing_assets:
if asset["destinationUrn"] not in assets:
assets.append(asset["destinationUrn"])
return assets
[docs]
@staticmethod
def add_assets_to_data_product(
data_product_id,
assets=list[str],
domain_urn: str | None = None,
):
"""
It updates a data product and adds assets to it.
:param data_product_id: the id of the data product used to create it. does not need to be a urn
:type data_product_id: _type_
:param assets: list of the assets. It should be generated using dataset_urn or be a valid urn, defaults to list[str]
:type assets: list[str], optional
:param domain_urn: the domain urn. It should be a urn generated using the `domain_urn` method of the `DataHubUtil` class. or a valid urn, defaults to None
:type domain_urn: str, optional
"""
from datahub.api.entities.dataproduct.dataproduct import DataProduct
if not len(assets):
print("No assets to add to data product")
return
if not domain_urn:
domain_urn = DataHubUtil.domain_urn()
graph = DataHubUtil.get_graph()
merged_assets = DataHubUtil.merge_data_product_assets(
data_product_id, assets, domain_urn
)
data_product_to_emit = DataProduct(
id=data_product_id,
assets=merged_assets,
domain=domain_urn,
)
for mcp in data_product_to_emit.generate_mcp(upsert=True):
graph.emit(mcp)
[docs]
@staticmethod
def dataset_exists(
dataset_urn: str,
) -> bool:
"""
Check if a dataset exists in the datahub
:param dataset_urn: The urn of the dataset
:type dataset_urn: str
:return: True if the dataset exists, False otherwise
:rtype: bool
"""
import requests
access_token = DataHubUtil._get_datahub_token()
return (
requests.head(
f"https://uxie.ready-internal.net/openapi/v2/entity/dataset/{dataset_urn}",
headers={"Authorization": f"Bearer {access_token}"},
).status_code
== 204
)
@staticmethod
def _add_tile_server_platform():
"""
Adds the TileServer GL platform to DataHub. You probably only need to run this once.
"""
DataHubUtil.add_new_platform(
"TileServer GL",
TILESERVER_PLATFORM,
"https://www.klokantech.com/img/projects/tileserver-gl-logo.png",
)
[docs]
@staticmethod
def get_schema_from_tilejson(
tile_json: dict,
):
"""
Get schema fields from a tile json.
:param tile_json: the tile json
:type tile_json: dict
:return: the schema fields
:rtype: list[SchemaFieldClass]
"""
from datahub.metadata.schema_classes import (
SchemaFieldClass,
SchemaFieldDataTypeClass,
StringTypeClass,
)
current_timestamp = DataHubUtil.get_time_stamp()
fields = []
for layer in tile_json.get("tilestats", {}).get("layers", []):
layer_id = layer.get("layer", "unknown")
for attribute in layer.get("attributes", []):
fields.append(
SchemaFieldClass(
fieldPath=f"{layer_id}@{attribute['attribute']}",
type=SchemaFieldDataTypeClass(type=StringTypeClass()),
nativeDataType=attribute["type"],
description=f"min:{attribute.get('min','')} , max:{attribute.get('max','')}, count:{attribute['count']}",
lastModified=current_timestamp,
)
)
return fields
[docs]
@staticmethod
def get_properties_from_tilejson(
tile_json: dict,
base_url: str,
) -> dict:
"""
Get properties from a tile json.
:param tile_json: the tile json
:type tile_json: dict
:param base_url: the base url of the tile server
:type base_url: str
:return: the properties
:rtype: dict
"""
if not len(tile_json.keys()):
return {}
tile_json_properties = {
item: str(tile_json[item])
for item in tile_json.keys()
if item not in ["tilestats", "vector_layers"]
}
tile_json_properties["tileJSON"] = f"{base_url}/data/{tile_json['id']}.json"
inspect_url = (
f"{base_url}/data/{tile_json['id']}/#{tile_json['center'][2]}/{tile_json['center'][1]}/{tile_json['center'][0]}"
if tile_json.get("center")
else f"{base_url}/data/{tile_json['id']}/"
)
tile_json_properties["inspect_url"] = inspect_url
layers = []
for layer in tile_json.get("tilestats", {}).get("layers", []):
layer_id = layer["layer"]
layers.append(layer_id)
tile_json_properties[f"{layer_id}.layer"] = layer_id
tile_json_properties[f"{layer_id}.count"] = str(layer["count"])
tile_json_properties[f"{layer_id}.geometry"] = str(layer["geometry"])
tile_json_properties[f"{layer_id}.attributeCount"] = str(
layer["attributeCount"]
)
tile_json_properties["layers"] = str(layers)
return tile_json_properties
[docs]
@staticmethod
def get_time_stamp(
user: str = INGESTION_USER,
):
"""
Get a time stamp for the current time.
:param user: the user who generated the time stamp. If you pass dag owner it should include the full name like `
:type user: str, optional
"""
import time
from datahub.metadata.schema_classes import (
AuditStampClass,
)
return AuditStampClass(
time=int(time.time() * 1000),
actor=f"urn:li:corpuser:{user}",
)
[docs]
@staticmethod
def add_tileset_schema_to_datahub(
tile_json,
entity_urn: str,
):
"""
Add tileset schema to DataHub.
:param tile_json: the tileset json
:type tile_json: dict
:param entity_urn: the tileset dataset urn
:type entity_urn: str
"""
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.emitter.mce_builder import make_data_platform_urn
from datahub.metadata.schema_classes import (
OtherSchemaClass,
SchemaMetadataClass,
)
current_timestamp = DataHubUtil.get_time_stamp()
fields = DataHubUtil.get_schema_from_tilejson(tile_json)
schema_event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
entityUrn=entity_urn,
aspect=SchemaMetadataClass(
schemaName="tileserver", # not used
platform=make_data_platform_urn(
TILESERVER_PLATFORM
), # 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__"),
lastModified=current_timestamp,
fields=fields,
),
)
DataHubUtil.get_emitter().emit(schema_event)
[docs]
@staticmethod
def add_tileset_properties_to_datahub(
tile_json,
entity_urn: str,
base_url: str,
):
"""
Add tileset properties to DataHub.
:param tile_json: the tileset json
:type tile_json: dict
:param entity_urn: the tileset dataset urn
:type entity_urn: str
:param base_url: the base url of the tile server
"""
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.metadata.schema_classes import (
DatasetPropertiesClass,
)
tile_json_properties = DataHubUtil.get_properties_from_tilejson(
tile_json, base_url
)
inspect_url = tile_json_properties["inspect_url"]
tile_json_url = tile_json_properties["tileJSON"]
properties_event: MetadataChangeProposalWrapper = MetadataChangeProposalWrapper(
entityUrn=entity_urn,
aspect=DatasetPropertiesClass(
description=tile_json.get("description", ""),
customProperties=tile_json_properties,
tags=["tileset", "pbf", "vector-tiles"],
externalUrl=inspect_url,
),
)
DataHubUtil.get_emitter().emit(properties_event)
DataHubUtil.add_link_to_dataset(
entity_urn, tile_json_url, "TileJSON Url", INGESTION_USER
)
DataHubUtil.add_link_to_dataset(
entity_urn, inspect_url, "Inspect Url", INGESTION_USER
)
print(f"Tile Inspection url: {inspect_url}")
print(f"Tile JSON url: {tile_json_url}")
[docs]
@staticmethod
def add_tile_set_to_datahub(
base_url: str,
tile_id: str,
entity_urn: str,
):
"""
Add a tileset to DataHub.
:param base_url: the base url of the tile server
:type base_url: str
:param tile_id: the tileset id
:type tile_id: str
"""
import requests
tile_json_url = f"{base_url}/data/{tile_id}.json"
print(f"Processing to load {tile_json_url} for datahub")
tile_json = requests.get(tile_json_url, verify=False).json()
DataHubUtil.add_tileset_schema_to_datahub(tile_json, entity_urn)
DataHubUtil.add_tileset_properties_to_datahub(tile_json, entity_urn, base_url)
[docs]
@staticmethod
def get_tileset_assertion_urn(
tile_id: str,
env: DataHubEnvType = "DEV",
):
"""
Get the tileset assertion urn.
:param tile_id: the tileset id
:type tile_id: str
:param env: the environment where the tileset is deployed, defaults to "DEV"
:type env: Literal["DEV", "STAGE", "PROD"], optional
"""
return f"urn:li:assertion:tileset_deployment_{tile_id}_{env}"
[docs]
@staticmethod
def add_tileset_assertions(
entity_urn: str,
assertion_urn: str,
assertion_type: str = "SUCCESS",
):
"""Add assertions to the tileset deployment.
:param entity_urn: the tileset dataset urn
:type entity_urn: str
:param assertion_urn: the assertion_urn id
:type assertion_urn: str
:param assertion_type: the assertion type, defaults to AssertionResultType.SUCCESS
:type assertion_type: AssertionResultType, optional
"""
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.metadata.com.linkedin.pegasus2avro.common import (
DataPlatformInstance,
)
import datahub.emitter.mce_builder as builder
import time
import uuid
from datahub.metadata.com.linkedin.pegasus2avro.assertion import (
AssertionInfo,
AssertionResult,
AssertionRunEvent,
AssertionRunStatus,
AssertionStdAggregation,
AssertionStdOperator,
AssertionType,
DatasetAssertionInfo,
DatasetAssertionScope,
)
assertion_tile_accessible = AssertionInfo(
type=AssertionType.DATASET,
datasetAssertion=DatasetAssertionInfo(
scope=DatasetAssertionScope.UNKNOWN,
operator=AssertionStdOperator._NATIVE_,
nativeType="Tile Deployment",
aggregation=AssertionStdAggregation._NATIVE_,
fields=[],
dataset=entity_urn,
),
customProperties={"suite_name": "demo_suite"},
)
# Construct a MetadataChangeProposalWrapper object.
assertion_tile_accessible_mcp = MetadataChangeProposalWrapper(
entityUrn=assertion_urn,
aspect=assertion_tile_accessible,
)
# Emit Assertion entity info aspect!
DataHubUtil.get_emitter().emit_mcp(assertion_tile_accessible_mcp)
# Construct an assertion platform object.
assertion_dataPlatformInstance = DataPlatformInstance(
platform=builder.make_data_platform_urn("tiling")
)
# Construct a MetadataChangeProposalWrapper object for assertion platform
assertion_dataPlatformInstance_mcp = MetadataChangeProposalWrapper(
entityUrn=assertion_urn,
aspect=assertion_dataPlatformInstance,
)
# Emit Assertion entity platform aspect!
DataHubUtil.get_emitter().emit(assertion_dataPlatformInstance_mcp)
# Construct batch assertion result object for partition 1 batch
assertionResult_maxVal_batch_partition1 = AssertionRunEvent(
timestampMillis=int(time.time() * 1000),
assertionUrn=assertion_urn,
asserteeUrn=entity_urn,
runId=str(uuid.uuid4()),
status=AssertionRunStatus.COMPLETE,
result=AssertionResult(type=assertion_type),
)
DataHubUtil.emit_assertion_result(
assertionResult_maxVal_batch_partition1,
)
[docs]
@staticmethod
def emit_assertion_result(assertionResult) -> None:
"""Emit assertion result to DataHub.
:param assertionResult: _description_
:type assertionResult: AssertionRunEvent
"""
from datahub.emitter.mcp import MetadataChangeProposalWrapper
dataset_assertionRunEvent_mcp = MetadataChangeProposalWrapper(
entityUrn=assertionResult.assertionUrn,
aspect=assertionResult,
)
# Emit BatchAssertion Result! (timeseries aspect)
DataHubUtil.get_emitter().emit_mcp(dataset_assertionRunEvent_mcp)
[docs]
@staticmethod
def create_data_product(
id: str,
display_name: str,
description,
assets: List[str] = [],
owners=[],
terms=[],
tags: Optional[Union[List[str], dict[str, str]]] = [],
properties: dict[str, str] = {},
external_url: str = None,
domain=None,
):
"""
Create a data product in DataHub.
:param id: the id of the data product used to create it. does not need to be a urn and should be unique per data product
:type id: str
:param display_name: the display name of the data product
:type display_name: str
:param domain: the domain of the data product. It should be a urn generated using the `domain_urn` method of the `DataHubUtil` class. or a valid urn. you can use `DataHubUtil.domain_urn()` to generate it.
:type domain: urn
:param description: the data product description
:type description: str
:param assets: an array of assets urns in str format. , defaults to []
:type assets: list, optional
:param owners: Urn of the owners, defaults to []
:type owners: list, optional
:param terms: _description_, defaults to []
:type terms: list, optional
:param tags: an array of tag urns, defaults to []
:type tags: list, optional
:param properties: a dict of properties, defaults to {}
:type properties: dict, optional
:param external_url: _description_, defaults to None
:type external_url: _type_, optional
"""
from datahub.api.entities.dataproduct.dataproduct import DataProduct
if isinstance(tags, dict):
tags = [DataHubUtil.create_tag(tag, tags[tag]) for tag in tags.keys()]
if not domain:
domain = DataHubUtil.domain_urn()
owner_object = []
for owner in owners:
if owner.get("email"):
owner["id"] = DataHubUtil.user_urn(owner["email"])
owner_object.append({"id": owner["id"], "type": owner["type"]})
else:
owner_object.append(owner)
merged_assets = DataHubUtil.merge_data_product_assets(id, assets, domain)
data_product = DataProduct(
id=id,
display_name=display_name,
domain=domain,
description=description,
assets=merged_assets,
owners=owner_object,
terms=terms,
tags=tags,
properties=properties,
external_url=external_url,
)
graph = DataHubUtil.get_graph()
for mcp in data_product.generate_mcp(upsert=True):
graph.emit(mcp)
[docs]
@staticmethod
def get_user_urn_from_dag(
context: Context,
) -> str:
"""Get a user urn from a dag context.
:param context: the username
:type username: str
"""
if len(context["dag"].default_args["email"]) > 0:
owner = context["dag"].default_args["email"][0]
user_urn = DataHubUtil.user_urn(owner)
return user_urn
return f"urn:li:corpuser:{INGESTION_USER}"
[docs]
@staticmethod
def remove_dataset(dataset_urn: str):
"""remove a dataset from DataHub.
:param dataset_urn: a dataset urn. It should be a urn generated using the `dataset_urn` method of the `DataHubUtil` class. or a valid urn
:type dataset_urn: str
"""
graph = DataHubUtil.get_graph()
# Soft-delete the dataset.
graph.delete_entity(urn=dataset_urn, hard=False)
[docs]
@staticmethod
def send_graphql_request(
query: str, variables: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
"""
Send a GraphQL request to DataHub.
:param query: GraphQL query string
:type query: str
:param variables: Variables for the GraphQL query
:type variables: Optional[Dict[str, Any]]
:return: Response from the GraphQL endpoint
:rtype: Dict[str, Any]
"""
import requests
from config import Config
# Get authentication token
access_token = DataHubUtil._get_datahub_token()
# Prepare the GraphQL endpoint
graphql_endpoint = f"{Config.DATAHUB_GMS_ENDPOINT}/api/graphql"
# Prepare headers
headers = {
"Authorization": f"Bearer {access_token}",
"Content-Type": "application/json",
}
# Prepare the request payload
payload = {"query": query, "variables": variables or {}}
try:
response = requests.post(graphql_endpoint, headers=headers, json=payload)
response.raise_for_status()
result = response.json()
if "errors" in result:
print(f"GraphQL errors: {result['errors']}")
return result
except requests.exceptions.RequestException as e:
print(f"Error sending GraphQL request: {e}")
raise
[docs]
@staticmethod
def get_queries_for_dataset(dataset_urn: str) -> List[Dict[str, Any]]:
"""
Get all queries associated with a dataset.
:param dataset_urn: The dataset URN to get queries for
:type dataset_urn: str
:return: List of queries with their details
:rtype: List[Dict[str, Any]]
"""
query = """
query listQueries($input: ListQueriesInput!) {
listQueries(input: $input) {
start
count
total
queries {
urn
properties {
name
description
statement {
value
language
}
}
subjects {
dataset {
urn
}
}
}
}
}
"""
variables = {
"input": {
"datasetUrn": dataset_urn,
"start": 0,
"count": 1000,
}
}
try:
result = DataHubUtil.send_graphql_request(query, variables)
if result and result.get("data", {}).get("listQueries", {}).get("queries"):
queries = []
for query_entity in result["data"]["listQueries"]["queries"]:
props = query_entity.get("properties", {})
if props:
queries.append(
{
"urn": query_entity.get("urn"),
"name": props.get("name"),
"description": props.get("description"),
"statement": props.get("statement", {}).get("value"),
}
)
return queries
else:
return []
except Exception as e:
print(f"Error getting queries for dataset: {e}")
return []
[docs]
@staticmethod
def query_exists_for_dataset(
dataset_urn: str, query_text: str = None, title: str = None
) -> bool:
"""
Check if a query already exists for a dataset.
Can check by either query text (exact match) or title.
:param dataset_urn: The dataset URN to check
:type dataset_urn: str
:param query_text: The SQL query text to check (optional)
:type query_text: str
:param title: The query title to check (optional)
:type title: str
:return: True if query exists, False otherwise
:rtype: bool
"""
existing_queries = DataHubUtil.get_queries_for_dataset(dataset_urn)
if not existing_queries:
return False
for query in existing_queries:
# Check by title if provided
if title and query.get("name") == title:
print(f"Query with title '{title}' already exists for {dataset_urn}")
return True
# Check by query text if provided (normalize whitespace for comparison)
if query_text and query.get("statement"):
# Normalize whitespace for comparison
normalized_existing = " ".join(query.get("statement", "").split())
normalized_new = " ".join(query_text.split())
if normalized_existing == normalized_new:
print(f"Query with same SQL already exists for {dataset_urn}")
return True
return False
[docs]
@staticmethod
def add_query_to_urn(
query_text: str,
dataset_urn: str,
title: str = "Sample Query",
description: str = "",
created_by: str = "ingestion",
check_duplicates: bool = True,
) -> bool:
"""
Add a sample query to a dataset URN using GraphQL.
Now includes duplicate checking by default.
:param query_text: The SQL query text to add
:type query_text: str
:param dataset_urn: The dataset URN to add the query to
:type dataset_urn: str
:param title: Title for the query
:type title: str
:param description: Description of the query
:type description: str
:param created_by: User who created the query
:type created_by: str
:param check_duplicates: Whether to check for duplicates before adding
:type check_duplicates: bool
:return: Success status
:rtype: bool
"""
# Check for duplicates if enabled
if check_duplicates:
if DataHubUtil.query_exists_for_dataset(
dataset_urn, query_text=query_text, title=title
):
print(f"Skipping duplicate query '{title}' for {dataset_urn}")
return True # Return True since the query already exists
# GraphQL mutation to add a query
mutation = """
mutation addQuery($input: CreateQueryInput!) {
createQuery(input: $input){
urn
}
}
"""
# Prepare variables
variables = {
"input": {
"properties": {
"name": title,
"description": description,
"statement": {"value": query_text, "language": "SQL"},
},
"subjects": [{"datasetUrn": dataset_urn}],
}
}
try:
result = DataHubUtil.send_graphql_request(mutation, variables)
if result.get("data", {}).get("createQuery"):
print(f"Successfully added query '{title}' to {dataset_urn}")
return True
else:
print(f"Failed to add query to {dataset_urn}")
return False
except Exception as e:
print(f"Error adding query to URN: {e}")
return False
[docs]
@staticmethod
def delete_query(query_urn: str) -> bool:
"""
Delete a query by URN.
:param query_urn: The URN of the query to delete
:type query_urn: str
:return: Success status
:rtype: bool
"""
mutation = """
mutation deleteQuery($urn: String!) {
deleteQuery(urn: $urn)
}
"""
variables = {"urn": query_urn}
try:
result = DataHubUtil.send_graphql_request(mutation, variables)
if result.get("data", {}).get("deleteQuery"):
print(f"Successfully deleted query: {query_urn}")
return True
else:
print(f"Failed to delete query: {query_urn}")
return False
except Exception as e:
print(f"Error deleting query: {e}")
return False
[docs]
@staticmethod
def set_column_documentation(
dataset_urn: str,
columns: List[Dict[str, Any]],
) -> Dict[str, Any]:
"""
Set documentation for multiple columns of a dataset.
:param dataset_urn: The dataset URN
:type dataset_urn: str
:param columns: List of column definitions, each with:
- name: Column name (required)
- description: Column description (required)
- field_type: Data type (optional, included in description)
- nullable: Whether nullable (optional)
- is_primary_key: Whether PK (optional)
- foreign_key_reference: FK reference like "Table.column" (optional)
- relationship_cardinality: FK cardinality like "1:1", "1:N", "N:M" (optional, only if foreign_key_reference is provided)
:type columns: List[Dict[str, Any]]
:return: Results dict with counts of successes and failures
:rtype: Dict[str, Any]
"""
mutation = """
mutation updateDescription($input: DescriptionUpdateInput!) {
updateDescription(input: $input)
}
"""
results: Dict[str, Any] = {"updated": 0, "failed": 0, "errors": []}
for col in columns:
col_name = col.get("name")
description = col.get("description", "")
metadata_parts = []
if col.get("field_type"):
metadata_parts.append(f"**Type:** `{col['field_type']}`")
if col.get("is_primary_key"):
metadata_parts.append("**Primary Key**")
if col.get("foreign_key_reference"):
fk_text = f"**Foreign Key:** → `{col['foreign_key_reference']}`"
if col.get("relationship_cardinality"):
fk_text += f" ({col['relationship_cardinality']})"
metadata_parts.append(fk_text)
if col.get("nullable"):
metadata_parts.append("*Nullable*")
full_description = description
if metadata_parts:
full_description = " | ".join(metadata_parts) + "\n\n" + description
variables = {
"input": {
"description": full_description,
"resourceUrn": dataset_urn,
"subResource": col_name,
"subResourceType": "DATASET_FIELD",
}
}
try:
result = DataHubUtil.send_graphql_request(mutation, variables)
if result.get("data", {}).get("updateDescription"):
results["updated"] += 1
print(f"Updated description for column: {col_name}")
else:
results["failed"] += 1
results["errors"].append(f"Failed to update {col_name}")
except Exception as e:
results["failed"] += 1
results["errors"].append(f"Error updating {col_name}: {e}")
return results
[docs]
@staticmethod
def search_datasets(
query: str = "*",
platform: Optional[str] = None,
is_production: bool = Config.PRODUCTION,
count: int = 100,
) -> List[Dict[str, Any]]:
"""
Search for datasets in DataHub.
:param query: Search query string (supports wildcards)
:type query: str
:param platform: Filter by platform (e.g., "postgres", "athena")
:type platform: Optional[str]
:param is_production: Whether to search PROD or DEV environment
:type is_production: bool
:param count: Maximum number of results to return
:type count: int
:return: List of datasets with urn, name, and description
:rtype: List[Dict[str, Any]]
"""
from datahub.metadata.schema_classes import FabricTypeClass
env = FabricTypeClass.PROD if is_production else FabricTypeClass.DEV
gql_query = """
query searchDatasets($input: SearchInput!) {
search(input: $input) {
searchResults {
entity {
urn
... on Dataset {
name
editableProperties {
description
}
properties {
description
}
}
}
}
}
}
"""
filters = [{"field": "origin", "values": [env]}]
if platform:
filters.append({"field": "platform", "values": [f"urn:li:dataPlatform:{platform}"]})
variables = {
"input": {
"type": "DATASET",
"query": query,
"start": 0,
"count": count,
"orFilters": [{"and": filters}],
}
}
result = DataHubUtil.send_graphql_request(gql_query, variables)
datasets = []
for item in result.get("data", {}).get("search", {}).get("searchResults", []):
entity = item.get("entity", {})
editable_desc = entity.get("editableProperties", {}) or {}
props_desc = entity.get("properties", {}) or {}
datasets.append({
"urn": entity.get("urn"),
"name": entity.get("name"),
"description": editable_desc.get("description") or props_desc.get("description") or "",
})
return datasets