Ensure kind matches when retrieving objects from the store

Author: dgarros

infrahub_sdk/exceptions.py

@@ -88,6 +88,10 @@ def __str__(self) -> str:
         """
 
 
+class NodeInvalidError(NodeNotFoundError):
+    pass
+
+
 class ResourceNotDefinedError(Error):
     """Raised when trying to access a resource that hasn't been defined."""
 

infrahub_sdk/node.py

@@ -814,6 +814,11 @@ def __repr__(self) -> str:
     def get_kind(self) -> str:
         return self._schema.kind
 
+    def get_all_kinds(self) -> list[str]:
+        if hasattr(self._schema, "inherit_from"):
+            return [self._schema.kind] + self._schema.inherit_from
+        return [self._schema.kind]
+
     def is_ip_prefix(self) -> bool:
         builtin_ipprefix_kind = "BuiltinIPPrefix"
         return self.get_kind() == builtin_ipprefix_kind or builtin_ipprefix_kind in self._schema.inherit_from  # type: ignore[union-attr]

infrahub_sdk/protocols_base.py

@@ -160,6 +160,8 @@ def get_human_friendly_id_as_string(self, include_kind: bool = False) -> str | N
 
     def get_kind(self) -> str: ...
 
+    def get_all_kinds(self) -> list[str]: ...
+
     def get_branch(self) -> str: ...
 
     def is_ip_prefix(self) -> bool: ...

infrahub_sdk/store.py

@@ -4,7 +4,7 @@
 from collections import defaultdict
 from typing import TYPE_CHECKING, Literal, overload
 
-from .exceptions import NodeNotFoundError
+from .exceptions import NodeInvalidError, NodeNotFoundError
 from .node import parse_human_friendly_id
 
 if TYPE_CHECKING:
@@ -53,62 +53,110 @@ def get(  # type: ignore[no-untyped-def]
         kind: type[SchemaType | SchemaTypeSync] | str | None = None,
         raise_when_missing: bool = True,
     ) -> InfrahubNode | InfrahubNodeSync | CoreNode | CoreNodeSync | None:
+        found_invalid = False
+
+        kind_name = get_schema_name(schema=kind)
+
         try:
-            return self._get_by_internal_id(key)
+            return self._get_by_internal_id(key, kind=kind_name)
+        except NodeInvalidError:
+            found_invalid = True
         except NodeNotFoundError:
             pass
 
         try:
-            return self._get_by_id(key)
+            return self._get_by_id(key, kind=kind_name)
+        except NodeInvalidError:
+            found_invalid = True
         except NodeNotFoundError:
             pass
 
         try:
-            return self._get_by_key(key)
+            return self._get_by_key(key, kind=kind_name)
+        except NodeInvalidError:
+            found_invalid = True
         except NodeNotFoundError:
             pass
 
-        kind_name = get_schema_name(schema=kind)
-
         try:
             return self._get_by_hfid(key, kind=kind_name)
         except NodeNotFoundError:
             pass
 
         if not raise_when_missing:
             return None
+
+        if kind and found_invalid:
+            raise NodeInvalidError(
+                node_type="n/a",
+                identifier={"key": [key]},
+                message=f"Found a node of a differentkind instead of {kind} for key {key!r} in the store ({self.branch_name})",
+            )
+
         raise NodeNotFoundError(
             node_type="n/a",
             identifier={"key": [key]},
             message=f"Unable to find the node {key!r} in the store ({self.branch_name})",
         )
 
-    def _get_by_internal_id(self, internal_id: str) -> InfrahubNode | InfrahubNodeSync | CoreNode | CoreNodeSync:
+    def _get_by_internal_id(
+        self, internal_id: str, kind: str | None = None
+    ) -> InfrahubNode | InfrahubNodeSync | CoreNode | CoreNodeSync:
         if internal_id not in self._objs:
             raise NodeNotFoundError(
                 node_type="n/a",
                 identifier={"internal_id": [internal_id]},
                 message=f"Unable to find the node {internal_id!r} in the store ({self.branch_name})",
             )
