Skip to content

Commit bf78dc2

Browse files
vaibhavatlanAryamanz29
authored andcommitted
Added generator files
1 parent fe70c10 commit bf78dc2

File tree

6 files changed

+226
-0
lines changed

6 files changed

+226
-0
lines changed
Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
2+
@overload
3+
@classmethod
4+
def creator(
5+
cls,
6+
*,
7+
name: str,
8+
ai_application_version: str,
9+
ai_application_development_stage: AIApplicationDevelopmentStage,
10+
) -> AIApplication: ...
11+
12+
@overload
13+
@classmethod
14+
def creator(
15+
cls,
16+
*,
17+
name: str,
18+
ai_application_version: str,
19+
ai_application_development_stage: AIApplicationDevelopmentStage,
20+
owner_groups: Set[str],
21+
owner_users: Set[str],
22+
models: List[AIModel],
23+
) -> AIApplication: ...
24+
25+
@classmethod
26+
@init_guid
27+
def creator(
28+
cls,
29+
*,
30+
name: str,
31+
ai_application_version: str,
32+
ai_application_development_stage: AIApplicationDevelopmentStage,
33+
owner_groups: Optional[Set[str]] = set(),
34+
owner_users: Optional[Set[str]] = set(),
35+
models: Optional[List[AIModel]] = [],
36+
) -> AIApplication:
37+
validate_required_fields(
38+
["name", "ai_application_version", "ai_application_development_stage"],
39+
[name, ai_application_version, ai_application_development_stage],
40+
)
41+
attributes = AIApplication.Attributes.creator(
42+
name=name,
43+
ai_application_version=ai_application_version,
44+
ai_application_development_stage=ai_application_development_stage,
45+
owner_groups=owner_groups,
46+
owner_users=owner_users,
47+
models=models,
48+
)
49+
return cls(attributes=attributes)
Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
2+
@overload
3+
@classmethod
4+
def creator(
5+
cls,
6+
*,
7+
name: str,
8+
ai_model_status: AIModelStatus,
9+
) -> AIModel: ...
10+
11+
@overload
12+
@classmethod
13+
def creator(
14+
cls,
15+
*,
16+
name: str,
17+
ai_model_status: AIModelStatus,
18+
owner_groups: Set[str],
19+
owner_users: Set[str],
20+
ai_model_version: str,
21+
) -> AIModel: ...
22+
23+
@classmethod
24+
@init_guid
25+
def creator(
26+
cls,
27+
*,
28+
name: str,
29+
ai_model_status: AIModelStatus,
30+
owner_groups: Optional[Set[str]] = set(),
31+
owner_users: Optional[Set[str]] = set(),
32+
ai_model_version: Optional[str] = None,
33+
) -> AIModel:
34+
validate_required_fields(
35+
["name", "ai_model_status"],
36+
[name, ai_model_status],
37+
)
38+
attributes = AIModel.Attributes.creator(
39+
name=name,
40+
ai_model_status=ai_model_status,
41+
owner_groups=owner_groups,
42+
owner_users=owner_users,
43+
ai_model_version=ai_model_version,
44+
)
45+
return cls(attributes=attributes)
46+
47+
@classmethod
48+
def processes_creator(
49+
cls,
50+
ai_model: AIModel,
51+
dataset_dict: Dict[AIDatasetType, list],
52+
) -> List[Process]:
53+
"""
54+
Creates a list of Process objects representing the relationships between an AI model and its datasets.
55+
56+
:param ai_model: the AI model for which to create processes
57+
:param dataset_dict: dictionary mapping AI dataset types to lists of assets
58+
:returns: list of Process objects representing the AI model's data lineage
59+
:raises ValueError: when the AI model is missing required attributes (guid or name)
60+
"""
61+
if not ai_model.guid or not ai_model.name:
62+
raise ValueError("AI model must have both guid and name attributes")
63+
process_list = []
64+
for key, value_list in dataset_dict.items():
65+
for value in value_list:
66+
asset_type = Asset._convert_to_real_type_(value)
67+
if key == AIDatasetType.OUTPUT:
68+
process_name = f"{ai_model.name} -> {value.name}"
69+
process_created = Process.creator(
70+
name=process_name,
71+
connection_qualified_name="default/ai/dataset",
72+
inputs=[AIModel.ref_by_guid(guid=ai_model.guid)],
73+
outputs=[asset_type.ref_by_guid(guid=value.guid)], # type: ignore
74+
extra_hash_params={key.value},
75+
)
76+
process_created.ai_dataset_type = key
77+
else:
78+
process_name = f"{value.name} -> {ai_model.name}"
79+
process_created = Process.creator(
80+
name=process_name,
81+
connection_qualified_name="default/ai/dataset",
82+
inputs=[asset_type.ref_by_guid(guid=value.guid)], # type: ignore
83+
outputs=[AIModel.ref_by_guid(guid=ai_model.guid)],
84+
extra_hash_params={key.value},
85+
)
86+
process_created.ai_dataset_type = key
87+
process_list.append(process_created)
88+
89+
return process_list
90+
91+
@classmethod
92+
def processes_batch_save(cls, client, process_list: List[Process]) -> List:
93+
"""
94+
Saves a list of Process objects to Atlan in batches to optimize performance.
95+
We save the processes in batches of 20.
96+
97+
:param client: Atlan client instance for making API calls
98+
:param process_list: list of Process objects to save
99+
:returns: list of API responses from each batch save operation
100+
"""
101+
batch_size = 20
102+
total_processes = len(process_list)
103+
responses = []
104+
105+
for i in range(0, total_processes, batch_size):
106+
batch = process_list[i : i + batch_size]
107+
response = client.asset.save(batch)
108+
responses.append(response)
109+
110+
return responses

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
outputs: List["Catalog"],
1010
process_id: Optional[str] = None,
1111
parent: Optional[Process] = None,
12+
extra_hash_params: Optional[Set[str]] = set(),
1213
) -> Process:
1314
return Process(
1415
attributes=Process.Attributes.create(
@@ -18,6 +19,7 @@
1819
inputs=inputs,
1920
outputs=outputs,
2021
parent=parent,
22+
extra_hash_params=extra_hash_params,
2123
)
2224
)
2325

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
2+
@classmethod
3+
@init_guid
4+
def creator(
5+
cls,
6+
*,
7+
name: str,
8+
ai_application_version: str,
9+
ai_application_development_stage: AIApplicationDevelopmentStage,
10+
owner_groups: Optional[Set[str]] = set(),
11+
owner_users: Optional[Set[str]] = set(),
12+
models: Optional[List[AIModel]] = [],
13+
) -> AIApplication.Attributes:
14+
validate_required_fields(
15+
["name", "ai_application_version", "ai_application_development_stage"],
16+
[name, ai_application_version, ai_application_development_stage],
17+
)
18+
name_camel_case = to_camel_case(name)
19+
return AIApplication.Attributes(
20+
name=name,
21+
qualified_name=f"default/ai/aiapplication/{name_camel_case}",
22+
connector_name=AtlanConnectorType.AI.value,
23+
ai_application_version=ai_application_version,
24+
ai_application_development_stage=ai_application_development_stage,
25+
owner_groups=owner_groups,
26+
owner_users=owner_users,
27+
models=models,
28+
certificate_status=CertificateStatus.DRAFT,
29+
asset_cover_image="/assets/default-product-cover-DeQonY47.webp",
30+
)
Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
2+
@classmethod
3+
@init_guid
4+
def creator(
5+
cls,
6+
*,
7+
name: str,
8+
ai_model_status: AIModelStatus,
9+
owner_groups: Optional[Set[str]] = set(),
10+
owner_users: Optional[Set[str]] = set(),
11+
ai_model_version: Optional[str] = None,
12+
) -> AIModel.Attributes:
13+
validate_required_fields(
14+
["name", "ai_model_status"],
15+
[name, ai_model_status],
16+
)
17+
name_camel_case = to_camel_case(name)
18+
return AIModel.Attributes(
19+
name=name,
20+
qualified_name=f"default/ai/aiapplication/{name_camel_case}",
21+
connector_name=AtlanConnectorType.AI.value,
22+
ai_model_status=ai_model_status,
23+
ai_model_version=ai_model_version,
24+
owner_groups=owner_groups,
25+
owner_users=owner_users,
26+
asset_cover_image="/assets/default-product-cover-DeQonY47.webp",
27+
)

