Repository: AlexsLemonade/refinebio

Issue:
Transcriptome indices won't make both long and short versions
### Context

I'm pretty sure this is an issue that should affect prod. I saw it locally as part of #2035. While I did get transcriptome indices made successfully, I didn't look hard enough and it appears that only one of the two jobs succeeded because after the first the input files are deleted and the second job doesn't have the files it needs.

### Problem or idea

This was working at some point. I can't figure out exactly where things went wrong, but my guess is that we moved the file deletion into the `end_job` and didn't properly make sure to protect the transcriptome-specific files at least until the second job ran.

### Solution or next step

I think we need to add a special case to `processor.utils.end_job` to make sure that there's no other processor jobs associated with the file before deleting it.



Discussion:

Resolving patch:
diff --git a/foreman/data_refinery_foreman/surveyor/transcriptome_index.py b/foreman/data_refinery_foreman/surveyor/transcriptome_index.py
--- a/foreman/data_refinery_foreman/surveyor/transcriptome_index.py
+++ b/foreman/data_refinery_foreman/surveyor/transcriptome_index.py
@@ -1,11 +1,12 @@
+import csv
 import re
+import shutil
 import urllib
 
 from abc import ABC
-from django.utils import timezone
 from typing import List, Dict
 
-from data_refinery_common.job_lookup import ProcessorPipeline, Downloaders
+from data_refinery_common.job_lookup import Downloaders
 from data_refinery_common.logging import get_and_configure_logger
 from data_refinery_common.models import (
     OriginalFile,
@@ -17,17 +18,23 @@
 logger = get_and_configure_logger(__name__)
 
 
+CHUNK_SIZE = 1024 * 256  # chunk_size is in bytes
+
+
 MAIN_DIVISION_URL_TEMPLATE = "https://rest.ensembl.org/info/species?content-type=application/json"
 DIVISION_URL_TEMPLATE = (
-    "https://rest.ensembl.org/info/genomes/division/{division}" "?content-type=application/json"
+    "https://rest.ensembl.org/info/genomes/division/{division}?content-type=application/json"
 )
 
+SPECIES_DETAIL_URL_TEMPLATE = (
+    "ftp://ftp.ensemblgenomes.org/pub/{short_division}/current/species_{division}.txt"
+)
 TRANSCRIPTOME_URL_TEMPLATE = (
-    "ftp://ftp.{url_root}/fasta/{species_sub_dir}/dna/"
+    "ftp://ftp.{url_root}/fasta/{collection}{species_sub_dir}/dna/"
     "{filename_species}.{assembly}.dna.{schema_type}.fa.gz"
 )
 GTF_URL_TEMPLATE = (
-    "ftp://ftp.{url_root}/gtf/{species_sub_dir}/"
+    "ftp://ftp.{url_root}/gtf/{collection}{species_sub_dir}/"
     "{filename_species}.{assembly}.{assembly_version}.gtf.gz"
 )
 
@@ -52,6 +59,49 @@
 DIVISION_RELEASE_URL = "https://rest.ensembl.org/info/eg_version?content-type=application/json"
 
 
+def get_strain_mapping_for_organism(
+    species_name: str, config_file="config/organism_strain_mapping.csv"
+) -> List[Dict]:
+    """Returns the row of the strain/organism mapping for the species_name
+    """
+    upper_name = species_name.upper()
+    with open(config_file) as csvfile:
+        reader = csv.DictReader(csvfile)
+        for row in reader:
+            if row["organism"] == upper_name:
+                return row
+
+    return None
+
+
+def get_species_detail_by_assembly(assembly: str, division: str) -> str:
+    """Returns additional detail about a species given an assembly and a division.
+
+    These details are necessary because the FTP directory for
+    EnsemblBacteria and EnsemblFungi have an additional level in their
+    paths that can only be determined by parsing this file. I found
+    this out via the Ensembl dev mailing list.
+    """
+    bacteria_species_detail_url = SPECIES_DETAIL_URL_TEMPLATE.format(
+        short_division=DIVISION_LOOKUP[division], division=division
+    )
+
+    urllib.request.urlcleanup()
+    collection_path = assembly + "_collection.tsv"
+    with open(collection_path, "wb") as collection_file:
+        with urllib.request.urlopen(bacteria_species_detail_url) as request:
+            shutil.copyfileobj(request, collection_file, CHUNK_SIZE)
+
+    # Ancient unresolved bug. WTF python: https://bugs.python.org/issue27973
+    urllib.request.urlcleanup()
+
+    with open(collection_path) as csvfile:
+        reader = csv.DictReader(csvfile, delimiter="\t")
+        for row in reader:
+            if row["assembly"] == assembly:
+                return row
+
+
 class EnsemblUrlBuilder(ABC):
     """Generates URLs for different divisions of Ensembl.
 
@@ -65,12 +115,19 @@ class EnsemblUrlBuilder(ABC):
     def __init__(self, species: Dict):
         """Species is a Dict containing parsed JSON from the Division API."""
         self.url_root = "ensemblgenomes.org/pub/release-{assembly_version}/{short_division}"
+        self.division = species["division"]
         self.short_division = DIVISION_LOOKUP[species["division"]]
-        self.assembly = species["assembly_name"].replace(" ", "_")
-        self.assembly_version = (
-            utils.requests_retry_session().get(DIVISION_RELEASE_URL).json()["version"]
-        )
 
+        mapping = get_strain_mapping_for_organism(species["name"])
+        if mapping:
+            self.assembly = mapping["assembly"]
+            self.strain = mapping["strain"]
+        else:
+            self.assembly = species["assembly_name"].replace(" ", "_")
+            self.strain = None
+
+        assembly_response = utils.requests_retry_session().get(DIVISION_RELEASE_URL)
+        self.assembly_version = assembly_response.json()["version"]
         self.species_sub_dir = species["name"]
         self.filename_species = species["name"].capitalize()
 
@@ -79,6 +136,9 @@ def __init__(self, species: Dict):
         self.scientific_name = species["name"].upper()
         self.taxonomy_id = species["taxonomy_id"]
 
+        # This field is only needed for EnsemblBacteria and EnsemblFungi.
+        self.collection = ""
+
     def build_transcriptome_url(self) -> str:
         url_root = self.url_root.format(
             assembly_version=self.assembly_version, short_division=self.short_division
@@ -86,6 +146,7 @@ def build_transcriptome_url(self) -> str:
         url = TRANSCRIPTOME_URL_TEMPLATE.format(
             url_root=url_root,
             species_sub_dir=self.species_sub_dir,
+            collection=self.collection,
             filename_species=self.filename_species,
             assembly=self.assembly,
             schema_type="primary_assembly",
@@ -98,7 +159,7 @@ def build_transcriptome_url(self) -> str:
             file_handle = urllib.request.urlopen(url)
             file_handle.close()
             urllib.request.urlcleanup()
-        except:
+        except Exception:
             url = url.replace("primary_assembly", "toplevel")
 
         return url
@@ -110,6 +171,7 @@ def build_gtf_url(self) -> str:
         return GTF_URL_TEMPLATE.format(
             url_root=url_root,
             species_sub_dir=self.species_sub_dir,
+            collection=self.collection,
             filename_species=self.filename_species,
             assembly=self.assembly,
             assembly_version=self.assembly_version,
@@ -131,6 +193,7 @@ def __init__(self, species: Dict):
         self.url_root = "ensembl.org/pub/release-{assembly_version}"
         self.short_division = None
         self.species_sub_dir = species["name"]
+        self.collection = ""
         self.filename_species = species["name"].capitalize()
         self.assembly = species["assembly"]
         self.assembly_version = (
@@ -150,11 +213,35 @@ class EnsemblProtistsUrlBuilder(EnsemblUrlBuilder):
 
     def __init__(self, species: Dict):
         super().__init__(species)
-        self.filename_species = species["species"].capitalize()
+        self.filename_species = species["name"].capitalize()
+
+
+class EnsemblBacteriaUrlBuilder(EnsemblUrlBuilder):
+    """The EnsemblBacteria URLs are extra tricky because they have an extra layer in them.
+
+    This requires parsing a special file to find out what collection
+    the species belongs to.
+    """
+
+    def __init__(self, species: Dict):
+        super().__init__(species)
+        species_detail = get_species_detail_by_assembly(self.assembly, self.division)
 
+        if species_detail:
+            self.species_sub_dir = species_detail["species"]
+            self.filename_species = species_detail["species"].capitalize()
+            collection_pattern = r"bacteria_.+_collection"
+            match = re.match(collection_pattern, species_detail["core_db"])
+            if match:
+                # Need to append a / to the collection because it's not
+                # present in all the routes so we don't want to put it in the
+                # template and end up with a // in the path if collection is
+                # blank.
+                self.collection = match.group(0) + "/"
 
-class EnsemblFungiUrlBuilder(EnsemblProtistsUrlBuilder):
-    """The EnsemblFungi URLs work the similarly to Protists division.
+
+class EnsemblFungiUrlBuilder(EnsemblUrlBuilder):
+    """The EnsemblFungi URLs work similarly to EnsemblBacteria.
 
     EnsemblFungi is special because there is an assembly_name TIGR
     which needs to be corrected to CADRE for some reason.
@@ -165,6 +252,20 @@ def __init__(self, species: Dict):
         if self.assembly == "TIGR":
             self.assembly = "CADRE"
 
+        species_detail = get_species_detail_by_assembly(self.assembly, self.division)
+
+        if species_detail:
+            self.species_sub_dir = species_detail["species"]
+            self.filename_species = species_detail["species"].capitalize()
+            collection_pattern = r"fungi_.+_collection"
+            match = re.match(collection_pattern, species_detail["core_db"])
+            if match:
+                # Need to append a / to the collection because it's not
+                # present in all the routes so we don't want to put it in the
+                # template and end up with a // in the path if collection is
+                # blank.
+                self.collection = match.group(0) + "/"
+
 
 def ensembl_url_builder_factory(species: Dict) -> EnsemblUrlBuilder:
     """Returns instance of EnsemblUrlBuilder or one of its subclasses.
@@ -177,6 +278,8 @@ def ensembl_url_builder_factory(species: Dict) -> EnsemblUrlBuilder:
         return EnsemblFungiUrlBuilder(species)
     elif species["division"] == "EnsemblVertebrates":
         return MainEnsemblUrlBuilder(species)
+    elif species["division"] == "EnsemblBacteria":
+        return EnsemblBacteriaUrlBuilder(species)
     else:
         return EnsemblUrlBuilder(species)
 
@@ -251,7 +354,7 @@ def survey(self, source_type=None) -> bool:
             species_files = self.discover_species()
         except Exception:
             logger.exception(
-                ("Exception caught while discovering species. " "Terminating survey job."),
+                "Exception caught while discovering species. Terminating survey job.",
                 survey_job=self.survey_job.id,
             )
             return False
@@ -263,7 +366,7 @@ def survey(self, source_type=None) -> bool:
                 )
         except Exception:
             logger.exception(
-                ("Failed to queue downloader jobs. " "Terminating survey job."),
+                "Failed to queue downloader jobs. Terminating survey job.",
                 survey_job=self.survey_job.id,
             )
             return False
