Make dataset collection reference non hardcoded (#956)

Author: bpblanken

v03_pipeline/lib/model/cached_reference_dataset_query.py

@@ -5,6 +5,9 @@
 
 from v03_pipeline.lib.model.dataset_type import DatasetType
 from v03_pipeline.lib.model.definitions import ReferenceGenome
+from v03_pipeline.lib.model.reference_dataset_collection import (
+    ReferenceDatasetCollection,
+)
 from v03_pipeline.lib.reference_data.queries import (
     clinvar_path_variants,
     gnomad_coding_and_noncoding_variants,
@@ -30,11 +33,13 @@ def dataset(self, dataset_type: DatasetType) -> str | None:
         }.get(self)
 
     @property
-    def query_raw_dataset(self) -> bool:
+    def reference_dataset_collection(self) -> ReferenceDatasetCollection:
         return {
-            CachedReferenceDatasetQuery.GNOMAD_CODING_AND_NONCODING_VARIANTS: True,
-            CachedReferenceDatasetQuery.GNOMAD_QC: True,
-        }.get(self, False)
+            CachedReferenceDatasetQuery.CLINVAR_PATH_VARIANTS: ReferenceDatasetCollection.COMBINED,
+            CachedReferenceDatasetQuery.GNOMAD_CODING_AND_NONCODING_VARIANTS: None,
+            CachedReferenceDatasetQuery.GNOMAD_QC: None,
+            CachedReferenceDatasetQuery.HIGH_AF_VARIANTS: ReferenceDatasetCollection.COMBINED,
+        }[self]
 
     @property
     def query(self) -> Callable[[hl.Table, ReferenceGenome], hl.Table]:

v03_pipeline/lib/tasks/reference_data/updated_cached_reference_dataset_query.py

@@ -72,7 +72,7 @@ def output(self) -> luigi.Target:
         )
 
     def requires(self) -> luigi.Task:
-        if self.crdq.query_raw_dataset:
+        if not self.crdq.reference_dataset_collection:
             return HailTableTask(
                 get_ht_path(
                     CONFIG[self.crdq.dataset(self.dataset_type)][
@@ -83,7 +83,7 @@ def requires(self) -> luigi.Task:
         # Special nested import to avoid a circular dependency issue
         # (ValidateCallset -> this file -> UpdatedReferenceDatasetCollection -> ValidateCallset)
         # The specific CRDQ referenced in ValidateCallset will never reach
-        # this line due to it being a "query_raw_dataset".  In theory this
+        # this line due to it being a raw dataset query.  In theory this
         # would be fixed by splitting the CRDQ into raw_dataset and non-raw_dataset
         # queries.
         from v03_pipeline.lib.tasks.reference_data.updated_reference_dataset_collection import (
@@ -93,12 +93,12 @@ def requires(self) -> luigi.Task:
         return UpdatedReferenceDatasetCollectionTask(
             self.reference_genome,
             self.dataset_type,
-            ReferenceDatasetCollection.COMBINED,
+            self.crdq.reference_dataset_collection,
         )
 
     def create_table(self) -> hl.Table:
         dataset: str = self.crdq.dataset(self.dataset_type)
-        if self.crdq.query_raw_dataset:
+        if not self.crdq.reference_dataset_collection:
             query_ht = import_ht_from_config_path(
                 CONFIG[dataset][self.reference_genome.v02_value],
                 dataset,
@@ -121,21 +121,21 @@ def create_table(self) -> hl.Table:
             paths=hl.Struct(
                 **{
                     dataset: query_ht.index_globals().path
-                    if self.crdq.query_raw_dataset
+                    if not self.crdq.reference_dataset_collection
                     else query_ht.index_globals().paths[dataset],
                 },
             ),
             versions=hl.Struct(
                 **{
                     dataset: query_ht.index_globals().version
-                    if self.crdq.query_raw_dataset
+                    if not self.crdq.reference_dataset_collection
                     else query_ht.index_globals().versions[dataset],
                 },
             ),
             enums=hl.Struct(
                 **{
                     dataset: query_ht.index_globals().enums
-                    if self.crdq.query_raw_dataset
+                    if not self.crdq.reference_dataset_collection
                     else query_ht.index_globals().enums[dataset],
                 },
             ),