pyatlan/generator/templates/methods/attribute/process.jinja2

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
outputs: List["Catalog"],
88
parent: Optional["Process"] = None,
99
process_id: Optional[str] = None,
10+
extra_hash_params: Optional[Set[str]] = set(),
1011
) -> str:
1112
def append_relationship(output: StringIO, relationship: Asset):
1213
if relationship.guid:
@@ -29,6 +30,11 @@
2930
append_relationship(buffer, parent)
3031
append_relationships(buffer, inputs)
3132
append_relationships(buffer, outputs)
33+
# Handles edge case where identical name, connection, input, and output caused hash collisions,
34+
# resulting in duplicate qualified names and backend skipping process creation.
35+
if extra_hash_params:
36+
for param in extra_hash_params:
37+
buffer.write(param)
3238
ret_value = hashlib.md5( # noqa: S303, S324
3339
buffer.getvalue().encode()
3440
).hexdigest()
@@ -45,6 +51,7 @@
4551
outputs: List["Catalog"],
4652
process_id: Optional[str] = None,
4753
parent: Optional[Process] = None,
54+
extra_hash_params: Optional[Set[str]] = set(),
4855
) -> Process.Attributes:
4956
qualified_name = Process.Attributes.generate_qualified_name(
5057
name=name,
@@ -53,6 +60,7 @@
5360
inputs=inputs,
5461
outputs=outputs,
5562
parent=parent,
63+
extra_hash_params=extra_hash_params,
5664
)
5765
connector_name = connection_qualified_name.split("/")[1]
5866
return Process.Attributes(

0 commit comments

Comments
 (0)