ops_utils.tdr_utils.tdr_bq_utils

  1"""Module for handling TDR interactions with BigQuery."""
  2
  3from ..bq_utils import BigQueryUtil
  4from .tdr_api_utils import TDR
  5
  6import logging
  7from typing import Optional, Any
  8
  9
 10class GetTdrAssetInfo:
 11    """Class to obtain TDR asset metadata (from dataset or snapshot) from BigQuery."""
 12
 13    def __init__(self, tdr: TDR, dataset_id: Optional[str] = None, snapshot_id: Optional[str] = None):
 14        """
 15        Initialize the GetTdrAssetInfo class.
 16
 17        **Args:**
 18        - tdr (`ops_utils.tdr_utils.tdr_api_utils.TDR`): TDR instance for interacting with the TDR API.
 19        - dataset_id (str, optional): ID of the dataset.
 20        - snapshot_id (str, optional): ID of the snapshot.
 21        """
 22        if not dataset_id and not snapshot_id:
 23            raise ValueError("Either dataset_id or snapshot_id must be provided.")
 24        self.tdr = tdr
 25        """@private"""
 26        self.dataset_id = dataset_id
 27        """@private"""
 28        self.snapshot_id = snapshot_id
 29        """@private"""
 30
 31    def _get_dataset_info(self) -> dict:
 32        """
 33        Retrieve dataset information from TDR.
 34
 35        **Returns:**
 36        - dict: A dictionary containing BigQuery project ID, schema, tables, and relationships.
 37        """
 38        dataset_info = self.tdr.get_dataset_info(
 39            dataset_id=self.dataset_id,  # type: ignore[arg-type]
 40            info_to_include=["SCHEMA", "ACCESS_INFORMATION"]
 41        ).json()
 42        return {
 43            "bq_project": dataset_info["accessInformation"]["bigQuery"]["projectId"],
 44            "bq_schema": dataset_info["accessInformation"]["bigQuery"]["datasetName"],
 45            "tables": dataset_info["schema"]["tables"],
 46            "relationships": dataset_info["schema"]["relationships"]
 47        }
 48
 49    def _get_snapshot_info(self) -> Optional[dict]:
 50        """
 51        Retrieve snapshot information from TDR.
 52
 53        **Returns:**
 54        - dict: A dictionary containing BigQuery project ID, schema, tables, and relationships.
 55        Returns None if the snapshot is not found or access is denied.
 56        """
 57        response = self.tdr.get_snapshot_info(
 58            snapshot_id=self.snapshot_id,  # type: ignore[arg-type]
 59            info_to_include=["TABLES", "RELATIONSHIPS", "ACCESS_INFORMATION"]
 60        )
 61        if response:
 62            snapshot_info = response.json()
 63            return {
 64                "bq_project": snapshot_info["accessInformation"]["bigQuery"]["projectId"],
 65                "bq_schema": snapshot_info["accessInformation"]["bigQuery"]["datasetName"],
 66                "tables": snapshot_info["tables"],
 67                "relationships": snapshot_info["relationships"]
 68            }
 69        return None
 70
 71    def run(self) -> Optional[dict]:
 72        """
 73        Execute the process to retrieve either dataset or snapshot information.
 74
 75        **Returns:**
 76        - dict (optional): A dictionary containing the relevant information based on whether `dataset_id` or
 77            `snapshot_id` is provided. Returns None if the snapshot is not found or access is denied.
 78        """
 79        if self.dataset_id:
 80            return self._get_dataset_info()
 81        return self._get_snapshot_info()
 82
 83
 84class TdrBq:
 85    """Class to interact with TDR BigQuery tables."""
 86
 87    def __init__(self, project_id: str, bq_schema: str):
 88        """
 89        Initialize the TdrBq class.
 90
 91        **Args:**
 92        - project_id (str): The Google Cloud project ID.
 93        - bq_schema (str): The BigQuery schema name.
 94        """
 95        self.project_id = project_id
 96        """@private"""
 97        self.bq_schema = bq_schema
 98        """@private"""
 99        self.bq_util = BigQueryUtil(project_id)
100        """@private"""
101
102    def check_permissions_for_dataset(self, raise_on_other_failure: bool) -> bool:
103        """
104        Check the permissions for accessing BigQuery for specific dataset.
105
106        **Args:**
107        - raise_on_other_failure (bool): Whether to raise an exception on other failures.
108
109        **Returns:**
110        - bool: `True` if permissions are sufficient, `False` otherwise.
111        """
112        query = f"""SELECT 1 FROM `{self.project_id}.{self.bq_schema}.INFORMATION_SCHEMA.TABLES`"""
113        return self.bq_util.check_permissions_for_query(
114            query=query,
115            raise_on_other_failure=raise_on_other_failure
116        )
117
118    def get_tdr_table_contents(self, exclude_datarepo_id: bool, table_name: str, to_dataframe: bool) -> Any:
119        """
120        Retrieve the contents of a TDR table from BigQuery.
121
122        **Args:**
123        - exclude_datarepo_id (bool): Whether to exclude the datarepo_row_id column.
124        - to_dataframe (bool): Whether to return the results as a DataFrame.
125
126        **Returns:**
127        - list[dict]: The contents of the table
128        """
129        if exclude_datarepo_id:
130            exclude_str = "EXCEPT (datarepo_row_id)"
131        else:
132            exclude_str = ""
133        query = f"""SELECT * {exclude_str} FROM `{self.project_id}.{self.bq_schema}.{table_name}`"""
134        logging.info(f"Getting contents of table {table_name} from BQ")
135        return self.bq_util.query_table(query=query, to_dataframe=to_dataframe)