Fix missing sample types in delete dags (#908)

* Fix missing sample types in delete dags

* Fix typo

Author: bpblanken

v03_pipeline/lib/misc/project_tables.py

@@ -0,0 +1,31 @@
+import os
+
+import hail as hl
+import hailtop.fs as hfs
+
+from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
+from v03_pipeline.lib.paths import project_table_path
+
+
+def get_valid_project_tables(
+    reference_genome: ReferenceGenome,
+    dataset_type: DatasetType,
+    project_guid: str,
+) -> set[tuple[hl.Table, SampleType]]:
+    project_tables = set()
+    for sample_type in SampleType:
+        project_ht_path = project_table_path(
+            reference_genome,
+            dataset_type,
+            sample_type,
+            project_guid,
+        )
+        if hfs.exists(project_ht_path) and hfs.exists(
+            os.path.join(project_ht_path, '_SUCCESS'),
+        ):
+            project_ht = hl.read_table(project_ht_path)
+            project_tables.add((project_ht, sample_type))
+    if len(project_tables) == 0:
+        msg = f'No project tables found for {project_guid}'
+        raise RuntimeError(msg)
+    return project_tables

v03_pipeline/lib/tasks/__init__.py

@@ -3,7 +3,7 @@
 from v03_pipeline.lib.tasks.delete_project_family_tables import (
     DeleteProjectFamilyTablesTask,
 )
-from v03_pipeline.lib.tasks.delete_project_table import DeleteProjectTableTask
+from v03_pipeline.lib.tasks.delete_project_tables import DeleteProjectTablesTask
 from v03_pipeline.lib.tasks.migrate_all_lookup_tables import MigrateAllLookupTablesTask
 from v03_pipeline.lib.tasks.migrate_all_variant_annotations_tables import (
     MigrateAllVariantAnnotationsTablesTask,
@@ -21,8 +21,8 @@
     UpdateLookupTableWithDeletedProjectTask,
 )
 from v03_pipeline.lib.tasks.update_project_table import UpdateProjectTableTask
-from v03_pipeline.lib.tasks.update_project_table_with_deleted_families import (
-    UpdateProjectTableWithDeletedFamiliesTask,
+from v03_pipeline.lib.tasks.update_project_tables_with_deleted_families import (
+    UpdateProjectTablesWithDeletedFamiliesTask,
 )
 from v03_pipeline.lib.tasks.update_variant_annotations_table_with_deleted_families import (
     UpdateVariantAnnotationsTableWithDeletedFamiliesTask,
@@ -42,11 +42,11 @@
     'DeleteFamilyTableTask',
     'DeleteFamilyTablesTask',
     'DeleteProjectFamilyTablesTask',
-    'DeleteProjectTableTask',
+    'DeleteProjectTablesTask',
     'MigrateAllLookupTablesTask',
     'MigrateAllVariantAnnotationsTablesTask',
     'UpdateProjectTableTask',
-    'UpdateProjectTableWithDeletedFamiliesTask',
+    'UpdateProjectTablesWithDeletedFamiliesTask',
     'UpdateLookupTableTask',
     'UpdateLookupTableWithDeletedProjectTask',
     'UpdateLookupTableWithDeletedFamiliesTask',

v03_pipeline/lib/tasks/delete_project_family_tables.py

@@ -1,15 +1,12 @@
 import hail as hl
-import hailtop.fs as hfs
 import luigi
 import luigi.util
 
-from v03_pipeline.lib.model import SampleType
-from v03_pipeline.lib.paths import project_table_path
+from v03_pipeline.lib.misc.project_tables import get_valid_project_tables
 from v03_pipeline.lib.tasks.base.base_loading_pipeline_params import (
     BaseLoadingPipelineParams,
 )
 from v03_pipeline.lib.tasks.delete_family_table import DeleteFamilyTableTask
-from v03_pipeline.lib.tasks.files import HailTableTask
 
 
 @luigi.util.inherits(BaseLoadingPipelineParams)
@@ -27,23 +24,11 @@ def complete(self) -> bool:
         )
 
     def run(self):
-        project_tables = set()
-        for sample_type in SampleType:
-            project_ht_path = project_table_path(
-                self.reference_genome,
-                self.dataset_type,
-                sample_type,
-                self.project_guid,
-            )
-            if hfs.exists(project_ht_path):
-                project_table_task: luigi.Target = yield HailTableTask(project_ht_path)
-                project_ht = hl.read_table(project_table_task.path)
-                project_tables.add((project_ht, sample_type))
-
-        if len(project_tables) == 0:
-            msg = f'No project tables found for {self.project_guid}'
-            raise RuntimeError(msg)
-
+        project_tables = get_valid_project_tables(
+            self.reference_genome,
+            self.dataset_type,
+            self.project_guid,
+        )
         for project_ht, sample_type in project_tables:
             family_guids = hl.eval(project_ht.globals.family_guids)
             for family_guid in family_guids:

v03_pipeline/lib/tasks/delete_project_tables.py

@@ -0,0 +1,34 @@
+import luigi
+
+from v03_pipeline.lib.model import SampleType
+from v03_pipeline.lib.tasks.base.base_loading_pipeline_params import (
+    BaseLoadingPipelineParams,
+)
+from v03_pipeline.lib.tasks.delete_project_table import DeleteProjectTableTask
+
+
+@luigi.util.inherits(BaseLoadingPipelineParams)
+class DeleteProjectTablesTask(luigi.Task):
+    project_guid = luigi.Parameter()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.dynamic_delete_project_table_tasks = set()
+
+    def complete(self) -> bool:
+        return len(self.dynamic_delete_project_table_tasks) >= 1 and all(
+            delete_project_table_task.complete()
+            for delete_project_table_task in self.dynamic_delete_project_table_tasks
+        )
+
+    def run(self):
+        for sample_type in SampleType:
+            self.dynamic_delete_project_table_tasks.add(
+                DeleteProjectTableTask(
+                    reference_genome=self.reference_genome,
+                    dataset_type=self.dataset_type,
+                    sample_type=sample_type,
+                    project_guid=self.project_guid,
+                ),
+            )
+        yield self.dynamic_delete_project_table_tasks

v03_pipeline/lib/tasks/update_project_tables_with_deleted_families.py

@@ -0,0 +1,44 @@
+import luigi
+import luigi.util
+
+from v03_pipeline.lib.misc.project_tables import get_valid_project_tables
+from v03_pipeline.lib.tasks.base.base_loading_pipeline_params import (
+    BaseLoadingPipelineParams,
+)
+from v03_pipeline.lib.tasks.update_project_table_with_deleted_families import (
+    UpdateProjectTableWithDeletedFamiliesTask,
+)
+
+
+@luigi.util.inherits(BaseLoadingPipelineParams)
+class UpdateProjectTablesWithDeletedFamiliesTask(luigi.Task):
+    project_guid = luigi.Parameter()
+    family_guids = luigi.ListParameter()
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.dynamic_update_project_table_tasks = set()
+
+    def complete(self) -> bool:
+        return len(self.dynamic_update_project_table_tasks) >= 1 and all(
+            update_project_table_task.complete()
+            for update_project_table_task in self.dynamic_update_project_table_tasks
+        )
+
+    def run(self):
+        project_tables = get_valid_project_tables(
+            self.reference_genome,
+            self.dataset_type,
+            self.project_guid,
+        )
+        for _, sample_type in project_tables:
+            self.dynamic_update_project_table_tasks.add(
+                UpdateProjectTableWithDeletedFamiliesTask(
+                    reference_genome=self.reference_genome,
+                    dataset_type=self.dataset_type,
+                    sample_type=sample_type,
+                    project_guid=self.project_guid,
+                    family_guids=self.family_guids,
+                ),
+            )
+        yield self.dynamic_update_project_table_tasks