Merge pull request #239 from atlanhq/DVX-165

DVX-165: Adds relationship append and removal support

pyatlan/generator/templates/imports.jinja2

@@ -49,6 +49,7 @@ from pyatlan.model.enums import (
     KafkaTopicCompressionType,
     MatillionJobType,
     OpenLineageRunState,
+    SaveSemantic,
     PersonaDomainAction,
     PersonaGlossaryAction,
     PersonaMetadataAction,

pyatlan/generator/templates/methods/asset/asset.jinja2

@@ -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

pyatlan/generator/templates/referenceable_attributes.jinja2

@@ -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."
+        ),
+    )

pyatlan/model/assets/asset00.py

@@ -33,6 +33,7 @@
     IconType,
     MatillionJobType,
     OpenLineageRunState,
+    SaveSemantic,
     SchemaRegistrySchemaCompatibility,
     SchemaRegistrySchemaType,
     SourceCostUnitType,
@@ -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"""
@@ -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

pyatlan/model/core.py

@@ -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
 
 
@@ -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

pyatlan/model/enums.py

@@ -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  **
 # **************************************