@@ -276,9 +379,37 @@ def discover_species(self):
         ).value
 
         logger.info(
-            "Surveying %s division of ensembl.", ensembl_division, survey_job=self.survey_job.id
+            "Surveying %s division of ensembl.", ensembl_division, survey_job=self.survey_job.id,
         )
 
+        try:
+            organism_name = SurveyJobKeyValue.objects.get(
+                survey_job_id=self.survey_job.id, key__exact="organism_name"
+            ).value
+            organism_name = organism_name.lower().replace(" ", "_")
+        except SurveyJobKeyValue.DoesNotExist:
+            organism_name = None
+
+        if ensembl_division in ["EnsemblFungi", "EnsemblBacteria"]:
+            if organism_name is None:
+                logger.error(
+                    "Organism name must be specified for Fungi and Bacteria divisions.",
+                    ensembl_division=ensembl_division,
+                    organism_name=organism_name,
+                )
+                return []
+            else:
+                if get_strain_mapping_for_organism(organism_name) is None:
+                    logger.error(
+                        (
+                            "Organism name must be listed in config/organism_strain_"
+                            "mappings.csv for Fungi and Bacteria divisions."
+                        ),
+                        ensembl_division=ensembl_division,
+                        organism_name=organism_name,
+                    )
+                    return []
+
         # The main division has a different base URL for its REST API.
         if ensembl_division == "Ensembl":
             r = utils.requests_retry_session().get(MAIN_DIVISION_URL_TEMPLATE)
