infrahub-server 1.4.13__py3-none-any.whl → 1.5.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (222) hide show
  1. infrahub/actions/tasks.py +208 -16
  2. infrahub/api/artifact.py +3 -0
  3. infrahub/api/diff/diff.py +1 -1
  4. infrahub/api/internal.py +2 -0
  5. infrahub/api/query.py +2 -0
  6. infrahub/api/schema.py +27 -3
  7. infrahub/auth.py +5 -5
  8. infrahub/cli/__init__.py +2 -0
  9. infrahub/cli/db.py +160 -157
  10. infrahub/cli/dev.py +118 -0
  11. infrahub/cli/upgrade.py +56 -9
  12. infrahub/computed_attribute/tasks.py +19 -7
  13. infrahub/config.py +7 -2
  14. infrahub/core/attribute.py +35 -24
  15. infrahub/core/branch/enums.py +1 -1
  16. infrahub/core/branch/models.py +9 -5
  17. infrahub/core/branch/needs_rebase_status.py +11 -0
  18. infrahub/core/branch/tasks.py +72 -10
  19. infrahub/core/changelog/models.py +2 -10
  20. infrahub/core/constants/__init__.py +4 -0
  21. infrahub/core/constants/infrahubkind.py +1 -0
  22. infrahub/core/convert_object_type/object_conversion.py +201 -0
  23. infrahub/core/convert_object_type/repository_conversion.py +89 -0
  24. infrahub/core/convert_object_type/schema_mapping.py +27 -3
  25. infrahub/core/diff/model/path.py +4 -0
  26. infrahub/core/diff/payload_builder.py +1 -1
  27. infrahub/core/diff/query/artifact.py +1 -0
  28. infrahub/core/diff/query/field_summary.py +1 -0
  29. infrahub/core/graph/__init__.py +1 -1
  30. infrahub/core/initialization.py +7 -4
  31. infrahub/core/manager.py +3 -81
  32. infrahub/core/migrations/__init__.py +3 -0
  33. infrahub/core/migrations/exceptions.py +4 -0
  34. infrahub/core/migrations/graph/__init__.py +11 -10
  35. infrahub/core/migrations/graph/load_schema_branch.py +21 -0
  36. infrahub/core/migrations/graph/m013_convert_git_password_credential.py +1 -1
  37. infrahub/core/migrations/graph/m037_index_attr_vals.py +11 -30
  38. infrahub/core/migrations/graph/m039_ipam_reconcile.py +9 -7
  39. infrahub/core/migrations/graph/m042_profile_attrs_in_db.py +147 -0
  40. infrahub/core/migrations/graph/m043_create_hfid_display_label_in_db.py +164 -0
  41. infrahub/core/migrations/graph/m044_backfill_hfid_display_label_in_db.py +864 -0
  42. infrahub/core/migrations/query/__init__.py +7 -8
  43. infrahub/core/migrations/query/attribute_add.py +8 -6
  44. infrahub/core/migrations/query/attribute_remove.py +134 -0
  45. infrahub/core/migrations/runner.py +54 -0
  46. infrahub/core/migrations/schema/attribute_kind_update.py +9 -3
  47. infrahub/core/migrations/schema/attribute_supports_profile.py +90 -0
  48. infrahub/core/migrations/schema/node_attribute_add.py +26 -5
  49. infrahub/core/migrations/schema/node_attribute_remove.py +13 -109
  50. infrahub/core/migrations/schema/node_kind_update.py +2 -1
  51. infrahub/core/migrations/schema/node_remove.py +2 -1
  52. infrahub/core/migrations/schema/placeholder_dummy.py +3 -2
  53. infrahub/core/migrations/shared.py +66 -19
  54. infrahub/core/models.py +2 -2
  55. infrahub/core/node/__init__.py +207 -54
  56. infrahub/core/node/create.py +53 -49
  57. infrahub/core/node/lock_utils.py +124 -0
  58. infrahub/core/node/node_property_attribute.py +230 -0
  59. infrahub/core/node/resource_manager/ip_address_pool.py +2 -1
  60. infrahub/core/node/resource_manager/ip_prefix_pool.py +2 -1
  61. infrahub/core/node/resource_manager/number_pool.py +2 -1
  62. infrahub/core/node/standard.py +1 -1
  63. infrahub/core/property.py +11 -0
  64. infrahub/core/protocols.py +8 -1
  65. infrahub/core/query/attribute.py +82 -15
  66. infrahub/core/query/ipam.py +16 -4
  67. infrahub/core/query/node.py +66 -188
  68. infrahub/core/query/relationship.py +44 -26
  69. infrahub/core/query/subquery.py +0 -8
  70. infrahub/core/relationship/model.py +69 -24
  71. infrahub/core/schema/__init__.py +56 -0
  72. infrahub/core/schema/attribute_schema.py +4 -2
  73. infrahub/core/schema/basenode_schema.py +42 -2
  74. infrahub/core/schema/definitions/core/__init__.py +2 -0
  75. infrahub/core/schema/definitions/core/check.py +1 -1
  76. infrahub/core/schema/definitions/core/generator.py +2 -0
  77. infrahub/core/schema/definitions/core/group.py +16 -2
  78. infrahub/core/schema/definitions/core/repository.py +7 -0
  79. infrahub/core/schema/definitions/core/transform.py +1 -1
  80. infrahub/core/schema/definitions/internal.py +12 -3
  81. infrahub/core/schema/generated/attribute_schema.py +2 -2
  82. infrahub/core/schema/generated/base_node_schema.py +6 -1
  83. infrahub/core/schema/manager.py +3 -0
  84. infrahub/core/schema/node_schema.py +1 -0
  85. infrahub/core/schema/relationship_schema.py +0 -1
  86. infrahub/core/schema/schema_branch.py +295 -10
  87. infrahub/core/schema/schema_branch_display.py +135 -0
  88. infrahub/core/schema/schema_branch_hfid.py +120 -0
  89. infrahub/core/validators/aggregated_checker.py +1 -1
  90. infrahub/database/graph.py +21 -0
  91. infrahub/display_labels/__init__.py +0 -0
  92. infrahub/display_labels/gather.py +48 -0
  93. infrahub/display_labels/models.py +240 -0
  94. infrahub/display_labels/tasks.py +192 -0
  95. infrahub/display_labels/triggers.py +22 -0
  96. infrahub/events/branch_action.py +27 -1
  97. infrahub/events/group_action.py +1 -1
  98. infrahub/events/node_action.py +1 -1
  99. infrahub/generators/constants.py +7 -0
  100. infrahub/generators/models.py +38 -12
  101. infrahub/generators/tasks.py +34 -16
  102. infrahub/git/base.py +38 -1
  103. infrahub/git/integrator.py +22 -14
  104. infrahub/graphql/api/dependencies.py +2 -4
  105. infrahub/graphql/api/endpoints.py +16 -6
  106. infrahub/graphql/app.py +2 -4
  107. infrahub/graphql/initialization.py +2 -3
  108. infrahub/graphql/manager.py +213 -137
  109. infrahub/graphql/middleware.py +12 -0
  110. infrahub/graphql/mutations/branch.py +16 -0
  111. infrahub/graphql/mutations/computed_attribute.py +110 -3
  112. infrahub/graphql/mutations/convert_object_type.py +44 -13
  113. infrahub/graphql/mutations/display_label.py +118 -0
  114. infrahub/graphql/mutations/generator.py +25 -7
  115. infrahub/graphql/mutations/hfid.py +125 -0
  116. infrahub/graphql/mutations/ipam.py +73 -41
  117. infrahub/graphql/mutations/main.py +61 -178
  118. infrahub/graphql/mutations/profile.py +195 -0
  119. infrahub/graphql/mutations/proposed_change.py +8 -1
  120. infrahub/graphql/mutations/relationship.py +2 -2
  121. infrahub/graphql/mutations/repository.py +22 -83
  122. infrahub/graphql/mutations/resource_manager.py +2 -2
  123. infrahub/graphql/mutations/webhook.py +1 -1
  124. infrahub/graphql/queries/resource_manager.py +1 -1
  125. infrahub/graphql/registry.py +173 -0
  126. infrahub/graphql/resolvers/resolver.py +2 -0
  127. infrahub/graphql/schema.py +8 -1
  128. infrahub/graphql/schema_sort.py +170 -0
  129. infrahub/graphql/types/branch.py +4 -1
  130. infrahub/graphql/types/enums.py +3 -0
  131. infrahub/groups/tasks.py +1 -1
  132. infrahub/hfid/__init__.py +0 -0
  133. infrahub/hfid/gather.py +48 -0
  134. infrahub/hfid/models.py +240 -0
  135. infrahub/hfid/tasks.py +191 -0
  136. infrahub/hfid/triggers.py +22 -0
  137. infrahub/lock.py +119 -42
  138. infrahub/locks/__init__.py +0 -0
  139. infrahub/locks/tasks.py +37 -0
  140. infrahub/patch/plan_writer.py +2 -2
  141. infrahub/permissions/constants.py +2 -0
  142. infrahub/profiles/__init__.py +0 -0
  143. infrahub/profiles/node_applier.py +101 -0
  144. infrahub/profiles/queries/__init__.py +0 -0
  145. infrahub/profiles/queries/get_profile_data.py +98 -0
  146. infrahub/profiles/tasks.py +63 -0
  147. infrahub/proposed_change/tasks.py +24 -5
  148. infrahub/repositories/__init__.py +0 -0
  149. infrahub/repositories/create_repository.py +113 -0
  150. infrahub/server.py +9 -1
  151. infrahub/services/__init__.py +8 -5
  152. infrahub/services/adapters/workflow/worker.py +5 -2
  153. infrahub/task_manager/event.py +5 -0
  154. infrahub/task_manager/models.py +7 -0
  155. infrahub/tasks/registry.py +6 -4
  156. infrahub/trigger/catalogue.py +4 -0
  157. infrahub/trigger/models.py +2 -0
  158. infrahub/trigger/setup.py +13 -4
  159. infrahub/trigger/tasks.py +6 -0
  160. infrahub/webhook/models.py +1 -1
  161. infrahub/workers/dependencies.py +3 -1
  162. infrahub/workers/infrahub_async.py +5 -1
  163. infrahub/workflows/catalogue.py +118 -3
  164. infrahub/workflows/initialization.py +21 -0
  165. infrahub/workflows/models.py +17 -2
  166. infrahub_sdk/branch.py +17 -8
  167. infrahub_sdk/checks.py +1 -1
  168. infrahub_sdk/client.py +376 -95
  169. infrahub_sdk/config.py +29 -2
  170. infrahub_sdk/convert_object_type.py +61 -0
  171. infrahub_sdk/ctl/branch.py +3 -0
  172. infrahub_sdk/ctl/check.py +2 -3
  173. infrahub_sdk/ctl/cli_commands.py +20 -12
  174. infrahub_sdk/ctl/config.py +8 -2
  175. infrahub_sdk/ctl/generator.py +6 -3
  176. infrahub_sdk/ctl/graphql.py +184 -0
  177. infrahub_sdk/ctl/repository.py +39 -1
  178. infrahub_sdk/ctl/schema.py +40 -10
  179. infrahub_sdk/ctl/task.py +110 -0
  180. infrahub_sdk/ctl/utils.py +4 -0
  181. infrahub_sdk/ctl/validate.py +5 -3
  182. infrahub_sdk/diff.py +4 -5
  183. infrahub_sdk/exceptions.py +2 -0
  184. infrahub_sdk/generator.py +7 -1
  185. infrahub_sdk/graphql/__init__.py +12 -0
  186. infrahub_sdk/graphql/constants.py +1 -0
  187. infrahub_sdk/graphql/plugin.py +85 -0
  188. infrahub_sdk/graphql/query.py +77 -0
  189. infrahub_sdk/{graphql.py → graphql/renderers.py} +88 -75
  190. infrahub_sdk/graphql/utils.py +40 -0
  191. infrahub_sdk/node/attribute.py +2 -0
  192. infrahub_sdk/node/node.py +28 -20
  193. infrahub_sdk/node/relationship.py +1 -3
  194. infrahub_sdk/playback.py +1 -2
  195. infrahub_sdk/protocols.py +54 -6
  196. infrahub_sdk/pytest_plugin/plugin.py +7 -4
  197. infrahub_sdk/pytest_plugin/utils.py +40 -0
  198. infrahub_sdk/repository.py +1 -2
  199. infrahub_sdk/schema/__init__.py +70 -4
  200. infrahub_sdk/schema/main.py +1 -0
  201. infrahub_sdk/schema/repository.py +8 -0
  202. infrahub_sdk/spec/models.py +7 -0
  203. infrahub_sdk/spec/object.py +54 -6
  204. infrahub_sdk/spec/processors/__init__.py +0 -0
  205. infrahub_sdk/spec/processors/data_processor.py +10 -0
  206. infrahub_sdk/spec/processors/factory.py +34 -0
  207. infrahub_sdk/spec/processors/range_expand_processor.py +56 -0
  208. infrahub_sdk/spec/range_expansion.py +118 -0
  209. infrahub_sdk/task/models.py +6 -4
  210. infrahub_sdk/timestamp.py +18 -6
  211. infrahub_sdk/transforms.py +1 -1
  212. {infrahub_server-1.4.13.dist-info → infrahub_server-1.5.0.dist-info}/METADATA +9 -10
  213. {infrahub_server-1.4.13.dist-info → infrahub_server-1.5.0.dist-info}/RECORD +221 -165
  214. infrahub_testcontainers/container.py +114 -2
  215. infrahub_testcontainers/docker-compose-cluster.test.yml +5 -0
  216. infrahub_testcontainers/docker-compose.test.yml +5 -0
  217. infrahub_testcontainers/models.py +2 -2
  218. infrahub_testcontainers/performance_test.py +4 -4
  219. infrahub/core/convert_object_type/conversion.py +0 -134
  220. {infrahub_server-1.4.13.dist-info → infrahub_server-1.5.0.dist-info}/LICENSE.txt +0 -0
  221. {infrahub_server-1.4.13.dist-info → infrahub_server-1.5.0.dist-info}/WHEEL +0 -0
  222. {infrahub_server-1.4.13.dist-info → infrahub_server-1.5.0.dist-info}/entry_points.txt +0 -0
@@ -0,0 +1,864 @@
1
+ from __future__ import annotations
2
+
3
+ from collections import defaultdict
4
+ from itertools import chain
5
+ from typing import TYPE_CHECKING, Any
6
+
7
+ import ujson
8
+ from rich.progress import Progress, TaskID
9
+
10
+ from infrahub.core import registry
11
+ from infrahub.core.branch import Branch
12
+ from infrahub.core.constants import GLOBAL_BRANCH_NAME, BranchSupportType, RelationshipDirection
13
+ from infrahub.core.initialization import get_root_node
14
+ from infrahub.core.migrations.shared import MigrationResult, get_migration_console
15
+ from infrahub.core.query import Query, QueryType
16
+ from infrahub.types import is_large_attribute_type
17
+
18
+ from ..shared import MigrationRequiringRebase
19
+ from .load_schema_branch import get_or_load_schema_branch
20
+
21
+ if TYPE_CHECKING:
22
+ from infrahub.core.schema import AttributeSchema, NodeSchema
23
+ from infrahub.core.schema.basenode_schema import SchemaAttributePath
24
+ from infrahub.core.schema.schema_branch import SchemaBranch
25
+ from infrahub.database import InfrahubDatabase
26
+
27
+
28
+ console = get_migration_console()
29
+
30
+
31
+ class DefaultBranchNodeCount(Query):
32
+ """
33
+ Get the number of Node vertices on the given branches that are not in the kinds_to_skip list
34
+ Only works for default and global branches. Non-default branches would only return a count of nodes
35
+ created on the given branches
36
+ """
37
+
38
+ name = "get_branch_node_count"
39
+ type = QueryType.READ
40
+
41
+ def __init__(self, kinds_to_skip: list[str], **kwargs: Any) -> None:
42
+ super().__init__(**kwargs)
43
+ self.kinds_to_skip = kinds_to_skip
44
+
45
+ async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None: # noqa: ARG002
46
+ self.params = {
47
+ "branch_names": [registry.default_branch, GLOBAL_BRANCH_NAME],
48
+ "kinds_to_skip": self.kinds_to_skip,
49
+ }
50
+ query = """
51
+ MATCH (n:Node)-[e:IS_PART_OF]->(:Root)
52
+ WHERE NOT n.kind IN $kinds_to_skip
53
+ AND e.branch IN $branch_names
54
+ AND e.status = "active"
55
+ AND e.to IS NULL
56
+ AND NOT exists((n)-[:IS_PART_OF {branch: e.branch, status: "deleted"}]->(:Root))
57
+ WITH count(*) AS num_nodes
58
+ """
59
+ self.add_to_query(query)
60
+ self.return_labels = ["num_nodes"]
61
+
62
+ def get_num_nodes(self) -> int:
63
+ result = self.get_result()
64
+ if not result:
65
+ return 0
66
+ return result.get_as_type(label="num_nodes", return_type=int)
67
+
68
+
69
+ class GetResultMapQuery(Query):
70
+ def get_result_map(self, schema_paths: list[SchemaAttributePath]) -> dict[str, list[str | None]]:
71
+ """
72
+ Get the values for the given schema paths for all the Nodes captured by this query
73
+ """
74
+ # the query results for attribute and schema paths are unordered
75
+ # so we make this list of keys for ordering the results from the query
76
+ schema_path_keys: list[tuple[str, RelationshipDirection, str] | str] = []
77
+ for schema_path in schema_paths:
78
+ if schema_path.is_type_attribute and schema_path.attribute_schema:
79
+ path_key: str | tuple[str, RelationshipDirection, str] = schema_path.attribute_schema.name
80
+ elif schema_path.is_type_relationship and schema_path.relationship_schema and schema_path.attribute_schema:
81
+ path_key = (
82
+ schema_path.relationship_schema.get_identifier(),
83
+ schema_path.relationship_schema.direction,
84
+ schema_path.attribute_schema.name,
85
+ )
86
+ schema_path_keys.append(path_key)
87
+
88
+ result_map: dict[str, list[str | None]] = {}
89
+ for result in self.get_results():
90
+ node_uuid = result.get_as_type(label="n_uuid", return_type=str)
91
+
92
+ # for each node, build a map of the schema path key to value so that they
93
+ # can be ordered correctly for the input `schema_paths`
94
+ schema_path_value_map: dict[str | tuple[str, RelationshipDirection, str], Any] = {}
95
+ attr_values_tuples: list[tuple[str, Any]] = result.get_as_type(label="attr_vals_list", return_type=list)
96
+ for attr_value_tuple in attr_values_tuples:
97
+ attr_name = attr_value_tuple[0]
98
+ attr_value = attr_value_tuple[1]
99
+ schema_path_value_map[attr_name] = attr_value
100
+
101
+ relationship_values_tuples: list[tuple[str, str, str, Any]] = result.get_as_type(
102
+ label="peer_attr_vals_list", return_type=list
103
+ )
104
+ for rel_value_tuple in relationship_values_tuples:
105
+ rel_name = rel_value_tuple[0]
106
+ direction_raw = rel_value_tuple[1]
107
+ direction = RelationshipDirection.BIDIR
108
+ match direction_raw:
109
+ case "outbound":
110
+ direction = RelationshipDirection.OUTBOUND
111
+ case "inbound":
112
+ direction = RelationshipDirection.INBOUND
113
+ peer_attr_name = rel_value_tuple[2]
114
+ peer_val = rel_value_tuple[3]
115
+ schema_path_value_map[rel_name, direction, peer_attr_name] = peer_val
116
+
117
+ schema_path_values: list[str | None] = []
118
+ for schema_path_key in schema_path_keys:
119
+ value = schema_path_value_map.get(schema_path_key)
120
+ schema_path_values.append(str(value) if value is not None else None)
121
+ result_map[node_uuid] = schema_path_values
122
+ return result_map
123
+
124
+
125
+ class GetPathDetailsBranchQuery(GetResultMapQuery):
126
+ name = "get_path_details_branch"
127
+ type = QueryType.READ
128
+ insert_limit = False
129
+
130
+ def __init__(
131
+ self, schema_kind: str, schema_paths: list[SchemaAttributePath], updates_only: bool = True, **kwargs: Any
132
+ ) -> None:
133
+ super().__init__(**kwargs)
134
+
135
+ if self.branch.name in [registry.default_branch, GLOBAL_BRANCH_NAME]:
136
+ raise ValueError("This query can only be used on non-default branches")
137
+ self.schema_kind = schema_kind
138
+ self.updates_only = updates_only
139
+ self.attribute_names = []
140
+ self.bidir_rel_attr_map: dict[str, list[str]] = defaultdict(list)
141
+ self.outbound_rel_attr_map: dict[str, list[str]] = defaultdict(list)
142
+ self.inbound_rel_attr_map: dict[str, list[str]] = defaultdict(list)
143
+ for schema_path in schema_paths:
144
+ if schema_path.is_type_attribute and schema_path.attribute_schema:
145
+ self.attribute_names.append(schema_path.attribute_schema.name)
146
+ elif schema_path.is_type_relationship and schema_path.relationship_schema and schema_path.attribute_schema:
147
+ key = schema_path.relationship_schema.get_identifier()
148
+ value = schema_path.attribute_schema.name
149
+ if schema_path.relationship_schema.direction is RelationshipDirection.BIDIR:
150
+ self.bidir_rel_attr_map[key].append(value)
151
+ elif schema_path.relationship_schema.direction is RelationshipDirection.OUTBOUND:
152
+ self.outbound_rel_attr_map[key].append(value)
153
+ elif schema_path.relationship_schema.direction is RelationshipDirection.INBOUND:
154
+ self.inbound_rel_attr_map[key].append(value)
155
+
156
+ async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None: # noqa: ARG002
157
+ branch_filter, branch_filter_params = self.branch.get_query_filter_path(at=self.at)
158
+ self.params.update(branch_filter_params)
159
+ self.params.update(
160
+ {
161
+ "branch_name": self.branch.name,
162
+ "attribute_names": self.attribute_names,
163
+ "outbound_rel_ids": list(self.outbound_rel_attr_map.keys()),
164
+ "inbound_rel_ids": list(self.inbound_rel_attr_map.keys()),
165
+ "bidirectional_rel_ids": list(self.bidir_rel_attr_map.keys()),
166
+ "outbound_rel_attr_map": self.outbound_rel_attr_map,
167
+ "inbound_rel_attr_map": self.inbound_rel_attr_map,
168
+ "bidirectional_rel_attr_map": self.bidir_rel_attr_map,
169
+ "offset": self.offset,
170
+ "limit": self.limit,
171
+ }
172
+ )
173
+ get_active_nodes_query = """
174
+ // ------------
175
+ // Get the active nodes of the given kind on the branches
176
+ // ------------
177
+ MATCH (n:%(schema_kind)s)-[r:IS_PART_OF]->(:Root)
178
+ WHERE %(branch_filter)s
179
+ WITH DISTINCT n
180
+ CALL (n) {
181
+ MATCH (n)-[r:IS_PART_OF]->(:Root)
182
+ WHERE %(branch_filter)s
183
+ RETURN r.status = "active" AS is_active
184
+ ORDER BY r.branch_level DESC, r.from DESC, r.status ASC
185
+ LIMIT 1
186
+ }
187
+ WITH n, is_active
188
+ WHERE is_active = TRUE
189
+ """ % {"schema_kind": self.schema_kind, "branch_filter": branch_filter}
190
+ self.add_to_query(get_active_nodes_query)
191
+
192
+ if self.updates_only:
193
+ updated_nodes_filter_query = """
194
+ // ------------
195
+ // filter to any nodes that might have changes on the branch we care about
196
+ // ------------
197
+ OPTIONAL MATCH (n)-[r1:HAS_ATTRIBUTE]->(attr:Attribute)-[r2:HAS_VALUE]->(attr_val:AttributeValue)
198
+ WHERE attr.name in $attribute_names
199
+ AND r2.branch = $branch_name
200
+ AND r2.status = "active"
201
+ AND r2.to IS NULL
202
+ WITH n, attr_val IS NOT NULL AS has_attr_update
203
+ OPTIONAL MATCH (n)-[r1:IS_RELATED]-(rel:Relationship)-[r2:IS_RELATED]-(peer:Node)-[r3:HAS_ATTRIBUTE]-(attr:Attribute)-[r4:HAS_VALUE]->(attr_val)
204
+ WHERE rel.name IN $bidirectional_rel_ids + $outbound_rel_ids + $inbound_rel_ids
205
+ AND (
206
+ attr.name IN $outbound_rel_attr_map[rel.name]
207
+ OR attr.name IN $inbound_rel_attr_map[rel.name]
208
+ OR attr.name IN $bidirectional_rel_attr_map[rel.name]
209
+ )
210
+ AND $branch_name IN [r1.branch, r2.branch, r3.branch, r4.branch]
211
+ WITH n, has_attr_update, attr_val IS NOT NULL AS has_rel_update
212
+ WITH n, any(x IN collect(has_attr_update OR has_rel_update) WHERE x = TRUE) AS has_update
213
+ WITH n, has_update
214
+ WHERE has_update = TRUE
215
+ """
216
+ self.add_to_query(updated_nodes_filter_query)
217
+
218
+ get_node_details_query = """
219
+ // ------------
220
+ // Order and limit the Nodes
221
+ // ------------
222
+ ORDER BY elementId(n)
223
+ SKIP toInteger($offset)
224
+ LIMIT toInteger($limit)
225
+ // ------------
226
+ // for every possibly updated node
227
+ // get all the attribute values on this branch
228
+ // ------------
229
+ OPTIONAL MATCH (n)-[r:HAS_ATTRIBUTE]->(attr:Attribute)
230
+ WHERE attr.name IN $attribute_names
231
+ WITH DISTINCT n, attr
232
+ CALL (n, attr) {
233
+ OPTIONAL MATCH (n)-[r1:HAS_ATTRIBUTE]->(attr)-[r2:HAS_VALUE]->(attr_val)
234
+ WHERE all(r in [r1, r2] WHERE %(branch_filter)s)
235
+ RETURN attr_val.value AS attr_value, r1.status = "active" AND r2.status = "active" AS is_active
236
+ ORDER BY r2.branch_level DESC, r2.from DESC, r2.status ASC, r1.branch_level DESC, r1.from DESC, r1.status ASC
237
+ LIMIT 1
238
+ }
239
+ WITH n, attr, attr_value
240
+ WHERE is_active = TRUE
241
+ WITH n, collect([attr.name, attr_value]) AS attr_vals_list
242
+ // ------------
243
+ // for every possibly updated node
244
+ // get all the relationships on this branch
245
+ // ------------
246
+ OPTIONAL MATCH (n)-[:IS_RELATED]-(rel:Relationship)
247
+ WHERE rel.name IN $bidirectional_rel_ids + $outbound_rel_ids + $inbound_rel_ids
248
+ WITH DISTINCT n, attr_vals_list, rel
249
+ CALL (n, rel) {
250
+ OPTIONAL MATCH (n)-[r1:IS_RELATED]-(rel)-[r2:IS_RELATED]-(peer:Node)
251
+ WHERE all(r in [r1, r2] WHERE %(branch_filter)s)
252
+ AND (
253
+ (startNode(r1) = n AND startNode(r2) = rel AND rel.name IN $outbound_rel_ids)
254
+ OR (startNode(r1) = rel AND startNode(r2) = peer AND rel.name IN $inbound_rel_ids)
255
+ OR (startNode(r1) = n AND startNode(r2) = peer AND rel.name IN $bidirectional_rel_ids)
256
+ )
257
+ RETURN
258
+ peer,
259
+ r1.status = "active" AND r2.status = "active" AS is_active,
260
+ CASE
261
+ WHEN startNode(r1) = n AND startNode(r2) = rel AND rel.name IN $outbound_rel_ids THEN "outbound"
262
+ WHEN startNode(r1) = rel AND startNode(r2) = peer AND rel.name IN $inbound_rel_ids THEN "inbound"
263
+ ELSE "bidir"
264
+ END AS direction
265
+ ORDER BY r2.branch_level DESC, r2.from DESC, r2.status ASC, r1.branch_level DESC, r1.from DESC, r1.status ASC
266
+ LIMIT 1
267
+ }
268
+ // ------------
269
+ // get the attribute values that we care about for each relationship
270
+ // ------------
271
+ WITH n, attr_vals_list, rel.name AS rel_name, direction, peer
272
+ WHERE is_active = TRUE OR rel_name IS NULL
273
+ WITH *, CASE
274
+ WHEN direction = "outbound" THEN $outbound_rel_attr_map[rel_name]
275
+ WHEN direction = "inbound" THEN $inbound_rel_attr_map[rel_name]
276
+ ELSE $bidirectional_rel_attr_map[rel_name]
277
+ END AS peer_attr_names
278
+ UNWIND COALESCE(peer_attr_names, [NULL]) AS peer_attr_name
279
+ CALL (rel_name, direction, peer, peer_attr_name){
280
+ OPTIONAL MATCH (peer)-[r1:HAS_ATTRIBUTE]->(attr:Attribute)-[r2:HAS_VALUE]->(attr_val)
281
+ WHERE attr.name = peer_attr_name
282
+ AND all(r in [r1, r2] WHERE %(branch_filter)s)
283
+ RETURN attr_val.value AS peer_attr_value, r1.status = "active" AND r2.status = "active" AS is_active
284
+ ORDER BY r2.branch_level DESC, r2.from DESC, r2.status ASC, r1.branch_level DESC, r1.from DESC, r1.status ASC
285
+ LIMIT 1
286
+ }
287
+ // ------------
288
+ // collect everything to return a pair of lists with each node UUID
289
+ // ------------
290
+ WITH DISTINCT n, attr_vals_list, rel_name, peer, direction, peer_attr_name, peer_attr_value
291
+ WITH n, attr_vals_list, collect([rel_name, direction, peer_attr_name, peer_attr_value]) AS peer_attr_vals_list
292
+ """ % {"branch_filter": branch_filter}
293
+ self.add_to_query(get_node_details_query)
294
+ self.return_labels = ["n.uuid AS n_uuid", "attr_vals_list", "peer_attr_vals_list"]
295
+
296
+
297
+ class GetPathDetailsDefaultBranch(GetResultMapQuery):
298
+ """
299
+ Get the values of the given schema paths for the given kind of node on the default and global branches
300
+ Supports limit and offset for pagination
301
+ """
302
+
303
+ name = "get_path_details_default_branch"
304
+ type = QueryType.READ
305
+ insert_limit = False
306
+
307
+ def __init__(self, schema_kind: str, schema_paths: list[SchemaAttributePath], **kwargs: Any) -> None:
308
+ super().__init__(**kwargs)
309
+
310
+ self.branch_names = [registry.default_branch, GLOBAL_BRANCH_NAME]
311
+ self.schema_kind = schema_kind
312
+ self.attribute_names = []
313
+ self.bidir_rel_attr_map: dict[str, list[str]] = defaultdict(list)
314
+ self.outbound_rel_attr_map: dict[str, list[str]] = defaultdict(list)
315
+ self.inbound_rel_attr_map: dict[str, list[str]] = defaultdict(list)
316
+ for schema_path in schema_paths:
317
+ if schema_path.is_type_attribute and schema_path.attribute_schema:
318
+ self.attribute_names.append(schema_path.attribute_schema.name)
319
+ elif schema_path.is_type_relationship and schema_path.relationship_schema and schema_path.attribute_schema:
320
+ key = schema_path.relationship_schema.get_identifier()
321
+ value = schema_path.attribute_schema.name
322
+ if schema_path.relationship_schema.direction is RelationshipDirection.BIDIR:
323
+ self.bidir_rel_attr_map[key].append(value)
324
+ elif schema_path.relationship_schema.direction is RelationshipDirection.OUTBOUND:
325
+ self.outbound_rel_attr_map[key].append(value)
326
+ elif schema_path.relationship_schema.direction is RelationshipDirection.INBOUND:
327
+ self.inbound_rel_attr_map[key].append(value)
328
+
329
+ async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None: # noqa: ARG002
330
+ self.params = {
331
+ "branch_names": self.branch_names,
332
+ "attribute_names": self.attribute_names,
333
+ "outbound_rel_ids": list(self.outbound_rel_attr_map.keys()),
334
+ "inbound_rel_ids": list(self.inbound_rel_attr_map.keys()),
335
+ "bidirectional_rel_ids": list(self.bidir_rel_attr_map.keys()),
336
+ "outbound_rel_attr_map": self.outbound_rel_attr_map,
337
+ "inbound_rel_attr_map": self.inbound_rel_attr_map,
338
+ "bidirectional_rel_attr_map": self.bidir_rel_attr_map,
339
+ "offset": self.offset,
340
+ "limit": self.limit,
341
+ }
342
+ get_details_query = """
343
+ // ------------
344
+ // Get the active nodes of the given kind on the branches
345
+ // ------------
346
+ MATCH (n:%(schema_kind)s)-[e:IS_PART_OF]->(:Root)
347
+ WHERE e.branch IN $branch_names
348
+ AND e.to IS NULL
349
+ AND e.status = "active"
350
+ // ------------
351
+ // Order and limit the Nodes
352
+ // ------------
353
+ WITH DISTINCT n
354
+ ORDER BY elementId(n)
355
+ SKIP toInteger($offset)
356
+ LIMIT toInteger($limit)
357
+ // ------------
358
+ // Get the values for the attribute schema paths of the Nodes, if any
359
+ // ------------
360
+ OPTIONAL MATCH (n)-[e:HAS_ATTRIBUTE]->(attr:Attribute)
361
+ WHERE attr.name IN $attribute_names
362
+ AND e.branch IN $branch_names
363
+ AND e.to IS NULL
364
+ AND e.status = "active"
365
+ WITH n, attr
366
+ OPTIONAL MATCH (attr)-[e:HAS_VALUE]->(attr_val:AttributeValue)
367
+ WHERE e.branch IN $branch_names
368
+ AND e.to IS NULL
369
+ AND e.status = "active"
370
+ WITH n, collect([attr.name, attr_val.value]) AS attr_vals_list
371
+ // ------------
372
+ // Get the values for the relationship schema paths of the Nodes, if any
373
+ // ------------
374
+ OPTIONAL MATCH (n)-[e1:IS_RELATED]-(rel:Relationship)-[e2:IS_RELATED]-(peer:Node)
375
+ WHERE rel.name IN $bidirectional_rel_ids + $outbound_rel_ids + $inbound_rel_ids
376
+ AND e1.branch IN $branch_names
377
+ AND e1.to IS NULL
378
+ AND e1.status = "active"
379
+ AND e2.branch IN $branch_names
380
+ AND e2.to IS NULL
381
+ AND e2.status = "active"
382
+ AND (
383
+ (startNode(e1) = n AND startNode(e2) = rel AND rel.name IN $outbound_rel_ids)
384
+ OR (startNode(e1) = rel AND startNode(e2) = peer AND rel.name IN $inbound_rel_ids)
385
+ OR (startNode(e1) = n AND startNode(e2) = peer AND rel.name IN $bidirectional_rel_ids)
386
+ )
387
+ WITH DISTINCT n, attr_vals_list, rel.name AS rel_name, peer, CASE
388
+ WHEN startNode(e1) = n AND startNode(e2) = rel AND rel.name IN $outbound_rel_ids THEN "outbound"
389
+ WHEN startNode(e1) = rel AND startNode(e2) = peer AND rel.name IN $inbound_rel_ids THEN "inbound"
390
+ ELSE "bidir"
391
+ END AS direction
392
+ OPTIONAL MATCH (peer)-[e1:HAS_ATTRIBUTE]->(attr:Attribute)-[e2:HAS_VALUE]->(peer_attr_val:AttributeValue)
393
+ WHERE (
394
+ (direction = "outbound" AND attr.name IN $outbound_rel_attr_map[rel_name])
395
+ OR (direction = "inbound" AND attr.name IN $inbound_rel_attr_map[rel_name])
396
+ OR (direction = "bidir" AND attr.name IN $bidirectional_rel_attr_map[rel_name])
397
+ )
398
+ AND e1.branch IN $branch_names
399
+ AND e1.to IS NULL
400
+ AND e1.status = "active"
401
+ AND e2.branch IN $branch_names
402
+ AND e2.to IS NULL
403
+ AND e2.status = "active"
404
+ // ------------
405
+ // collect everything to return a pair of lists with each node UUID
406
+ // ------------
407
+ WITH DISTINCT n, attr_vals_list, rel_name, peer, direction, attr.name AS peer_attr_name, peer_attr_val.value AS peer_val
408
+ WITH n, attr_vals_list, collect([rel_name, direction, peer_attr_name, peer_val]) AS peer_attr_vals_list
409
+ """ % {"schema_kind": self.schema_kind}
410
+ self.add_to_query(get_details_query)
411
+ self.return_labels = ["n.uuid AS n_uuid", "attr_vals_list", "peer_attr_vals_list"]
412
+
413
+
414
+ class UpdateAttributeValuesQuery(Query):
415
+ """
416
+ Update the values of the given attribute schema for the input Node-id-to-value map
417
+ Includes special handling for updating large-type attributes b/c they are not indexed and will be slow to update
418
+ on large data sets
419
+ """
420
+
421
+ name = "update_attribute_values"
422
+ type = QueryType.WRITE
423
+ insert_return = False
424
+
425
+ def __init__(self, attribute_schema: AttributeSchema, values_by_id_map: dict[str, Any], **kwargs: Any) -> None:
426
+ super().__init__(**kwargs)
427
+ self.attribute_name = attribute_schema.name
428
+ self.is_large_type_attribute = is_large_attribute_type(attribute_schema.kind)
429
+ self.is_branch_agnostic = attribute_schema.get_branch() is BranchSupportType.AGNOSTIC
430
+ self.values_by_id_map = values_by_id_map
431
+
432
+ async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None: # noqa: ARG002
433
+ self.params = {
434
+ "node_uuids": list(self.values_by_id_map.keys()),
435
+ "attribute_name": self.attribute_name,
436
+ "values_by_id": self.values_by_id_map,
437
+ "default_branch": registry.default_branch,
438
+ "global_branch": GLOBAL_BRANCH_NAME,
439
+ "branch": GLOBAL_BRANCH_NAME if self.is_branch_agnostic else self.branch.name,
440
+ "branch_level": 1 if self.is_branch_agnostic else self.branch.hierarchy_level,
441
+ "at": self.at.to_string(),
442
+ }
443
+ branch_filter, branch_filter_params = self.branch.get_query_filter_path(at=self.at)
444
+ self.params.update(branch_filter_params)
445
+
446
+ if self.is_large_type_attribute:
447
+ # we make our own index of value to database ID, creating any vertices that are missing
448
+ # the mapping is a list of tuples instead of an actual mapping b/c creating an actual map is not possible
449
+ # without apoc functions
450
+ all_distinct_values = list(set(self.values_by_id_map.values()))
451
+ diy_index_query = """
452
+ MATCH (av:AttributeValue&!AttributeValueIndexed {is_default: false})
453
+ WHERE av.value IN $all_distinct_values
454
+ WITH collect([av.value, elementId(av)]) AS value_id_pairs, collect(av.value) AS found_values
455
+ WITH value_id_pairs, found_values,
456
+ reduce(
457
+ missing_distinct_values = [], value IN $all_distinct_values |
458
+ CASE
459
+ WHEN value IN found_values THEN missing_distinct_values
460
+ ELSE missing_distinct_values + [value]
461
+ END
462
+ ) AS missing_distinct_values
463
+ CALL (missing_distinct_values) {
464
+ UNWIND missing_distinct_values AS missing_value
465
+ CREATE (av:AttributeValue {is_default: false, value: missing_value})
466
+ RETURN collect([av.value, elementId(av)]) AS created_value_id_pairs
467
+ }
468
+ WITH value_id_pairs + created_value_id_pairs AS value_id_pairs
469
+ """
470
+ self.params["all_distinct_values"] = all_distinct_values
471
+ else:
472
+ # if this is not a large-type attribute, then just set the map to be empty
473
+ diy_index_query = """WITH [] AS value_id_pairs"""
474
+
475
+ self.add_to_query(diy_index_query)
476
+
477
+ if self.branch.name in [registry.default_branch, GLOBAL_BRANCH_NAME]:
478
+ update_value_query = """
479
+ // ------------
480
+ // Find the Nodes and Attributes we need to update
481
+ // ------------
482
+ MATCH (n:Node)-[e:IS_PART_OF]->(:Root)
483
+ WHERE n.uuid IN $node_uuids
484
+ AND e.branch IN [$default_branch, $global_branch]
485
+ AND e.to IS NULL
486
+ AND e.status = "active"
487
+ WITH DISTINCT n, value_id_pairs
488
+ MATCH (n)-[e:HAS_ATTRIBUTE]->(attr:Attribute {name: $attribute_name})
489
+ WHERE e.branch IN [$default_branch, $global_branch]
490
+ AND e.to IS NULL
491
+ AND e.status = "active"
492
+ // ------------
493
+ // If the attribute has an existing value on the branch, then set the to time on it
494
+ // but only if the value is different from the new value
495
+ // ------------
496
+ WITH DISTINCT n, attr, value_id_pairs
497
+ CALL (attr) {
498
+ OPTIONAL MATCH (attr)-[e:HAS_VALUE]->(existing_av)
499
+ WHERE e.branch IN [$default_branch, $global_branch]
500
+ AND e.to IS NULL
501
+ AND e.status = "active"
502
+ RETURN existing_av, e AS existing_has_value
503
+ }
504
+ CALL (existing_has_value) {
505
+ WITH existing_has_value
506
+ WHERE existing_has_value IS NOT NULL
507
+ SET existing_has_value.to = $at
508
+ }
509
+ """
510
+ else:
511
+ update_value_query = """
512
+ // ------------
513
+ // Find the Nodes and Attributes we need to update
514
+ // ------------
515
+ MATCH (n:Node)
516
+ WHERE n.uuid IN $node_uuids
517
+ CALL (n) {
518
+ MATCH (n)-[r:IS_PART_OF]->(:Root)
519
+ WHERE %(branch_filter)s
520
+ RETURN r.status = "active" AS is_active
521
+ ORDER BY r.branch_level DESC, r.from DESC, r.status ASC
522
+ LIMIT 1
523
+ }
524
+ WITH n, value_id_pairs, is_active
525
+ WHERE is_active = TRUE
526
+ WITH DISTINCT n, value_id_pairs
527
+ CALL (n) {
528
+ MATCH (n)-[r:HAS_ATTRIBUTE]->(attr:Attribute {name: $attribute_name})
529
+ WHERE %(branch_filter)s
530
+ RETURN attr, r.status = "active" AS is_active
531
+ ORDER BY r.branch_level DESC, r.from DESC, r.status ASC
532
+ LIMIT 1
533
+ }
534
+ WITH DISTINCT n, attr, value_id_pairs, is_active
535
+ WHERE is_active = TRUE
536
+ // ------------
537
+ // If the attribute has an existing value on the branch, then set the to time on it
538
+ // but only if the value is different from the new value
539
+ // ------------
540
+ CALL (n, attr) {
541
+ OPTIONAL MATCH (attr)-[r:HAS_VALUE]->(existing_av)
542
+ WHERE %(branch_filter)s
543
+ WITH r, existing_av
544
+ ORDER BY r.branch_level DESC, r.from DESC, r.status ASC
545
+ LIMIT 1
546
+ WITH CASE
547
+ WHEN existing_av.value <> $values_by_id[n.uuid]
548
+ AND r.status = "active"
549
+ AND r.branch = $branch
550
+ THEN [r, existing_av]
551
+ ELSE [NULL, NULL]
552
+ END AS existing_details
553
+ RETURN existing_details[0] AS existing_has_value, existing_details[1] AS existing_av
554
+ }
555
+ CALL (existing_has_value) {
556
+ WITH existing_has_value
557
+ WHERE existing_has_value IS NOT NULL
558
+ SET existing_has_value.to = $at
559
+ }
560
+ """ % {"branch_filter": branch_filter}
561
+ self.add_to_query(update_value_query)
562
+
563
+ if self.is_large_type_attribute:
564
+ # use the index we created at the start to get the database ID of the AttributeValue vertex
565
+ # and then link the Attribute to the AttributeValue
566
+ set_value_query = """
567
+ // ------------
568
+ // only make updates if the existing value is not the same as the new value
569
+ // ------------
570
+ WITH attr, existing_av, value_id_pairs, $values_by_id[n.uuid] AS required_value
571
+ WHERE existing_av.value <> required_value
572
+ OR existing_av IS NULL
573
+ WITH attr, value_id_pairs, required_value,
574
+ reduce(av_vertex_id = NULL, pair IN value_id_pairs |
575
+ CASE
576
+ WHEN av_vertex_id IS NOT NULL THEN av_vertex_id
577
+ WHEN pair[0] = required_value THEN pair[1]
578
+ ELSE av_vertex_id
579
+ END
580
+ ) AS av_vertex_id
581
+ MATCH (av:AttributeValue)
582
+ WHERE elementId(av) = av_vertex_id
583
+ CREATE (attr)-[r:HAS_VALUE { branch: $branch, branch_level: $branch_level, status: "active", from: $at }]->(av)
584
+ """
585
+ else:
586
+ # if not a large-type attribute, then we can just use the regular MERGE clause
587
+ # that makes use of the index on AttributeValueIndexed
588
+ set_value_query = """
589
+ // ------------
590
+ // only make updates if the existing value is not the same as the new value
591
+ // ------------
592
+ WITH n, attr, existing_av, value_id_pairs, $values_by_id[n.uuid] AS required_value
593
+ WHERE existing_av.value <> required_value
594
+ OR existing_av IS NULL
595
+ CALL (n, attr) {
596
+ MERGE (av:AttributeValue&AttributeValueIndexed {is_default: false, value: $values_by_id[n.uuid]} )
597
+ WITH av, attr
598
+ LIMIT 1
599
+ CREATE (attr)-[r:HAS_VALUE { branch: $branch, branch_level: $branch_level, status: "active", from: $at }]->(av)
600
+ }
601
+ """
602
+ self.add_to_query(set_value_query)
603
+
604
+
605
+ class Migration044(MigrationRequiringRebase):
606
+ """
607
+ Backfill `human_friendly_id` and `display_label` attributes for nodes with schemas that define them.
608
+ """
609
+
610
+ name: str = "044_backfill_hfid_display_label_in_db"
611
+ minimum_version: int = 43
612
+ update_batch_size: int = 1000
613
+ # skip these b/c the attributes on these schema-related nodes are used to define the values included in
614
+ # the human_friendly_id and display_label attributes on instances of these schema, so should not be updated
615
+ kinds_to_skip: list[str] = ["SchemaNode", "SchemaAttribute", "SchemaRelationship", "SchemaGeneric"]
616
+
617
+ async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult: # noqa: ARG002
618
+ return MigrationResult()
619
+
620
+ async def _do_one_schema_all(
621
+ self,
622
+ db: InfrahubDatabase,
623
+ branch: Branch,
624
+ schema: NodeSchema,
625
+ schema_branch: SchemaBranch,
626
+ attribute_schema_map: dict[AttributeSchema, AttributeSchema],
627
+ progress: Progress | None = None,
628
+ update_task: TaskID | None = None,
629
+ ) -> None:
630
+ print(f"Processing {schema.kind}...", end="")
631
+
632
+ schema_paths_by_name: dict[str, list[SchemaAttributePath]] = {}
633
+ for source_attribute_schema in attribute_schema_map.keys():
634
+ node_schema_property = getattr(schema, source_attribute_schema.name)
635
+ if not node_schema_property:
636
+ continue
637
+ if isinstance(node_schema_property, list):
638
+ schema_paths_by_name[source_attribute_schema.name] = [
639
+ schema.parse_schema_path(path=str(path), schema=schema_branch) for path in node_schema_property
640
+ ]
641
+ else:
642
+ schema_paths_by_name[source_attribute_schema.name] = [
643
+ schema.parse_schema_path(path=str(node_schema_property), schema=schema_branch)
644
+ ]
645
+ all_schema_paths = list(chain(*schema_paths_by_name.values()))
646
+ offset = 0
647
+
648
+ # loop until we get no results from the get_details_query
649
+ while True:
650
+ if branch.is_default:
651
+ get_details_query: GetResultMapQuery = await GetPathDetailsDefaultBranch.init(
652
+ db=db,
653
+ schema_kind=schema.kind,
654
+ schema_paths=all_schema_paths,
655
+ offset=offset,
656
+ limit=self.update_batch_size,
657
+ )
658
+ else:
659
+ get_details_query = await GetPathDetailsBranchQuery.init(
660
+ db=db,
661
+ branch=branch,
662
+ schema_kind=schema.kind,
663
+ schema_paths=all_schema_paths,
664
+ updates_only=False,
665
+ offset=offset,
666
+ limit=self.update_batch_size,
667
+ )
668
+ await get_details_query.execute(db=db)
669
+
670
+ num_updates = 0
671
+ for source_attribute_schema, destination_attribute_schema in attribute_schema_map.items():
672
+ schema_paths = schema_paths_by_name[source_attribute_schema.name]
673
+ schema_path_values_map = get_details_query.get_result_map(schema_paths)
674
+ num_updates = max(num_updates, len(schema_path_values_map))
675
+ formatted_schema_path_values_map = {}
676
+ for k, v in schema_path_values_map.items():
677
+ if not v:
678
+ continue
679
+ if destination_attribute_schema.kind == "List":
680
+ formatted_schema_path_values_map[k] = ujson.dumps(v)
681
+ else:
682
+ formatted_schema_path_values_map[k] = " ".join(item for item in v if item is not None)
683
+
684
+ if not formatted_schema_path_values_map:
685
+ continue
686
+
687
+ update_display_label_query = await UpdateAttributeValuesQuery.init(
688
+ db=db,
689
+ branch=branch,
690
+ attribute_schema=destination_attribute_schema,
691
+ values_by_id_map=formatted_schema_path_values_map,
692
+ )
693
+ await update_display_label_query.execute(db=db)
694
+
695
+ if progress is not None and update_task is not None:
696
+ progress.update(update_task, advance=num_updates)
697
+
698
+ if num_updates == 0:
699
+ break
700
+
701
+ offset += self.update_batch_size
702
+
703
+ print("done")
704
+
705
+ async def execute(self, db: InfrahubDatabase) -> MigrationResult:
706
+ root_node = await get_root_node(db=db, initialize=False)
707
+ default_branch_name = root_node.default_branch
708
+ default_branch = await Branch.get_by_name(db=db, name=default_branch_name)
709
+
710
+ main_schema_branch = await get_or_load_schema_branch(db=db, branch=default_branch)
711
+
712
+ total_nodes_query = await DefaultBranchNodeCount.init(db=db, kinds_to_skip=self.kinds_to_skip)
713
+ await total_nodes_query.execute(db=db)
714
+ total_nodes_count = total_nodes_query.get_num_nodes()
715
+
716
+ base_node_schema = main_schema_branch.get("SchemaNode", duplicate=False)
717
+ display_label_attribute_schema = base_node_schema.get_attribute("display_label")
718
+ display_labels_attribute_schema = base_node_schema.get_attribute("display_labels")
719
+ hfid_attribute_schema = base_node_schema.get_attribute("human_friendly_id")
720
+
721
+ try:
722
+ with Progress(console=console) as progress:
723
+ update_task = progress.add_task(
724
+ f"Set display_label and human_friendly_id for {total_nodes_count} nodes on default branch",
725
+ total=total_nodes_count,
726
+ )
727
+ for node_schema_name in main_schema_branch.node_names:
728
+ if node_schema_name in self.kinds_to_skip:
729
+ continue
730
+
731
+ node_schema = main_schema_branch.get_node(name=node_schema_name, duplicate=False)
732
+ attribute_schema_map = {}
733
+ if node_schema.display_labels:
734
+ attribute_schema_map[display_labels_attribute_schema] = display_label_attribute_schema
735
+ if node_schema.human_friendly_id:
736
+ attribute_schema_map[hfid_attribute_schema] = hfid_attribute_schema
737
+ if not attribute_schema_map:
738
+ continue
739
+
740
+ await self._do_one_schema_all(
741
+ db=db,
742
+ branch=default_branch,
743
+ schema=node_schema,
744
+ schema_branch=main_schema_branch,
745
+ attribute_schema_map=attribute_schema_map,
746
+ progress=progress,
747
+ update_task=update_task,
748
+ )
749
+
750
+ except Exception as exc:
751
+ return MigrationResult(errors=[str(exc)])
752
+ return MigrationResult()
753
+
754
+ async def _do_one_schema_branch(
755
+ self,
756
+ db: InfrahubDatabase,
757
+ branch: Branch,
758
+ schema: NodeSchema,
759
+ schema_branch: SchemaBranch,
760
+ source_attribute_schema: AttributeSchema,
761
+ destination_attribute_schema: AttributeSchema,
762
+ ) -> None:
763
+ print(f"Processing {schema.kind}.{destination_attribute_schema.name} for {branch.name}...", end="")
764
+
765
+ schema_property = getattr(schema, source_attribute_schema.name)
766
+ if isinstance(schema_property, list):
767
+ schema_paths = [
768
+ schema.parse_schema_path(path=str(path_part), schema=schema_branch) for path_part in schema_property
769
+ ]
770
+ else:
771
+ schema_paths = [schema.parse_schema_path(path=str(schema_property), schema=schema_branch)]
772
+
773
+ offset = 0
774
+
775
+ while True:
776
+ # loop until we get no results from the get_details_query
777
+ get_details_query = await GetPathDetailsBranchQuery.init(
778
+ db=db,
779
+ branch=branch,
780
+ schema_kind=schema.kind,
781
+ schema_paths=schema_paths,
782
+ offset=offset,
783
+ limit=self.update_batch_size,
784
+ )
785
+ await get_details_query.execute(db=db)
786
+
787
+ schema_path_values_map = get_details_query.get_result_map(schema_paths)
788
+ if not schema_path_values_map:
789
+ print("done")
790
+ break
791
+ formatted_schema_path_values_map = {}
792
+ for k, v in schema_path_values_map.items():
793
+ if not v:
794
+ continue
795
+ if destination_attribute_schema.kind == "List":
796
+ formatted_v = ujson.dumps(v)
797
+ else:
798
+ formatted_v = " ".join(item for item in v if item is not None)
799
+ formatted_schema_path_values_map[k] = formatted_v
800
+
801
+ update_attr_values_query = await UpdateAttributeValuesQuery.init(
802
+ db=db,
803
+ branch=branch,
804
+ attribute_schema=destination_attribute_schema,
805
+ values_by_id_map=formatted_schema_path_values_map,
806
+ )
807
+ await update_attr_values_query.execute(db=db)
808
+
809
+ offset += self.update_batch_size
810
+
811
+ async def execute_against_branch(self, db: InfrahubDatabase, branch: Branch) -> MigrationResult:
812
+ default_branch = await Branch.get_by_name(db=db, name=registry.default_branch)
813
+ main_schema_branch = await get_or_load_schema_branch(db=db, branch=default_branch)
814
+ schema_branch = await get_or_load_schema_branch(db=db, branch=branch)
815
+
816
+ base_node_schema = schema_branch.get("SchemaNode", duplicate=False)
817
+ display_label_attribute_schema = base_node_schema.get_attribute("display_label")
818
+ display_labels_attribute_schema = base_node_schema.get_attribute("display_labels")
819
+ hfid_attribute_schema = base_node_schema.get_attribute("human_friendly_id")
820
+
821
+ try:
822
+ for node_schema_name in schema_branch.node_names:
823
+ if node_schema_name in self.kinds_to_skip:
824
+ continue
825
+
826
+ node_schema = schema_branch.get_node(name=node_schema_name, duplicate=False)
827
+ default_node_schema = main_schema_branch.get_node(name=node_schema_name, duplicate=False)
828
+ schemas_for_universal_update_map = {}
829
+ schemas_for_targeted_update_map = {}
830
+ if default_node_schema.display_label != node_schema.display_label:
831
+ schemas_for_universal_update_map[display_labels_attribute_schema] = display_label_attribute_schema
832
+ elif node_schema.display_labels:
833
+ schemas_for_targeted_update_map[display_labels_attribute_schema] = display_label_attribute_schema
834
+
835
+ if default_node_schema.human_friendly_id != node_schema.human_friendly_id:
836
+ schemas_for_universal_update_map[hfid_attribute_schema] = hfid_attribute_schema
837
+ elif node_schema.human_friendly_id:
838
+ schemas_for_targeted_update_map[hfid_attribute_schema] = hfid_attribute_schema
839
+
840
+ if schemas_for_universal_update_map:
841
+ await self._do_one_schema_all(
842
+ db=db,
843
+ branch=branch,
844
+ schema=node_schema,
845
+ schema_branch=schema_branch,
846
+ attribute_schema_map=schemas_for_universal_update_map,
847
+ )
848
+
849
+ if not schemas_for_targeted_update_map:
850
+ continue
851
+
852
+ for source_attribute_schema, destination_attribute_schema in schemas_for_targeted_update_map.items():
853
+ await self._do_one_schema_branch(
854
+ db=db,
855
+ branch=branch,
856
+ schema=node_schema,
857
+ schema_branch=schema_branch,
858
+ source_attribute_schema=source_attribute_schema,
859
+ destination_attribute_schema=destination_attribute_schema,
860
+ )
861
+
862
+ except Exception as exc:
863
+ return MigrationResult(errors=[str(exc)])
864
+ return MigrationResult()