From 73d438a51e237d4a4435cf3022a0b4f47cb337be Mon Sep 17 00:00:00 2001
From: glrs <5999366+glrs@users.noreply.github.com>
Date: Wed, 20 Nov 2024 11:08:59 +0100
Subject: [PATCH] Update to interact with YggDB and follow AbstractProject
 changes

---
 lib/realms/smartseq3/ss3_project.py | 15 ++++++++++-----
 lib/realms/smartseq3/ss3_sample.py  |  9 ++++++++-
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/lib/realms/smartseq3/ss3_project.py b/lib/realms/smartseq3/ss3_project.py
index 2e07147..e15f8dc 100644
--- a/lib/realms/smartseq3/ss3_project.py
+++ b/lib/realms/smartseq3/ss3_project.py
@@ -4,10 +4,9 @@
 from lib.base.abstract_project import AbstractProject
 from lib.core_utils.config_loader import ConfigLoader
 from lib.core_utils.logging_utils import custom_logger
-from lib.module_utils.ngi_report_generator import generate_ngi_report
 
 # from datetime import datetime
-# from lib.couchdb.manager import YggdrasilDBManager
+from lib.module_utils.ngi_report_generator import generate_ngi_report
 from lib.realms.smartseq3.ss3_sample import SS3Sample
 
 logging = custom_logger("SS3 Project")
@@ -35,12 +34,14 @@ def __init__(self, doc, yggdrasil_db_manager):
         Args:
             doc (dict): Document containing project metadata.
         """
+        super().__init__(doc, yggdrasil_db_manager)
         self.doc = doc
         self.ydm = yggdrasil_db_manager
-        self.proceed = self._check_required_fields()
+        self.proceed = self.check_required_fields()
 
         # TODO: What if I return None if not self.proceed?
         if self.proceed:
+            self.initialize_project_in_db()
             self.project_info = self._extract_project_info()
             self.project_dir = self.ensure_project_directory()
             self.project_info["project_dir"] = self.project_dir
@@ -74,7 +75,7 @@ def _extract_project_info(self):
                 {}
             )  # Return an empty dict or some default values to allow continuation
 
-    def _check_required_fields(self):
+    def check_required_fields(self):
         """
         Checks if the document contains all required fields.
 
@@ -154,6 +155,8 @@ async def launch(self):
             logging.warning("No samples found for processing. Returning...")
             return
 
+        self.add_samples_to_project_in_db()
+
         # Pre-process samples
         pre_tasks = [sample.pre_process() for sample in self.samples]
         await asyncio.gather(*pre_tasks)
@@ -201,7 +204,9 @@ def extract_samples(self):
         samples = []
 
         for sample_id, sample_data in self.doc.get("samples", {}).items():
-            sample = SS3Sample(sample_id, sample_data, self.project_info, self.config)
+            sample = SS3Sample(
+                sample_id, sample_data, self.project_info, self.config, self.ydm
+            )
 
             if sample.flowcell_id:
                 samples.append(sample)
diff --git a/lib/realms/smartseq3/ss3_sample.py b/lib/realms/smartseq3/ss3_sample.py
index cde756a..1a33cd9 100644
--- a/lib/realms/smartseq3/ss3_sample.py
+++ b/lib/realms/smartseq3/ss3_sample.py
@@ -30,7 +30,9 @@ class SS3Sample(AbstractSample):
         file_handler (SampleFileHandler): Handler for sample files.
     """
 
-    def __init__(self, sample_id, sample_data, project_info, config):
+    def __init__(
+        self, sample_id, sample_data, project_info, config, yggdrasil_db_manager
+    ):
         """
         Initialize a SmartSeq3 sample instance.
 
@@ -44,6 +46,8 @@ def __init__(self, sample_id, sample_data, project_info, config):
         self._id = sample_id
         self.sample_data = sample_data
         self.project_info = project_info
+        # TODO: ensure project_id is always available
+        self.project_id = self.project_info.get("project_id", "")
 
         # Initialize barcode
         self.barcode = self.get_barcode()
@@ -52,6 +56,7 @@ def __init__(self, sample_id, sample_data, project_info, config):
         self.flowcell_id = self._get_latest_flowcell()
 
         self.config = config
+        self.ydm = yggdrasil_db_manager
         # self.job_id = None
 
         # TODO: Currently not used much, but should be used if we write to a database
@@ -79,6 +84,8 @@ def status(self):
     @status.setter
     def status(self, value):
         self._status = value
+        # Update the status in the database
+        self.ydm.update_sample_status(self.project_id, self.id, value)
 
     async def pre_process(self):
         """Pre-process the sample by collecting metadata and creating YAML files."""