Merge pull request #5030 from broadinstitute/datset-type-specific-discovery-genes

Dataset type specific discovery genes

Author: hanars

seqr/views/apis/project_api_tests.py

@@ -390,7 +390,7 @@ def test_project_families(self):
         gene_ids = self._assert_expected_project_families(url, response_keys)
         self.assertSetEqual(gene_ids, {'ENSG00000135953', 'ENSG00000240361'})
 
-    def _assert_expected_project_families(self, url, response_keys):
+    def _assert_expected_project_families(self, url, response_keys, no_discovery_tags=False):
         response = self.client.get(url)
         self.assertEqual(response.status_code, 200)
 
@@ -436,12 +436,14 @@ def _assert_expected_project_families(self, url, response_keys):
 
         self.assertListEqual(family_3['discoveryTags'], [])
         self.assertListEqual(empty_family['discoveryTags'], [])
-        self.assertListEqual(family_1['discoveryTags'], [{
+        family_1_tags = [] if no_discovery_tags else [{
             'transcripts': {'ENSG00000135953': [mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY, mock.ANY]},
             'mainTranscriptId': 'ENST00000258436',
             'selectedMainTranscriptId': None,
-        }])
-        self.assertListEqual(response_json['familiesByGuid']['F000002_2']['discoveryTags'], [self.DISCOVERY_TAG])
+        }]
+        self.assertListEqual(family_1['discoveryTags'], family_1_tags)
+        family_2_tags = [] if no_discovery_tags else [self.DISCOVERY_TAG]
+        self.assertListEqual(response_json['familiesByGuid']['F000002_2']['discoveryTags'], family_2_tags)
         no_discovery_families = set(response_json['familiesByGuid'].keys()) - {'F000001_1', 'F000002_2'}
         self.assertSetEqual({
             len(response_json['familiesByGuid'][family_guid]['discoveryTags']) for family_guid in no_discovery_families
@@ -802,8 +804,8 @@ def _assert_expected_project_families(self, *args, **kwargs):
         self.reset_logs()
         connections['clickhouse'].close()
         self.DISCOVERY_TAG = {**DISCOVERY_TAG, 'transcripts': {}}
-        no_clickhouse_gene_ids = super()._assert_expected_project_families(*args, **kwargs)
-        self.assertSetEqual(no_clickhouse_gene_ids, {'ENSG00000135953'})
+        no_clickhouse_gene_ids = super()._assert_expected_project_families(*args, **kwargs, no_discovery_tags=True)
+        self.assertSetEqual(no_clickhouse_gene_ids, set())
         self.assert_json_logs(None, [
             ("Error loading discovery genes from clickhouse: An error occurred in the current transaction. You can't execute queries until the end of the 'atomic' block.", {
                 'severity': 'ERROR',

seqr/views/utils/project_context_utils.py

@@ -1,9 +1,11 @@
 from collections import defaultdict
 from django.db.models import Count, Q, F, prefetch_related_objects
+from django.contrib.postgres.aggregates import ArrayAgg
+from django.db.models.functions import JSONObject
 
-from clickhouse_search.search import get_transcripts_by_key
+from clickhouse_search.search import get_transcripts_by_key, get_annotations_queryset
 from seqr.models import Individual, IgvSample, AnalysisGroup, DynamicAnalysisGroup, LocusList, VariantTagType,\
-    VariantFunctionalData, FamilyNote, SavedVariant, VariantTag, VariantNote
+    VariantFunctionalData, FamilyNote, SavedVariant, VariantTag, VariantNote, Sample
 from seqr.utils.gene_utils import get_genes
 from seqr.utils.logging_utils import SeqrLogger
 from seqr.utils.search.utils import backend_specific_call
@@ -119,38 +121,60 @@ def families_discovery_tags(families, genome_version, project=None):
     families_by_guid = {f['familyGuid']: dict(discoveryTags=[], **f) for f in families}
 
     family_filter = {'family__project': project} if project else {'family__guid__in': families_by_guid.keys()}
-    discovery_tags_by_key = {v.pop('key'): v for v in SavedVariant.objects.filter(
+    discovery_variants = SavedVariant.objects.filter(
         varianttag__variant_tag_type__category='CMG Discovery Tags', **family_filter,
-    ).values(
-        'key', 'family__guid', selectedMainTranscriptId=F('selected_main_transcript_id'),
-        transcripts=F('saved_variant_json__transcripts'), mainTranscriptId=F('saved_variant_json__mainTranscriptId'),
-    )}
+    )
     try:
-        backend_specific_call(lambda *args: None, _add_clickhouse_transcripts)(discovery_tags_by_key, genome_version)
+        discovery_tags = backend_specific_call(_get_no_key_tags, _get_clickhouse_tags)(
+            discovery_variants, genome_version=genome_version,
+        )
     except Exception as e:
         logger.error(f'Error loading discovery genes from clickhouse: {e}', None)
+        discovery_tags = []
 
     gene_ids = set()
-    for tag in discovery_tags_by_key.values():
+    for tag in discovery_tags:
         tag['transcripts'] = tag.get('transcripts') or {}
         gene_ids.update(list(tag['transcripts'].keys()))
-        families_by_guid[tag.pop('family__guid')]['discoveryTags'].append(tag)
+        families_by_guid[tag.pop('family_guid')]['discoveryTags'].append(tag)
 
     return {
         'familiesByGuid': families_by_guid,
         'genesById': get_genes(gene_ids),
     }
 
 
-def _add_clickhouse_transcripts(discovery_tags_by_key, genome_version):
-    transcripts_by_key = get_transcripts_by_key(genome_version, discovery_tags_by_key.keys())
-    for key, tag in discovery_tags_by_key.items():
-        if key in transcripts_by_key:
+def _get_no_key_tags(discovery_variants, **kwargs):
+    return discovery_variants.values(
+        family_guid=F('family__guid'), selectedMainTranscriptId=F('selected_main_transcript_id'),
+        transcripts=F('saved_variant_json__transcripts'), mainTranscriptId=F('saved_variant_json__mainTranscriptId'),
+    )
+
+
+def _get_clickhouse_tags(discovery_variants, genome_version):
+    discovery_tags = list(_get_no_key_tags(discovery_variants.filter(key__isnull=True)))
+
+    tags_by_dataset_type = discovery_variants.filter(key__isnull=False).values('dataset_type').annotate(
+        keys=ArrayAgg('key', distinct=True),
+        tags=ArrayAgg(JSONObject(key='key', family_guid='family__guid', selectedMainTranscriptId='selected_main_transcript_id')),
+    )
+
+    for dataset_type, keys, tags in tags_by_dataset_type.values_list('dataset_type', 'keys', 'tags'):
+        if dataset_type == Sample.DATASET_TYPE_VARIANT_CALLS:
+            transcripts_by_key = get_transcripts_by_key(genome_version, keys)
+        else:
+            qs = get_annotations_queryset(genome_version, dataset_type, keys)
+            transcripts_by_key = dict(qs.values_list('key', qs.transcript_field))
+        for tag in tags:
+            key = tag.pop('key')
             tag['transcripts'] = transcripts_by_key[key]
-            tag['mainTranscriptId'] = next(
+            tag['mainTranscriptId'] = next((
                 t['transcriptId'] for gene_transcripts in tag['transcripts'].values() for t in gene_transcripts
-                if t['transcriptRank'] == 0
-            )
+                if t.get('transcriptRank') == 0
+            ), None)
+            discovery_tags.append(tag)
+
+    return discovery_tags
 
 
 MME_TAG_NAME = 'MME Submission'