Merge pull request #67 from broadinstitute/sn_POD-2822_interact_with_dev

Adding ability to interact with dev

Author: snovod

VERSION.txt

@@ -1,2 +1,4 @@
-10.6.3
-- gcp cloud function checks for errors
+11.0.0
+- Adding ability to create snapshots
+- Be able to pass in environment variables to terra and TDR
+- Be able to pass in service account json for Token and GCP util

ops_utils/bq_utils.py

@@ -2,7 +2,7 @@
 import logging
 from google.cloud import bigquery
 from google.api_core.exceptions import Forbidden
-from typing import Optional
+from typing import Optional, Any
 
 
 class BigQueryUtil:
@@ -69,7 +69,7 @@ def upload_data_to_table(self, table_id: str, rows: list[dict], delete_existing_
         new_rows = destination_table.num_rows
         logging.info(f"Table now contains {new_rows} rows after upload")
 
-    def query_table(self, query: str, to_dataframe: bool = False) -> list[dict]:
+    def query_table(self, query: str, to_dataframe: bool = False) -> Any:
         """
         Execute a SQL query on a BigQuery table and returns the results.
 
@@ -83,7 +83,7 @@ def query_table(self, query: str, to_dataframe: bool = False) -> list[dict]:
         query_job = self.client.query(query)
         if to_dataframe:
             return query_job.result().to_dataframe()
-        return [row for row in query_job.result()]
+        return [dict(row) for row in query_job.result()]
 
     def check_permissions_to_project(self, raise_on_other_failure: bool = True) -> bool:
         """

ops_utils/gcp_utils.py

@@ -2,6 +2,7 @@
 import os
 import logging
 import io
+import json
 import hashlib
 import base64
 import subprocess
@@ -10,6 +11,9 @@
 from mimetypes import guess_type
 from typing import Optional, Any
 from google.cloud.storage.blob import Blob
+from google.oauth2 import service_account
+from google.cloud import storage
+from google.auth import default
 
 from .vars import ARG_DEFAULTS
 from .thread_pool_executor_util import MultiThreadedJobs
@@ -27,21 +31,42 @@
 class GCPCloudFunctions:
     """Class to handle GCP Cloud Functions."""
 
-    def __init__(self, project: Optional[str] = None) -> None:
+    def __init__(
+            self,
+            project: Optional[str] = None,
+            service_account_json: Optional[str] = None
+    ) -> None:
         """
         Initialize the GCPCloudFunctions class.
 
-        Authenticates using the default credentials and sets up the Storage Client.
-        Uses the `project_id` if provided, otherwise utilizes the default project set.
+        Authenticates using service account JSON if provided or default credentials,
+        and sets up the Storage Client.
 
-        **Args:**
-        - project (str, optional): The GCP project ID
-        """
-        from google.cloud import storage  # type: ignore[attr-defined]
-        from google.auth import default
-        credentials, default_project = default()
+        Args:
+            project: Optional[str] = None
+                The GCP project ID. If not provided, will use project from service account or default.
+            service_account_json: Optional[str] = None
+                Path to service account JSON key file. If provided, will use these credentials.
+        """
+        # Initialize credentials and project
+        credentials = None
+        default_project = None
+
+        if service_account_json:
+            credentials = service_account.Credentials.from_service_account_file(service_account_json)
+            # Extract project from service account if not specified
+            if not project:
+                with open(service_account_json, 'r') as f:
+                    sa_info = json.load(f)
+                    project = sa_info.get('project_id')
+        else:
+            # Use default credentials
+            credentials, default_project = default()
+
+        # Set project if not already set
         if not project:
             project = default_project
+
         self.client = storage.Client(credentials=credentials, project=project)
         """@private"""
 

ops_utils/tdr_utils/tdr_api_utils.py

@@ -18,17 +18,24 @@
 class TDR:
     """Class to interact with the Terra Data Repository (TDR) API."""
 
-    TDR_LINK = "https://data.terra.bio/api/repository/v1"
+    PROD_LINK = "https://data.terra.bio/api/repository/v1"
+    DEV_LINK = "https://jade.datarepo-dev.broadinstitute.org/api/repository/v1"
     """(str): The base URL for the TDR API."""
 
-    def __init__(self, request_util: RunRequest):
+    def __init__(self, request_util: RunRequest, env: str = 'prod'):
         """
         Initialize the TDR class (A class to interact with the Terra Data Repository (TDR) API).
 
         **Args:**
         - request_util (`ops_utils.request_util.RunRequest`): Utility for making HTTP requests.
         """
         self.request_util = request_util
+        if env.lower() == 'prod':
+            self.tdr_link = self.PROD_LINK
+        elif env.lower() == 'dev':
+            self.tdr_link = self.DEV_LINK
+        else:
+            raise RuntimeError(f"Unsupported environment: {env}. Must be 'prod' or 'dev'.")
         """@private"""
 
     @staticmethod
@@ -89,7 +96,7 @@ def get_dataset_files(
         **Returns:**
         - list[dict]: A list of dictionaries containing the metadata of the files in the dataset.
         """
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}/files"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}/files"
         logging.info(f"Getting all files in dataset {dataset_id}")
         return self._get_response_from_batched_endpoint(uri=uri, limit=limit)
 
