infrahub-server 1.6.0b0__py3-none-any.whl → 1.6.1__py3-none-any.whl

infrahub/core/node/__init__.py

@@ -170,7 +170,7 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
         return self._human_friendly_id is not None
 
     async def add_human_friendly_id(self, db: InfrahubDatabase) -> None:
-        if not self._schema.human_friendly_id or self._human_friendly_id:
+        if self._human_friendly_id:
             return
 
         self._human_friendly_id = HumanFriendlyIdentifier(
@@ -179,11 +179,8 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
         await self._human_friendly_id.compute(db=db, node=self)
 
     async def get_display_label(self, db: InfrahubDatabase) -> str:
-        if self._display_label:
-            if isinstance(self._display_label._value, str):
-                return self._display_label._value
-            if self._display_label._value:
-                return self._display_label._value.value
+        if self._display_label and (value := self._display_label.get_value(node=self, at=self._at)):
+            return value
 
         return await self.render_display_label(db=db)
 
@@ -191,7 +188,7 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
         return self._display_label is not None
 
     async def add_display_label(self, db: InfrahubDatabase) -> None:
-        if not self._schema.display_label or self._display_label:
+        if self._display_label:
             return
 
         self._display_label = DisplayLabel(node_schema=self._schema, template=self._schema.display_label)
@@ -467,15 +464,21 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
         for attribute_name in template._attributes:
             if attribute_name in list(fields) + [OBJECT_TEMPLATE_NAME_ATTR]:
                 continue
-            attr_value = getattr(template, attribute_name).value
+            attr = getattr(template, attribute_name)
+            attr_value = attr.value
             if attr_value is not None:
-                fields[attribute_name] = {"value": attr_value, "source": template.id}
+                # Preserve is_from_profile flag when copying from template
+                field_data = {"value": attr_value, "source": attr.source_id or template.id}
+                if attr.is_from_profile:
+                    field_data["is_from_profile"] = True
+                fields[attribute_name] = field_data
 
         for relationship_name in template._relationships:
             relationship_schema = template._schema.get_relationship(name=relationship_name)
             if (
                 relationship_name in list(fields)
-                or relationship_schema.kind not in [RelationshipKind.ATTRIBUTE, RelationshipKind.GENERIC]
+                or relationship_schema.kind
+                not in [RelationshipKind.ATTRIBUTE, RelationshipKind.GENERIC, RelationshipKind.PROFILE]
                 or relationship_name == OBJECT_TEMPLATE_RELATIONSHIP_NAME
             ):
                 continue
@@ -570,7 +573,7 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
                     self,
                     rel_schema.name,
                     await generator_method(
-                        db=db, name=rel_schema.name, schema=rel_schema, data=fields.get(rel_schema.name, None)
+                        db=db, name=rel_schema.name, schema=rel_schema, data=fields.get(rel_schema.name)
                     ),
                 )
             except ValidationError as exc:
@@ -600,7 +603,7 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
                     self,
                     attr_schema.name,
                     await generator_method(
-                        db=db, name=attr_schema.name, schema=attr_schema, data=fields.get(attr_schema.name, None)
+                        db=db, name=attr_schema.name, schema=attr_schema, data=fields.get(attr_schema.name)
                     ),
                 )
                 if not self._existing:
@@ -1088,7 +1091,7 @@ class Node(BaseNode, metaclass=BaseNodeMeta):
 
             if key in self._relationships:
                 rel: RelationshipManager = getattr(self, key)
-                changed |= await rel.update(db=db, data=value)
+                changed |= await rel.update(db=db, data=value, process_delete=process_pools)
 
         return changed
 

infrahub/core/node/create.py