-        return self._objs[internal_id]
 
-    def _get_by_key(self, key: str) -> InfrahubNode | InfrahubNodeSync | CoreNode | CoreNodeSync:
+        node = self._objs[internal_id]
+        if kind and kind not in node.get_all_kinds():
+            raise NodeInvalidError(
+                node_type=kind,
+                identifier={"internal_id": [internal_id]},
+                message=f"Found a node of kind {node.get_kind()} instead of {kind} for internal_id {internal_id!r} in the store ({self.branch_name})",
+            )
+
+        return node
+
+    def _get_by_key(
+        self, key: str, kind: str | None = None
+    ) -> InfrahubNode | InfrahubNodeSync | CoreNode | CoreNodeSync:
         if key not in self._keys:
             raise NodeNotFoundError(
                 node_type="n/a",
                 identifier={"key": [key]},
                 message=f"Unable to find the node {key!r} in the store ({self.branch_name})",
             )
-        return self._get_by_internal_id(self._keys[key])
 
-    def _get_by_id(self, id: str) -> InfrahubNode | InfrahubNodeSync | CoreNode | CoreNodeSync:
+        node = self._get_by_internal_id(self._keys[key])
+
+        if kind and node.get_kind() != kind:
+            raise NodeInvalidError(
+                node_type=kind,
+                identifier={"key": [key]},
+                message=f"Found a node of kind {node.get_kind()} instead of {kind} for key {key!r} in the store ({self.branch_name})",
+            )
+
+        return node
+
+    def _get_by_id(self, id: str, kind: str | None = None) -> InfrahubNode | InfrahubNodeSync | CoreNode | CoreNodeSync:
         if id not in self._uuids:
             raise NodeNotFoundError(
                 node_type="n/a",
                 identifier={"id": [id]},
                 message=f"Unable to find the node {id!r} in the store ({self.branch_name})",
             )
-        return self._get_by_internal_id(self._uuids[id])
+
+        node = self._get_by_internal_id(self._uuids[id])
+        if kind and kind not in node.get_all_kinds():
+            raise NodeInvalidError(
+                node_type=kind,
+                identifier={"id": [id]},
+                message=f"Found a node of kind {node.get_kind()} instead of {kind} for id {id!r} in the store ({self.branch_name})",
+            )
+
+        return node
 
     def _get_by_hfid(
         self, hfid: str, kind: str | None = None

tests/unit/sdk/test_store.py

@@ -1,6 +1,6 @@
 import pytest
 
-from infrahub_sdk.exceptions import NodeNotFoundError
+from infrahub_sdk.exceptions import NodeInvalidError, NodeNotFoundError
 from infrahub_sdk.node import InfrahubNode, InfrahubNodeSync
 from infrahub_sdk.store import NodeStore, NodeStoreSync
 
@@ -91,6 +91,8 @@ def test_node_store_get(client_type, clients, location_schema):
     assert store.get(kind="BuiltinLocation", key="mykey").id == node.id
     assert store.get(key="mykey").id == node.id
 
+    assert store.get(kind="BuiltinTest", key="mykey", raise_when_missing=False) is None
+
     assert store.get(kind="BuiltinLocation", key="anotherkey", raise_when_missing=False) is None
     assert store.get(key="anotherkey", raise_when_missing=False) is None
 
@@ -101,6 +103,9 @@ def test_node_store_get(client_type, clients, location_schema):
     with pytest.raises(NodeNotFoundError):
         store.get(key="mykey", branch="mybranch")
 
+    with pytest.raises(NodeInvalidError):
+        store.get(kind="BuiltinTest", key="mykey")
+
     store.set(key="mykey", node=node, branch="mybranch")
     assert store.get(key="mykey", branch="mybranch").id == node.id