Remove shortuuid and default branch value

Author: dgarros

infrahub_sdk/client.py

@@ -281,7 +281,7 @@ def _initialize(self) -> None:
         self.schema = InfrahubSchema(self)
         self.branch = InfrahubBranchManager(self)
         self.object_store = ObjectStore(self)
-        self.store = NodeStore()
+        self.store = NodeStore(default_branch=self.default_branch)
         self.task = InfrahubTaskManager(self)
         self.concurrent_execution_limit = asyncio.Semaphore(self.max_concurrent_execution)
         self._request_method: AsyncRequester = self.config.requester or self._default_request_method
@@ -1529,7 +1529,7 @@ def _initialize(self) -> None:
         self.schema = InfrahubSchemaSync(self)
         self.branch = InfrahubBranchManagerSync(self)
         self.object_store = ObjectStoreSync(self)
-        self.store = NodeStoreSync()
+        self.store = NodeStoreSync(default_branch=self.default_branch)
         self.task = InfrahubTaskManagerSync(self)
         self._request_method: SyncRequester = self.config.sync_requester or self._default_request_method
         self.group_context = InfrahubGroupContextSync(self)

infrahub_sdk/exceptions.py

@@ -69,12 +69,12 @@ def __init__(self, message: str | None = None):
 class NodeNotFoundError(Error):
     def __init__(
         self,
-        node_type: str,
         identifier: Mapping[str, list[str]],
         message: str = "Unable to find the node in the database.",
         branch_name: str | None = None,
+        node_type: str | None = None,
     ):
-        self.node_type = node_type
+        self.node_type = node_type or "unknown"
         self.identifier = identifier
         self.branch_name = branch_name
 

infrahub_sdk/node.py

@@ -6,8 +6,6 @@
 from copy import copy
 from typing import TYPE_CHECKING, Any, Callable, Union, get_args
 
-import shortuuid
-
 from .constants import InfrahubClientMode
 from .exceptions import (
     Error,
@@ -17,7 +15,7 @@
 )
 from .graphql import Mutation, Query
 from .schema import GenericSchemaAPI, RelationshipCardinality, RelationshipKind
-from .utils import compare_lists, get_flat_value
+from .utils import compare_lists, generate_short_id, get_flat_value
 from .uuidt import UUIDT
 
 if TYPE_CHECKING:
@@ -693,7 +691,7 @@ def __init__(self, schema: MainSchemaTypesAPI, branch: str, data: dict | None =
         # Generate a unique ID only to be used inside the SDK
         # The format if this ID is purposely different from the ID used by the API
         # This is done to avoid confusion and potential conflicts between the IDs
-        self._internal_id = shortuuid.uuid()
+        self._internal_id = generate_short_id()
 
         self.id = data.get("id", None) if isinstance(data, dict) else None
         self.display_label: str | None = data.get("display_label", None) if isinstance(data, dict) else None

infrahub_sdk/store.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import warnings
-from collections import defaultdict
 from typing import TYPE_CHECKING, Literal, overload
 
 from .exceptions import NodeInvalidError, NodeNotFoundError
@@ -28,7 +27,7 @@ def __init__(self, name: str) -> None:
         self.branch_name = name
 
         self._objs: dict[str, InfrahubNode | InfrahubNodeSync | CoreNode | CoreNodeSync] = {}
-        self._hfids: dict[str, dict[tuple, str]] = defaultdict(dict)
+        self._hfids: dict[str, dict[tuple, str]] = {}
         self._keys: dict[str, str] = {}
         self._uuids: dict[str, str] = {}
 
@@ -45,9 +44,12 @@ def set(self, node: InfrahubNode | InfrahubNodeSync | CoreNode | CoreNodeSync, k
             self._uuids[node.id] = node._internal_id
 
         if hfid := node.get_human_friendly_id():
-            self._hfids[node.get_kind()][tuple(hfid)] = node._internal_id
+            for kind in node.get_all_kinds():
+                if kind not in self._hfids:
+                    self._hfids[kind] = {}
+                self._hfids[kind][tuple(hfid)] = node._internal_id
 
-    def get(  # type: ignore[no-untyped-def]
+    def get(
         self,
         key: str,
         kind: type[SchemaType | SchemaTypeSync] | str | None = None,
@@ -88,13 +90,11 @@ def get(  # type: ignore[no-untyped-def]
 
         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})",
+                message=f"Found a node of a different kind 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})",
         )
@@ -104,7 +104,6 @@ def _get_by_internal_id(
     ) -> 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})",
             )
@@ -124,7 +123,6 @@ def _get_by_key(
     ) -> 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})",
             )
@@ -143,7 +141,6 @@ def _get_by_key(
     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})",
             )
