ops_utils.tdr_utils.tdr_api_utils

   1"""Utility classes for interacting with TDR API."""
   2
   3import json
   4import logging
   5import requests
   6from typing import Any, Optional, Union
   7from pydantic import ValidationError
   8
   9from ..request_util import GET, POST, DELETE, PUT, RunRequest
  10from ..tdr_api_schema.create_dataset_schema import CreateDatasetSchema
  11from ..tdr_api_schema.update_dataset_schema import UpdateSchema
  12from .tdr_job_utils import MonitorTDRJob, SubmitAndMonitorMultipleJobs
  13from ..vars import ARG_DEFAULTS, GCP, APPLICATION_JSON
  14
  15
  16class TDR:
  17    """Class to interact with the Terra Data Repository (TDR) API."""
  18
  19    PROD_LINK = "https://data.terra.bio/api/repository/v1"
  20    DEV_LINK = "https://jade.datarepo-dev.broadinstitute.org/api/repository/v1"
  21    """(str): The base URL for the TDR API."""
  22
  23    def __init__(self, request_util: RunRequest, env: str = 'prod', dry_run: bool = False):
  24        """
  25        Initialize the TDR class (A class to interact with the Terra Data Repository (TDR) API).
  26
  27        **Args:**
  28        - request_util (`ops_utils.request_util.RunRequest`): Utility for making HTTP requests.
  29        """
  30        self.request_util = request_util
  31        # NOTE: dry_run is not fully implemented in this class, only in delete_files_and_snapshots
  32        self.dry_run = dry_run
  33        if env.lower() == 'prod':
  34            self.tdr_link = self.PROD_LINK
  35        elif env.lower() == 'dev':
  36            self.tdr_link = self.DEV_LINK
  37        else:
  38            raise RuntimeError(f"Unsupported environment: {env}. Must be 'prod' or 'dev'.")
  39        """@private"""
  40
  41    @staticmethod
  42    def _check_policy(policy: str) -> None:
  43        """
  44        Check if the policy is valid.
  45
  46        **Args:**
  47        - policy (str): The role to check.
  48
  49        **Raises:**
  50        - ValueError: If the policy is not one of the allowed options.
  51        """
  52        if policy not in ["steward", "custodian", "snapshot_creator"]:
  53            raise ValueError(f"Policy {policy} is not valid. Must be steward, custodian, or snapshot_creator")
  54
  55    def get_dataset_files(
  56            self,
  57            dataset_id: str,
  58            limit: int = ARG_DEFAULTS['batch_size_to_list_files']  # type: ignore[assignment]
  59    ) -> list[dict]:
  60        """
  61        Get all files in a dataset.
  62
  63        Returns json like below
  64
  65            {
  66                "fileId": "68ba8bfc-1d84-4ef3-99b8-cf1754d5rrrr",
  67                "collectionId": "b20b6024-5943-4c23-82e7-9c24f545fuy7",
  68                "path": "/path/set/in/ingest.csv",
  69                "size": 1722,
  70                "checksums": [
  71                    {
  72                        "checksum": "82f7e79v",
  73                        "type": "crc32c"
  74                    },
  75                    {
  76                        "checksum": "fff973507e30b74fa47a3d6830b84a90",
  77                        "type": "md5"
  78                    }
  79                ],
  80                "created": "2024-13-11T15:01:00.256Z",
  81                "description": null,
  82                "fileType": "file",
  83                "fileDetail": {
  84                    "datasetId": "b20b6024-5943-4c23-82e7-9c24f5456444",
  85                    "mimeType": null,
  86                    "accessUrl": "gs://datarepo-bucket/path/to/actual/file.csv",
  87                    "loadTag": "RP_3333-RP_3333"
  88                },
  89                "directoryDetail": null
  90            }
  91
  92        **Args:**
  93        - dataset_id (str): The ID of the dataset.
  94        - limit (int, optional): The maximum number of records to retrieve per batch. Defaults to `20000`.
  95
  96        **Returns:**
  97        - list[dict]: A list of dictionaries containing the metadata of the files in the dataset.
  98        """
  99        uri = f"{self.tdr_link}/datasets/{dataset_id}/files"
 100        logging.info(f"Getting all files in dataset {dataset_id}")
 101        return self._get_response_from_batched_endpoint(uri=uri, limit=limit)
 102
 103    def create_file_dict(
 104            self,
 105            dataset_id: str,
 106            limit: int = ARG_DEFAULTS['batch_size_to_list_files']  # type: ignore[assignment]
 107    ) -> dict:
 108        """
 109        Create a dictionary of all files in a dataset where the key is the file UUID.
 110
 111        **Args:**
 112        - dataset_id (str): The ID of the dataset.
 113        - limit (int, optional): The maximum number of records to retrieve per batch. Defaults to `20000`.
 114
 115        **Returns:**
 116        - dict: A dictionary where the key is the file UUID and the value is the file metadata.
 117        """
 118        return {
 119            file_dict["fileId"]: file_dict
 120            for file_dict in self.get_dataset_files(dataset_id=dataset_id, limit=limit)
 121        }
 122
 123    def create_file_uuid_dict_for_ingest_for_experimental_self_hosted_dataset(
 124            self,
 125            dataset_id: str,
 126            limit: int = ARG_DEFAULTS['batch_size_to_list_files']  # type: ignore[assignment]
 127    ) -> dict:
 128        """
 129        Create a dictionary of all files in a dataset where the key is the file 'path' and the value is the file UUID.
 130
 131        This assumes that the TDR 'path' is original path of the file in the cloud storage with `gs://` stripped out.
 132
 133        This will ONLY work if dataset was created with `experimentalSelfHosted = True`
 134
 135        **Args:**
 136        - dataset_id (str): The ID of the dataset.
 137        - limit (int, optional): The maximum number of records to retrieve per batch. Defaults to `20000`.
 138
 139        **Returns:**
 140        - dict: A dictionary where the key is the file UUID and the value is the file path.
 141        """
 142        return {
 143            file_dict['fileDetail']['accessUrl']: file_dict['fileId']
 144            for file_dict in self.get_dataset_files(dataset_id=dataset_id, limit=limit)
 145        }
 146
 147    def delete_file(self, file_id: str, dataset_id: str) -> requests.Response:
 148        """
 149        Delete a file from a dataset.
 150
 151        **Args:**
 152        - file_id (str): The ID of the file to be deleted.
 153        - dataset_id (str): The ID of the dataset.
 154
 155        **Returns:**
 156        - requests.Response: The response from the request.
 157        """
 158        uri = f"{self.tdr_link}/datasets/{dataset_id}/files/{file_id}"
 159        logging.info(f"Submitting delete job for file {file_id}")
 160        return self.request_util.run_request(uri=uri, method=DELETE)
 161
 162    def delete_files(
 163            self,
 164            file_ids: list[str],
 165            dataset_id: str,
 166            batch_size_to_delete_files: int = ARG_DEFAULTS["batch_size_to_delete_files"],  # type: ignore[assignment]
 167            check_interval: int = 15) -> None:
 168        """
 169        Delete multiple files from a dataset in batches and monitor delete jobs until completion for each batch.
 170
 171        **Args:**
 172        - file_ids (list[str]): A list of file IDs to be deleted.
 173        - dataset_id (str): The ID of the dataset.
 174        - batch_size_to_delete_files (int, optional): The number of files to delete per batch. Defaults to `200`.
 175        - check_interval (int, optional): The interval in seconds to wait between status checks. Defaults to `15`.
 176        """
 177        SubmitAndMonitorMultipleJobs(
 178            tdr=self,
 179            job_function=self.delete_file,
 180            job_args_list=[(file_id, dataset_id) for file_id in file_ids],
 181            batch_size=batch_size_to_delete_files,
 182            check_interval=check_interval
 183        ).run()
 184
 185    def _delete_snapshots_for_files(self, dataset_id: str, file_ids: set[str]) -> None:
 186        """Delete snapshots that reference any of the provided file IDs."""
 187        snapshots_resp = self.get_dataset_snapshots(dataset_id=dataset_id)
 188        snapshot_items = snapshots_resp.json().get('items', [])
 189        snapshots_to_delete = []
 190        logging.info(
 191            "Checking %d snapshots for references",
 192            len(snapshot_items),
 193        )
 194        for snap in snapshot_items:
 195            snap_id = snap.get('id')
 196            if not snap_id:
 197                continue
 198            snap_files = self.get_files_from_snapshot(snapshot_id=snap_id)
 199            snap_file_ids = {
 200                fd.get('fileId') for fd in snap_files if fd.get('fileId')
 201            }
 202            # Use set intersection to check for any matching file IDs
 203            if snap_file_ids & file_ids:
 204                snapshots_to_delete.append(snap_id)
 205        if snapshots_to_delete:
 206            self.delete_snapshots(snapshot_ids=snapshots_to_delete)
 207        else:
 208            logging.info("No snapshots reference the provided file ids")
 209
 210    def _dry_run_msg(self) -> str:
 211        return '[Dry run] ' if self.dry_run else ''
 212
 213    def delete_files_and_snapshots(self, dataset_id: str, file_ids: set[str]) -> None:
 214        """Delete files from a dataset by their IDs, handling snapshots."""
 215        self._delete_snapshots_for_files(dataset_id=dataset_id, file_ids=file_ids)
 216
 217        logging.info(
 218            f"{self._dry_run_msg()}Submitting delete request for {len(file_ids)} files in "
 219            f"dataset {dataset_id}")
 220        if not self.dry_run:
 221            self.delete_files(
 222                file_ids=list(file_ids),
 223                dataset_id=dataset_id
 224            )
 225
 226    def add_user_to_dataset(self, dataset_id: str, user: str, policy: str) -> requests.Response:
 227        """
 228        Add a user to a dataset with a specified policy.
 229
 230        **Args:**
 231        - dataset_id (str): The ID of the dataset.
 232        - user (str): The email of the user to be added.
 233        - policy (str): The policy to be assigned to the user.
 234                Must be one of `steward`, `custodian`, or `snapshot_creator`.
 235
 236        **Returns:**
 237        - requests.Response: The response from the request.
 238
 239        **Raises:**
 240        - ValueError: If the policy is not valid.
 241        """
 242        self._check_policy(policy)
 243        uri = f"{self.tdr_link}/datasets/{dataset_id}/policies/{policy}/members"
 244        member_dict = {"email": user}
 245        logging.info(f"Adding user {user} to dataset {dataset_id} with policy {policy}")
 246        return self.request_util.run_request(
 247            uri=uri,
 248            method=POST,
 249            data=json.dumps(member_dict),
 250            content_type=APPLICATION_JSON
 251        )
 252
 253    def remove_user_from_dataset(self, dataset_id: str, user: str, policy: str) -> requests.Response:
 254        """
 255        Remove a user from a dataset.
 256
 257        **Args:**
 258        - dataset_id (str): The ID of the dataset.
 259        - user (str): The email of the user to be removed.
 260        - policy (str): The policy to be removed from the user.
 261                Must be one of `steward`, `custodian`, or `snapshot_creator`.
 262
 263        **Returns:**
 264        - requests.Response: The response from the request.
 265
 266        **Raises:**
 267        - ValueError: If the policy is not valid.
 268        """
 269        self._check_policy(policy)
 270        uri = f"{self.tdr_link}/datasets/{dataset_id}/policies/{policy}/members/{user}"
 271        logging.info(f"Removing user {user} from dataset {dataset_id} with policy {policy}")
 272        return self.request_util.run_request(uri=uri, method=DELETE)
 273
 274    def delete_dataset(self, dataset_id: str) -> None:
 275        """
 276        Delete a dataset and monitors the job until completion.
 277
 278        **Args:**
 279            dataset_id (str): The ID of the dataset to be deleted.
 280        """
 281        uri = f"{self.tdr_link}/datasets/{dataset_id}"
 282        logging.info(f"Deleting dataset {dataset_id}")
 283        response = self.request_util.run_request(uri=uri, method=DELETE)
 284        job_id = response.json()['id']
 285        MonitorTDRJob(tdr=self, job_id=job_id, check_interval=30, return_json=False).run()
 286
 287    def make_snapshot_public(self, snapshot_id: str) -> requests.Response:
 288        """
 289        Make a snapshot public.
 290
 291        **Args:**
 292        - snapshot_id (str): The ID of the snapshot to be made public.
 293
 294        **Returns:**
 295        - requests.Response: The response from the request.
 296        """
 297        uri = f"{self.tdr_link}/snapshots/{snapshot_id}/public"
 298        logging.info(f"Making snapshot {snapshot_id} public")
 299        return self.request_util.run_request(uri=uri, method=PUT, content_type=APPLICATION_JSON, data="true")
 300
 301    def get_snapshot_info(
 302            self,
 303            snapshot_id: str,
 304            continue_not_found: bool = False,
 305            info_to_include: Optional[list[str]] = None
 306    ) -> Optional[requests.Response]:
 307        """
 308        Get information about a snapshot.
 309
 310        **Args:**
 311        - snapshot_id (str): The ID of the snapshot.
 312        - continue_not_found (bool, optional): Whether to accept a `404` response. Defaults to `False`.
 313        - info_to_include (list[str], optional): A list of additional information to include. Defaults to None.
 314                Options are: `SOURCES`, `TABLES`, `RELATIONSHIPS`, `ACCESS_INFORMATION`, `PROFILE`, `PROPERTIES`,
 315                `DATA_PROJECT`,`CREATION_INFORMATION`, `DUOS`
 316
 317        **Returns:**
 318        - requests.Response (optional): The response from the request (returns None if the snapshot is not
 319         found or access is denied).
 320        """
 321        acceptable_return_code = [404, 403] if continue_not_found else []
 322        acceptable_include_info = [
 323            "SOURCES",
 324            "TABLES",
 325            "RELATIONSHIPS",
 326            "ACCESS_INFORMATION",
 327            "PROFILE",
 328            "PROPERTIES",
 329            "DATA_PROJECT",
 330            "CREATION_INFORMATION",
 331            "DUOS"
 332        ]
 333        if info_to_include:
 334            if not all(info in acceptable_include_info for info in info_to_include):
 335                raise ValueError(f"info_to_include must be a subset of {acceptable_include_info}")
 336            include_string = '&include='.join(info_to_include)
 337        else:
 338            include_string = ""
 339        uri = f"{self.tdr_link}/snapshots/{snapshot_id}?include={include_string}"
 340        response = self.request_util.run_request(
 341            uri=uri,
 342            method=GET,
 343            accept_return_codes=acceptable_return_code
 344        )
 345        if response.status_code == 404:
 346            logging.warning(f"Snapshot {snapshot_id} not found")
 347            return None
 348        if response.status_code == 403:
 349            logging.warning(f"Access denied for snapshot {snapshot_id}")
 350            return None
 351        return response
 352
 353    def delete_snapshots(
 354            self,
 355            snapshot_ids: list[str],
 356            batch_size: int = 25,
 357            check_interval: int = 10,
 358            verbose: bool = False) -> None:
 359        """
 360        Delete multiple snapshots from a dataset in batches and monitor delete jobs until completion for each batch.
 361
 362        **Args:**
 363        - snapshot_ids (list[str]): A list of snapshot IDs to be deleted.
 364        - batch_size (int, optional): The number of snapshots to delete per batch. Defaults to `25`.
 365        - check_interval (int, optional): The interval in seconds to wait between status checks. Defaults to `10`.
 366        - verbose (bool, optional): Whether to log detailed information about each job. Defaults to `False`.
 367        """
 368        logging.info(f"{self._dry_run_msg()}Deleting {len(snapshot_ids)} snapshots")
 369        if not self.dry_run:
 370            SubmitAndMonitorMultipleJobs(
 371                tdr=self,
 372                job_function=self.delete_snapshot,
 373                job_args_list=[(snapshot_id,) for snapshot_id in snapshot_ids],
 374                batch_size=batch_size,
 375                check_interval=check_interval,
 376                verbose=verbose
 377            ).run()
 378
 379    def delete_snapshot(self, snapshot_id: str) -> requests.Response:
 380        """
 381        Delete a snapshot.
 382
 383        **Args:**
 384        - snapshot_id (str): The ID of the snapshot to be deleted.
 385
 386        **Returns:**
 387        - requests.Response: The response from the request.
 388        """
 389        uri = f"{self.tdr_link}/snapshots/{snapshot_id}"
 390        logging.info(f"Deleting snapshot {snapshot_id}")
 391        return self.request_util.run_request(uri=uri, method=DELETE)
 392
 393    def _yield_existing_datasets(
 394            self, filter: Optional[str] = None, batch_size: int = 100, direction: str = "asc"
 395    ) -> Any:
 396        """
 397        Get all datasets in TDR, optionally filtered by dataset name.
 398
 399        **Args:**
 400            filter (Optional[str]): A filter string to match dataset names. Defaults to None.
 401            batch_size (int): The number of datasets to retrieve per batch. Defaults to 100.
 402            direction (str): The direction to sort the datasets by creation date. Defaults to "asc".
 403
 404        Yields:
 405            Any: A generator yielding datasets.
 406        """
 407        offset = 0
 408        if filter:
 409            filter_str = f"&filter={filter}"
 410            log_message = f"Searching for datasets with filter {filter} in batches of {batch_size}"
 411        else:
 412            filter_str = ""
 413            log_message = f"Searching for all datasets in batches of {batch_size}"
 414        logging.info(log_message)
 415        while True:
 416            uri = f"{self.tdr_link}/datasets?offset={offset}&limit={batch_size}&sort=created_date&direction={direction}{filter_str}"  # noqa: E501
 417            response = self.request_util.run_request(uri=uri, method=GET)
 418            datasets = response.json()["items"]
 419            if not datasets:
 420                break
 421            for dataset in datasets:
 422                yield dataset
 423            offset += batch_size
 424            break
 425
 426    def check_if_dataset_exists(self, dataset_name: str, billing_profile: Optional[str] = None) -> list[dict]:
 427        """
 428        Check if a dataset exists by name and optionally by billing profile.
 429
 430        **Args:**
 431        - dataset_name (str): The name of the dataset to check.
 432        - billing_profile (str, optional): The billing profile ID to match. Defaults to None.
 433
 434        **Returns:**
 435        - list[dict]: A list of matching datasets.
 436        """
 437        matching_datasets = []
 438        for dataset in self._yield_existing_datasets(filter=dataset_name):
 439            # Search uses wildcard so could grab more datasets where dataset_name is substring
 440            if dataset_name == dataset["name"]:
 441                if billing_profile:
 442                    if dataset["defaultProfileId"] == billing_profile:
 443                        logging.info(
 444                            f"Dataset {dataset['name']} already exists under billing profile {billing_profile}")
 445                        dataset_id = dataset["id"]
 446                        logging.info(f"Dataset ID: {dataset_id}")
 447                        matching_datasets.append(dataset)
 448                    else:
 449                        logging.warning(
 450                            f"Dataset {dataset['name']} exists but is under {dataset['defaultProfileId']} " +
 451                            f"and not under billing profile {billing_profile}"
 452                        )
 453                        # Datasets names need to be unique regardless of billing profile, so raise an error if
 454                        # a dataset with the same name is found but is not under the requested billing profile
 455                        raise ValueError(
 456                            f"Dataset {dataset_name} already exists but is not under billing profile {billing_profile}")
 457                else:
 458                    matching_datasets.append(dataset)
 459        return matching_datasets
 460
 461    def get_dataset_info(self, dataset_id: str, info_to_include: Optional[list[str]] = None) -> requests.Response:
 462        """
 463        Get information about a dataset.
 464
 465        **Args:**
 466        - dataset_id (str): The ID of the dataset.
 467        - info_to_include (list[str], optional): A list of additional information to include. Valid options include:
 468        `SCHEMA`, `ACCESS_INFORMATION`, `PROFILE`, `PROPERTIES`, `DATA_PROJECT`, `STORAGE`, `SNAPSHOT_BUILDER_SETTING`.
 469        Defaults to None.
 470
 471        **Returns:**
 472        - requests.Response: The response from the request.
 473
 474        **Raises:**
 475        - ValueError: If `info_to_include` contains invalid information types.
 476        """
 477        acceptable_include_info = [
 478            "SCHEMA",
 479            "ACCESS_INFORMATION",
 480            "PROFILE",
 481            "PROPERTIES",
 482            "DATA_PROJECT",
 483            "STORAGE",
 484            "SNAPSHOT_BUILDER_SETTING"
 485        ]
 486        if info_to_include:
 487            if not all(info in acceptable_include_info for info in info_to_include):
 488                raise ValueError(f"info_to_include must be a subset of {acceptable_include_info}")
 489            include_string = '&include='.join(info_to_include)
 490        else:
 491            include_string = ""
 492        uri = f"{self.tdr_link}/datasets/{dataset_id}?include={include_string}"
 493        return self.request_util.run_request(uri=uri, method=GET)
 494
 495    def get_table_schema_info(
 496            self,
 497            dataset_id: str,
 498            table_name: str,
 499            dataset_info: Optional[dict] = None
 500    ) -> Union[dict, None]:
 501        """
 502        Get schema information for a specific table within a dataset.
 503
 504        **Args:**
 505        - dataset_id (str): The ID of the dataset.
 506        - table_name (str): The name of the table.
 507        - dataset_info (dict, optional): The dataset information if already retrieved. Defaults to None.
 508
 509        **Returns:**
 510        - Union[dict, None]: A dictionary containing the table schema information, or None if the table is not found.
 511        """
 512        if not dataset_info:
 513            dataset_info = self.get_dataset_info(dataset_id=dataset_id, info_to_include=["SCHEMA"]).json()
 514        for table in dataset_info["schema"]["tables"]:  # type: ignore[index]
 515            if table["name"] == table_name:
 516                return table
 517        return None
 518
 519    def get_job_result(self, job_id: str, expect_failure: bool = False) -> requests.Response:
 520        """
 521        Retrieve the result of a job.
 522
 523        **Args:**
 524        - job_id (str): The ID of the job.
 525        - expect_failure (bool, optional): Whether the job is expected to fail. Defaults to `False`.
 526
 527        **Returns:**
 528        - requests.Response: The response from the request.
 529        """
 530        uri = f"{self.tdr_link}/jobs/{job_id}/result"
 531        # If job is expected to fail, accept any return code
 532        acceptable_return_code = list(range(100, 600)) if expect_failure else []
 533        return self.request_util.run_request(uri=uri, method=GET, accept_return_codes=acceptable_return_code)
 534
 535    def ingest_to_dataset(self, dataset_id: str, data: dict) -> requests.Response:
 536        """
 537        Load data into a TDR dataset.
 538
 539        **Args:**
 540        - dataset_id (str): The ID of the dataset.
 541        - data (dict): The data to be ingested.
 542
 543        **Returns:**
 544        - requests.Response: The response from the request.
 545        """
 546        uri = f"{self.tdr_link}/datasets/{dataset_id}/ingest"
 547        logging.info(
 548            "If recently added TDR SA to source bucket/dataset/workspace and you receive a 400/403 error, " +
 549            "it can sometimes take up to 12/24 hours for permissions to propagate. Try rerunning the script later.")
 550        return self.request_util.run_request(
 551            uri=uri,
 552            method=POST,
 553            content_type=APPLICATION_JSON,
 554            data=data
 555        )
 556
 557    def file_ingest_to_dataset(
 558            self,
 559            dataset_id: str,
 560            profile_id: str,
 561            file_list: list[dict],
 562            load_tag: str = "file_ingest_load_tag"
 563    ) -> dict:
 564        """
 565        Load files into a TDR dataset.
 566
 567        **Args:**
 568        - dataset_id (str): The ID of the dataset.
 569        - profile_id (str): The billing profile ID.
 570        - file_list (list[dict]): The list of files to be ingested.
 571        - load_tag (str): The tag to be used in the ingest job. Defaults to `file_ingest_load_tag`.
 572
 573        **Returns:**
 574        - dict: A dictionary containing the response from the ingest operation job monitoring.
 575        """
 576        uri = f"{self.tdr_link}/datasets/{dataset_id}/files/bulk/array"
 577        data = {
 578            "profileId": profile_id,
 579            "loadTag": f"{load_tag}",
 580            "maxFailedFileLoads": 0,
 581            "loadArray": file_list
 582        }
 583
 584        response = self.request_util.run_request(
 585            uri=uri,
 586            method=POST,
 587            content_type=APPLICATION_JSON,
 588            data=json.dumps(data)
 589        )
 590        job_id = response.json()['id']
 591        job_results = MonitorTDRJob(tdr=self, job_id=job_id, check_interval=30, return_json=True).run()
 592        return job_results  # type: ignore[return-value]
 593
 594    def get_dataset_table_metrics(
 595            self, dataset_id: str, target_table_name: str, query_limit: int = 1000
 596    ) -> list[dict]:
 597        """
 598        Retrieve all metrics for a specific table within a dataset.
 599
 600        **Args:**
 601        - dataset_id (str): The ID of the dataset.
 602        - target_table_name (str): The name of the target table.
 603        - query_limit (int, optional): The maximum number of records to retrieve per batch. Defaults to `1000`.
 604
 605        **Returns:**
 606        - list[dict]: A list of dictionaries containing the metrics for the specified table.
 607        """
 608        return [
 609            metric
 610            for metric in self._yield_dataset_metrics(
 611                dataset_id=dataset_id,
 612                target_table_name=target_table_name,
 613                query_limit=query_limit
 614            )
 615        ]
 616
 617    def _yield_dataset_metrics(self, dataset_id: str, target_table_name: str, query_limit: int = 1000) -> Any:
 618        """
 619        Yield all entity metrics from a dataset.
 620
 621        **Args:**
 622            dataset_id (str): The ID of the dataset.
 623            target_table_name (str): The name of the target table.
 624            query_limit (int, optional): The maximum number of records to retrieve per batch. Defaults to 1000.
 625
 626        Yields:
 627            Any: A generator yielding dictionaries containing the metrics for the specified table.
 628        """
 629        search_request = {
 630            "offset": 0,
 631            "limit": query_limit,
 632            "sort": "datarepo_row_id"
 633        }
 634        uri = f"{self.tdr_link}/datasets/{dataset_id}/data/{target_table_name}"
 635        while True:
 636            batch_number = int((search_request["offset"] / query_limit)) + 1  # type: ignore[operator]
 637            response = self.request_util.run_request(
 638                uri=uri,
 639                method=POST,
 640                content_type=APPLICATION_JSON,
 641                data=json.dumps(search_request)
 642            )
 643            if not response or not response.json()["result"]:
 644                break
 645            logging.info(
 646                f"Downloading batch {batch_number} of max {query_limit} records from {target_table_name} table " +
 647                f"dataset {dataset_id}"
 648            )
 649            for record in response.json()["result"]:
 650                yield record
 651            search_request["offset"] += query_limit  # type: ignore[operator]
 652
 653    def get_dataset_sample_ids(self, dataset_id: str, target_table_name: str, entity_id: str) -> list[str]:
 654        """
 655        Get existing IDs from a dataset.
 656
 657        **Args:**
 658        - dataset_id (str): The ID of the dataset.
 659        - target_table_name (str): The name of the target table.
 660        - entity_id (str): The entity ID to retrieve.
 661
 662        **Returns:**
 663        - list[str]: A list of entity IDs from the specified table.
 664        """
 665        dataset_metadata = self._yield_dataset_metrics(dataset_id=dataset_id, target_table_name=target_table_name)
 666        return [str(sample_dict[entity_id]) for sample_dict in dataset_metadata]
 667
 668    def get_job_status(self, job_id: str) -> requests.Response:
 669        """
 670        Retrieve the status of a job.
 671
 672        **Args:**
 673        - job_id (str): The ID of the job.
 674
 675        **Returns:**
 676        - requests.Response: The response from the request.
 677        """
 678        uri = f"{self.tdr_link}/jobs/{job_id}"
 679        return self.request_util.run_request(uri=uri, method=GET)
 680
 681    def get_dataset_file_uuids_from_metadata(self, dataset_id: str) -> list[str]:
 682        """
 683        Get all file UUIDs from the metadata of a dataset.
 684
 685        **Args:**
 686        - dataset_id (str): The ID of the dataset.
 687
 688        **Returns:**
 689        - list[str]: A list of file UUIDs from the dataset metadata.
 690        """
 691        dataset_info = self.get_dataset_info(dataset_id=dataset_id, info_to_include=["SCHEMA"]).json()
 692        all_metadata_file_uuids = []
 693        tables = 0
 694        for table in dataset_info["schema"]["tables"]:
 695            tables += 1
 696            table_name = table["name"]
 697            logging.info(f"Getting all file information for {table_name}")
 698            # Get just columns where datatype is fileref
 699            file_columns = [column["name"] for column in table["columns"] if column["datatype"] == "fileref"]
 700            dataset_metrics = self.get_dataset_table_metrics(dataset_id=dataset_id, target_table_name=table_name)
 701            # Get unique list of file uuids
 702            file_uuids = list(
 703                set(
 704                    [
 705                        value for metric in dataset_metrics for key, value in metric.items() if key in file_columns
 706                    ]
 707                )
 708            )
 709            logging.info(f"Got {len(file_uuids)} file uuids from table '{table_name}'")
 710            all_metadata_file_uuids.extend(file_uuids)
 711            # Make full list unique
 712            all_metadata_file_uuids = list(set(all_metadata_file_uuids))
 713        logging.info(f"Got {len(all_metadata_file_uuids)} file uuids from {tables} total table(s)")
 714        return all_metadata_file_uuids
 715
 716    def soft_delete_entries(
 717            self,
 718            dataset_id: str,
 719            table_name: str,
 720            datarepo_row_ids: list[str],
 721            check_intervals: int = 15
 722    ) -> Optional[dict]:
 723        """
 724        Soft delete specific records from a table.
 725
 726        **Args:**
 727        - dataset_id (str): The ID of the dataset.
 728        - table_name (str): The name of the target table.
 729        - datarepo_row_ids (list[str]): A list of row IDs to be deleted.
 730        - check_intervals (int, optional): The interval in seconds to wait between status checks. Defaults to `15`.
 731
 732        **Returns:**
 733        - dict (optional): A dictionary containing the response from the soft delete operation job
 734        monitoring. Returns None if no row IDs are provided.
 735        """
 736        if not datarepo_row_ids:
 737            logging.info(f"No records found to soft delete in table {table_name}")
 738            return None
 739        logging.info(f"Soft deleting {len(datarepo_row_ids)} records from table {table_name}")
 740        uri = f"{self.tdr_link}/datasets/{dataset_id}/deletes"
 741        payload = {
 742            "deleteType": "soft",
 743            "specType": "jsonArray",
 744            "tables": [
 745                {
 746                    "tableName": table_name,
 747                    "jsonArraySpec": {
 748                        "rowIds": datarepo_row_ids
 749                    }
 750                }
 751            ]
 752        }
 753        response = self.request_util.run_request(
 754            method=POST,
 755            uri=uri,
 756            data=json.dumps(payload),
 757            content_type=APPLICATION_JSON
 758        )
 759        job_id = response.json()["id"]
 760        return MonitorTDRJob(tdr=self, job_id=job_id, check_interval=check_intervals, return_json=False).run()
 761
 762    def soft_delete_all_table_entries(
 763            self,
 764            dataset_id: str,
 765            table_name: str,
 766            query_limit: int = 1000,
 767            check_intervals: int = 15
 768    ) -> Optional[dict]:
 769        """
 770        Soft deletes all records in a table.
 771
 772        **Args:**
 773        - dataset_id (str): The ID of the dataset.
 774        - table_name (str): The name of the target table.
 775        - query_limit (int, optional): The maximum number of records to retrieve per batch. Defaults to `1000`.
 776        - check_intervals (int, optional): The interval in seconds to wait between status checks. Defaults to `15`.
 777
 778        **Returns:**
 779        - dict (optional): A dictionary containing the response from the soft delete operation job monitoring. Returns
 780        None if no row IDs are provided.
 781        """
 782        dataset_metrics = self.get_dataset_table_metrics(
 783            dataset_id=dataset_id, target_table_name=table_name, query_limit=query_limit
 784        )
 785        row_ids = [metric["datarepo_row_id"] for metric in dataset_metrics]
 786        return self.soft_delete_entries(
 787            dataset_id=dataset_id,
 788            table_name=table_name,
 789            datarepo_row_ids=row_ids,
 790            check_intervals=check_intervals
 791        )
 792
 793    def get_or_create_dataset(
 794            self,
 795            dataset_name: str,
 796            billing_profile: str,
 797            schema: dict,
 798            description: str,
 799            relationships: Optional[list[dict]] = None,
 800            delete_existing: bool = False,
 801            continue_if_exists: bool = False,
 802            additional_properties_dict: Optional[dict] = None
 803    ) -> str:
 804        """
 805        Get or create a dataset.
 806
 807        **Args:**
 808        - dataset_name (str): The name of the dataset.
 809        - billing_profile (str): The billing profile ID.
 810        - schema (dict): The schema of the dataset.
 811        - description (str): The description of the dataset.
 812        - relationships (Optional[list[dict]], optional): A list of relationships to add to the dataset schema.
 813                Defaults to None.
 814        - additional_properties_dict (Optional[dict], optional): Additional properties
 815                for the dataset. Defaults to None.
 816        - delete_existing (bool, optional): Whether to delete the existing dataset if found.
 817                Defaults to `False`.
 818        - continue_if_exists (bool, optional): Whether to continue if the dataset already exists.
 819                Defaults to `False`.
 820
 821        **Returns:**
 822        - str: The ID of the dataset.
 823
 824        **Raises:**
 825        - ValueError: If multiple datasets with the same name are found under the billing profile.
 826        """
 827        existing_datasets = self.check_if_dataset_exists(dataset_name, billing_profile)
 828        if existing_datasets:
 829            if not continue_if_exists:
 830                raise ValueError(
 831                    f"Run with continue_if_exists=True to use the existing dataset {dataset_name}"
 832                )
 833            # If delete_existing is True, delete the existing dataset and set existing_datasets to an empty list
 834            if delete_existing:
 835                logging.info(f"Deleting existing dataset {dataset_name}")
 836                self.delete_dataset(existing_datasets[0]["id"])
 837                existing_datasets = []
 838            # If not delete_existing and continue_if_exists then grab existing datasets id
 839            else:
 840                dataset_id = existing_datasets[0]["id"]
 841        if not existing_datasets:
 842            logging.info("Did not find existing dataset")
 843            # Create dataset
 844            dataset_id = self.create_dataset(
 845                schema=schema,
 846                dataset_name=dataset_name,
 847                description=description,
 848                profile_id=billing_profile,
 849                additional_dataset_properties=additional_properties_dict
 850            )
 851        return dataset_id
 852
 853    def create_dataset(  # type: ignore[return]
 854            self,
 855            schema: dict,
 856            dataset_name: str,
 857            description: str,
 858            profile_id: str,
 859            additional_dataset_properties: Optional[dict] = None
 860    ) -> Optional[str]:
 861        """
 862        Create a new dataset.
 863
 864        **Args:**
 865        - schema (dict): The schema of the dataset.
 866        - dataset_name (str): The name of the dataset.
 867        - description (str): The description of the dataset.
 868        - profile_id (str): The billing profile ID.
 869        - additional_dataset_properties (Optional[dict], optional): Additional
 870                properties for the dataset. Defaults to None.
 871
 872        **Returns:**
 873        - Optional[str]: The ID of the created dataset, or None if creation failed.
 874
 875        Raises:
 876        - ValueError: If the schema validation fails.
 877        """
 878        dataset_properties = {
 879            "name": dataset_name,
 880            "description": description,
 881            "defaultProfileId": profile_id,
 882            "region": "us-central1",
 883            "cloudPlatform": GCP,
 884            "schema": schema
 885        }
 886
 887        if additional_dataset_properties:
 888            dataset_properties.update(additional_dataset_properties)
 889        try:
 890            CreateDatasetSchema(**dataset_properties)  # type: ignore[arg-type]
 891        except ValidationError as e:
 892            raise ValueError(f"Schema validation error: {e}")
 893        uri = f"{self.tdr_link}/datasets"
 894        logging.info(f"Creating dataset {dataset_name} under billing profile {profile_id}")
 895        response = self.request_util.run_request(
 896            method=POST,
 897            uri=uri,
 898            data=json.dumps(dataset_properties),
 899            content_type=APPLICATION_JSON
 900        )
 901        job_id = response.json()["id"]
 902        job_results = MonitorTDRJob(tdr=self, job_id=job_id, check_interval=30, return_json=True).run()
 903        dataset_id = job_results["id"]  # type: ignore[index]
 904        logging.info(f"Successfully created dataset {dataset_name}: {dataset_id}")
 905        return dataset_id
 906
 907    def update_dataset_schema(  # type: ignore[return]
 908            self,
 909            dataset_id: str,
 910            update_note: str,
 911            tables_to_add: Optional[list[dict]] = None,
 912            relationships_to_add: Optional[list[dict]] = None,
 913            columns_to_add: Optional[list[dict]] = None
 914    ) -> Optional[str]:
 915        """
 916        Update the schema of a dataset.
 917
 918        **Args:**
 919        - dataset_id (str): The ID of the dataset.
 920        - update_note (str): A note describing the update.
 921        - tables_to_add (list[dict], optional): A list of tables to add. Defaults to None.
 922        - relationships_to_add (list[dict], optional): A list of relationships to add. Defaults to None.
 923        - columns_to_add (list[dict], optional): A list of columns to add. Defaults to None.
 924
 925        **Returns:**
 926        - Optional[str]: The ID of the updated dataset, or None if the update failed.
 927
 928        **Raises:**
 929        - ValueError: If the schema validation fails.
 930        """
 931        uri = f"{self.tdr_link}/datasets/{dataset_id}/updateSchema"
 932        request_body: dict = {"description": f"{update_note}", "changes": {}}
 933        if tables_to_add:
 934            request_body["changes"]["addTables"] = tables_to_add
 935        if relationships_to_add:
 936            request_body["changes"]["addRelationships"] = relationships_to_add
 937        if columns_to_add:
 938            request_body["changes"]["addColumns"] = columns_to_add
 939        try:
 940            UpdateSchema(**request_body)
 941        except ValidationError as e:
 942            raise ValueError(f"Schema validation error: {e}")
 943
 944        response = self.request_util.run_request(
 945            uri=uri,
 946            method=POST,
 947            content_type=APPLICATION_JSON,
 948            data=json.dumps(request_body)
 949        )
 950        job_id = response.json()["id"]
 951        job_results = MonitorTDRJob(tdr=self, job_id=job_id, check_interval=30, return_json=True).run()
 952        dataset_id = job_results["id"]  # type: ignore[index]
 953        logging.info(f"Successfully ran schema updates in dataset {dataset_id}")
 954        return dataset_id
 955
 956    def _get_response_from_batched_endpoint(self, uri: str, limit: int = 1000) -> list[dict]:
 957        """
 958        Get response from a batched endpoint.
 959
 960        Helper method for all GET endpoints that require batching.
 961
 962        Given the URI and the limit (optional), will
 963        loop through batches until all metadata is retrieved.
 964
 965        **Args:**
 966        - uri (str): The base URI for the endpoint (without query params for offset or limit).
 967        - limit (int, optional): The maximum number of records to retrieve per batch. Defaults to `1000`.
 968
 969        **Returns:**
 970        - list[dict]: A list of dictionaries containing the metadata retrieved from the endpoint.
 971        """
 972        batch = 1
 973        offset = 0
 974        metadata: list = []
 975        while True:
 976            logging.info(f"Retrieving {(batch - 1) * limit} to {batch * limit} records in metadata")
 977            response_json = self.request_util.run_request(uri=f"{uri}?offset={offset}&limit={limit}", method=GET).json()
 978
 979            # If no more files, break the loop
 980            if not response_json:
 981                logging.info(f"No more results to retrieve, found {len(metadata)} total records")
 982                break
 983
 984            metadata.extend(response_json)
 985            if len(response_json) < limit:
 986                logging.info(f"Retrieved final batch of results, found {len(metadata)} total records")
 987                break
 988
 989            # Increment the offset by limit for the next page
 990            offset += limit
 991            batch += 1
 992        return metadata
 993
 994    def get_files_from_snapshot(self, snapshot_id: str, limit: int = 1000) -> list[dict]:
 995        """
 996        Return all the metadata about files in a given snapshot.
 997
 998        Not all files can be returned at once, so the API
 999        is used repeatedly until all "batches" have been returned.
1000
1001        **Args:**
1002        - snapshot_id (str): The ID of the snapshot.
1003        - limit (int, optional): The maximum number of records to retrieve per batch. Defaults to `1000`.
1004
1005        **Returns:**
1006        - list[dict]: A list of dictionaries containing the metadata of the files in the snapshot.
1007        """
1008        uri = f"{self.tdr_link}/snapshots/{snapshot_id}/files"
1009        return self._get_response_from_batched_endpoint(uri=uri, limit=limit)
1010
1011    def get_dataset_snapshots(self, dataset_id: str) -> requests.Response:
1012        """
1013        Return snapshots belonging to specified dataset.
1014
1015        **Args:**
1016        - dataset_id: uuid of dataset to query.
1017
1018        **Returns:**
1019        - requests.Response: The response from the request.
1020        """
1021        uri = f"{self.tdr_link}/snapshots?datasetIds={dataset_id}"
1022        return self.request_util.run_request(
1023            uri=uri,
1024            method=GET
1025        )
1026
1027    def create_snapshot(
1028            self,
1029            snapshot_name: str,
1030            description: str,
1031            dataset_name: str,
1032            snapshot_mode: str,  # byFullView is entire dataset
1033            profile_id: str,
1034            stewards: Optional[list[str]] = [],
1035            readers: Optional[list[str]] = [],
1036            consent_code: Optional[str] = None,
1037            duos_id: Optional[str] = None,
1038            data_access_control_groups: Optional[list[str]] = None,
1039    ) -> None:
1040        """
1041        Create a snapshot in TDR.
1042
1043        **Returns:**
1044        - requests.Response: The response from the request.
1045        """
1046        uri = f"{self.tdr_link}/snapshots"
1047        payload = {
1048            "name": snapshot_name,
1049            "description": description,
1050            "contents": [
1051                {
1052                    "datasetName": dataset_name,
1053                    "mode": snapshot_mode,
1054                }
1055            ],
1056            "policies": {
1057                "stewards": stewards,
1058                "readers": readers,
1059            },
1060            "profileId": profile_id,
1061            "globalFileIds": True,
1062        }
1063        if consent_code:
1064            payload["consentCode"] = consent_code
1065        if duos_id:
1066            payload["duosId"] = duos_id
1067        if data_access_control_groups:
1068            payload["dataAccessControlGroups"] = data_access_control_groups
1069        logging.info(f"Creating snapshot {snapshot_name} in dataset {dataset_name}")
1070        response = self.request_util.run_request(
1071            uri=uri,
1072            method=POST,
1073            content_type=APPLICATION_JSON,
1074            data=json.dumps(payload)
1075        )
1076        job_id = response.json()["id"]
1077        job_results = MonitorTDRJob(tdr=self, job_id=job_id, check_interval=30, return_json=True).run()
1078        snapshot_id = job_results["id"]  # type: ignore[index]
1079        logging.info(f"Successfully created snapshot {snapshot_name} - {snapshot_id}")
1080
1081
1082class FilterOutSampleIdsAlreadyInDataset:
1083    """Class to filter ingest metrics to remove sample IDs that already exist in the dataset."""
1084
1085    def __init__(
1086            self,
1087            ingest_metrics: list[dict],
1088            dataset_id: str,
1089            tdr: TDR,
1090            target_table_name: str,
1091            filter_entity_id: str
1092    ):
1093        """
1094        Initialize the FilterOutSampleIdsAlreadyInDataset class.
1095
1096        **Args:**
1097        - ingest_metrics (list[dict]): The metrics to be ingested.
1098        - dataset_id (str): The ID of the dataset.
1099        - tdr (`ops_utils.tdr_utils.tdr_utils.TDR`): The TDR instance
1100        - target_table_name (str): The name of the target table.
1101        - filter_entity_id (str): The entity ID to filter on.
1102        """
1103        self.ingest_metrics = ingest_metrics
1104        """@private"""
1105        self.tdr = tdr
1106        """@private"""
1107        self.dataset_id = dataset_id
1108        """@private"""
1109        self.target_table_name = target_table_name
1110        """@private"""
1111        self.filter_entity_id = filter_entity_id
1112        """@private"""
1113
1114    def run(self) -> list[dict]:
1115        """
1116        Run the filter process to remove sample IDs that already exist in the dataset.
1117
1118        **Returns:**
1119        - list[dict]: The filtered ingest metrics.
1120        """
1121        # Get all sample ids that already exist in dataset
1122        logging.info(
1123            f"Getting all {self.filter_entity_id} that already exist in table {self.target_table_name} in "
1124            f"dataset {self.dataset_id}"
1125        )
1126
1127        dataset_sample_ids = self.tdr.get_dataset_sample_ids(
1128            dataset_id=self.dataset_id,
1129            target_table_name=self.target_table_name,
1130            entity_id=self.filter_entity_id
1131        )
1132        # Filter out rows that already exist in dataset
1133        filtered_ingest_metrics = [
1134            row
1135            for row in self.ingest_metrics
1136            if str(row[self.filter_entity_id]) not in dataset_sample_ids
1137        ]
1138        if len(filtered_ingest_metrics) < len(self.ingest_metrics):
1139            logging.info(
1140                f"Filtered out {len(self.ingest_metrics) - len(filtered_ingest_metrics)} rows that already exist in "
1141                f"dataset. There is {len(filtered_ingest_metrics)} rows left to ingest"
1142            )
1143
1144            if filtered_ingest_metrics:
1145                return filtered_ingest_metrics
1146            else:
1147                logging.info("All rows filtered out as they all exist in dataset, nothing to ingest")
1148                return []
1149        else:
1150            logging.info("No rows were filtered out as they all do not exist in dataset")
1151            return filtered_ingest_metrics