@@ -287,28 +418,29 @@ def discover_species(self):
             # distinguish between a singlular species and multiple species.
             specieses = r.json()["species"]
         else:
-            r = utils.requests_retry_session().get(
-                DIVISION_URL_TEMPLATE.format(division=ensembl_division)
-            )
+            formatted_division_url = DIVISION_URL_TEMPLATE.format(division=ensembl_division)
+            r = utils.requests_retry_session().get(formatted_division_url)
             specieses = r.json()
 
-        try:
-            organism_name = SurveyJobKeyValue.objects.get(
-                survey_job_id=self.survey_job.id, key__exact="organism_name"
-            ).value
-            organism_name = organism_name.lower().replace(" ", "_")
-        except SurveyJobKeyValue.DoesNotExist:
-            organism_name = None
-
         all_new_species = []
         if organism_name:
             for species in specieses:
                 # This key varies based on whether the division is the
                 # main one or not... why couldn't they just make them
                 # consistent?
-                if ("species" in species and species["species"] == organism_name) or (
-                    "name" in species and species["name"] == organism_name
+                if ("species" in species and organism_name in species["species"]) or (
+                    "name" in species and organism_name in species["name"]
                 ):
+                    # Fungi have a strain identifier in their
+                    # names. This is different than everything else,
+                    # so we're going to handle this special case by
+                    # just overwriting this. This is okay because we
+                    # just have to discover one species for the
+                    # organism, and then our strain mapping will make
+                    # sure we use the correct strain and assembly.
+                    if ensembl_division == "EnsemblFungi" and organism_name != species["name"]:
+                        species["name"] = organism_name
+
                     all_new_species.append(self._generate_files(species))
                     break
         else:
diff --git a/workers/data_refinery_workers/processors/utils.py b/workers/data_refinery_workers/processors/utils.py
--- a/workers/data_refinery_workers/processors/utils.py
+++ b/workers/data_refinery_workers/processors/utils.py
@@ -320,7 +320,8 @@ def end_job(job_context: Dict, abort=False):
         # Cleanup Original Files
         if "original_files" in job_context:
             for original_file in job_context["original_files"]:
-                original_file.delete_local_file()
+                if original_file.needs_processing(job.id):
+                    original_file.delete_local_file()
 
     # If the pipeline includes any steps, save it.
     if "pipeline" in job_context:


Test patch:
diff --git a/foreman/data_refinery_foreman/surveyor/test_transcriptome_index.py b/foreman/data_refinery_foreman/surveyor/test_transcriptome_index.py
--- a/foreman/data_refinery_foreman/surveyor/test_transcriptome_index.py
+++ b/foreman/data_refinery_foreman/surveyor/test_transcriptome_index.py
@@ -1,29 +1,23 @@
-import os
-import json
 import urllib
-from unittest.mock import Mock, patch, call
-from unittest import skip
+from unittest.mock import patch, call
 from django.test import TestCase
-from urllib.request import URLError
 from data_refinery_common.job_lookup import Downloaders
-from data_refinery_common.models import DownloaderJob, SurveyJob, SurveyJobKeyValue, Organism
+from data_refinery_common.models import DownloaderJob, SurveyJob, SurveyJobKeyValue
 from data_refinery_foreman.surveyor.transcriptome_index import TranscriptomeIndexSurveyor
 
 
 class SurveyTestCase(TestCase):
-    def setUp(self):
+    @patch("data_refinery_foreman.surveyor.external_source.message_queue.send_job")
+    def test_survey(self, mock_send_job):
         survey_job = SurveyJob(source_type="TRANSCRIPTOME_INDEX")
         survey_job.save()
-        self.survey_job = survey_job
 
         key_value_pair = SurveyJobKeyValue(
             survey_job=survey_job, key="ensembl_division", value="EnsemblPlants"
         )
         key_value_pair.save()
 
-    @patch("data_refinery_foreman.surveyor.external_source.message_queue.send_job")
-    def test_survey(self, mock_send_job):
-        surveyor = TranscriptomeIndexSurveyor(self.survey_job)
+        surveyor = TranscriptomeIndexSurveyor(survey_job)
         surveyor.survey(source_type="TRANSCRIPTOME_INDEX")
 
         downloader_jobs = DownloaderJob.objects.order_by("id").all()
@@ -41,7 +35,6 @@ def test_correct_index_location(self):
         """
         survey_job = SurveyJob(source_type="TRANSCRIPTOME_INDEX")
         survey_job.save()
-        self.survey_job = survey_job
 
         key_value_pair = SurveyJobKeyValue(
             survey_job=survey_job, key="ensembl_division", value="Ensembl"
@@ -53,7 +46,7 @@ def test_correct_index_location(self):
         )
         key_value_pair.save()
 
-        surveyor = TranscriptomeIndexSurveyor(self.survey_job)
+        surveyor = TranscriptomeIndexSurveyor(survey_job)
         files = surveyor.discover_species()[0]
 
         for file in files:
@@ -66,7 +59,6 @@ def test_correct_index_location_metazoa(self):
         """
         survey_job = SurveyJob(source_type="TRANSCRIPTOME_INDEX")
         survey_job.save()
-        self.survey_job = survey_job
 
         key_value_pair = SurveyJobKeyValue(
             survey_job=survey_job, key="ensembl_division", value="EnsemblMetazoa"
@@ -78,8 +70,146 @@ def test_correct_index_location_metazoa(self):
         )
         key_value_pair.save()
 
-        surveyor = TranscriptomeIndexSurveyor(self.survey_job)
+        surveyor = TranscriptomeIndexSurveyor(survey_job)
         files = surveyor.discover_species()[0]
 
         for file in files:
             urllib.request.urlopen(file.source_url)
+
+    def test_single_plant(self):
+        """ Tests that the files returned actually exist.
+
+        Tests the Metazoa division instead of the main division.
+        """
+        survey_job = SurveyJob(source_type="TRANSCRIPTOME_INDEX")
+        survey_job.save()
+
+        key_value_pair = SurveyJobKeyValue(
+            survey_job=survey_job, key="ensembl_division", value="EnsemblPlants"
+        )
+        key_value_pair.save()
+
+        key_value_pair = SurveyJobKeyValue(
+            survey_job=survey_job, key="organism_name", value="Arabidopsis thaliana"
+        )
+        key_value_pair.save()
+
+        surveyor = TranscriptomeIndexSurveyor(survey_job)
+        files = surveyor.discover_species()[0]
+
+        for file in files:
+            urllib.request.urlopen(file.source_url)
+
+    def test_correct_index_location_protist(self):
+        """ Tests that the files returned actually exist.
+
+        Tests the Metazoa division instead of the main division.
+        """
+        survey_job = SurveyJob(source_type="TRANSCRIPTOME_INDEX")
+        survey_job.save()
+
+        key_value_pair = SurveyJobKeyValue(
+            survey_job=survey_job, key="ensembl_division", value="EnsemblProtists"
+        )
+        key_value_pair.save()
+
+        key_value_pair = SurveyJobKeyValue(
+            survey_job=survey_job, key="organism_name", value="Leishmania major"
+        )
+        key_value_pair.save()
+
+        surveyor = TranscriptomeIndexSurveyor(survey_job)
+        files = surveyor.discover_species()[0]
+
+        for file in files:
+            urllib.request.urlopen(file.source_url)
+
+    @patch("data_refinery_foreman.surveyor.external_source.message_queue.send_job")
+    def test_survey_fungi(self, mock_send_job):
+        survey_job = SurveyJob(source_type="TRANSCRIPTOME_INDEX")
+        survey_job.save()
+
+        key_value_pair = SurveyJobKeyValue(
+            survey_job=survey_job, key="ensembl_division", value="EnsemblFungi"
+        )
+        key_value_pair.save()
+
+        key_value_pair = SurveyJobKeyValue(
+            survey_job=survey_job, key="organism_name", value="CANDIDA_ALBICANS"
+        )
+        key_value_pair.save()
+
+        surveyor = TranscriptomeIndexSurveyor(survey_job)
+        surveyor.survey(source_type="TRANSCRIPTOME_INDEX")
+
+        downloader_jobs = DownloaderJob.objects.order_by("id").all()
+        self.assertEqual(downloader_jobs.count(), 1)
+        send_job_calls = []
+        for downloader_job in downloader_jobs:
+            send_job_calls.append(call(Downloaders.TRANSCRIPTOME_INDEX, downloader_job))
+
+        mock_send_job.assert_has_calls(send_job_calls)
+
+    @patch("data_refinery_foreman.surveyor.external_source.message_queue.send_job")
+    def test_survey_bacteria(self, mock_send_job):
+        survey_job = SurveyJob(source_type="TRANSCRIPTOME_INDEX")
+        survey_job.save()
+
+        key_value_pair = SurveyJobKeyValue(
+            survey_job=survey_job, key="ensembl_division", value="EnsemblBacteria"
+        )
+        key_value_pair.save()
+
+        key_value_pair = SurveyJobKeyValue(
+            survey_job=survey_job, key="organism_name", value="PSEUDOMONAS_AERUGINOSA"
+        )
+        key_value_pair.save()
+
+        surveyor = TranscriptomeIndexSurveyor(survey_job)
+        surveyor.survey(source_type="TRANSCRIPTOME_INDEX")
+
+        downloader_jobs = DownloaderJob.objects.order_by("id").all()
+        self.assertEqual(downloader_jobs.count(), 1)
+        send_job_calls = []
+        for downloader_job in downloader_jobs:
+            send_job_calls.append(call(Downloaders.TRANSCRIPTOME_INDEX, downloader_job))
+
+        mock_send_job.assert_has_calls(send_job_calls)
+
+    @patch("data_refinery_foreman.surveyor.external_source.message_queue.send_job")
+    def test_survey_bacteria_none(self, mock_send_job):
+        """When surveying fungi an organism_name must be supplied."""
+        survey_job = SurveyJob(source_type="TRANSCRIPTOME_INDEX")
+        survey_job.save()
+
+        key_value_pair = SurveyJobKeyValue(
+            survey_job=survey_job, key="ensembl_division", value="EnsemblBacteria"
+        )
+        key_value_pair.save()
+
+        surveyor = TranscriptomeIndexSurveyor(survey_job)
+        surveyor.survey(source_type="TRANSCRIPTOME_INDEX")
+
+        downloader_jobs = DownloaderJob.objects.order_by("id").all()
+        self.assertEqual(downloader_jobs.count(), 0)
+
+        mock_send_job.assert_not_called()
+
+    @patch("data_refinery_foreman.surveyor.external_source.message_queue.send_job")
+    def test_survey_fungi_none(self, mock_send_job):
+        """When surveying fungi an organism_name must be supplied."""
+        survey_job = SurveyJob(source_type="TRANSCRIPTOME_INDEX")
+        survey_job.save()
+
+        key_value_pair = SurveyJobKeyValue(
+            survey_job=survey_job, key="ensembl_division", value="EnsemblFungi"
+        )
+        key_value_pair.save()
+
+        surveyor = TranscriptomeIndexSurveyor(survey_job)
+        surveyor.survey(source_type="TRANSCRIPTOME_INDEX")
+
+        downloader_jobs = DownloaderJob.objects.order_by("id").all()
+        self.assertEqual(downloader_jobs.count(), 0)
+
+        mock_send_job.assert_not_called()