feat(nanodb): Implement update_one method for DocumentCollection

resolves #68

Author: asaf

packages/nanodb/src/flux0_nanodb/api.py

@@ -1,11 +1,18 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Generic, Mapping, Optional, Sequence, Tuple, Type
+from typing import Generic, List, Mapping, Optional, Sequence, Tuple, Type
 
 from flux0_nanodb.projection import Projection
 from flux0_nanodb.query import QueryFilter
-from flux0_nanodb.types import DeleteResult, InsertOneResult, SortingOrder, TDocument
+from flux0_nanodb.types import (
+    DeleteResult,
+    InsertOneResult,
+    JSONPatchOperation,
+    SortingOrder,
+    TDocument,
+    UpdateOneResult,
+)
 
 
 class DocumentDatabase(ABC):
@@ -64,6 +71,24 @@ async def insert_one(self, document: TDocument) -> InsertOneResult:
         """
         pass
 
+    @abstractmethod
+    async def update_one(
+        self, filters: QueryFilter, patch: List[JSONPatchOperation], upsert: bool = False
+    ) -> UpdateOneResult:
+        """
+        Apply a JSON Patch (RFC 6902) to a single document that matches the provided filters.
+        If upsert is True and no document matches, insert a new document.
+
+        Parameters:
+            filters (QueryFilter): Query to match a document.
+            patch (List[JSONPatchOperation]): JSON Patch operations in a type-safe structured format.
+            upsert (bool): If True, insert a new document if no match is found.
+
+        Returns:
+            UpdateOneResult: Metadata about the update operation.
+        """
+        pass
+
     @abstractmethod
     async def delete_one(self, filters: QueryFilter) -> DeleteResult[TDocument]:
         """

packages/nanodb/src/flux0_nanodb/common.py

@@ -1,12 +1,28 @@
-from typing import Any, Mapping, Type, get_type_hints
+from typing import Any, Dict, List, Mapping, Type, get_type_hints
+
+from flux0_nanodb.types import JSONPatchOperation
 
 
 def validate_is_total(document: Mapping[str, Any], schema: Type[Mapping[str, Any]]) -> None:
-    required_keys = get_type_hints(schema).keys()
+    # Use the __required_keys__ attribute if it exists, otherwise fall back to all type hints.
+    required_keys = getattr(schema, "__required_keys__", get_type_hints(schema).keys())
     missing_keys = [key for key in required_keys if key not in document]
 
     if missing_keys:
         raise TypeError(
             f"TypedDict '{schema.__qualname__}' is missing required keys: {missing_keys}. "
             f"Expected at least the keys: {list(required_keys)}."
         )
+
+
+def convert_patch(patch: List[JSONPatchOperation]) -> List[Dict[str, Any]]:
+    converted = []
+    for op in patch:
+        # Create a shallow copy to avoid mutating the original.
+        op_copy = dict(op)
+        # For move and copy operations, rename "from_" to "from" as expected by jsonpatch.
+        if op_copy.get("op") in ("move", "copy"):
+            if "from_" in op_copy:
+                op_copy["from"] = op_copy.pop("from_")
+        converted.append(op_copy)
+    return converted

packages/nanodb/src/flux0_nanodb/memory.py

@@ -1,15 +1,19 @@
-from typing import Any, Mapping, Optional, Protocol, Sequence, Tuple, Type, cast
+from typing import Any, List, Mapping, Optional, Protocol, Sequence, Tuple, Type, cast
+
+import jsonpatch
 
 from flux0_nanodb.api import DocumentCollection, DocumentDatabase
-from flux0_nanodb.common import validate_is_total
+from flux0_nanodb.common import convert_patch, validate_is_total
 from flux0_nanodb.projection import Projection, apply_projection
 from flux0_nanodb.query import QueryFilter, matches_query
 from flux0_nanodb.types import (
     DeleteResult,
     DocumentID,
     InsertOneResult,
+    JSONPatchOperation,
     SortingOrder,
     TDocument,
+    UpdateOneResult,
 )
 
 