@@ -58,18 +58,34 @@ async def extract_peer_data(
             except ValueError:
                 pass
 
-        obj_peer_data[attr_name] = {"value": template_attr.value, "source": template_peer.id}
+        # If the template attribute comes from a profile, preserve the profile as the source
+        # Otherwise, use the template itself as the source
+        source_id = template_attr.source_id or template_peer.id
+        attr_data = {"value": template_attr.value, "source": source_id}
+        if template_attr.is_from_profile:
+            attr_data["is_from_profile"] = True
+        obj_peer_data[attr_name] = attr_data
 
     for rel in template_peer.get_schema().relationship_names:
         rel_manager: RelationshipManager = getattr(template_peer, rel)
-
-        if rel_manager.schema.name not in obj_peer_schema.relationship_names:
+        if (
+            rel_manager.schema.kind
+            not in [
+                RelationshipKind.COMPONENT,
+                RelationshipKind.PARENT,
+                RelationshipKind.PROFILE,
+                RelationshipKind.ATTRIBUTE,
+            ]
+            or rel_manager.schema.name not in obj_peer_schema.relationship_names
+        ):
             continue
 
         peers_map = await rel_manager.get_peers(db=db)
-        if rel_manager.schema.kind in [RelationshipKind.COMPONENT, RelationshipKind.PARENT] and list(
-            peers_map.keys()
-        ) == [current_template.id]:
+        if rel_manager.schema.kind in [
+            RelationshipKind.COMPONENT,
+            RelationshipKind.PARENT,
+            RelationshipKind.PROFILE,
+        ] and list(peers_map.keys()) == [current_template.id]:
             obj_peer_data[rel] = {"id": parent_obj.id}
             continue
 
@@ -80,7 +96,13 @@ async def extract_peer_data(
                 continue
             rel_peer_ids.append({"id": peer_id})
 
-        obj_peer_data[rel] = rel_peer_ids
+        # Only set the relationship data if there are actual peers to set
+        if rel_peer_ids:
+            obj_peer_data[rel] = rel_peer_ids
+
+        if rel_manager.schema.kind == RelationshipKind.PROFILE:
+            profiles = list(await rel_manager.get_peers(db=db))
+            obj_peer_data[rel] = profiles
 
     return obj_peer_data
 
@@ -125,6 +147,12 @@ async def handle_template_relationships(
             await constraint_runner.check(node=obj_peer, field_filters=list(obj_peer_data))
             await obj_peer.save(db=db)
 
+            template_profile_ids = await get_profile_ids(db=db, obj=template_relationship_peer)
+            if template_profile_ids:
+                node_profiles_applier = NodeProfilesApplier(db=db, branch=branch)
+                await node_profiles_applier.apply_profiles(node=obj_peer)
+                await obj_peer.save(db=db)
+
             await handle_template_relationships(
                 db=db,
                 branch=branch,
@@ -136,7 +164,7 @@ async def handle_template_relationships(
             )
 
 
-async def get_profile_ids(db: InfrahubDatabase, obj: Node) -> set[str]:
+async def get_profile_ids(db: InfrahubDatabase, obj: Node | CoreObjectTemplate) -> set[str]:
     if not hasattr(obj, "profiles"):
         return set()
     profile_rels = await obj.profiles.get_relationships(db=db)

infrahub/core/node/proposed_change.py

@@ -1,11 +1,13 @@
-from typing import cast
+from typing import TYPE_CHECKING, cast
 
 from infrahub.core.constants.infrahubkind import THREADCOMMENT
 from infrahub.core.manager import NodeManager
 from infrahub.core.node import Node
-from infrahub.core.protocols import CoreProposedChange as CoreProposedChangeProtocol
 from infrahub.database import InfrahubDatabase
 
+if TYPE_CHECKING:
+    from infrahub.core.protocols import CoreProposedChange as CoreProposedChangeProtocol
+
 
 class CoreProposedChange(Node):
     async def to_graphql(
@@ -29,7 +31,7 @@ class CoreProposedChange(Node):
         if fields:
             if "total_comments" in fields:
                 total_comments = 0
-                proposed_change = cast(CoreProposedChangeProtocol, self)
+                proposed_change = cast("CoreProposedChangeProtocol", self)
                 change_comments = await proposed_change.comments.get_relationships(db=db)
                 total_comments += len(change_comments)
 

infrahub/core/node/standard.py

@@ -72,7 +72,7 @@ class StandardNode(BaseModel):
         response: dict[str, Any] = {"id": self.uuid}
 
         for field_name in fields.keys():
-            if field_name in ["id"]:
+            if field_name == "id":
                 continue
             if field_name == "__typename":
                 response[field_name] = self.get_type()

infrahub/core/protocols.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Optional, Protocol, runtime_checkable
+from typing import TYPE_CHECKING
 
 from infrahub.core.protocols_base import CoreNode
 
@@ -16,22 +16,16 @@ if TYPE_CHECKING:
         DateTime,
         DateTimeOptional,
         Dropdown,
-        DropdownOptional,
         HashedPassword,
-        HashedPasswordOptional,
         Integer,
         IntegerOptional,
         IPHost,
-        IPHostOptional,
         IPNetwork,
-        IPNetworkOptional,
         JSONAttribute,
         JSONAttributeOptional,
-        ListAttribute,
         ListAttributeOptional,
         String,
         StringOptional,
-        URLOptional,
     )
     from infrahub.core.relationship import RelationshipManager
 

infrahub/core/query/attribute.py

@@ -373,7 +373,7 @@ async def default_attribute_query_filter(
         if property_name not in [v.value for v in NodeProperty]:
             raise ValueError(f"filter {filter_name}: {filter_value}, {property_name} is not a valid property")
 
-        if property_attr not in ["id"]:
+        if property_attr != "id":
             raise ValueError(f"filter {filter_name}: {filter_value}, {property_attr} is supported")
 
         clean_filter_name = f"{property_name}_{property_attr}"

infrahub/core/query/node.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 from collections import defaultdict
 from copy import copy
 from dataclasses import dataclass
@@ -12,6 +13,7 @@ from infrahub.core import registry
 from infrahub.core.constants import (
     GLOBAL_BRANCH_NAME,
     PROFILE_NODE_RELATIONSHIP_IDENTIFIER,
+    PROFILE_TEMPLATE_RELATIONSHIP_IDENTIFIER,
     AttributeDBNodeType,
     RelationshipDirection,
     RelationshipHierarchyDirection,
@@ -623,7 +625,8 @@ class NodeListGetAttributeQuery(Query):
 
     async def query_init(self, db: InfrahubDatabase, **kwargs) -> None:  # noqa: ARG002
         self.params["ids"] = self.ids
-        self.params["profile_relationship_name"] = PROFILE_NODE_RELATIONSHIP_IDENTIFIER
+        self.params["profile_node_relationship_name"] = PROFILE_NODE_RELATIONSHIP_IDENTIFIER
+        self.params["profile_template_relationship_name"] = PROFILE_TEMPLATE_RELATIONSHIP_IDENTIFIER
 
         branch_filter, branch_params = self.branch.get_query_filter_path(
             at=self.at, branch_agnostic=self.branch_agnostic
@@ -632,7 +635,10 @@ class NodeListGetAttributeQuery(Query):
 
         query = """
         MATCH (n:Node) WHERE n.uuid IN $ids
-        WITH n, exists((n)-[:IS_RELATED]-(:Relationship {name: $profile_relationship_name})) AS might_use_profile
+        WITH n, (
+            exists((n)-[:IS_RELATED]-(:Relationship {name: $profile_node_relationship_name})) OR
+            exists((n)-[:IS_RELATED]-(:Relationship {name: $profile_template_relationship_name}))
+        ) AS might_use_profile
         MATCH (n)-[:HAS_ATTRIBUTE]-(a:Attribute)
         """
         if self.fields:
@@ -1134,10 +1140,8 @@ class NodeGetListQuery(Query):
             self._variables_to_track.append(variable)
 
     def _untrack_variable(self, variable: str) -> None:
-        try:
+        with contextlib.suppress(ValueError):
             self._variables_to_track.remove(variable)
-        except ValueError:
-            ...
 
     def _get_tracked_variables(self) -> list[str]:
         return self._variables_to_track

infrahub/core/relationship/model.py

@@ -912,6 +912,8 @@ class RelationshipManager:
         db: InfrahubDatabase,
         peer_type: type[PeerType],
         branch_agnostic: bool = ...,
+        include_source: bool = ...,
+        include_owner: bool = ...,
     ) -> Mapping[str, PeerType]: ...
 
     @overload
@@ -920,6 +922,8 @@ class RelationshipManager:
         db: InfrahubDatabase,
         peer_type: None = None,
         branch_agnostic: bool = ...,
+        include_source: bool = ...,
+        include_owner: bool = ...,
     ) -> Mapping[str, Node]: ...
 
     async def get_peers(
@@ -927,11 +931,18 @@ class RelationshipManager:
         db: InfrahubDatabase,
         peer_type: type[PeerType] | None = None,  # noqa: ARG002
         branch_agnostic: bool = False,
+        include_source: bool = False,
+        include_owner: bool = False,
     ) -> Mapping[str, Node | PeerType]:
         rels = await self.get_relationships(db=db, branch_agnostic=branch_agnostic)
         peer_ids = [rel.peer_id for rel in rels if rel.peer_id]
         nodes = await registry.manager.get_many(
-            db=db, ids=peer_ids, branch=self.branch, branch_agnostic=branch_agnostic
+            db=db,
+            ids=peer_ids,
+            branch=self.branch,
+            branch_agnostic=branch_agnostic,
+            include_source=include_source,
+            include_owner=include_owner,
         )
         return nodes
 
@@ -1061,7 +1072,12 @@ class RelationshipManager:
 
         return self._relationships.as_list()
 
-    async def update(self, data: list[str | Node] | dict[str, Any] | str | Node | None, db: InfrahubDatabase) -> bool:
+    async def update(
+        self,
+        data: list[str | Node] | dict[str, Any] | str | Node | None,
+        db: InfrahubDatabase,
+        process_delete: bool = True,
+    ) -> bool:
         """Replace and Update the list of relationships with this one."""
         if not isinstance(data, list):
             list_data: Sequence[str | Node | dict[str, Any] | None] = [data]
@@ -1087,8 +1103,9 @@ class RelationshipManager:
 
             if item is None:
                 if previous_relationships:
-                    for rel in previous_relationships.values():
-                        await rel.delete(db=db)
+                    if process_delete:
+                        for rel in previous_relationships.values():
+                            await rel.delete(db=db)
                     changed = True
                 continue
 

infrahub/core/schema/generic_schema.py

@@ -51,4 +51,4 @@ class GenericSchema(GeneratedGenericSchema):
     def _get_field_names_for_diff(self) -> list[str]:
         """Exclude used_by from the diff for generic nodes"""
         fields = super()._get_field_names_for_diff()
-        return [field for field in fields if field not in ["used_by"]]
+        return [field for field in fields if field != "used_by"]

infrahub/core/schema/manager.py

@@ -774,10 +774,15 @@ class SchemaManager(NodeManager):
         """Return non active branches that were purged."""
 
         hashes_to_keep: set[str] = set()
+        branch_processed: set[str] = set()
         for active_branch in active_branches:
-            if branch := self._branches.get(active_branch):
-                nodes = branch.get_all(include_internal=True, duplicate=False)
-                hashes_to_keep.update([node.get_hash() for node in nodes.values()])
+            branch_hash = self._branch_hash_by_name.get(active_branch)
+            if not branch_hash or branch_hash not in branch_processed:
+                if branch_hash:
+                    branch_processed.add(branch_hash)
+                if branch := self._branches.get(active_branch):
+                    nodes = branch.get_all(include_internal=True, duplicate=False)
+                    hashes_to_keep.update([node.get_hash() for node in nodes.values()])
 
         removed_branches: list[str] = []
         for branch_name in list(self._branches.keys()):

infrahub/core/schema/schema_branch.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import copy
 import hashlib
 import keyword
@@ -19,6 +20,7 @@ from infrahub.core.constants import (
     OBJECT_TEMPLATE_NAME_ATTR,
     OBJECT_TEMPLATE_RELATIONSHIP_NAME,
     PROFILE_NODE_RELATIONSHIP_IDENTIFIER,
+    PROFILE_TEMPLATE_RELATIONSHIP_IDENTIFIER,
     RESERVED_ATTR_GEN_NAMES,
     RESERVED_ATTR_REL_NAMES,
     RESTRICTED_NAMESPACES,
@@ -72,6 +74,16 @@ from .schema_branch_hfid import HFIDs
 log = get_logger()
 
 
+profiles_rel_settings: dict[str, Any] = {
+    "name": "profiles",
+    "identifier": PROFILE_NODE_RELATIONSHIP_IDENTIFIER,
+    "peer": InfrahubKind.PROFILE,
+    "kind": RelationshipKind.PROFILE,
+    "cardinality": RelationshipCardinality.MANY,
+    "branch": BranchSupportType.AWARE,
+}
+
+
 class SchemaBranch:
     def __init__(
         self,
@@ -334,10 +346,8 @@ class SchemaBranch:
             )
 
         schema: MainSchemaTypes | None = None
-        try:
+        with contextlib.suppress(KeyError):
             schema = self._cache[key]
-        except KeyError:
-            pass
 
         if not schema:
             raise ValueError(f"Schema {name!r} on branch {self.name} has incorrect hash: {key!r}")
@@ -1116,7 +1126,7 @@ class SchemaBranch:
                     ) from None
 
             for rel in node.relationships:
-                if rel.peer in [InfrahubKind.GENERICGROUP]:
+                if rel.peer == InfrahubKind.GENERICGROUP:
                     continue
                 if not self.has(rel.peer) or self.get(rel.peer, duplicate=False).state == HashableModelState.ABSENT:
                     raise ValueError(
@@ -2163,10 +2173,8 @@ class SchemaBranch:
                 or not node.generate_profile
                 or node.state == HashableModelState.ABSENT
             ):
-                try:
+                with contextlib.suppress(SchemaNotFoundError):
                     self.delete(name=self._get_profile_kind(node_kind=node.kind))
-                except SchemaNotFoundError:
-                    ...
                 continue
 
             profile = self.generate_profile_from_node(node=node)
@@ -2212,15 +2220,6 @@ class SchemaBranch:
             ):
                 continue
 
-            profiles_rel_settings: dict[str, Any] = {
-                "name": "profiles",
-                "identifier": PROFILE_NODE_RELATIONSHIP_IDENTIFIER,
-                "peer": InfrahubKind.PROFILE,
-                "kind": RelationshipKind.PROFILE,
-                "cardinality": RelationshipCardinality.MANY,
-                "branch": BranchSupportType.AWARE,
-            }
-
             # Add relationship between node and profile
             if "profiles" not in node.relationship_names:
                 node_schema = self.get(name=node_name, duplicate=True)
@@ -2285,6 +2284,18 @@ class SchemaBranch:
                 )
             ],
         )
+        if f"Template{node.kind}" in self.all_names:
+            template = self.get(name=f"Template{node.kind}", duplicate=False)
+            profile.relationships.append(
+                RelationshipSchema(
+                    name="related_templates",
+                    identifier=PROFILE_TEMPLATE_RELATIONSHIP_IDENTIFIER,
+                    peer=template.kind,
+                    kind=RelationshipKind.PROFILE,
+                    cardinality=RelationshipCardinality.MANY,
+                    branch=BranchSupportType.AWARE,
+                )
+            )
 
         for node_attr in node.attributes:
             if not node_attr.support_profiles:
@@ -2415,6 +2426,14 @@ class SchemaBranch:
                 template_schema.human_friendly_id = [parent_hfid] + template_schema.human_friendly_id
                 template_schema.uniqueness_constraints[0].append(relationship.name)
 
+        if getattr(node, "generate_profile", False):
+            if "profiles" not in [r.name for r in template_schema.relationships]:
+                settings = dict(profiles_rel_settings)
+                settings["identifier"] = PROFILE_TEMPLATE_RELATIONSHIP_IDENTIFIER
+                template_schema.relationships.append(RelationshipSchema(**settings))
+
+        self.set(name=template_schema.kind, schema=template_schema)
+
     def generate_object_template_from_node(
         self, node: NodeSchema | GenericSchema, need_templates: set[NodeSchema | GenericSchema]
     ) -> TemplateSchema | GenericSchema:

infrahub/core/validators/attribute/choices.py

@@ -4,13 +4,13 @@ from typing import TYPE_CHECKING, Any, cast
 
 from infrahub.core.constants import NULL_VALUE, PathType
 from infrahub.core.path import DataPath, GroupedDataPaths
-from infrahub.core.schema.generic_schema import GenericSchema
 
 from ..interface import ConstraintCheckerInterface
 from ..shared import AttributeSchemaValidatorQuery
 
 if TYPE_CHECKING:
     from infrahub.core.branch import Branch
+    from infrahub.core.schema.generic_schema import GenericSchema
     from infrahub.database import InfrahubDatabase
 
     from ..model import SchemaConstraintValidatorRequest
@@ -106,7 +106,7 @@ class AttributeChoicesChecker(ConstraintCheckerInterface):
         # skip inheriting schemas that override the attribute being checked
         excluded_kinds: list[str] = []
         if request.node_schema.is_generic_schema:
-            request.node_schema = cast(GenericSchema, request.node_schema)
+            request.node_schema = cast("GenericSchema", request.node_schema)
             for inheriting_kind in request.node_schema.used_by:
                 inheriting_schema = request.schema_branch.get_node(name=inheriting_kind, duplicate=False)
                 inheriting_schema_attribute = inheriting_schema.get_attribute(name=request.schema_path.field_name)

infrahub/core/validators/determiner.py

@@ -1,3 +1,4 @@
+import contextlib
 from typing import TYPE_CHECKING, Any
 
 from infrahub.core.constants import RelationshipKind, SchemaPathType
@@ -84,14 +85,10 @@ class ConstraintValidatorDeterminer:
         constraints: list[SchemaUpdateConstraintInfo] = []
         schemas = list(self.schema_branch.get_all(duplicate=False).values())
         # added here to check their uniqueness constraints
-        try:
+        with contextlib.suppress(SchemaNotFoundError):
             schemas.append(self.schema_branch.get_node(name="SchemaAttribute", duplicate=False))
-        except SchemaNotFoundError:
-            pass
-        try:
+        with contextlib.suppress(SchemaNotFoundError):
             schemas.append(self.schema_branch.get_node(name="SchemaRelationship", duplicate=False))
-        except SchemaNotFoundError:
-            pass
         for schema in schemas:
             constraints.extend(await self._get_property_constraints_for_one_schema(schema=schema))
         return constraints

infrahub/database/__init__.py

@@ -356,7 +356,7 @@ class InfrahubDatabase:
                     type
                     and type == QueryType.READ
                     and runtime not in [Neo4jRuntime.DEFAULT, Neo4jRuntime.UNDEFINED]
-                    and not (self.is_transaction and runtime in [Neo4jRuntime.PARALLEL])
+                    and not (self.is_transaction and runtime == Neo4jRuntime.PARALLEL)
                 ):
                     query = f"CYPHER runtime = {runtime.value}\n" + query
                 else:

infrahub/git/base.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import contextlib
 import shutil
 from abc import ABC, abstractmethod
 from pathlib import Path
@@ -748,10 +749,8 @@ class InfrahubRepositoryBase(BaseModel, ABC):
         for short_name, branch_data in branches.items():
             branch = None
 
-            try:
+            with contextlib.suppress(BranchNotFoundError):
                 branch = registry.get_branch_from_registry(branch=short_name)
-            except BranchNotFoundError:
-                ...
 
             branch_exists_import_sync_condition = branch and (
                 branch.name not in {registry.default_branch, self.default_branch}

infrahub/git/models.py

@@ -1,6 +1,8 @@
 from pydantic import BaseModel, ConfigDict, Field
 
 from infrahub.context import InfrahubContext
+from infrahub.core.node import Node
+from infrahub.core.protocols import CoreReadOnlyRepository, CoreRepository
 from infrahub.message_bus.types import ProposedChangeBranchDiff
 
 
@@ -201,11 +203,22 @@ class RepositoryBranchInfo(BaseModel):
 
 
 class RepositoryData(BaseModel):
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
     repository_id: str = Field(..., description="Id of the repository")
     repository_name: str = Field(..., description="Name of the repository")
+    repository: CoreRepository | CoreReadOnlyRepository | Node = Field(
+        ..., description="InfrahubNode representing a Repository"
+    )
     branches: dict[str, str] = Field(
         ...,
         description="Dictionary with the name of the branch as the key and the active commit id as the value",
     )
 
     branch_info: dict[str, RepositoryBranchInfo] = Field(default_factory=dict)
+
+    def get_staging_branch(self) -> str | None:
+        for branch, info in self.branch_info.items():
+            if info.internal_status == "staging":
+                return branch
+        return None