Add batch feature to sync client

Author: gmazoyer

infrahub_sdk/batch.py

@@ -1,9 +1,10 @@
 import asyncio
 from collections.abc import AsyncGenerator, Awaitable
+from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
-from typing import Any, Callable, Optional
+from typing import Any, Callable, Generator, Optional
 
-from .node import InfrahubNode
+from .node import InfrahubNode, InfrahubNodeSync
 
 
 @dataclass
@@ -14,13 +15,32 @@ class BatchTask:
     node: Optional[Any] = None
 
 
+@dataclass
+class BatchTaskSync:
+    task: Callable[..., Any]
+    args: tuple[Any, ...]
+    kwargs: dict[str, Any]
+    node: Optional[InfrahubNodeSync] = None
+
+    def execute(self, return_exceptions: bool = False) -> tuple[Optional[InfrahubNodeSync], Any]:
+        """Executes the stored task."""
+        result = None
+        try:
+            result = self.task(*self.args, **self.kwargs)
+        except Exception as exc:
+            if return_exceptions:
+                return (self.node, exc)
+            raise exc
+
+        return self.node, result
+
+
 async def execute_batch_task_in_pool(
     task: BatchTask, semaphore: asyncio.Semaphore, return_exceptions: bool = False
 ) -> tuple[Optional[InfrahubNode], Any]:
     async with semaphore:
         try:
             result = await task.task(*task.args, **task.kwargs)
-
         except Exception as exc:  # pylint: disable=broad-exception-caught
             if return_exceptions:
                 return (task.node, exc)
@@ -64,3 +84,26 @@ async def execute(self) -> AsyncGenerator:
             if isinstance(result, Exception) and not self.return_exceptions:
                 raise result
             yield node, result
+
+
+class InfrahubBatchSync:
+    def __init__(self, max_concurrent_execution: int = 5, return_exceptions: bool = False):
+        self._tasks: list[BatchTaskSync] = []
+        self.max_concurrent_execution = max_concurrent_execution
+        self.return_exceptions = return_exceptions
+
+    @property
+    def num_tasks(self) -> int:
+        return len(self._tasks)
+
+    def add(self, *args: Any, task: Callable[..., Any], node: Optional[Any] = None, **kwargs: Any) -> None:
+        self._tasks.append(BatchTaskSync(task=task, node=node, args=args, kwargs=kwargs))
+
+    def execute(self) -> Generator[tuple[Optional[InfrahubNodeSync], Any], None, None]:
+        with ThreadPoolExecutor(max_workers=self.max_concurrent_execution) as executor:
+            futures = [executor.submit(task.execute) for task in self._tasks]
+            for future in futures:
+                node, result = future.result()
+                if isinstance(result, Exception) and not self.return_exceptions:
+                    raise result
+                yield node, result

infrahub_sdk/client.py

@@ -23,7 +23,7 @@
 import ujson
 from typing_extensions import Self
 
-from .batch import InfrahubBatch
+from .batch import InfrahubBatch, InfrahubBatchSync
 from .branch import (
     BranchData,
     InfrahubBranchManager,
@@ -1454,9 +1454,6 @@ def delete(self, kind: Union[str, type[SchemaTypeSync]], id: str, branch: Option
         node = InfrahubNodeSync(client=self, schema=schema, branch=branch, data={"id": id})
         node.delete()
 
-    def create_batch(self, return_exceptions: bool = False) -> InfrahubBatch:
-        raise NotImplementedError("This method hasn't been implemented in the sync client yet.")
-
     def clone(self) -> InfrahubClientSync:
         """Return a cloned version of the client using the same configuration"""
         return InfrahubClientSync(config=self.config)
@@ -1955,6 +1952,11 @@ def get(
 
         return results[0]
 
+    def create_batch(self, return_exceptions: bool = False) -> InfrahubBatchSync:
+        return InfrahubBatchSync(
+            max_concurrent_execution=self.max_concurrent_execution, return_exceptions=return_exceptions
+        )
+
     def get_list_repositories(
         self, branches: Optional[dict[str, BranchData]] = None, kind: str = "CoreGenericRepository"
     ) -> dict[str, RepositoryData]: