Skip to content

Commit

Permalink
Merge pull request #61 from atlanhq/ACTIV-595
Browse files Browse the repository at this point in the history 
Refactor usage of custom metadata
  • Loading branch information
@cmgrote
cmgrote authored Jun 5, 2023
2 parents 672a812 + 355a1bd commit f4b15e1
Show file tree
Hide file tree
Showing 13 changed files with 636 additions and 623 deletions.
76 changes: 16 additions & 60 deletions pyatlan/cache/custom_metadata_cache.py
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,8 @@
# SPDX-License-Identifier: Apache-2.0
# Copyright 2022 Atlan Pte. Ltd.
import json
from typing import Any, Optional
from typing import Optional

from pyatlan.client.atlan import AtlanClient
from pyatlan.error import InvalidRequestError, LogicError, NotFoundError
from pyatlan.model.core import CustomMetadata
from pyatlan.model.enums import AtlanTypeCategory
from pyatlan.model.typedef import AttributeDef, CustomMetadataDef

Expand Down Expand Up @@ -35,7 +32,7 @@ class CustomMetadataCache:

@classmethod
def refresh_cache(cls) -> None:
from pyatlan.model.core import CustomMetadata, to_snake_case
from pyatlan.client.atlan import AtlanClient

client = AtlanClient.get_default_client()
if client is None:
Expand All @@ -56,39 +53,21 @@ def refresh_cache(cls) -> None:
cls.map_name_to_id[type_name] = type_id
cls.map_attr_id_to_name[type_id] = {}
cls.map_attr_name_to_id[type_id] = {}
meta_name = cm.display_name.replace(" ", "")
attribute_class_name = f"Attributes_{meta_name}"
attrib_type = type(attribute_class_name, (CustomMetadata,), {})
attrib_type._meta_data_type_id = type_id # type: ignore
attrib_type._meta_data_type_name = type_name # type: ignore
cls.map_id_to_type[type_id] = attrib_type
applicable_types: set[str] = set()
if cm.attribute_defs:
for attr in cm.attribute_defs:
if attr.options and attr.options.custom_applicable_entity_types:
applicable_types.update(
json.loads(attr.options.custom_applicable_entity_types)
)
attr_id = str(attr.name)
attr_name = str(attr.display_name)
# Use a renamed attribute everywhere
attr_renamed = to_snake_case(attr_name)
cls.map_attr_id_to_name[type_id][attr_id] = attr_renamed
cls.map_attr_id_to_name[type_id][attr_id] = attr_name
if attr.options and attr.options.is_archived:
cls.archived_attr_ids[attr_id] = attr_renamed
elif attr_renamed in cls.map_attr_name_to_id[type_id]:
cls.archived_attr_ids[attr_id] = attr_name
elif attr_name in cls.map_attr_name_to_id[type_id]:
raise LogicError(
f"Multiple custom attributes with exactly the same name ({attr_renamed}) "
f"Multiple custom attributes with exactly the same name ({attr_name}) "
f"found for: {type_name}",
code="ATLAN-PYTHON-500-100",
)
else:
setattr(attrib_type, attr_renamed, Synonym(attr_id))
cls.map_attr_name_to_id[type_id][attr_renamed] = attr_id
for asset_type in applicable_types:
if asset_type not in cls.types_by_asset:
cls.types_by_asset[asset_type] = set()
cls.types_by_asset[asset_type].add(attrib_type)
cls.map_attr_name_to_id[type_id][attr_name] = attr_id

@classmethod
def get_id_for_name(cls, name: str) -> str:
Expand Down Expand Up @@ -202,18 +181,22 @@ def get_attr_id_for_name(cls, set_name: str, attr_name: str) -> str:
)

@classmethod
def get_attr_name_for_id(cls, set_id: str, attr_id: str) -> Optional[str]:
def get_attr_name_for_id(cls, set_id: str, attr_id: str) -> str:
"""
Translate the provided human-readable custom metadata set and attribute names to the Atlan-internal ID string
for the attribute.
Given the Atlan-internal ID stringfor the set and the Atlan-internal ID for the attribute return the
human-readable custom metadata name for the attribute.
"""
if sub_map := cls.map_attr_id_to_name.get(set_id):
if attr_name := sub_map.get(attr_id):
return attr_name
cls.refresh_cache()
if sub_map := cls.map_attr_id_to_name.get(set_id):
return sub_map.get(attr_id)
return None
if attr_name := sub_map.get(attr_id):
return attr_name
raise NotFoundError(
message=f"Custom metadata property with ID {attr_id} does not exist in the custom metadata {set_id}.",
code="ATLAN-PYTHON-404-009",
)

@classmethod
def _get_attributes_for_search_results(cls, set_id: str) -> Optional[list[str]]:
Expand All @@ -234,33 +217,6 @@ def get_attributes_for_search_results(cls, set_name: str) -> Optional[list[str]]
return cls._get_attributes_for_search_results(set_id)
return None

@classmethod
def get_custom_metadata(
cls,
name: str,
asset_type: type,
business_attributes: Optional[dict[str, Any]] = None,
) -> CustomMetadata:
type_name = asset_type.__name__
ba_id = cls.get_id_for_name(name)
if ba_id is None:
raise ValueError(f"No custom metadata with the name: {name} exist")
for a_type in CustomMetadataCache.types_by_asset[type_name]:
if (
hasattr(a_type, "_meta_data_type_name")
and a_type._meta_data_type_name == name
):
break
else:
raise ValueError(f"Custom metadata {name} is not applicable to {type_name}")
if ba_type := CustomMetadataCache.get_type_for_id(ba_id):
return (
ba_type(business_attributes[ba_id])
if business_attributes and ba_id in business_attributes
else ba_type()
)
raise ValueError(f"Custom metadata {name} is not applicable to {type_name}")

@classmethod
def get_custom_metadata_def(cls, name: str) -> CustomMetadataDef:
"""
Expand Down
75 changes: 35 additions & 40 deletions pyatlan/client/atlan.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,9 +76,8 @@
Classification,
ClassificationName,
Classifications,
CustomMetadata,
CustomMetadataReqest,
)
from pyatlan.model.custom_metadata import CustomMetadataDict, CustomMetadataRequest
from pyatlan.model.enums import (
AtlanConnectorType,
AtlanDeleteType,
Expand Down Expand Up @@ -1015,60 +1014,56 @@ def remove_announcement(
return self._update_asset_by_attribute(asset, asset_type, qualified_name)

def update_custom_metadata_attributes(
self, guid: str, custom_metadata: CustomMetadata
self, guid: str, custom_metadata: CustomMetadataDict
):
custom_metadata_request = CustomMetadataReqest(__root__=custom_metadata)
custom_metadata_request = CustomMetadataRequest.create(
custom_metadata_dict=custom_metadata
)
self._call_api(
ADD_BUSINESS_ATTRIBUTE_BY_ID.format_path(
{"entity_guid": guid, "bm_id": custom_metadata._meta_data_type_id}
{
"entity_guid": guid,
"bm_id": custom_metadata_request.custom_metadata_set_id,
}
),
None,
custom_metadata_request,
)

def replace_custom_metadata(self, guid: str, custom_metadata: CustomMetadata):
from pyatlan.cache.custom_metadata_cache import CustomMetadataCache

# Iterate through the custom metadata provided and explicitly set every
# single attribute, so that they are all serialized out (forcing removal
# of any empty ones)
for k, v in custom_metadata.items():
# Need to translate the hashed-string key here back to an attribute name
attr_name = str(
CustomMetadataCache.get_attr_name_for_id(
set_id=custom_metadata._meta_data_type_id, attr_id=k
)
)
if not v:
setattr(custom_metadata, attr_name, None)
else:
setattr(custom_metadata, attr_name, v)
custom_metadata_request = CustomMetadataReqest(__root__=custom_metadata)
def replace_custom_metadata(self, guid: str, custom_metadata: CustomMetadataDict):
# clear unset attributes so that they are removed
custom_metadata.clear_unset()
custom_metadata_request = CustomMetadataRequest.create(
custom_metadata_dict=custom_metadata
)
self._call_api(
ADD_BUSINESS_ATTRIBUTE_BY_ID.format_path(
{"entity_guid": guid, "bm_id": custom_metadata._meta_data_type_id}
{
"entity_guid": guid,
"bm_id": custom_metadata_request.custom_metadata_set_id,
}
),
None,
custom_metadata_request,
)

def remove_custom_metadata(self, guid: str, cm_name: str):
from pyatlan.cache.custom_metadata_cache import CustomMetadataCache

# Ensure the custom metadata exists first - let this throw an error if not
if cm_id := CustomMetadataCache.get_id_for_name(cm_name):
# Initialize a dict of empty attributes for the custom metadata, and then
# send that so that they are removed accordingly
if cm_type := CustomMetadataCache.get_type_for_id(cm_id):
custom_metadata = cm_type()
custom_metadata_request = CustomMetadataReqest(__root__=custom_metadata)
self._call_api(
ADD_BUSINESS_ATTRIBUTE_BY_ID.format_path(
{"entity_guid": guid, "bm_id": cm_id}
),
None,
custom_metadata_request,
)
custom_metadata = CustomMetadataDict(name=cm_name)
# invoke clear_all so all attributes are set to None and consequently removed
custom_metadata.clear_all()
custom_metadata_request = CustomMetadataRequest.create(
custom_metadata_dict=custom_metadata
)
self._call_api(
ADD_BUSINESS_ATTRIBUTE_BY_ID.format_path(
{
"entity_guid": guid,
"bm_id": custom_metadata_request.custom_metadata_set_id,
}
),
None,
custom_metadata_request,
)

@validate_arguments()
def append_terms(
Expand Down
69 changes: 18 additions & 51 deletions pyatlan/generator/templates/entity.jinja2
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,15 +10,10 @@ from io import StringIO
from typing import Any, ClassVar, Dict, List, Optional, TypeVar
from urllib.parse import quote, unquote

from pydantic import Field, StrictStr, root_validator, validator
from pydantic import Field, PrivateAttr, StrictStr, root_validator, validator

from pyatlan.model.core import (
Announcement,
AtlanObject,
Classification,
CustomMetadata,
Meaning,
)
from pyatlan.model.core import Announcement, AtlanObject, Classification, Meaning
from pyatlan.model.custom_metadata import CustomMetadataDict, CustomMetadataProxy
from pyatlan.model.enums import (
ADLSAccessTier,
ADLSAccountStatus,
Expand Down Expand Up @@ -90,9 +85,16 @@ SelfAsset = TypeVar("SelfAsset", bound="Asset")
class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes %}, type_name='{{ entity_def.name }}'{% endif %}):
"""Description"""
{% if entity_def.name == "Referenceable" %}
def __init__(__pydantic_self__, **data:Any)->None:
def __init__(__pydantic_self__, **data: Any) -> None:
super().__init__(**data)
__pydantic_self__.__fields_set__.update(["attributes", "type_name"])
__pydantic_self__._metadata_proxy = CustomMetadataProxy(
__pydantic_self__.business_attributes
)

def json(self, *args, **kwargs) -> str:
self.business_attributes = self._metadata_proxy.business_attributes
return super().json(**kwargs)
{% endif %}
def __setattr__(self, name, value):
if name in {{ entity_def.name }}._convience_properties:
Expand All @@ -117,6 +119,7 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
def validate_required(self):
pass

_metadata_proxy: CustomMetadataProxy = PrivateAttr()
attributes: '{{entity_def.name}}.Attributes' = Field(
default_factory = lambda : {{entity_def.name}}.Attributes(),
description='Map of attributes in the instance and their values. The specific keys of this map will vary '
Expand Down Expand Up @@ -219,51 +222,15 @@ class {{ entity_def.name }}({{super_classes[0]}} {%- if "Asset" in super_classes
if not self.create_time or self.created_by:
self.attributes.validate_required()

def get_custom_metadata(self, name: str) -> CustomMetadata:
from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
def get_custom_metadata(self, name: str) -> CustomMetadataDict:
return self._metadata_proxy.get_custom_metadata(name=name)

ba_id = CustomMetadataCache.get_id_for_name(name)
if ba_id is None:
raise ValueError(f"No custom metadata with the name: {name} exist")
for a_type in CustomMetadataCache.types_by_asset[self.type_name]:
if (
hasattr(a_type, "_meta_data_type_name")
and a_type._meta_data_type_name == name
):
break
else:
raise ValueError(
f"Custom metadata attributes {name} are not applicable to {self.type_name}"
)
if ba_type := CustomMetadataCache.get_type_for_id(ba_id):
return (
ba_type(self.business_attributes[ba_id])
if self.business_attributes and ba_id in self.business_attributes
else ba_type()
)
else:
raise ValueError(
f"Custom metadata attributes {name} are not applicable to {self.type_name}"
)
def set_custom_metadata(self, custom_metadata: CustomMetadataDict):
return self._metadata_proxy.set_custom_metadata(custom_metadata=custom_metadata)

def set_custom_metadata(self, custom_metadata: CustomMetadata) -> None:
from pyatlan.cache.custom_metadata_cache import CustomMetadataCache
def flush_custom_metadata(self):
self.business_attributes = self._metadata_proxy.business_attributes

if not isinstance(custom_metadata, CustomMetadata):
raise ValueError(
"business_attributes must be an instance of CustomMetadata"
)
if (
type(custom_metadata)
not in CustomMetadataCache.types_by_asset[self.type_name]
):
raise ValueError(
f"Business attributes {custom_metadata._meta_data_type_name} are not applicable to {self.type_name}"
)
ba_dict = dict(custom_metadata)
if not self.business_attributes:
self.business_attributes = {}
self.business_attributes[custom_metadata._meta_data_type_id] = ba_dict

{%- else %}
{%- if entity_def.name == "Asset" %}
Expand Down
2 changes: 1 addition & 1 deletion pyatlan/generator/templates/structs.jinja2
View file
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,6 @@ class {{struct.name}}(AtlanObject):
{% endif %}
{%- for attribute_def in struct.attribute_defs %}
{%- set type = attribute_def.type_name | get_type %}
{{attribute_def.name | to_snake_case }}: {% if attribute_def.is_optional %}Optional[{% endif %}{{type}}{% if attribute_def.is_optional %}]{% endif %} = Field(None, description='' , alias='{{attribute_def.name}}')
{{attribute_def.name | to_snake_case }}: {% if attribute_def.is_optional %}Optional[{% endif %}'{{type}}'{% if attribute_def.is_optional %}]{% endif %} = Field(None, description='' , alias='{{attribute_def.name}}')
{%- endfor %}
{% endfor %}
Loading

0 comments on commit f4b15e1

Please sign in to comment.