diff --git a/lib/couchdb/document.py b/lib/couchdb/document.py index 61efdef..373c3c4 100644 --- a/lib/couchdb/document.py +++ b/lib/couchdb/document.py @@ -80,9 +80,11 @@ def to_dict(self) -> Dict[str, Any]: def add_sample( self, sample_id: str, - lib_prep_option: str, + # lib_prep_option: str, status: str = "pending", flowcell_ids_processed_for: Optional[List[str]] = None, + start_time: Optional[str] = None, + end_time: Optional[str] = None, ) -> None: """Adds a new sample to the document. @@ -91,16 +93,38 @@ def add_sample( lib_prep_option (str): The library preparation option. status (str, optional): The status of the sample. Defaults to "pending". flowcell_ids_processed_for (List[str], optional): Flowcell IDs the sample has been processed for. + start_time (str, optional): Start time of the sample processing. + end_time (str, optional): End time of the sample processing. """ - sample = { - "sample_id": sample_id, - "status": status, - "lib_prep_option": lib_prep_option, - "start_time": "", - "end_time": "", - "flowcell_ids_processed_for": flowcell_ids_processed_for or [], - } - self.samples.append(sample) + existing_sample = self.get_sample(sample_id) + if existing_sample: + # Update existing sample + existing_sample["status"] = status + if flowcell_ids_processed_for: + existing_sample["flowcell_ids_processed_for"].extend( + flowcell_ids_processed_for + ) + # Remove duplicates + existing_sample["flowcell_ids_processed_for"] = list( + set(existing_sample["flowcell_ids_processed_for"]) + ) + if start_time: + existing_sample["start_time"] = start_time + if end_time: + existing_sample["end_time"] = end_time + # logging.debug(f"Updated sample: {existing_sample}") + else: + # Add new sample + sample = { + "sample_id": sample_id, + "status": status, + # "lib_prep_option": lib_prep_option, + "start_time": start_time or "", + "end_time": end_time or "", + "flowcell_ids_processed_for": flowcell_ids_processed_for or [], + } + self.samples.append(sample) + # logging.debug(f"Added sample: {sample}") def update_sample_status(self, sample_id: str, status: str) -> None: """Updates the status of a specific sample. @@ -132,7 +156,7 @@ def update_sample_status(self, sample_id: str, status: str) -> None: self.check_project_completion() def get_sample(self, sample_id: str) -> Optional[Dict[str, Any]]: - """Retrieves a specific sample by its ID. + """Retrieves a specific sample from the samples list by its ID. Args: sample_id (str): The sample ID to retrieve. diff --git a/lib/couchdb/manager.py b/lib/couchdb/manager.py index ed24be0..24b95b2 100644 --- a/lib/couchdb/manager.py +++ b/lib/couchdb/manager.py @@ -255,6 +255,33 @@ def get_document_by_project_id(self, project_id: str) -> Optional[Dict[str, Any] logging.error(f"Error accessing project: {e}") return None + def add_sample( + self, + project_id: str, + sample_id: str, + # lib_prep_option: str, + status: str = "pending", + ) -> None: + """Adds a sample to a project and saves the document. + + Args: + project_id (str): The project ID. + sample_id (str): The sample ID. + lib_prep_option (str): The library preparation option. + status (str, optional): The status of the sample. Defaults to "pending". + """ + try: + document_dict = self.get_document_by_project_id(project_id) + if document_dict: + ygg_doc = YggdrasilDocument.from_dict(document_dict) + ygg_doc.add_sample(sample_id, status) + self.save_document(ygg_doc) + logging.info(f"Updated project {project_id} with sample.") + else: + logging.error(f"Project {project_id} does not exist in YggdrasilDB.") + except Exception as e: + logging.error(f"Error adding sample: {e}") + def update_sample_status( self, project_id: str, sample_id: str, status: str ) -> None: