ADD validator - RM from dataset

Author: eboileau

Diff for: server/src/scimodom/api/management.py

@@ -9,8 +9,8 @@
 from scimodom.config import get_config
 from scimodom.services.assembly import LiftOverError
 
-from scimodom.services.dataset import (
-    get_dataset_service,
+from scimodom.services.dataset import get_dataset_service
+from scimodom.services.validator import (
     SelectionNotFoundError,
     DatasetImportError,
     DatasetHeaderError,

Diff for: server/src/scimodom/cli/assembly.py

@@ -2,6 +2,7 @@
 
 from scimodom.services.assembly import (
     AssemblyVersionError,
+    AssemblyAbortedError,
     get_assembly_service,
 )
 
@@ -28,20 +29,20 @@ def add_assembly(**kwargs) -> None:
             return
         try:
             assembly_service.prepare_assembly_for_version(assembly_id)
-        except AssemblyVersionError:
+        except FileExistsError:
             click.secho(
-                "Cannot create assembly for this database version... Aborting!",
+                "Assembly directory already exists... Aborting!",
                 fg="red",
             )
             return
-        except FileExistsError:
+        except AssemblyVersionError as exc:
             click.secho(
-                "Assembly directory already exists... Aborting!",
+                f"Cannot create assembly for this database version: {exc}",
                 fg="red",
             )
             return
-        except Exception as exc:
-            click.secho(f"Failed to prepare assembly: {exc}. Aborting!", fg="red")
+        except AssemblyAbortedError as exc:
+            click.secho(f"Failed to prepare assembly: {exc}", fg="red")
             return
         click.secho("... done!", fg="green")
     else:
@@ -56,16 +57,8 @@ def add_assembly(**kwargs) -> None:
         if c not in ["y", "Y"]:
             return
         try:
-            assembly_id = assembly_service.add_assembly(taxa_id, assembly_name)
-        except FileExistsError:
-            click.secho(
-                "Directory exists, but not assembly... check for data corruption. Aborting!",
-                fg="red",
-            )
-            return
-        except Exception as exc:
-            click.secho(
-                f"Failed to add alternative assembly: {exc}. Aborting!", fg="red"
-            )
+            assembly_id = assembly_service.get_assembly_by_name(taxa_id, assembly_name)
+        except AssemblyAbortedError as exc:
+            click.secho(f"Failed to add alternative assembly: {exc}", fg="red")
             return
         click.secho(f"... done! Assembly ID is {assembly_id}.", fg="green")

Diff for: server/src/scimodom/cli/utilities.py

@@ -51,8 +51,8 @@ def validate_dataset_title(ctx, param, value):
 def add_assembly_to_template_if_none(organism, assembly_service):
     click.secho("Checking if assembly ID is defined...", fg="green")
     if organism.assembly_id is None:
-        assembly_id = assembly_service.add_assembly(
-            organism.taxa_id, organism.assembly_name
+        assembly_id = assembly_service.get_assembly_by_name(
+            organism.taxa_id, organism.assembly_name, fail_safe=False
         )
         click.secho(
             f"Updating project metadata template with assembly ID '{assembly_id}'... ",

Diff for: server/src/scimodom/services/assembly.py

@@ -2,7 +2,7 @@
 import logging
 from functools import cache
 from posixpath import join as urljoin
-from typing import Any, Sequence
+from typing import Any, Sequence, TextIO
 
 from sqlalchemy import select, func
 from sqlalchemy.exc import NoResultFound
@@ -25,13 +25,20 @@
 
 
 class AssemblyNotFoundError(Exception):
-    """Exception handling for a non-existing Assembly."""
+    """Exception for handling a non-existing Assembly."""
 
     pass
 
 
 class AssemblyVersionError(Exception):
-    """Exception handling for Assembly version mismatch."""
+    """Exception for handling an Assembly version mismatch."""
+
+    pass
+
+
+class AssemblyAbortedError(Exception):
+    """Exception for handling general errors associated with
+    preparing/adding assemblies e.g. request streaming, etc."""
 
     pass
 
@@ -86,6 +93,36 @@ def get_assembly_by_id(self, assembly_id: int) -> Assembly:
         except NoResultFound:
             raise AssemblyNotFoundError(f"No such assembly with ID: {assembly_id}.")
 
+    def get_assembly_by_name(
+        self, taxa_id: int, assembly_name: str, fail_safe: bool = True
+    ) -> Assembly:
+        """Retrieve assembly by name. If not found, add it to the database,
+        unless fail_safe is False, in which case raises AssemblyNotFoundError.
+
+        :param taxa_id: Taxonomy ID
+        :type taxa_id: int
+        :param assembly_name: Assembly name
+        :type assembly_name: str
+        :param fail_safe: If True (default), add assembly if not found.
+        :type fail_safe: bool
+        :returns: Newly created or existing assembly
+        :rtype: Assembly
+
+        :raises: AssemblyNotFoundError
+        """
+        try:
+            assembly = self._session.execute(
+                select(Assembly).filter_by(taxa_id=taxa_id, name=assembly_name)
+            ).scalar_one()
+            return assembly
+        except NoResultFound:
+            if fail_safe:
+                return self._add_assembly(taxa_id, assembly_name)
+            else:
+                raise AssemblyNotFoundError(
+                    f"No such assembly '{assembly_name}' for organism '{taxa_id}'."
+                )
+
     def get_assemblies_by_taxa(self, taxa_id: int) -> Sequence[Assembly]:
         """Retrieve all assemblies for a given organism.
 
@@ -158,13 +195,13 @@ def _yield_chroms():
 
         return list(_yield_chroms())
 
-    def liftover(
+    def create_lifted_file(
         self,
         assembly: Assembly,
         raw_file: str,
         unmapped_file: str | None = None,
         threshold: float = ImportLimits.LIFTOVER.max,
-    ) -> str:
+    ) -> TextIO:
         """Liftover records to current assembly.
 
         :param assembly: Assembly instance
@@ -175,8 +212,8 @@ def liftover(
         :type unmapped_file: str | None
         :param threshold: Threshold for raising LiftOverError
         :type threshold: float
-        :returns: Files pointing to the liftedOver features
-        :rtype: str
+        :returns: File handle pointing to the liftedOver features
+        :rtype: TextIO
         """
         if self.is_latest_assembly(assembly):
             raise AssemblyVersionError("Cannot liftover for latest assembly.")
@@ -207,60 +244,7 @@ def liftover(
                 f"{unmapped_lines} records could not be mapped... "
                 "Contact the system administrator if you have questions."
             )
-        return lifted_file
-
-    def add_assembly(self, taxa_id: int, assembly_name: str) -> int:
-        """Add an alternative assembly to the database.
-
-        If assembly exists, nothing is done.
-
-        :param taxa_id: Taxonomy ID
-        :type taxa_id: int
-        :param assembly_name: Assembly name
-        :type assembly_name: str
-        :returns: Newly created or existing assembly ID
-        :rtype: int
-        """
-        try:
-            assembly = self._session.execute(
-                select(Assembly).filter_by(taxa_id=taxa_id, name=assembly_name)
-            ).scalar_one()
-            return assembly.id
-        except NoResultFound:
-            pass
-
-        if self._file_service.check_if_assembly_exists(taxa_id, assembly_name):
-            raise FileExistsError(
-                f"Directory exists, but assembly '{assembly_name}' does not exist!"
-            )
-
-        chain_file_name = self._get_chain_file_name(
-            assembly_name, self.get_name_for_version(taxa_id)
-        )
-        url = self._get_ensembl_chain_file_url(taxa_id, chain_file_name)
-
-        logger.info(f"Setting up a new assembly for {assembly_name}...")
-        try:
-            with self._file_service.create_chain_file(
-                taxa_id, chain_file_name, assembly_name
-            ) as fh:
-                self._web_service.stream_request_to_file(url, fh)
-            version_nums = (
-                self._session.execute(select(func.distinct(Assembly.version)))
-                .scalars()
-                .all()
-            )
-            version_num = gen_short_uuid(Identifiers.ASSEMBLY.length, version_nums)
-            assembly = Assembly(
-                name=assembly_name, taxa_id=taxa_id, version=version_num
-            )
-            self._session.add(assembly)
-            self._session.commit()
-            return assembly.id
-        except Exception:
-            self._session.rollback()
-            self._file_service.delete_assembly(taxa_id, assembly_name)
-            raise
+        return self._file_service.open_file_for_reading(lifted_file)
 
     def prepare_assembly_for_version(self, assembly_id: int) -> None:
         """Prepare directories and files for the latest version.
@@ -288,14 +272,56 @@ def prepare_assembly_for_version(self, assembly_id: int) -> None:
         try:
             self._handle_gene_build(assembly)
             self._handle_release(assembly)
-        except Exception:
+        except AssemblyVersionError:
             self._file_service.delete_assembly(assembly.taxa_id, assembly.name)
             raise
+        except Exception as exc:
+            self._file_service.delete_assembly(assembly.taxa_id, assembly.name)
+            raise AssemblyAbortedError(
+                f"Adding assembly for ID '{assembly_id}' aborted."
+            ) from exc
 
     @staticmethod
     def _get_chain_file_name(source_assembly_name, target_assembly_name):
         return f"{source_assembly_name}_to_{target_assembly_name}.chain.gz"
 
+    def _add_assembly(self, taxa_id: int, assembly_name: str) -> Assembly:
+        if self._file_service.check_if_assembly_exists(taxa_id, assembly_name):
+            raise AssemblyAbortedError(
+                f"Suspected incomplete or inconsistent data: files were found on the "
+                f"system for '{assembly_name}', but assembly does not exist in the database."
+            )
+
+        chain_file_name = self._get_chain_file_name(
+            assembly_name, self.get_name_for_version(taxa_id)
+        )
+        url = self._get_ensembl_chain_file_url(taxa_id, chain_file_name)
+
+        logger.info(f"Setting up a new assembly for {assembly_name}...")
+        try:
+            with self._file_service.create_chain_file(
+                taxa_id, chain_file_name, assembly_name
+            ) as fh:
+                self._web_service.stream_request_to_file(url, fh)
+            version_nums = (
+                self._session.execute(select(func.distinct(Assembly.version)))
+                .scalars()
+                .all()
+            )
+            version_num = gen_short_uuid(Identifiers.ASSEMBLY.length, version_nums)
+            assembly = Assembly(
+                name=assembly_name, taxa_id=taxa_id, version=version_num
+            )
+            self._session.add(assembly)
+            self._session.commit()
+            return assembly
+        except Exception as exc:
+            self._session.rollback()
+            self._file_service.delete_assembly(taxa_id, assembly_name)
+            raise AssemblyAbortedError(
+                f"Adding assembly for '{assembly_name}' aborted."
+            ) from exc
+
     def _get_ensembl_chain_file_url(self, taxa_id: int, chain_file_name):
         return urljoin(
             Ensembl.FTP.value,

Diff for: server/src/scimodom/services/data.py

@@ -17,7 +17,7 @@ class NoDataRecords(Exception):
 
 class DataService:
     def __init__(self, session: Session):
-        self._db_session = session
+        self._session = session
 
     def get_by_dataset(
         self, datasets: Union[str, Dataset, List[Union[str, Dataset]]]
@@ -30,7 +30,7 @@ def get_by_dataset(
             .where(Data.dataset_id.in_(dataset_ids))
         )
         count = 0
-        for record in self._db_session.execute(query).all():
+        for record in self._session.execute(query).all():
             count += 1
             yield record[0]
         if count == 0: