Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

DVX-165: Adds relationship append and removal support #239

Merged
merged 5 commits into from
Feb 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions pyatlan/generator/templates/imports.jinja2
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ from pyatlan.model.enums import (
KafkaTopicCompressionType,
MatillionJobType,
OpenLineageRunState,
SaveSemantic,
PersonaDomainAction,
PersonaGlossaryAction,
PersonaMetadataAction,
Expand Down
12 changes: 10 additions & 2 deletions pyatlan/generator/templates/methods/asset/asset.jinja2
Original file line number Diff line number Diff line change
Expand Up @@ -39,17 +39,25 @@
return cls(attributes=cls.Attributes(qualified_name=qualified_name, name=name))

@classmethod
def ref_by_guid(cls: type[SelfAsset], guid: str) -> SelfAsset:
def ref_by_guid(
cls: type[SelfAsset], guid: str, semantic: SaveSemantic = SaveSemantic.REPLACE
) -> SelfAsset:
retval: SelfAsset = cls(attributes=cls.Attributes())
retval.guid = guid
retval.semantic = semantic
return retval

@classmethod
def ref_by_qualified_name(cls: type[SelfAsset], qualified_name: str) -> SelfAsset:
def ref_by_qualified_name(
cls: type[SelfAsset],
qualified_name: str,
semantic: SaveSemantic = SaveSemantic.REPLACE,
) -> SelfAsset:
ret_value: SelfAsset = cls(
attributes=cls.Attributes(name="", qualified_name=qualified_name)
)
ret_value.unique_attributes = {"qualifiedName": qualified_name}
ret_value.semantic = semantic
return ret_value

@classmethod
Expand Down
18 changes: 18 additions & 0 deletions pyatlan/generator/templates/referenceable_attributes.jinja2
Original file line number Diff line number Diff line change
Expand Up @@ -176,3 +176,21 @@
pending_tasks: Optional[list[str]] = Field(None)

unique_attributes: Optional[dict[str, Any]] = Field(None)

append_relationship_attributes: Optional[dict[str, Any]] = Field(
None,
alias="appendRelationshipAttributes",
description="Map of append relationship attributes.",
)
remove_relationship_attributes: Optional[dict[str, Any]] = Field(
None,
alias="removeRelationshipAttributes",
description="Map of remove relationship attributes.",
)
semantic: Optional[SaveSemantic] = Field(
exclude=True,
description=(
"Semantic for how this relationship should be saved, "
"if used in an asset request on which `.save()` is called."
),
)
31 changes: 29 additions & 2 deletions pyatlan/model/assets/asset00.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
IconType,
MatillionJobType,
OpenLineageRunState,
SaveSemantic,
SchemaRegistrySchemaCompatibility,
SchemaRegistrySchemaType,
SourceCostUnitType,
Expand Down Expand Up @@ -336,6 +337,24 @@ def validate_required(self):

unique_attributes: Optional[dict[str, Any]] = Field(None)

append_relationship_attributes: Optional[dict[str, Any]] = Field(
None,
alias="appendRelationshipAttributes",
description="Map of append relationship attributes.",
)
remove_relationship_attributes: Optional[dict[str, Any]] = Field(
None,
alias="removeRelationshipAttributes",
description="Map of remove relationship attributes.",
)
semantic: Optional[SaveSemantic] = Field(
exclude=True,
description=(
"Semantic for how this relationship should be saved, "
"if used in an asset request on which `.save()` is called."
),
)


class Asset(Referenceable):
"""Description"""
Expand Down Expand Up @@ -386,17 +405,25 @@ def create_for_modification(
return cls(attributes=cls.Attributes(qualified_name=qualified_name, name=name))

@classmethod
def ref_by_guid(cls: type[SelfAsset], guid: str) -> SelfAsset:
def ref_by_guid(
cls: type[SelfAsset], guid: str, semantic: SaveSemantic = SaveSemantic.REPLACE
) -> SelfAsset:
retval: SelfAsset = cls(attributes=cls.Attributes())
retval.guid = guid
retval.semantic = semantic
return retval

@classmethod
def ref_by_qualified_name(cls: type[SelfAsset], qualified_name: str) -> SelfAsset:
def ref_by_qualified_name(
cls: type[SelfAsset],
qualified_name: str,
semantic: SaveSemantic = SaveSemantic.REPLACE,
) -> SelfAsset:
ret_value: SelfAsset = cls(
attributes=cls.Attributes(name="", qualified_name=qualified_name)
)
ret_value.unique_attributes = {"qualifiedName": qualified_name}
ret_value.semantic = semantic
return ret_value

@classmethod
Expand Down
85 changes: 80 additions & 5 deletions pyatlan/model/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
from pydantic.generics import GenericModel

from pyatlan.model.constants import DELETED_, DELETED_SENTINEL
from pyatlan.model.enums import AnnouncementType, EntityStatus
from pyatlan.model.enums import AnnouncementType, EntityStatus, SaveSemantic
from pyatlan.model.structs import SourceTagAttachment


Expand Down Expand Up @@ -233,9 +233,84 @@ class BulkRequest(AtlanObject, GenericModel, Generic[T]):
entities: list[T]

@validator("entities", each_item=True)
def flush_custom_metadata(cls, v):
def process_attributes_and_flush_cm(cls, asset):
from pyatlan.model.assets import Asset

if isinstance(v, Asset):
v.flush_custom_metadata()
return v
if not isinstance(asset, Asset):
return asset

# Manually need to set these to "None" so that we can exclude
# them from the request playload when they're not set by the user
asset.remove_relationship_attributes = None
asset.append_relationship_attributes = None
for attribute in asset.attributes:
asset = cls.process_relationship_attributes(asset, attribute)

# Flush custom metadata
asset.flush_custom_metadata()
return asset.__class__(
**asset.dict(by_alias=True, exclude_unset=True, exclude_none=True)
)

@classmethod
def process_relationship_attributes(cls, asset, attribute):
from pyatlan.model.assets import Asset

append_attributes = []
remove_attributes = []
replace_attributes = []

attribute_name, attribute_value = attribute[0], getattr(
asset, attribute[0], None
)

# Process list of relationship attributes
if attribute_value and isinstance(attribute_value, list):
for value in attribute_value:
if value and isinstance(value, Asset):
if value.semantic == SaveSemantic.REMOVE:
remove_attributes.append(value)
elif value.semantic == SaveSemantic.APPEND:
append_attributes.append(value)
else:
replace_attributes.append(value)

# Update asset based on processed relationship attributes
if remove_attributes:
asset.remove_relationship_attributes = {
to_camel_case(attribute_name): remove_attributes
}
if append_attributes:
asset.append_relationship_attributes = {
to_camel_case(attribute_name): append_attributes
}
if replace_attributes:
setattr(asset, attribute_name, replace_attributes)

# If only remove or append attributes are present without any replace attributes,
# set the attribute to `None` to exclude it from the bulk request payload
# This avoids including unwanted replace attributes that could alter the request behavior
if (remove_attributes or append_attributes) and not replace_attributes:
setattr(asset, attribute_name, None)

# Process single relationship attribute
elif attribute_value and isinstance(attribute_value, Asset):
if attribute_value.semantic == SaveSemantic.REMOVE:
# Set the replace attribute to "None" so that we exclude it
# from the request payload's "attributes" property
# We only want to pass this attribute under
# "remove_relationship_attributes," not both
setattr(asset, attribute_name, None)
asset.remove_relationship_attributes = {
to_camel_case(attribute_name): attribute_value
}
elif attribute_value.semantic == SaveSemantic.APPEND:
# Set the replace attribute to "None" so that we exclude it
# from the request payload's "attributes" property
# We only want to pass this attribute under
# "append_relationship_attributes," not both
setattr(asset, attribute_name, None)
asset.append_relationship_attributes = {
to_camel_case(attribute_name): attribute_value
}
return asset
6 changes: 6 additions & 0 deletions pyatlan/model/enums.py
Original file line number Diff line number Diff line change
Expand Up @@ -1929,6 +1929,12 @@ class QueryStatus(str, Enum):
ERROR = "error"


class SaveSemantic(Enum):
REPLACE = "REPLACE"
APPEND = "APPEND"
REMOVE = "REMOVE"


# **************************************
# CODE BELOW IS GENERATED NOT MODIFY **
# **************************************
Expand Down
Loading
Loading