Skip to content

Commit 7d81c20

Browse files
Aryamanz29Aryamanz29
committed
APP-6912: Added support for creating user-defined relationships
1 parent 212ce28 commit 7d81c20

File tree

4 files changed

+181
-1
lines changed

4 files changed

+181
-1
lines changed

‎pyatlan/model/assets/core/referenceable.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99

1010
from pydantic.v1 import Field, PrivateAttr, root_validator
1111

12+
from pyatlan.model.assets.relations import RelationshipAttributes
1213
from pyatlan.model.core import AtlanObject, AtlanTag, Meaning
1314
from pyatlan.model.custom_metadata import CustomMetadataDict, CustomMetadataProxy
1415
from pyatlan.model.enums import EntityStatus, SaveSemantic
@@ -186,6 +187,11 @@ class Attributes(AtlanObject):
186187
meanings: Optional[List[AtlasGlossaryTerm]] = Field(
187188
default=None, description=""
188189
) # relationship
190+
relationship_attributes: Optional[RelationshipAttributes] = Field(
191+
default=None,
192+
description="Map of relationships for the entity. The specific keys of this map will vary by type, "
193+
"so are described in the sub-types of this schema.",
194+
)
189195

190196
def validate_required(self):
191197
pass
@@ -303,7 +309,7 @@ def validate_required(self):
303309
labels: Optional[List[str]] = Field(
304310
default=None, description="Arbitrary textual labels for the asset."
305311
)
306-
relationship_attributes: Optional[Dict[str, Any]] = Field(
312+
relationship_attributes: Optional[RelationshipAttributes] = Field(
307313
default=None,
308314
description="Map of relationships for the entity. The specific keys of this map will vary by type, "
309315
"so are described in the sub-types of this schema.",

‎pyatlan/model/assets/relations/__init__.py

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
import lazy_loader as lazy
2+
3+
__PYATLAN_ASSET_RELATIONS__ = {
4+
"relationship_attributes": ["RelationshipAttributes"],
5+
"user_def_relationship": ["UserDefRelationship"],
6+
}
7+
8+
lazy_loader = lazy.attach(__name__, submod_attrs=__PYATLAN_ASSET_RELATIONS__)
9+
__getattr__, __dir__, __all__ = lazy_loader

‎pyatlan/model/assets/relations/relationship_attributes.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
from __future__ import annotations
2+
3+
import importlib
4+
from typing import Optional
5+
6+
from pydantic.v1 import Field
7+
8+
from pyatlan.model.core import AtlanObject
9+
10+
11+
class RelationshipAttributes(AtlanObject):
12+
type_name: Optional[str] = Field(
13+
default=None,
14+
description="Name of the relationship type that defines the relationship.",
15+
)
16+
17+
@classmethod
18+
def __get_validators__(cls):
19+
yield cls._parse_relationship_attributes
20+
21+
@classmethod
22+
def _parse_relationship_attributes(cls, data):
23+
if isinstance(data, RelationshipAttributes):
24+
return data
25+
26+
if isinstance(data, list): # Recursively process lists
27+
return [cls._parse_relationship_attributes(item) for item in data]
28+
29+
type_name = (
30+
data.get("type_name") if "type_name" in data else data.get("typeName")
31+
)
32+
relationship_attribute_cls = getattr(
33+
importlib.import_module("pyatlan.model.assets.relations"),
34+
type_name,
35+
RelationshipAttributes,
36+
)
37+
return relationship_attribute_cls(**data)

‎pyatlan/model/assets/relations/user_def_relationship.py

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
from __future__ import annotations
2+
3+
from typing import Any, Optional
4+
5+
from pydantic.v1 import Field, validator
6+
7+
from pyatlan.model.assets import Asset
8+
from pyatlan.model.assets.relations import RelationshipAttributes
9+
from pyatlan.model.core import AtlanObject
10+
from pyatlan.model.enums import SaveSemantic
11+
12+
13+
class UserDefRelationship(RelationshipAttributes):
14+
type_name: str = Field(
15+
allow_mutation=False,
16+
default="UserDefRelationship",
17+
description="Name of the relationship type that defines the relationship.",
18+
)
19+
attributes: UserDefRelationship.Attributes = Field(
20+
default_factory=lambda: UserDefRelationship.Attributes(),
21+
description="Map of attributes in the instance and their values",
22+
)
23+
24+
class Attributes(AtlanObject):
25+
from_type_label: Optional[str] = Field(
26+
default=None,
27+
alias="fromTypeLabel",
28+
description="Name for the relationship when referring from endDef2 asset to endDef1 asset.",
29+
)
30+
to_type_label: Optional[str] = Field(
31+
default=None,
32+
alias="toTypeLabel",
33+
description="Name for the relationship when referring from endDef1 asset to endDef2 asset.",
34+
)
35+
36+
def __init__(__pydantic_self__, **data: Any) -> None:
37+
if "attributes" not in data:
38+
data = {"attributes": data}
39+
super().__init__(**data)
40+
__pydantic_self__.__fields_set__.update(["attributes", "type_name"])
41+
42+
class UserDefRelationshipFrom(Asset):
43+
type_name: str = Field(
44+
default="UserDefRelationship",
45+
description="Name of the relationship type that defines the relationship.",
46+
)
47+
relationship_type: str = Field(
48+
default="UserDefRelationship",
49+
const=True,
50+
description="Fixed typeName for UserDefRelationship.",
51+
)
52+
relationship_attributes: UserDefRelationship = Field(
53+
default=None,
54+
description="Attributes of the UserDefRelationship.",
55+
)
56+
57+
@validator("type_name")
58+
def validate_type_name(cls, v):
59+
return v
60+
61+
def __init__(__pydantic_self__, **data: Any) -> None:
62+
super().__init__(**data)
63+
__pydantic_self__.__fields_set__.update(["type_name", "relationship_type"])
64+
65+
class UserDefRelationshipTo(Asset):
66+
type_name: str = Field(
67+
default="UserDefRelationship",
68+
description="Name of the relationship type that defines the relationship.",
69+
)
70+
relationship_type: str = Field(
71+
default="UserDefRelationship",
72+
description="Fixed typeName for UserDefRelationship.",
73+
)
74+
relationship_attributes: UserDefRelationship = Field(
75+
default=None,
76+
description="Attributes of the UserDefRelationship.",
77+
)
78+
79+
@validator("type_name")
80+
def validate_type_name(cls, v):
81+
return v
82+
83+
def __init__(__pydantic_self__, **data: Any) -> None:
84+
super().__init__(**data)
85+
__pydantic_self__.__fields_set__.update(["type_name", "relationship_type"])
86+
87+
def user_def_relationship_to(
88+
self, related: Asset, semantic: SaveSemantic = SaveSemantic.REPLACE
89+
) -> UserDefRelationship.UserDefRelationshipTo:
90+
if related.guid:
91+
return UserDefRelationship.UserDefRelationshipTo._create_ref(
92+
type_name=related.type_name,
93+
guid=related.guid,
94+
semantic=semantic,
95+
relationship_attributes=self,
96+
)
97+
98+
# If the related asset does not have a GUID, we use qualifiedName
99+
return UserDefRelationship.UserDefRelationshipTo._create_ref(
100+
type_name=related.type_name,
101+
unique_attributes={"qualifiedName": related.qualified_name},
102+
semantic=semantic,
103+
relationship_attributes=self,
104+
)
105+
106+
def user_def_relationship_from(
107+
self, related: Asset, semantic: SaveSemantic = SaveSemantic.REPLACE
108+
) -> UserDefRelationship.UserDefRelationshipFrom:
109+
if related.guid:
110+
return UserDefRelationship.UserDefRelationshipfrom._create_ref(
111+
type_name=related.type_name,
112+
guid=related.guid,
113+
semantic=semantic,
114+
relationship_attributes=self,
115+
)
116+
117+
# If the related asset does not have a GUID, we use qualifiedName
118+
return UserDefRelationship.UserDefRelationshipfrom._create_ref(
119+
type_name=related.type_name,
120+
unique_attributes={"qualifiedName": related.qualified_name},
121+
semantic=semantic,
122+
relationship_attributes=self,
123+
)
124+
125+
126+
UserDefRelationship.UserDefRelationshipTo.update_forward_refs()
127+
UserDefRelationship.UserDefRelationshipFrom.update_forward_refs()
128+
UserDefRelationship.update_forward_refs()

0 commit comments

Comments
 (0)