Refactor Dataset.session_list() to match MATLAB API and improve error handling

src/ndi/database_fun.py

@@ -390,7 +390,7 @@ def copy_session_to_dataset(
 
     # Check for already-copied sessions
     try:
-        refs, session_ids = ndi_dataset_obj.session_list()
+        refs, session_ids, *_ = ndi_dataset_obj.session_list()
         session_id = ndi_session_obj.id()
         if session_id in session_ids:
             return (

src/ndi/dataset/_dataset.py

@@ -108,16 +108,23 @@ def add_linked_session(self, session: Any) -> Dataset:
 
         Returns:
             self for chaining
+
+        Raises:
+            ValueError: If the session is already part of this dataset
         """
-        # Check if already linked
         existing = self._find_session_doc(session.id())
         if existing is not None:
-            return self
+            raise ValueError(
+                f"Session with id {session.id()} is already part of " f"dataset {self.id()}."
+            )
 
         # Create session_in_a_dataset document
         doc = self._create_session_doc(session, is_linked=True)
         self._session.database_add(doc)
 
+        # Cache the session object
+        self._session_cache[session.id()] = session
+
         return self
 
     def add_ingested_session(self, session: Any) -> Dataset:
@@ -132,11 +139,25 @@ def add_ingested_session(self, session: Any) -> Dataset:
 
         Returns:
             self for chaining
+
+        Raises:
+            ValueError: If the session is already part of this dataset
+            ValueError: If the session is not fully ingested
         """
-        # Check if already present
         existing = self._find_session_doc(session.id())
         if existing is not None:
-            return self
+            raise ValueError(
+                f"Session with id {session.id()} is already part of " f"dataset {self.id()}."
+            )
+
+        # Check if session is fully ingested
+        if hasattr(session, "is_fully_ingested") and not session.is_fully_ingested():
+            raise ValueError(
+                f"Session with id {session.id()} and reference "
+                f"{session.reference} is not yet fully ingested. "
+                f"It must be fully ingested before it can be added "
+                f"in ingested form to a dataset."
+            )
 
         # Copy all documents from source session into the dataset's database.
         # We bypass session.database_add() because it enforces session_id ==
@@ -160,25 +181,38 @@ def add_ingested_session(self, session: Any) -> Dataset:
     def unlink_session(
         self,
         session_id: str,
-        remove_documents: bool = False,
+        are_you_sure: bool = False,
     ) -> Dataset:
         """
-        Remove a session from this dataset.
+        Unlink a linked session from this dataset.
+
+        The session must be a linked session (not ingested). Use
+        delete_ingested_session() for ingested sessions.
 
         Args:
             session_id: ID of the session to unlink
-            remove_documents: If True, also remove ingested documents
+            are_you_sure: Must be True to proceed
 
         Returns:
             self for chaining
+
+        Raises:
+            ValueError: If not confirmed, session not found, or session is ingested
         """
+        if not are_you_sure:
+            raise ValueError("Must set are_you_sure=True to unlink a session.")
+
         doc = self._find_session_doc(session_id)
         if doc is None:
-            return self
+            raise ValueError(f"Session with ID {session_id} not found in " f"dataset {self.id()}.")
 
-        # Optionally remove ingested documents
-        if remove_documents:
-            self._remove_session_documents(session_id)
+        props = doc.document_properties.get("session_in_a_dataset", {})
+        if not props.get("is_linked", False):
+            raise ValueError(
+                f"Session with ID {session_id} is an INGESTED session, "
+                f"not a linked session. Cannot unlink. Use "
+                f"delete_ingested_session() instead."
+            )
 
         # Remove the session_in_a_dataset document
         self._session.database_rm(doc)
@@ -224,38 +258,43 @@ def open_session(self, session_id: str) -> Any | None:
 
         return session
 
-    def session_list(self) -> tuple[list[str], list[dict[str, Any]]]:
+    def session_list(
+        self,
+    ) -> tuple[list[str], list[str], list[str], str]:
         """
         List all sessions in this dataset.
 
         Returns:
-            A tuple of (session_references, session_list) where:
-                - session_references: List of session reference strings
-                - session_list: List of dicts with keys:
-                    - session_id: Session identifier
-                    - session_reference: Session reference name
-                    - is_linked: True if linked, False if ingested
-                    - document_id: ID of the session_in_a_dataset document
+            A tuple of (ref_list, id_list, session_doc_ids, dataset_session_doc_id):
+                - ref_list: List of session reference strings
+                - id_list: List of session ID strings
+                - session_doc_ids: List of document IDs for the
+                  session_in_a_dataset documents
+                - dataset_session_doc_id: Document ID of the dataset's
+                  own session document (empty string if not found)
         """
         q = Query("").isa("session_in_a_dataset")
         docs = self._session.database_search(q)
 
-        session_references = []
-        session_list = []
+        ref_list: list[str] = []
+        id_list: list[str] = []
+        session_doc_ids: list[str] = []
         for doc in docs:
             props = doc.document_properties.get("session_in_a_dataset", {})
-            ref = props.get("session_reference", "")
-            session_references.append(ref)
-            session_list.append(
-                {
-                    "session_id": props.get("session_id", ""),
-                    "session_reference": ref,
-                    "is_linked": bool(props.get("is_linked", False)),
-                    "document_id": doc.id,
-                }
-            )
+            ref_list.append(props.get("session_reference", ""))
+            id_list.append(props.get("session_id", ""))
+            session_doc_ids.append(doc.id)
+
+        # Find the dataset's own session document
+        dataset_session_doc_id = ""
+        q_ds = Query("").isa("session") & (Query("base.session_id") == self.id())
+        ds_docs = self._session.database_search(q_ds)
+        if len(ds_docs) == 1:
+            dataset_session_doc_id = ds_docs[0].id
+        elif len(ds_docs) > 1:
+            raise ValueError("More than 1 session document for the dataset session found.")
 
-        return session_references, session_list
+        return ref_list, id_list, session_doc_ids, dataset_session_doc_id
 
     # =========================================================================
     # Database Operations (delegated to internal session)
@@ -276,14 +315,34 @@ def database_rm(
         return self
 
     def database_search(self, query: Query) -> list[Document]:
-        """Search the dataset database.
+        """Search the dataset database and all linked sessions.
 
         Unlike Session.database_search(), this does NOT filter by session_id
         because a dataset stores documents from multiple ingested sessions.
+        Results from linked sessions are also included.
         """
         if self._session._database is None:
-            return []
-        return self._session._database.search(query)
+            results: list[Document] = []
+        else:
+            results = list(self._session._database.search(query))
+
+        # Also search linked sessions
+        self._open_linked_sessions()
+        q = Query("").isa("session_in_a_dataset")
+        info_docs = self._session.database_search(q)
+        for info_doc in info_docs:
+            props = info_doc.document_properties.get("session_in_a_dataset", {})
+            if props.get("is_linked", False):
+                sid = props.get("session_id", "")
+                session = self._session_cache.get(sid)
+                if session is not None:
+                    try:
+                        linked_results = session.database_search(query)
+                        results.extend(linked_results)
+                    except Exception:
+                        pass
+
+        return results
 
     def database_openbinarydoc(
         self,
@@ -324,11 +383,14 @@ def delete_ingested_session(
 
         doc = self._find_session_doc(session_id)
         if doc is None:
-            return self
+            raise ValueError(f"Session {session_id} not found in dataset.")
 
         props = doc.document_properties.get("session_in_a_dataset", {})
         if props.get("is_linked", False):
-            raise ValueError("Cannot delete a linked session. Use unlink_session() instead.")
+            raise ValueError(
+                f"Session {session_id} is a linked session, not an "
+                f"ingested one. Use unlink_session() instead."
+            )
 
         # Remove all documents from this session
         self._remove_session_documents(session_id)
@@ -432,6 +494,17 @@ def _remove_session_documents(self, session_id: str) -> None:
             except Exception as exc:
                 logger.warning("Failed to remove document %s: %s", doc.id, exc)
 
+    def _open_linked_sessions(self) -> None:
+        """Ensure all linked sessions are open and cached."""
+        q = Query("").isa("session_in_a_dataset")
+        docs = self._session.database_search(q)
+        for doc in docs:
+            props = doc.document_properties.get("session_in_a_dataset", {})
+            if props.get("is_linked", False):
+                sid = props.get("session_id", "")
+                if sid and sid not in self._session_cache:
+                    self.open_session(sid)
+
     def _recreate_linked_session(self, props: dict[str, Any]) -> Any | None:
         """Recreate a linked session from stored creator args."""
         creator = props.get("session_creator", "")
@@ -465,5 +538,5 @@ def _recreate_linked_session(self, props: dict[str, Any]) -> Any | None:
 
     def __repr__(self) -> str:
         """String representation."""
-        sessions = self.session_list()
-        return f"Dataset('{self._reference}', sessions={len(sessions)})"
+        refs, _ids, _doc_ids, _ds_doc_id = self.session_list()
+        return f"Dataset('{self._reference}', sessions={len(refs)})"