@@ -170,7 +167,7 @@ def _get_by_hfid(
             node_hfid = [hfid]
 
         exception_to_raise_if_not_found = NodeNotFoundError(
-            node_type=node_kind or "unknown",
+            node_type=node_kind,
             identifier={"hfid": node_hfid},
             message=f"Unable to find the node {hfid!r} in the store ({self.branch_name})",
         )
@@ -192,17 +189,27 @@ class NodeStoreBase:
     we need to save them in order to reuse them later to associate them with another node for example.
     """
 
-    def __init__(self, default_branch: str = "main") -> None:
+    def __init__(self, default_branch: str | None = None) -> None:
         self._branches: dict[str, NodeStoreBranch] = {}
         self._default_branch = default_branch
 
+    def _get_branch(self, branch: str | None = None) -> str:
+        branch = branch or self._default_branch
+
+        if branch is None:
+            raise ValueError(
+                "A Branch must be provided to use the store, either as a parameter or by setting the default branch on the store"
+            )
+
+        return branch
+
     def _set(
         self,
         node: InfrahubNode | InfrahubNodeSync | SchemaType | SchemaTypeSync,
         key: str | None = None,
         branch: str | None = None,
     ) -> None:
-        branch = branch or node.get_branch() or self._default_branch
+        branch = self._get_branch(branch or node.get_branch())
 
         if branch not in self._branches:
             self._branches[branch] = NodeStoreBranch(name=branch)
@@ -216,15 +223,15 @@ def _get(  # type: ignore[no-untyped-def]
         raise_when_missing: bool = True,
         branch: str | None = None,
     ):
-        branch = branch or self._default_branch
+        branch = self._get_branch(branch)
 
         if branch not in self._branches:
             self._branches[branch] = NodeStoreBranch(name=branch)
 
         return self._branches[branch].get(key=key, kind=kind, raise_when_missing=raise_when_missing)
 
     def count(self, branch: str | None = None) -> int:
-        branch = branch or self._default_branch
+        branch = self._get_branch(branch)
 
         if branch not in self._branches:
             return 0

infrahub_sdk/utils.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import base64
 import hashlib
 import json
+import uuid
 from itertools import groupby
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
@@ -25,6 +27,11 @@
     from whenever import TimeDelta
 
 
+def generate_short_id() -> str:
+    """Generate a short unique ID"""
+    return base64.urlsafe_b64encode(uuid.uuid4().bytes).rstrip(b"=").decode("ascii").lower()
+
+
 def base36encode(number: int) -> str:
     if not isinstance(number, (int)):
         raise TypeError("number must be an integer")

poetry.lock

@@ -43,7 +43,6 @@ pyyaml = { version = "^6", optional = true }
 eval-type-backport = { version = "^0.2.2", python = "~3.9" }
 dulwich = "^0.21.4"
 whenever = "0.7.2"
-shortuuid = "^1.0.13"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "*"

pyproject.toml

@@ -11,11 +11,11 @@
 def test_node_store_set(client_type, clients, schema_with_hfid):
     if client_type == "standard":
         client = clients.standard
-        store = NodeStore()
+        store = NodeStore(default_branch="main")
         node_class = InfrahubNode
     else:
         client = clients.sync
-        store = NodeStoreSync()
+        store = NodeStoreSync(default_branch="main")
         node_class = InfrahubNodeSync
 
     data = {
@@ -36,11 +36,11 @@ def test_node_store_set(client_type, clients, schema_with_hfid):
 def test_node_store_set_no_hfid(client_type, clients, location_schema):
     if client_type == "standard":
         client = clients.standard
-        store = NodeStore()
+        store = NodeStore(default_branch="main")
         node_class = InfrahubNode
     else:
         client = clients.sync
-        store = NodeStoreSync()
+        store = NodeStoreSync(default_branch="main")
         node_class = InfrahubNodeSync
 
     data = {
@@ -70,11 +70,11 @@ def test_node_store_set_no_hfid(client_type, clients, location_schema):
 def test_node_store_get(client_type, clients, location_schema):
     if client_type == "standard":
         client = clients.standard
-        store = NodeStore()
+        store = NodeStore(default_branch="main")
         node_class = InfrahubNode
     else:
         client = clients.sync
-        store = NodeStoreSync()
+        store = NodeStoreSync(default_branch="main")
         node_class = InfrahubNodeSync
 
     data = {
@@ -114,11 +114,11 @@ def test_node_store_get(client_type, clients, location_schema):
 def test_node_store_get_with_hfid(client_type, clients, schema_with_hfid):
     if client_type == "standard":
         client = clients.standard
-        store = NodeStore()
+        store = NodeStore(default_branch="main")
         node_class = InfrahubNode
     else:
         client = clients.sync
-        store = NodeStoreSync()
+        store = NodeStoreSync(default_branch="main")
         node_class = InfrahubNodeSync
 
     data = {

tests/unit/sdk/test_store.py

@@ -18,6 +18,7 @@
     dict_hash,
     duplicates,
     extract_fields,
+    generate_short_id,
     get_flat_value,
     is_valid_url,
     is_valid_uuid,
@@ -26,6 +27,12 @@
 )
 
 
+def test_generate_short_id():
+    assert len(generate_short_id()) == 22
+    assert isinstance(generate_short_id(), str)
+    assert generate_short_id() != generate_short_id()
+
+
 def test_is_valid_uuid():
     assert is_valid_uuid(uuid.uuid4()) is True
     assert is_valid_uuid(uuid.UUID("ba0aecd9-546a-4d77-9187-23e17a20633e")) is True