@@ -153,12 +160,12 @@ def get_sas_token(self, snapshot_id: str = "", dataset_id: str = "") -> dict:
         - ValueError: If neither `snapshot_id` nor `dataset_id` is provided.
         """
         if snapshot_id:
-            uri = f"{self.TDR_LINK}/snapshots/{snapshot_id}?include=ACCESS_INFORMATION"
+            uri = f"{self.tdr_link}/snapshots/{snapshot_id}?include=ACCESS_INFORMATION"
             response = self.request_util.run_request(uri=uri, method=GET)
             snapshot_info = json.loads(response.text)
             sas_token = snapshot_info["accessInformation"]["parquet"]["sasToken"]
         elif dataset_id:
-            uri = f"{self.TDR_LINK}/datasets/{dataset_id}?include=ACCESS_INFORMATION"
+            uri = f"{self.tdr_link}/datasets/{dataset_id}?include=ACCESS_INFORMATION"
             response = self.request_util.run_request(uri=uri, method=GET)
             snapshot_info = json.loads(response.text)
             sas_token = snapshot_info["accessInformation"]["parquet"]["sasToken"]
@@ -182,7 +189,7 @@ def delete_file(self, file_id: str, dataset_id: str) -> requests.Response:
         **Returns:**
         - requests.Response: The response from the request.
         """
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}/files/{file_id}"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}/files/{file_id}"
         logging.info(f"Submitting delete job for file {file_id}")
         return self.request_util.run_request(uri=uri, method=DELETE)
 
@@ -226,7 +233,7 @@ def add_user_to_dataset(self, dataset_id: str, user: str, policy: str) -> reques
         - ValueError: If the policy is not valid.
         """
         self._check_policy(policy)
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}/policies/{policy}/members"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}/policies/{policy}/members"
         member_dict = {"email": user}
         logging.info(f"Adding user {user} to dataset {dataset_id} with policy {policy}")
         return self.request_util.run_request(
@@ -253,7 +260,7 @@ def remove_user_from_dataset(self, dataset_id: str, user: str, policy: str) -> r
         - ValueError: If the policy is not valid.
         """
         self._check_policy(policy)
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}/policies/{policy}/members/{user}"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}/policies/{policy}/members/{user}"
         logging.info(f"Removing user {user} from dataset {dataset_id} with policy {policy}")
         return self.request_util.run_request(uri=uri, method=DELETE)
 
@@ -264,7 +271,7 @@ def delete_dataset(self, dataset_id: str) -> None:
         **Args:**
             dataset_id (str): The ID of the dataset to be deleted.
         """
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}"
         logging.info(f"Deleting dataset {dataset_id}")
         response = self.request_util.run_request(uri=uri, method=DELETE)
         job_id = response.json()['id']
@@ -308,7 +315,7 @@ def get_snapshot_info(
             include_string = '&include='.join(info_to_include)
         else:
             include_string = ""
-        uri = f"{self.TDR_LINK}/snapshots/{snapshot_id}?include={include_string}"
+        uri = f"{self.tdr_link}/snapshots/{snapshot_id}?include={include_string}"
         response = self.request_util.run_request(
             uri=uri,
             method=GET,
@@ -356,7 +363,7 @@ def delete_snapshot(self, snapshot_id: str) -> requests.Response:
         **Returns:**
         - requests.Response: The response from the request.
         """