@@ -73,6 +77,39 @@ async def insert_one(self, document: TDocument) -> InsertOneResult:
             raise ValueError("Document is missing an 'id' field")
         return InsertOneResult(acknowledged=True, inserted_id=inserted_id)
 
+    async def update_one(
+        self, filters: QueryFilter, patch: List[JSONPatchOperation], upsert: bool = False
+    ) -> UpdateOneResult:
+        standard_patch = convert_patch(patch)
+        # Look for an existing document matching the filters.
+        for i, doc in enumerate(self._documents):
+            if matches_query(filters, doc):
+                try:
+                    updated_doc = jsonpatch.apply_patch(doc, standard_patch, in_place=False)
+                except jsonpatch.JsonPatchException as e:
+                    raise ValueError("Invalid JSON patch") from e
+                # validate_is_total(updated_doc, self._schema)
+                self._documents[i] = cast(TDocument, updated_doc)
+                return UpdateOneResult(
+                    acknowledged=True, matched_count=1, modified_count=1, upserted_id=None
+                )
+        # No matching document found.
+        if upsert:
+            try:
+                new_doc = jsonpatch.apply_patch({}, standard_patch, in_place=False)
+            except jsonpatch.JsonPatchException as e:
+                raise ValueError("Invalid JSON patch for upsert") from e
+            if "id" not in new_doc:
+                raise ValueError("Upserted document is missing an 'id' field")
+            validate_is_total(new_doc, self._schema)
+            self._documents.append(cast(TDocument, new_doc))
+            return UpdateOneResult(
+                acknowledged=True, matched_count=0, modified_count=0, upserted_id=new_doc["id"]
+            )
+        return UpdateOneResult(
+            acknowledged=True, matched_count=0, modified_count=0, upserted_id=None
+        )
+
     async def delete_one(self, filters: QueryFilter) -> DeleteResult[TDocument]:
         for i, doc in enumerate(self._documents):
             if matches_query(filters, doc):

packages/nanodb/src/flux0_nanodb/types.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from enum import Enum
-from typing import Generic, NewType, Optional, TypedDict, TypeVar
+from typing import Any, Generic, Literal, NewType, Optional, TypedDict, TypeVar, Union
 
 DocumentID = NewType("DocumentID", str)
 DocumentVersion = NewType("DocumentVersion", str)
@@ -26,9 +26,57 @@ class InsertOneResult:
     inserted_id: DocumentID  # Mimicking MongoDB’s inserted_id field.
 
 
+@dataclass(frozen=True)
+class UpdateOneResult:
+    acknowledged: bool
+    matched_count: int
+    modified_count: int
+    upserted_id: Optional[DocumentID]
+
+
 # MongoDB-like result structure for delete operations.
 @dataclass(frozen=True)
 class DeleteResult(Generic[TDocument]):
     acknowledged: bool
     deleted_count: int
     deleted_document: Optional[TDocument]
+
+
+# Define a type-safe JSON Patch operation
+class AddOp(TypedDict):
+    op: Literal["add"]
+    path: str
+    value: Any
+
+
+class RemoveOp(TypedDict):
+    op: Literal["remove"]
+    path: str
+
+
+class ReplaceOp(TypedDict):
+    op: Literal["replace"]
+    path: str
+    value: Any
+
+
+class MoveOp(TypedDict):
+    op: Literal["move"]
+    from_: str  # note: using from_ instead of 'from' because it's a keyword
+    path: str
+
+
+class CopyOp(TypedDict):
+    op: Literal["copy"]
+    from_: str
+    path: str
+
+
+class TestOp(TypedDict):
+    op: Literal["test"]
+    path: str
+    value: Any
+
+
+# Union type for JSON patch operations.
+JSONPatchOperation = Union[AddOp, RemoveOp, ReplaceOp, MoveOp, CopyOp, TestOp]