Add helper functions for querying Terra Data Repository (#998)

* Add service account credentialing

* ruff

* First pass

* tests passing

* add coverage of bigquery test

* change function names

* use generators everywhere

* bq requirement

* resolver

* Update sample id name

* Build Sex Check Table from TDR Metrics (#999)

Author: bpblanken

requirements.in

@@ -4,3 +4,4 @@ gnomad==0.6.4
 aiofiles==24.1.0
 pydantic==2.8.2
 google-cloud-dataproc==5.14.0
+google-cloud-bigquery==3.27.0

requirements.txt

@@ -97,17 +97,30 @@ frozenlist==1.5.0
 gnomad==0.6.4
     # via -r requirements.in
 google-api-core[grpc]==2.22.0
-    # via google-cloud-dataproc
+    # via
+    #   google-cloud-bigquery
+    #   google-cloud-core
+    #   google-cloud-dataproc
 google-auth==2.35.0
     # via
     #   google-api-core
     #   google-auth-oauthlib
+    #   google-cloud-bigquery
+    #   google-cloud-core
     #   google-cloud-dataproc
     #   hail
 google-auth-oauthlib==0.8.0
     # via hail
+google-cloud-bigquery==3.27.0
+    # via -r requirements.in
+google-cloud-core==2.4.1
+    # via google-cloud-bigquery
 google-cloud-dataproc==5.14.0
     # via -r requirements.in
+google-crc32c==1.6.0
+    # via google-resumable-media
+google-resumable-media==2.7.2
+    # via google-cloud-bigquery
 googleapis-common-protos[grpc]==1.65.0
     # via
     #   google-api-core
@@ -197,6 +210,7 @@ orjson==3.10.10
 packaging==24.1
     # via
     #   bokeh
+    #   google-cloud-bigquery
     #   plotly
 pandas==2.2.3
     # via
@@ -256,16 +270,15 @@ pygments==2.18.0
     #   ipython
     #   rich
 pyjwt[crypto]==2.9.0
-    # via
-    #   msal
-    #   pyjwt
+    # via msal
 pyspark==3.5.3
     # via hail
 python-daemon==3.1.0
     # via luigi
 python-dateutil==2.9.0.post0
     # via
     #   botocore
+    #   google-cloud-bigquery
     #   luigi
     #   pandas
 python-json-logger==2.0.7
@@ -282,6 +295,7 @@ requests==2.32.3
     # via
     #   azure-core
     #   google-api-core
+    #   google-cloud-bigquery
     #   hail
     #   msal
     #   msrest

v03_pipeline/lib/misc/allele_registry.py

@@ -8,9 +8,9 @@
 import hailtop.fs as hfs
 import requests
 from requests import HTTPError
-from requests.adapters import HTTPAdapter, Retry
 
 from v03_pipeline.lib.logger import get_logger
+from v03_pipeline.lib.misc.requests import requests_retry_session
 from v03_pipeline.lib.model import Env, ReferenceGenome
 
 MAX_VARIANTS_PER_REQUEST = 1000000
@@ -96,13 +96,7 @@ def register_alleles(
     logger.info('Calling the ClinGen Allele Registry')
     with hfs.open(formatted_vcf_file_name, 'r') as vcf_in:
         data = vcf_in.read()
-        s = requests.Session()
-        retries = Retry(
-            total=5,
-            backoff_factor=1,
-            status_forcelist=[500, 502, 503, 504],
-        )
-        s.mount('https://', HTTPAdapter(max_retries=retries))
+        s = requests_retry_session()
         res = s.put(
             url=build_url(base_url, reference_genome),
             data=data,

v03_pipeline/lib/misc/requests.py

@@ -0,0 +1,14 @@
+import requests
+from requests.adapters import HTTPAdapter, Retry
+
+
+def requests_retry_session():
+    s = requests.Session()
+    retries = Retry(
+        total=5,
+        backoff_factor=1,
+        status_forcelist=[500, 502, 503, 504],
+    )
+    s.mount('http://', HTTPAdapter(max_retries=retries))
+    s.mount('https://', HTTPAdapter(max_retries=retries))
+    return s

v03_pipeline/lib/misc/terra_data_repository.py

@@ -0,0 +1,68 @@
+import os
+import re
+from collections.abc import Generator
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+import google.cloud.bigquery
+from google.cloud import bigquery
+
+from v03_pipeline.lib.misc.gcp import get_service_account_credentials
+from v03_pipeline.lib.misc.requests import requests_retry_session
+
+BIGQUERY_METRICS = [
+    'collaborator_sample_id',
+    'predicted_sex',
+]
+BIGQUERY_RESOURCE = 'bigquery'
+TABLE_NAME_VALIDATION_REGEX = r'datarepo-\w+.datarepo_\w+'
+TDR_ROOT_URL = 'https://data.terra.bio/api/repository/v1/'
+
+
+def _tdr_request(resource: str) -> dict:
+    service_account_token = get_service_account_credentials().token
+    s = requests_retry_session()
+    res = s.get(
+        url=os.path.join(TDR_ROOT_URL, resource),
+        headers={'Authorization': f'Bearer {service_account_token}'},
+        timeout=10,
+    )
+    res.raise_for_status()
+    return res.json()
+
+
+def _get_dataset_ids() -> list[str]:
+    res_body = _tdr_request('datasets')
+    items = res_body['items']
+    for item in items:
+        if not any(x['cloudResource'] == BIGQUERY_RESOURCE for x in item['storage']):
+            # Hard failure on purpose to prompt manual investigation.
+            msg = 'Datasets without bigquery sources are unsupported'
+            raise ValueError(msg)
+    return [x['id'] for x in items]
+
+
+def gen_bq_table_names() -> Generator[str]:
+    with ThreadPoolExecutor(max_workers=5) as executor:
+        futures = [
+            executor.submit(
+                _tdr_request,
+                f'datasets/{dataset_id}?include=ACCESS_INFORMATION',
+            )
+            for dataset_id in _get_dataset_ids()
+        ]
+        for future in as_completed(futures):
+            result = future.result()
+            yield f"{result['accessInformation']['bigQuery']['projectId']}.{result['accessInformation']['bigQuery']['datasetName']}"
+
+
+def bq_metrics_query(bq_table_name: str) -> google.cloud.bigquery.table.RowIterator:
+    if not re.match(TABLE_NAME_VALIDATION_REGEX, bq_table_name):
+        msg = f'{bq_table_name} does not match expected pattern'
+        raise ValueError(msg)
+    client = bigquery.Client()
+    return client.query_and_wait(
+        f"""
+        SELECT {','.join(BIGQUERY_METRICS)}
+        FROM `{bq_table_name}.sample`
+    """,  # noqa: S608
+    )