digitalkin 0.2.7__py3-none-any.whl → 0.2.8__py3-none-any.whl

digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.2.7"
+    __version__ = "0.2.8"

digitalkin/services/storage/default_storage.py

@@ -5,6 +5,7 @@ import json
 import logging
 import tempfile
 from pathlib import Path
+from typing import Any
 
 from pydantic import BaseModel
 
@@ -14,115 +15,91 @@ logger = logging.getLogger(__name__)
 
 
 class DefaultStorage(StorageStrategy):
-    """This class implements the default storage strategy with file persistence.
+    """Persist records in a local JSON file for quick local development.
 
-    The storage persists data in a local JSON file, enabling data retention
-    across multiple requests or module instances.
+    File format: a JSON object of
+      { "<collection>:<record_id>": { ... StorageRecord fields ... },
     """
 
+    @staticmethod
+    def _json_default(o: Any) -> str:  # noqa: ANN401
+        """JSON serializer for non-standard types (datetime → ISO).
+
+        Args:
+            o: The object to serialize
+
+        Returns:
+            str: The serialized object
+
+        Raises:
+            TypeError: If the object is not serializable
+        """
+        if isinstance(o, datetime.datetime):
+            return o.isoformat()
+        msg = f"Type {o.__class__.__name__} not serializable"
+        raise TypeError(msg)
+
     def _load_from_file(self) -> dict[str, StorageRecord]:
         """Load storage data from the file.
 
         Returns:
             A dictionary containing the loaded storage records
         """
-        file_path = Path(self.storage_file_path)
-        if not file_path.exists():
+        if not self.storage_file.exists():
             return {}
 
         try:
-            records = {}
-            file_content = json.loads(file_path.read_text(encoding="utf-8"))
-
-            for key, record_dict in file_content.items():
-                # Get the stored record data
-                name = record_dict.get("name", "")
-                model_class = self.config.get(name)
-
-                if not model_class:
-                    logger.warning("No model found for record %s", name)
+            raw = json.loads(self.storage_file.read_text(encoding="utf-8"))
+            out: dict[str, StorageRecord] = {}
+
+            for key, rd in raw.items():
+                # rd is a dict with the StorageRecord fields
+                model_cls = self.config.get(rd["collection"])
+                if not model_cls:
+                    logger.warning("No model for collection '%s'", rd["collection"])
                     continue
-
-                # Create a model instance from the stored data
-                data_dict = record_dict.get("data", {})
-                try:
-                    data_model = model_class.model_validate(data_dict)
-                except Exception:
-                    logger.exception("Failed to validate data for record %s", name)
-                    continue
-
-                # Create a StorageRecord object
-                record = StorageRecord(
-                    mission_id=record_dict.get("mission_id", ""),
-                    name=name,
+                data_model = model_cls.model_validate(rd["data"])
+                rec = StorageRecord(
+                    mission_id=rd["mission_id"],
+                    collection=rd["collection"],
+                    record_id=rd["record_id"],
                     data=data_model,
-                    data_type=DataType[record_dict.get("data_type", "OUTPUT")],
+                    data_type=DataType[rd["data_type"]],
+                    creation_date=datetime.datetime.fromisoformat(rd["creation_date"])
+                    if rd.get("creation_date")
+                    else None,
+                    update_date=datetime.datetime.fromisoformat(rd["update_date"]) if rd.get("update_date") else None,
                 )
-
-                # Set dates if they exist
-                if "creation_date" in record_dict:
-                    record.creation_date = datetime.datetime.fromisoformat(record_dict["creation_date"])
-                if "update_date" in record_dict:
-                    record.update_date = datetime.datetime.fromisoformat(record_dict["update_date"])
-
-                records[key] = record
-        except json.JSONDecodeError:
-            logger.exception("Error decoding JSON from file")
-            return {}
-        except FileNotFoundError:
-            logger.info("Storage file not found, starting with empty storage")
-            return {}
+                out[key] = rec
         except Exception:
-            logger.exception("Unexpected error loading storage")
+            logger.exception("Failed to load default storage file")
             return {}
-        return records
+        return out
 
     def _save_to_file(self) -> None:
-        """Save storage data to the file using a safe write pattern."""
-        # Usage of pathlib for file operations
-        file_path = Path(self.storage_file_path)
-        file_path.parent.mkdir(parents=True, exist_ok=True)
-
-        try:
-            # Convert storage to a serializable format
-            serializable_data = {}
-            for key, record in self.storage.items():
-                record_dict = {
-                    "mission_id": record.mission_id,
-                    "name": record.name,
-                    "data_type": record.data_type.name,  # Convert enum to string
-                    "data": record.data.model_dump(),  # Convert Pydantic model to dict
-                }
-
-                # Handle dates (convert to ISO format strings)
-                if record.creation_date:
-                    record_dict["creation_date"] = record.creation_date.isoformat()
-                if record.update_date:
-                    record_dict["update_date"] = record.update_date.isoformat()
-
-                serializable_data[key] = record_dict
-
-            # usage of NamedTemporaryFile for atomic writes
-            with tempfile.NamedTemporaryFile(
-                mode="w", encoding="utf-8", dir=str(file_path.parent), delete=False, suffix=".tmp"
-            ) as temp_file:
-                json.dump(serializable_data, temp_file, indent=2)
-                temp_path = temp_file.name
-
-            # Creation of a backup if the file already exists
-            if file_path.exists():
-                backup_path = f"{self.storage_file_path}.bak"
-                file_path.replace(backup_path)
-
-            # Remplacement du fichier (opération atomique) avec pathlib
-            Path(temp_path).replace(str(file_path))
-
-        except PermissionError:
-            logger.exception("Permission denied when saving to file")
-        except OSError:
-            logger.exception("OS error when saving to file")
-        except Exception:
-            logger.exception("Unexpected error saving storage")
+        """Atomically write `self.storage` back to disk as JSON."""
+        self.storage_file.parent.mkdir(parents=True, exist_ok=True)
+        with tempfile.NamedTemporaryFile(
+            mode="w", encoding="utf-8", delete=False, dir=str(self.storage_file.parent), suffix=".tmp"
+        ) as temp:
+            try:
+                # Convert storage to a serializable format
+                serial: dict[str, dict] = {}
+                for key, record in self.storage.items():
+                    serial[key] = {
+                        "mission_id": record.mission_id,
+                        "collection": record.collection,
+                        "record_id": record.record_id,
+                        "data_type": record.data_type.name,
+                        "data": record.data.model_dump(),
+                        "creation_date": record.creation_date.isoformat() if record.creation_date else None,
+                        "update_date": record.update_date.isoformat() if record.update_date else None,
+                    }
+                json.dump(serial, temp, indent=2, default=self._json_default)
+                temp.flush()
+                Path(temp.name).replace(self.storage_file)
+            except Exception:
+                logger.exception("Unexpected error saving storage")
 
     def _store(self, record: StorageRecord) -> StorageRecord:
         """Store a new record in the database and persist to file.
@@ -136,73 +113,97 @@ class DefaultStorage(StorageStrategy):
         Raises:
             ValueError: If the record already exists
         """
-        name = record.name
-        if name in self.storage:
-            msg = f"Record with name {name} already exists"
+        key = f"{record.collection}:{record.record_id}"
+        if key in self.storage:
+            msg = f"Document {key!r} already exists"
             raise ValueError(msg)
-        self.storage[name] = record
-        self.storage[name].creation_date = datetime.datetime.now(datetime.timezone.utc)
-        self.storage[name].update_date = datetime.datetime.now(datetime.timezone.utc)
-
-        # Persist to file
+        now = datetime.datetime.now(datetime.timezone.utc)
+        record.creation_date = now
+        record.update_date = now
+        self.storage[key] = record
         self._save_to_file()
+        logger.debug("Created %s", key)
+        return record
 
-        logger.info("CREATE %s:%s successful", name, record)
-        return self.storage[name]
-
-    def _read(self, name: str) -> StorageRecord | None:
+    def _read(self, collection: str, record_id: str) -> StorageRecord | None:
         """Get records from the database.
 
         Args:
-            name: The unique name to retrieve data for
+            collection: The unique name to retrieve data for
+            record_id: The unique ID of the record
 
         Returns:
             StorageRecord: The corresponding record
         """
-        logger.info("GET record linked to the key = %s", name)
-        if name not in self.storage:
-            logger.info("GET key = %s: DOESN'T EXIST", name)
-            return None
-        return self.storage[name]
+        key = f"{collection}:{record_id}"
+        return self.storage.get(key)
 
-    def _modify(self, name: str, data: BaseModel) -> StorageRecord | None:
+    def _update(self, collection: str, record_id: str, data: BaseModel) -> StorageRecord | None:
         """Update records in the database and persist to file.
 
         Args:
-            name: The unique name to store the data under
+            collection: The unique name to retrieve data for
+            record_id: The unique ID of the record
             data: The data to modify
 
         Returns:
             StorageRecord: The modified record
         """
-        if name not in self.storage:
-            logger.info("UPDATE key = %s: DOESN'T EXIST", name)
+        key = f"{collection}:{record_id}"
+        rec = self.storage.get(key)
+        if not rec:
             return None
-        self.storage[name].data = data
-        self.storage[name].update_date = datetime.datetime.now(datetime.timezone.utc)
-
-        # Persist to file
+        rec.data = data
+        rec.update_date = datetime.datetime.now(datetime.timezone.utc)
         self._save_to_file()
+        logger.debug("Modified %s", key)
+        return rec
 
-        return self.storage[name]
-
-    def _remove(self, name: str) -> bool:
+    def _remove(self, collection: str, record_id: str) -> bool:
         """Delete records from the database and update file.
 
         Args:
-            name: The unique name to remove a record
+            collection: The unique name to retrieve data for
+            record_id: The unique ID of the record
 
         Returns:
             bool: True if the record was removed, False otherwise
         """
-        if name not in self.storage:
-            logger.info("DELETE key = %s: DOESN'T EXIST", name)
+        key = f"{collection}:{record_id}"
+        if key not in self.storage:
             return False
-        del self.storage[name]
-
-        # Persist to file
+        del self.storage[key]
         self._save_to_file()
+        logger.debug("Removed %s", key)
+        return True
+
+    def _list(self, collection: str) -> list[StorageRecord]:
+        """Implements StorageStrategy._list.
+
+        Args:
+            collection: The unique name to retrieve data for
 
+        Returns:
+            A list of storage records
+        """
+        prefix = f"{collection}:"
+        return [r for k, r in self.storage.items() if k.startswith(prefix)]
+
+    def _remove_collection(self, collection: str) -> bool:
+        """Implements StorageStrategy._remove_collection.
+
+        Args:
+            collection: The unique name to retrieve data for
+
+        Returns:
+            bool: True if the collection was removed, False otherwise
+        """
+        prefix = f"{collection}:"
+        to_delete = [k for k in self.storage if k.startswith(prefix)]
+        for k in to_delete:
+            del self.storage[k]
+        self._save_to_file()
+        logger.debug("Removed collection %s (%d docs)", collection, len(to_delete))
         return True
 
     def __init__(
@@ -215,4 +216,5 @@ class DefaultStorage(StorageStrategy):
         """Initialize the storage."""
         super().__init__(mission_id=mission_id, config=config)
         self.storage_file_path = f"{self.mission_id}_{storage_file_path}.json"
+        self.storage_file = Path(self.storage_file_path)
         self.storage = self._load_from_file()

digitalkin/services/storage/grpc_storage.py

@@ -17,19 +17,36 @@ logger = logging.getLogger(__name__)
 class GrpcStorage(StorageStrategy, GrpcClientWrapper):
     """This class implements the default storage strategy."""
 
-    def __init__(
-        self,
-        mission_id: str,
-        config: dict[str, type[BaseModel]],
-        client_config: ClientConfig,
-        **kwargs,  # noqa: ANN003, ARG002
-    ) -> None:
-        """Initialize the storage."""
-        super().__init__(mission_id=mission_id, config=config)
+    def _build_record_from_proto(self, proto: data_pb2.StorageRecord) -> StorageRecord:
+        """Convert a protobuf StorageRecord message into our Pydantic model.
 
-        channel = self._init_channel(client_config)
-        self.stub = storage_service_pb2_grpc.StorageServiceStub(channel)
-        logger.info("Channel client 'storage' initialized succesfully")
+        Args:
+            proto: gRPC StorageRecord
+
+        Returns:
+            A fully validated StorageRecord.
+        """
+        raw = json_format.MessageToDict(
+            proto,
+            preserving_proto_field_name=True,
+            always_print_fields_with_no_presence=True,
+        )
+        mission = raw["mission_id"]
+        coll = raw["collection"]
+        rid = raw["record_id"]
+        dtype = DataType[raw["data_type"]]
+        payload = raw.get("data", {})
+
+        validated = self._validate_data(coll, payload)
+        return StorageRecord(
+            mission_id=mission,
+            collection=coll,
+            record_id=rid,
+            data=validated,
+            data_type=dtype,
+            creation_date=raw.get("creation_date"),
+            update_date=raw.get("update_date"),
+        )
 
     def _store(self, record: StorageRecord) -> StorageRecord:
         """Create a new record in the database.
@@ -44,113 +61,137 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
             StorageServiceError: If there is an error while storing the record
         """
         try:
-            # Create a Struct for the data
             data_struct = Struct()
             data_struct.update(record.data.model_dump())
-
-            request = data_pb2.StoreRecordRequest(
+            req = data_pb2.StoreRecordRequest(
                 data=data_struct,
                 mission_id=record.mission_id,
-                name=record.name,
+                collection=record.collection,
+                record_id=record.record_id,
                 data_type=record.data_type.name,
             )
-            return self.exec_grpc_query("StoreRecord", request)
-        except Exception:
-            msg = f"Error while storing record {record.name}"
-            logger.exception(msg)
-            raise StorageServiceError(msg)
+            resp = self.exec_grpc_query("StoreRecord", req)
+            return self._build_record_from_proto(resp.stored_data)
+        except Exception as e:
+            logger.exception("gRPC StoreRecord failed for %s:%s", record.collection, record.record_id)
+            raise StorageServiceError(str(e)) from e
 
-    def _read(self, name: str) -> StorageRecord | None:
-        """Get records from the database.
+    def _read(self, collection: str, record_id: str) -> StorageRecord | None:
+        """Fetch a single document by collection + record_id.
 
         Returns:
-            list[StorageData]: The list of records
+            StorageData: The record
         """
         try:
-            request = data_pb2.ReadRecordRequest(mission_id=self.mission_id, name=name)
-            response: data_pb2.ReadRecordResponse = self.exec_grpc_query("ReadRecord", request)
-            response_dict = json_format.MessageToDict(
-                response.stored_data,
-                preserving_proto_field_name=True,
-                always_print_fields_with_no_presence=True,
-            )
-            return StorageRecord(
-                mission_id=response_dict["mission_id"],
-                name=response_dict["name"],
-                data_type=response_dict["data_type"],
-                data=self._validate_data(name, response_dict["data"]),
+            req = data_pb2.ReadRecordRequest(
+                mission_id=self.mission_id,
+                collection=collection,
+                record_id=record_id,
             )
+            resp = self.exec_grpc_query("ReadRecord", req)
+            return self._build_record_from_proto(resp.stored_data)
         except Exception:
-            msg = f"Error while reading record {name}"
-            logger.exception(msg)
+            logger.exception("gRPC ReadRecord failed for %s:%s", collection, record_id)
             return None
 
-    def _modify(self, name: str, data: BaseModel) -> StorageRecord | None:
-        """Update records in the database.
+    def _update(self, collection: str, record_id: str, data: BaseModel) -> StorageRecord | None:
+        """Overwrite a document via gRPC.
+
+        Args:
+            collection: The unique name for the record type
+            record_id: The unique ID for the record
+            data: The validated data model
 
         Returns:
-            int: The number of records updated
+            StorageRecord: The updated record
         """
         try:
-            # Create a Struct for the data
-            data_struct = Struct()
-            data_struct.update(data.model_dump())
-
-            request = data_pb2.ModifyRecordRequest(data=data_struct, mission_id=self.mission_id, name=name)
-            response: data_pb2.ModifyRecordResponse = self.exec_grpc_query("ModifyRecord", request)
-            return self._build_record_from_proto(response.stored_data, name)
+            struct = Struct()
+            struct.update(data.model_dump())
+            req = data_pb2.UpdateRecordRequest(
+                data=struct,
+                mission_id=self.mission_id,
+                collection=collection,
+                record_id=record_id,
+            )
+            resp = self.exec_grpc_query("ModifyRecord", req)
+            return self._build_record_from_proto(resp.stored_data)
         except Exception:
-            msg = f"Error while modifing record {name}"
-            logger.exception(msg)
+            logger.exception("gRPC ModifyRecord failed for %s:%s", collection, record_id)
             return None
 
-    def _remove(self, name: str) -> bool:
-        """Delete records from the database.
+    def _remove(self, collection: str, record_id: str) -> bool:
+        """Delete a document via gRPC.
+
+        Args:
+            collection: The unique name for the record type
+            record_id: The unique ID for the record
 
         Returns:
-            int: The number of records deleted
+            bool: True if the record was deleted, False otherwise
         """
         try:
-            request = data_pb2.RemoveRecordRequest(
+            req = data_pb2.RemoveRecordRequest(
                 mission_id=self.mission_id,
-                name=name,
+                collection=collection,
+                record_id=record_id,
             )
-            self.exec_grpc_query("RemoveRecord", request)
+            self.exec_grpc_query("RemoveRecord", req)
         except Exception:
-            msg = f"Error while removin record {name}"
-            logger.exception(msg)
+            logger.exception("gRPC RemoveRecord failed for %s:%s", collection, record_id)
             return False
         return True
 
-    def _build_record_from_proto(self, stored_data: data_pb2.StorageRecord, default_name: str) -> StorageRecord:
-        """Helper to construire un StorageRecord complet à partir du message gRPC.
+    def _list(self, collection: str) -> list[StorageRecord]:
+        """List all documents in a collection via gRPC.
 
         Args:
-            stored_data: Le message gRPC contenant les données stockées.
-            default_name: Le nom par défaut à utiliser si le nom n'est pas présent dans les données.
+            collection: The unique name for the record type
 
         Returns:
-            StorageRecord: Un objet StorageRecord construit à partir des données stockées.
+            list[StorageRecord]: A list of storage records
         """
-        # Converti en dict, avec tous les champs même s'ils sont absents
-        raw = json_format.MessageToDict(
-            stored_data,
-            preserving_proto_field_name=True,
-            always_print_fields_with_no_presence=True,
-        )
-        # On récupère ou on complète les champs obligatoires
-        name = raw.get("name", default_name)
-        dtype = raw.get("data_type", DataType.OUTPUT.name)
-        payload = raw.get("data", {})
+        try:
+            req = data_pb2.ListRecordsRequest(
+                mission_id=self.mission_id,
+                collection=collection,
+            )
+            resp = self.exec_grpc_query("ListRecords", req)
+            return [self._build_record_from_proto(r) for r in resp.records]
+        except Exception:
+            logger.exception("gRPC ListRecords failed for %s", collection)
+            return []
 
-        # Valide le modèle pydantic pour le champ `data`
-        validated = self._validate_data(name, payload)
+    def _remove_collection(self, collection: str) -> bool:
+        """Delete an entire collection via gRPC.
 
-        return StorageRecord(
-            mission_id=self.mission_id,
-            name=name,
-            data_type=DataType[dtype],
-            data=validated,
-            creation_date=raw.get("creation_date"),
-            update_date=raw.get("update_date"),
-        )
+        Args:
+            collection: The unique name for the record type
+
+        Returns:
+            bool: True if the collection was deleted, False otherwise
+        """
+        try:
+            req = data_pb2.RemoveCollectionRequest(
+                mission_id=self.mission_id,
+                collection=collection,
+            )
+            self.exec_grpc_query("RemoveCollection", req)
+        except Exception:
+            logger.exception("gRPC RemoveCollection failed for %s", collection)
+            return False
+        return True
+
+    def __init__(
+        self,
+        mission_id: str,
+        config: dict[str, type[BaseModel]],
+        client_config: ClientConfig,
+        **kwargs,  # noqa: ANN003, ARG002
+    ) -> None:
+        """Initialize the storage."""
+        super().__init__(mission_id=mission_id, config=config)
+
+        channel = self._init_channel(client_config)
+        self.stub = storage_service_pb2_grpc.StorageServiceStub(channel)
+        logger.debug("Channel client 'storage' initialized succesfully")

digitalkin/services/storage/storage_strategy.py

@@ -4,6 +4,7 @@ import datetime
 from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Any, Literal, TypeGuard
+from uuid import uuid4
 
 from pydantic import BaseModel, Field
 
@@ -24,35 +25,69 @@ class DataType(Enum):
 
 
 class StorageRecord(BaseModel):
-    """Container for stored records with metadata."""
+    """A single record stored in a collection, with metadata."""
 
-    # Metadata
-    mission_id: str = Field(description="The ID of the mission this record is associated with")
-    name: str = Field(description="The name of the record")
-    creation_date: datetime.datetime | None = Field(default=None, description="The date the record was created")
-    update_date: datetime.datetime | None = Field(default=None, description="The date the record was last updated")
-    data_type: DataType = Field(default=DataType.OUTPUT, description="The type of data stored")
-    # Actual data payload
-    data: BaseModel = Field(description="The data stored in the record")
+    mission_id: str = Field(..., description="ID of the mission (bucket) this doc belongs to")
+    collection: str = Field(..., description="Logical collection name")
+    record_id: str = Field(..., description="Unique ID of this record in its collection")
+    data_type: DataType = Field(default=DataType.OUTPUT, description="Category of the data of this record")
+    data: BaseModel = Field(..., description="The typed payload of this record")
+    creation_date: datetime.datetime | None = Field(default=None, description="When this record was first created")
+    update_date: datetime.datetime | None = Field(default=None, description="When this record was last modified")
 
 
 class StorageStrategy(BaseStrategy, ABC):
-    """Abstract base class for storage strategies.
+    """Define CRUD + list/remove-collection against a collection/record store."""
 
-    This strategy defines how data is stored and retrieved, with
-    type validation through registered Pydantic models.
-    """
+    def _validate_data(self, collection: str, data: dict[str, Any]) -> BaseModel:
+        """Validate data against the model schema for the given key.
 
-    def __init__(self, mission_id: str, config: dict[str, type[BaseModel]]) -> None:
-        """Initialize the storage strategy.
+        Args:
+            collection: The unique name for the record type
+            data: The data to validate
+
+        Returns:
+            A validated model instance
+
+        Raises:
+            ValueError: If the key has no associated model or validation fails
+        """
+        model_cls = self.config.get(collection)
+        if not model_cls:
+            msg = f"No schema registered for collection '{collection}'"
+            raise ValueError(msg)
+
+        try:
+            return model_cls.model_validate(data)
+        except Exception as e:
+            msg = f"Validation failed for '{collection}': {e!s}"
+            raise ValueError(msg) from e
+
+    def _create_storage_record(
+        self,
+        collection: str,
+        record_id: str,
+        validated_data: BaseModel,
+        data_type: DataType,
+    ) -> StorageRecord:
+        """Create a storage record with metadata.
 
         Args:
-            mission_id: The ID of the mission this strategy is associated with
-            config: A dictionary mapping names to Pydantic model classes
+            collection: The unique name for the record type
+            record_id: The unique ID for the record
+            validated_data: The validated data model
+            data_type: The type of data
+
+        Returns:
+            A complete storage record with metadata
         """
-        super().__init__(mission_id)
-        # Schema configuration mapping keys to model classes
-        self.config: dict[str, type[BaseModel]] = config
+        return StorageRecord(
+            mission_id=self.mission_id,
+            collection=collection,
+            record_id=record_id,
+            data=validated_data,
+            data_type=data_type,
+        )
 
     @staticmethod
     def _is_valid_data_type_name(value: str) -> TypeGuard[str]:
@@ -69,145 +104,162 @@ class StorageStrategy(BaseStrategy, ABC):
             The ID of the created record
         """
 
-    def store(
-        self,
-        name: str,
-        data: dict[str, Any],
-        data_type: Literal["OUTPUT", "VIEW", "LOGS", "OTHER"] = "OUTPUT",
-    ) -> StorageRecord:
-        """Store a new record in the storage.
+    @abstractmethod
+    def _read(self, collection: str, record_id: str) -> StorageRecord | None:
+        """Get records from storage by key.
 
         Args:
-            name: The unique name to store the data under
-            data: The data to store
-            data_type: The type of data being stored (default: OUTPUT)
+            collection: The unique name to retrieve data for
+            record_id: The unique ID of the record
 
         Returns:
-            The ID of the created record
+            A storage record with validated data
+        """
 
-        Raises:
-            ValueError: If the data type is invalid or if validation fails
+    @abstractmethod
+    def _update(self, collection: str, record_id: str, data: BaseModel) -> StorageRecord | None:
+        """Overwrite an existing record's payload.
+
+        Args:
+            collection: The unique name for the record type
+            record_id: The unique ID of the record
+            data: The new data to store
+
+        Returns:
+            StorageRecord: The modified record
         """
-        if not self._is_valid_data_type_name(data_type):
-            msg = f"Invalid data type '{data_type}'. Must be one of {list(DataType.__members__.keys())}"
-            raise ValueError(msg)
-        data_type_enum = DataType[data_type]
-        validated_data = self._validate_data(name, {**data, "mission_id": self.mission_id})
-        record = self._create_storage_record(name, validated_data, data_type_enum)
-        return self._store(record)
 
     @abstractmethod
-    def _read(self, name: str) -> StorageRecord | None:
-        """Get records from storage by key.
+    def _remove(self, collection: str, record_id: str) -> bool:
+        """Delete a record from the storage.
 
         Args:
-            name: The unique name to retrieve data for
+            collection: The unique name for the record type
+            record_id: The unique ID of the record
 
         Returns:
-            A storage record with validated data
+            True if the deletion was successful, False otherwise
         """
 
-    def read(self, name: str) -> StorageRecord | None:
-        """Get records from storage by key.
+    @abstractmethod
+    def _list(self, collection: str) -> list[StorageRecord]:
+        """List all records in a collection.
 
         Args:
-            name: The unique name to retrieve data for
+            collection: The unique name for the record type
 
         Returns:
-            A storage record with validated data
+            A list of storage records
         """
-        return self._read(name)
 
     @abstractmethod
-    def _modify(self, name: str, data: BaseModel) -> StorageRecord | None:
-        """Update a record in the storage.
+    def _remove_collection(self, collection: str) -> bool:
+        """Delete all records in a collection.
 
         Args:
-            name: The unique name for the record type
-            data: The new data to store
+            collection: The unique name for the record type
 
         Returns:
-            StorageRecord: The modified record
+            True if the deletion was successful, False otherwise
         """
 
-    def modify(self, name: str, data: dict[str, Any]) -> StorageRecord | None:
-        """Update a record in the storage (overwrite all the data).
+    def __init__(self, mission_id: str, config: dict[str, type[BaseModel]]) -> None:
+        """Initialize the storage strategy.
 
         Args:
-            name: The unique name for the record type
-            data: The new data to store
+            mission_id: The ID of the mission this strategy is associated with
+            config: A dictionary mapping names to Pydantic model classes
+        """
+        super().__init__(mission_id)
+        # Schema configuration mapping keys to model classes
+        self.config: dict[str, type[BaseModel]] = config
+
+    def store(
+        self,
+        collection: str,
+        record_id: str | None,
+        data: dict[str, Any],
+        data_type: Literal["OUTPUT", "VIEW", "LOGS", "OTHER"] = "OUTPUT",
+    ) -> StorageRecord:
+        """Store a new record in the storage.
+
+        Args:
+            collection: The unique name for the record type
+            record_id: The unique ID for the record (optional)
+            data: The data to store
+            data_type: The type of data being stored (default: OUTPUT)
 
         Returns:
-            StorageRecord: The modified record
+            The ID of the created record
+
+        Raises:
+            ValueError: If the data type is invalid or if validation fails
         """
-        validated_data = self._validate_data(name, data)
-        return self._modify(name, validated_data)
+        if not self._is_valid_data_type_name(data_type):
+            msg = f"Invalid data type '{data_type}'. Must be one of {list(DataType.__members__.keys())}"
+            raise ValueError(msg)
+        record_id = record_id or uuid4().hex
+        data_type_enum = DataType[data_type]
+        validated_data = self._validate_data(record_id, {**data, "mission_id": self.mission_id})
+        record = self._create_storage_record(collection, record_id, validated_data, data_type_enum)
+        return self._store(record)
 
-    @abstractmethod
-    def _remove(self, name: str) -> bool:
-        """Delete a record from the storage.
+    def read(self, collection: str, record_id: str) -> StorageRecord | None:
+        """Get records from storage by key.
 
         Args:
-            name: The unique name for the record type
+            collection: The unique name to retrieve data for
+            record_id: The unique ID of the record
 
         Returns:
-            True if the deletion was successful, False otherwise
+            A storage record with validated data
+        """
+        return self._read(collection, record_id)
+
+    def update(self, collection: str, record_id: str, data: dict[str, Any]) -> StorageRecord | None:
+        """Validate & overwrite an existing record.
+
+        Args:
+            collection: The unique name for the record type
+            record_id: The unique ID of the record
+            data: The new data to store
+
+        Returns:
+            StorageRecord: The modified record
         """
+        validated_data = self._validate_data(record_id, data)
+        return self._update(collection, record_id, validated_data)
 
-    def remove(self, name: str) -> bool:
+    def remove(self, collection: str, record_id: str) -> bool:
         """Delete a record from the storage.
 
         Args:
-            name: The unique name for the record type
+            collection: The unique name for the record type
+            record_id: The unique ID of the record
 
         Returns:
             True if the deletion was successful, False otherwise
         """
-        return self._remove(name)
+        return self._remove(collection, record_id)
 
-    def _validate_data(self, name: str, data: dict[str, Any]) -> BaseModel:
-        """Validate data against the model schema for the given key.
+    def list(self, collection: str) -> list[StorageRecord]:
+        """Get all records within a collection.
 
         Args:
-            name: The unique name to get the model type for
-            data: The data to validate
+            collection: The unique name for the record type
 
         Returns:
-            A validated model instance
-
-        Raises:
-            ValueError: If the key has no associated model or validation fails
+            A list of storage records
         """
-        model_cls = self.config.get(name)
-        if not model_cls:
-            msg = f"No model schema defined for name: {name}"
-            raise ValueError(msg)
+        return self._list(collection)
 
-        try:
-            return model_cls.model_validate(data)
-        except Exception as e:
-            msg = f"Data validation failed for key '{name}': {e!s}"
-            raise ValueError(msg) from e
-
-    def _create_storage_record(
-        self,
-        name: str,
-        validated_data: BaseModel,
-        data_type: DataType,
-    ) -> StorageRecord:
-        """Create a storage record with metadata.
+    def remove_collection(self, collection: str) -> bool:
+        """Wipe a record clean.
 
         Args:
-            name: The unique name for the record
-            validated_data: The validated data model
-            data_type: The type of data
+            collection: The unique name for the record type
 
         Returns:
-            A complete storage record with metadata
+            True if the deletion was successful, False otherwise
         """
-        return StorageRecord(
-            mission_id=self.mission_id,
-            name=name,
-            data=validated_data,
-            data_type=data_type,
-        )
+        return self._remove_collection(collection)

{digitalkin-0.2.7.dist-info → digitalkin-0.2.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.2.7
+Version: 0.2.8
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -452,20 +452,20 @@ Classifier: License :: Other/Proprietary License
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: digitalkin-proto>=0.1.7
+Requires-Dist: digitalkin-proto>=0.1.8
 Requires-Dist: grpcio-health-checking>=1.71.0
 Requires-Dist: grpcio-reflection>=1.71.0
 Requires-Dist: grpcio-status>=1.71.0
-Requires-Dist: openai>=1.70.0
-Requires-Dist: pydantic>=2.11.3
+Requires-Dist: openai>=1.76.2
+Requires-Dist: pydantic>=2.11.4
 Provides-Extra: dev
 Requires-Dist: pytest>=8.3.4; extra == "dev"
 Requires-Dist: pytest-asyncio>=0.26.0; extra == "dev"
 Requires-Dist: pytest-cov>=6.1.0; extra == "dev"
-Requires-Dist: typos>=1.31.1; extra == "dev"
-Requires-Dist: ruff>=0.11.2; extra == "dev"
+Requires-Dist: typos>=1.31.2; extra == "dev"
+Requires-Dist: ruff>=0.11.7; extra == "dev"
 Requires-Dist: mypy>=1.15.0; extra == "dev"
-Requires-Dist: pyright>=1.1.398; extra == "dev"
+Requires-Dist: pyright>=1.1.400; extra == "dev"
 Requires-Dist: pre-commit>=4.2.0; extra == "dev"
 Requires-Dist: bump2version>=1.0.1; extra == "dev"
 Requires-Dist: build>=1.2.2; extra == "dev"

{digitalkin-0.2.7.dist-info → digitalkin-0.2.8.dist-info}/RECORD

@@ -7,7 +7,7 @@ base_server/mock/__init__.py,sha256=YZFT-F1l_TpvJYuIPX-7kTeE1CfOjhx9YmNRXVoi-jQ,
 base_server/mock/mock_pb2.py,sha256=sETakcS3PAAm4E-hTCV1jIVaQTPEAIoVVHupB8Z_k7Y,1843
 base_server/mock/mock_pb2_grpc.py,sha256=BbOT70H6q3laKgkHfOx1QdfmCS_HxCY4wCOX84YAdG4,3180
 digitalkin/__init__.py,sha256=7LLBAba0th-3SGqcpqFO-lopWdUkVLKzLZiMtB-mW3M,162
-digitalkin/__version__.py,sha256=sq3ZYGvIa5rczwcukAFcGJdR0MZ8h0cIPYmDlhbjOP4,190
+digitalkin/__version__.py,sha256=hDLYpdit14gwOhkz3Wcb9mCY0_Pbep3DECM9bgKB1pY,190
 digitalkin/logger.py,sha256=9cDgyJV2QXXT8F--xRODFlZyDgjuTTXNdpCU3GdqCsk,382
 digitalkin/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 digitalkin/grpc_servers/__init__.py,sha256=0cJBlwipSmFdXkyH3T0i6OJ1WpAtNsZgYX7JaSnkbtg,804
@@ -62,18 +62,18 @@ digitalkin/services/snapshot/__init__.py,sha256=Uzlnzo0CYlSpVsdiI37hW7xQk8hu3YA1
 digitalkin/services/snapshot/default_snapshot.py,sha256=Mb8QwWRsHh9I_tN0ln_ZiFa1QCZxOVWmuVLemQOTWpc,1058
 digitalkin/services/snapshot/snapshot_strategy.py,sha256=B1TU3V_k9A-OdqBkdyc41-ihnrW5Btcwd1KyQdHT46A,898
 digitalkin/services/storage/__init__.py,sha256=T-ocYLLphudkQgzvG47jBOm5GQsRFRIGA88y7Ur4akg,341
-digitalkin/services/storage/default_storage.py,sha256=bHNPm8nLvytKqKP2ntLkikvqH1qDKOwofrStVJH6PJg,7765
-digitalkin/services/storage/grpc_storage.py,sha256=eSadiI3JuveItx8LdiC3GbjbjnB8GmPqTVxw1uA1-9E,5809
-digitalkin/services/storage/storage_strategy.py,sha256=vGo4aYkEp_GZV11m7vd-xY_Z3gVa5K0gMTzbj2Au_3o,6600
+digitalkin/services/storage/default_storage.py,sha256=qzLPrND92NR9hB7Ok6BF3Yxot14Efa_CHIvVte6kLsU,7817
+digitalkin/services/storage/grpc_storage.py,sha256=uFWxTsh2WXx2MY9W_dLSgpymZaLamN81RhyTDlhl_K8,6937
+digitalkin/services/storage/storage_strategy.py,sha256=Om05zryRfTwlI6l_H_eKM_6ijhuRQbV5Yltn2cBhh2A,8633
 digitalkin/utils/__init__.py,sha256=sJnY-ZUgsjMfojAjONC1VN14mhgIDnzyOlGkw21rRnM,28
 digitalkin/utils/arg_parser.py,sha256=3YyI6oZhhrlTmPTrzlwpQzbCNWDFAT3pggcLxNtJoc0,4388
 digitalkin/utils/llm_ready_schema.py,sha256=JjMug_lrQllqFoanaC091VgOqwAd-_YzcpqFlS7p778,2375
-digitalkin-0.2.7.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
+digitalkin-0.2.8.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
 modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 modules/minimal_llm_module.py,sha256=W-E3OrRbAsRJ6hvSeTU8pzmacdJC_PbcWfDapRv5A1A,5617
-modules/storage_module.py,sha256=bu52lW4RFcWB8VFDhrpBFfCaTSkVL6so3zrkfW4LO9E,6270
+modules/storage_module.py,sha256=528tfWyRw3q5h0lUlP9Or8E7m3AnnyXezXgwisXd8BI,6399
 modules/text_transform_module.py,sha256=1KaA7abwxltKKtbmiW1rkkIK3BTYFPegUq54px0LOQs,7277
-digitalkin-0.2.7.dist-info/METADATA,sha256=1-9MQvIcY6e4S5Cfzea9Nk2lHFx7coGQZ935DzJ00zg,29125
-digitalkin-0.2.7.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
-digitalkin-0.2.7.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
-digitalkin-0.2.7.dist-info/RECORD,,
+digitalkin-0.2.8.dist-info/METADATA,sha256=8lmICvmb-B3owpEGkyFk0vmIuVNabxkYORbaoWS5fqU,29125
+digitalkin-0.2.8.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+digitalkin-0.2.8.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
+digitalkin-0.2.8.dist-info/RECORD,,

{digitalkin-0.2.7.dist-info → digitalkin-0.2.8.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.0.0)
+Generator: setuptools (80.3.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

modules/storage_module.py

@@ -119,7 +119,12 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
         )
 
         # Store the output data in storage
-        storage_id = self.storage.store("example_outputs", output_data.model_dump(), data_type="OUTPUT")
+        storage_id = self.storage.store(
+            collection="example",
+            record_id=f"example_outputs",
+            data=output_data.model_dump(),
+            data_type="OUTPUT"
+        )
 
         logger.info("Stored output data with ID: %s", storage_id)
 
@@ -159,7 +164,7 @@ async def test_module() -> None:
 
     # Check the storage
     if module.status == ModuleStatus.STOPPED:
-        result: StorageRecord = module.storage.read("example_outputs")
+        result: StorageRecord = module.storage.read("example", "example_outputs")
         if result:
             pass
 
@@ -170,10 +175,10 @@ def test_storage_directly() -> None:
     storage = ServicesConfig().storage(mission_id="test-mission", config={"test_table": ExampleStorage})
 
     # Create a test record
-    storage.store("test_table", {"test_key": "test_value"}, "OUTPUT")
+    storage.store("example", "test_table", {"test_key": "test_value"}, "OUTPUT")
 
     # Retrieve the record
-    retrieved = storage.read("test_table")
+    retrieved = storage.read("example", "test_table")
 
     if retrieved:
         pass