-        uri = f"{self.TDR_LINK}/snapshots/{snapshot_id}"
+        uri = f"{self.tdr_link}/snapshots/{snapshot_id}"
         logging.info(f"Deleting snapshot {snapshot_id}")
         return self.request_util.run_request(uri=uri, method=DELETE)
 
@@ -383,7 +390,7 @@ def _yield_existing_datasets(
             log_message = f"Searching for all datasets in batches of {batch_size}"
         logging.info(log_message)
         while True:
-            uri = f"{self.TDR_LINK}/datasets?offset={offset}&limit={batch_size}&sort=created_date&direction={direction}{filter_str}"  # noqa: E501
+            uri = f"{self.tdr_link}/datasets?offset={offset}&limit={batch_size}&sort=created_date&direction={direction}{filter_str}"  # noqa: E501
             response = self.request_util.run_request(uri=uri, method=GET)
             datasets = response.json()["items"]
             if not datasets:
@@ -459,7 +466,7 @@ def get_dataset_info(self, dataset_id: str, info_to_include: Optional[list[str]]
             include_string = '&include='.join(info_to_include)
         else:
             include_string = ""
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}?include={include_string}"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}?include={include_string}"
         return self.request_util.run_request(uri=uri, method=GET)
 
     def get_table_schema_info(
@@ -497,7 +504,7 @@ def get_job_result(self, job_id: str, expect_failure: bool = False) -> requests.
         **Returns:**
         - requests.Response: The response from the request.
         """
-        uri = f"{self.TDR_LINK}/jobs/{job_id}/result"
+        uri = f"{self.tdr_link}/jobs/{job_id}/result"
         # If job is expected to fail, accept any return code
         acceptable_return_code = list(range(100, 600)) if expect_failure else []
         return self.request_util.run_request(uri=uri, method=GET, accept_return_codes=acceptable_return_code)
@@ -513,7 +520,7 @@ def ingest_to_dataset(self, dataset_id: str, data: dict) -> requests.Response:
         **Returns:**
         - requests.Response: The response from the request.
         """
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}/ingest"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}/ingest"
         logging.info(
             "If recently added TDR SA to source bucket/dataset/workspace and you receive a 400/403 error, " +
             "it can sometimes take up to 12/24 hours for permissions to propagate. Try rerunning the script later.")
@@ -543,7 +550,7 @@ def file_ingest_to_dataset(
         **Returns:**
         - dict: A dictionary containing the response from the ingest operation job monitoring.
         """
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}/files/bulk/array"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}/files/bulk/array"
         data = {
             "profileId": profile_id,
             "loadTag": f"{load_tag}",
@@ -601,7 +608,7 @@ def _yield_dataset_metrics(self, dataset_id: str, target_table_name: str, query_
             "limit": query_limit,
             "sort": "datarepo_row_id"
         }
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}/data/{target_table_name}"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}/data/{target_table_name}"
         while True:
             batch_number = int((search_request["offset"] / query_limit)) + 1  # type: ignore[operator]
             response = self.request_util.run_request(
@@ -645,7 +652,7 @@ def get_job_status(self, job_id: str) -> requests.Response:
         **Returns:**
         - requests.Response: The response from the request.
         """
-        uri = f"{self.TDR_LINK}/jobs/{job_id}"
+        uri = f"{self.tdr_link}/jobs/{job_id}"
         return self.request_util.run_request(uri=uri, method=GET)
 
     def get_dataset_file_uuids_from_metadata(self, dataset_id: str) -> list[str]:
@@ -707,7 +714,7 @@ def soft_delete_entries(
             logging.info(f"No records found to soft delete in table {table_name}")
             return None
         logging.info(f"Soft deleting {len(datarepo_row_ids)} records from table {table_name}")
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}/deletes"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}/deletes"
         payload = {
             "deleteType": "soft",
             "specType": "jsonArray",
@@ -766,6 +773,7 @@ def get_or_create_dataset(
             billing_profile: str,
             schema: dict,
             description: str,
+            relationships: Optional[list[dict]] = None,
             delete_existing: bool = False,
             continue_if_exists: bool = False,
             additional_properties_dict: Optional[dict] = None
@@ -778,6 +786,8 @@ def get_or_create_dataset(
         - billing_profile (str): The billing profile ID.
         - schema (dict): The schema of the dataset.
         - description (str): The description of the dataset.
+        - relationships (Optional[list[dict]], optional): A list of relationships to add to the dataset schema.
+                Defaults to None.
         - additional_properties_dict (Optional[dict], optional): Additional properties
                 for the dataset. Defaults to None.
         - delete_existing (bool, optional): Whether to delete the existing dataset if found.
@@ -857,7 +867,7 @@ def create_dataset(  # type: ignore[return]
             CreateDatasetSchema(**dataset_properties)  # type: ignore[arg-type]
         except ValidationError as e:
             raise ValueError(f"Schema validation error: {e}")
-        uri = f"{self.TDR_LINK}/datasets"
+        uri = f"{self.tdr_link}/datasets"
         logging.info(f"Creating dataset {dataset_name} under billing profile {profile_id}")
         response = self.request_util.run_request(
             method=POST,
@@ -895,7 +905,7 @@ def update_dataset_schema(  # type: ignore[return]
         **Raises:**
         - ValueError: If the schema validation fails.
         """
-        uri = f"{self.TDR_LINK}/datasets/{dataset_id}/updateSchema"
+        uri = f"{self.tdr_link}/datasets/{dataset_id}/updateSchema"
         request_body: dict = {"description": f"{update_note}", "changes": {}}
         if tables_to_add:
             request_body["changes"]["addTables"] = tables_to_add
@@ -968,7 +978,7 @@ def get_files_from_snapshot(self, snapshot_id: str, limit: int = 1000) -> list[d
         **Returns:**
         - list[dict]: A list of dictionaries containing the metadata of the files in the snapshot.
         """
-        uri = f"{self.TDR_LINK}/snapshots/{snapshot_id}/files"
+        uri = f"{self.tdr_link}/snapshots/{snapshot_id}/files"
         return self._get_response_from_batched_endpoint(uri=uri, limit=limit)
 
     def get_dataset_snapshots(self, dataset_id: str) -> requests.Response:
@@ -981,12 +991,66 @@ def get_dataset_snapshots(self, dataset_id: str) -> requests.Response:
         **Returns:**
         - requests.Response: The response from the request.
         """
-        uri = f"{self.TDR_LINK}/snapshots?datasetIds={dataset_id}"
+        uri = f"{self.tdr_link}/snapshots?datasetIds={dataset_id}"
         return self.request_util.run_request(
             uri=uri,
             method=GET
         )
 
+    def create_snapshot(
+            self,
+            snapshot_name: str,
+            description: str,
+            dataset_name: str,
+            snapshot_mode: str,  # byFullView is entire dataset
+            profile_id: str,
+            stewards: Optional[list[str]] = [],
+            readers: Optional[list[str]] = [],
+            consent_code: Optional[str] = None,
+            duos_id: Optional[str] = None,
+            data_access_control_groups: Optional[list[str]] = None,
+    ) -> None:
+        """
+        Create a snapshot in TDR.
+
+        **Returns:**
+        - requests.Response: The response from the request.
+        """
+        uri = f"{self.tdr_link}/snapshots"
+        payload = {
+            "name": snapshot_name,
+            "description": description,
+            "contents": [
+                {
+                    "datasetName": dataset_name,
+                    "mode": snapshot_mode,
+                }
+            ],
+            "policies": {
+                "stewards": stewards,
+                "readers": readers,
+            },
+            "profileId": profile_id,
+            "globalFileIds": True,
+        }
+        if consent_code:
+            payload["consentCode"] = consent_code
+        if duos_id:
+            payload["duosId"] = duos_id
+        if data_access_control_groups:
+            payload["dataAccessControlGroups"] = data_access_control_groups
+        logging.info(f"Creating snapshot {snapshot_name} in dataset {dataset_name}")
+        response = self.request_util.run_request(
+            uri=uri,
+            method=POST,
+            content_type="application/json",
+            data=json.dumps(payload)
+        )
+        job_id = response.json()["id"]
+        job_results = MonitorTDRJob(tdr=self, job_id=job_id, check_interval=30, return_json=True).run()
+        snapshot_id = job_results["id"]  # type: ignore[index]
+        logging.info(f"Successfully created snapshot {snapshot_name} - {snapshot_id}")
+
 
 class FilterOutSampleIdsAlreadyInDataset:
     """Class to filter ingest metrics to remove sample IDs that already exist in the dataset."""