common/py_libs/bq_helper.py

# Copyright 2023 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""Library for BigQuery related functions."""

from collections import abc
from enum import Enum
import logging
import pathlib
import typing

from google.api_core import retry
from google.cloud import bigquery
from google.cloud.exceptions import BadRequest
from google.cloud.exceptions import Conflict
from google.cloud.exceptions import Forbidden
from google.cloud.exceptions import NotFound

from common.py_libs import constants
from common.py_libs import cortex_exceptions as exc

from common.py_libs.bq_materializer import add_cluster_to_table_def
from common.py_libs.bq_materializer import add_partition_to_table_def

logger = logging.getLogger(__name__)



def execute_sql_file(bq_client: bigquery.Client,
                     sql_file: typing.Union[str, pathlib.Path]) -> None:
    """Executes a Bigquery sql file."""
    # TODO: Convert sql_file from str to Path.
    with open(sql_file, mode="r", encoding="utf-8") as sqlf:
        sql_str = sqlf.read()
        logging.debug("Executing SQL: %s", sql_str)
        try:
            query_job = bq_client.query(sql_str)
            # Let's wait for query to complete.
            _ = query_job.result()
        except BadRequest:
            logging.error("Error when executing SQL:\n%s", sql_str)
            raise


def table_exists(bq_client: bigquery.Client, full_table_name: str) -> bool:
    """Checks if a given table exists in BigQuery."""
    try:
        bq_client.get_table(full_table_name)
        return True
    except NotFound:
        return False


def dataset_exists(bq_client: bigquery.Client, full_dataset_name: str) -> bool:
    """Checks if a given dataset exists in BigQuery."""
    try:
        bq_client.get_dataset(full_dataset_name)
        return True
    except NotFound:
        return False


DatasetExistence = Enum("DatasetExistence",
                        ["NOT_EXISTS",
                         "EXISTS_IN_LOCATION",
                         "EXISTS_IN_ANOTHER_LOCATION"])

def dataset_exists_in_location(bq_client: bigquery.Client,
                               full_dataset_name: str,
                               location: str) -> DatasetExistence:
    """Checks if a given dataset exists in BigQuery in a location."""
    try:
        dataset = bq_client.get_dataset(full_dataset_name)
        return (DatasetExistence.EXISTS_IN_LOCATION
                  if dataset.location.lower() == location.lower() # type: ignore
                  else DatasetExistence.EXISTS_IN_ANOTHER_LOCATION)
    except NotFound:
        return DatasetExistence.NOT_EXISTS


def _wait_for_bq_jobs(jobs: typing.List[typing.Union[bigquery.CopyJob,
                                                  bigquery.LoadJob]],
                      continue_if_failed: bool):
    """Waits for BigQuery jobs to finish."""

    # We can simply wait for them in a for loop
    # because we need all of them to finish.
    for job in jobs:
        try:
            job.result(retry=retry.Retry(deadline=60))
            logging.info("✅ Table %s has been loaded.", job.destination)
        except Conflict:
            logging.warning("⚠️ Table %s already exists. Skipping it.",
                        job.destination)
        except Exception:  # pylint: disable=bare-except, broad-exception-caught
            logging.error(
                "⛔️ Failed to load table %s.\n",
                job.destination,
                exc_info=True)
            if not continue_if_failed:
                raise


def load_tables(bq_client: bigquery.Client,
                sources: typing.Union[str, typing.List[str]],
                target_tables: typing.Union[str, typing.List[str]],
                location: str,
                continue_if_failed: bool = False,
                skip_existing_tables: bool = False,
                write_disposition: str = bigquery.WriteDisposition.WRITE_EMPTY,
                parallel_jobs: int = 5):
    """Loads data to multiple BigQuery tables.

    Args:
        bq_client (bigquery.Client): BigQuery client to use.
        sources (str | list[str]): data source URI or name.
                               Supported sources:
                                - BigQuery table name as project.dataset.table
                                - Any URI supported by load_table_from_uri
                                  for avro, csv, json and parquet files.
        target_tables (str | list[str]): full target tables names as
                                      "project.dataset.table".
        location (str): BigQuery location.
        continue_if_failed (bool): continue loading tables if some jobs fail.
        skip_existing_tables (bool): Skip tables that already exist.
                                    Defaults to False.
        write_disposition (bigquery.WriteDisposition): write disposition,
                          Defaults to WRITE_EMPTY (skip if has data).
        parallel_jobs (int): maximum number of parallel jobs. Defaults to 5.

     Raises:
        ValueError: If the number of source URIs is not equal to the number
                    of target tables, or if an unsupported source format
                    is provided.
    """

    if not isinstance(sources, abc.Sequence):
        sources = [sources]
    if not isinstance(target_tables, abc.Sequence):
        target_tables = [target_tables]
    if len(target_tables) != len(sources):
        raise ValueError(("Number of source URIs must be equal to "
                          "number of target tables."))

    jobs = []
    for index, source in enumerate(sources):
        target = target_tables[index]
        logging.info("Loading table %s from %s.", target, source)

        if skip_existing_tables and table_exists(bq_client, target):
            logging.warning("⚠️ Table %s already exists. Skipping it.", target)
            continue

        if "://" in source:
            ext = source.split(".")[-1].lower()
            if ext == "avro":
                source_format = bigquery.SourceFormat.AVRO
            elif ext == "parquet":
                source_format = bigquery.SourceFormat.PARQUET
            elif ext == "csv":
                source_format = bigquery.SourceFormat.CSV
            elif ext == "json":
                source_format = bigquery.SourceFormat.NEWLINE_DELIMITED_JSON
            else:
                raise ValueError((f"Extension `{ext}` "
                                 "is an unsupported source format."))
            job_config = bigquery.LoadJobConfig(
                autodetect=True,
                source_format=source_format,
                write_disposition=write_disposition,
            )
            load_job = bq_client.load_table_from_uri(
                source_uris=source,
                destination=target,
                job_config=job_config,
                location=location,
                retry=retry.Retry(deadline=60),
            )
        else:
            job_config = bigquery.CopyJobConfig(
                write_disposition=write_disposition)
            load_job = bq_client.copy_table(source,
                                            target,
                                            location=location,
                                            job_config=job_config,
                                            retry=retry.Retry(deadline=60))
        jobs.append(load_job)
        # If reached parallel_jobs number, wait for them to finish.
        if len(jobs) >= parallel_jobs:
            _wait_for_bq_jobs(jobs, continue_if_failed)
            jobs.clear()

    # Wait for the rest of jobs to finish.
    _wait_for_bq_jobs(jobs, continue_if_failed)


def create_dataset(bq_client: bigquery.Client,
                   dataset_name: str,
                   location: str,
                   suppress_success_logging: bool = False) -> None:
    """Creates a BigQuery dataset."""
    dataset_ref = bigquery.Dataset(dataset_name)
    dataset_ref.location = location
    try:
        bq_client.create_dataset(dataset_ref, timeout=30)
        if not suppress_success_logging:
            logging.info("✅ Dataset %s has been created in %s.",
                        dataset_name,
                        location)
    except Conflict:
        logging.warning("⚠️ Dataset %s already exists in %s. Skipping it.",
                        dataset_name,
                        location)
    except Exception:
        logging.error("⛔️ Failed to create dataset %s in %s.", dataset_name,
                      location,
                      exc_info=True)
        raise


def create_table(bq_client: bigquery.Client,
                 full_table_name: str,
                 schema_tuples_list: list[tuple[str, str]],
                 exists_ok=False) -> None:
    """Creates a BigQuery table based on given schema."""
    create_table_from_schema(
        bq_client,
        full_table_name,
        [bigquery.SchemaField(t[0], t[1]) for t in schema_tuples_list],
        None,
        None,
        exists_ok
    )


def create_table_from_schema(bq_client: bigquery.Client,
                             full_table_name: str,
                             schema: typing.List[bigquery.SchemaField],
                             partition_details=None,
                             cluster_details=None,
                             exists_ok=False) -> None:
    """Creates a table in BigQuery. Skips creation if it exists.

    Args:
        bq_client (bigquery.Client): BQ client.
        full_table_name (str): Full table name (project.dataset.table).
        schema (list[bigquery.SchemaField]): BQ schema.
        partition_details (dict): Partition details.
        cluster_details (dict): Clustering details.
        exists_ok (Optional[bool]):
            Allow the table to exist prior to this call. Defaults to False.
    """

    logging.info("Creating table %s", full_table_name)

    project, dataset_id, table_id = full_table_name.split(".")

    table_ref = bigquery.TableReference(
        bigquery.DatasetReference(project, dataset_id), table_id)

    table = bigquery.Table(
        table_ref,
        schema=schema)

    if partition_details:
        table = add_partition_to_table_def(table, partition_details)

    if cluster_details:
        table = add_cluster_to_table_def(table, cluster_details)

    bq_client.create_table(table, exists_ok=exists_ok)


def get_table_list(bq_client: bigquery.Client,
                   project_id: str,
                   dataset_name: str) -> typing.List[str]:
    ds_ref = bigquery.DatasetReference(project_id, dataset_name)
    return [t.table_id for t in bq_client.list_tables(ds_ref)]


def delete_table(bq_client: bigquery.Client, full_table_name: str) -> None:
    """ Calls the BQ API to delete the table, returns nothing """
    logger.info("Deleting table `%s`.", full_table_name)
    bq_client.delete_table(full_table_name, not_found_ok=True)


def copy_dataset(bq_client: bigquery.Client,
                 source_project: str,
                 source_dataset: str,
                 target_project: str,
                 target_dataset: str,
                 location: str,
                 skip_existing_tables: bool = False,
                 write_disposition: str = (
                    bigquery.WriteDisposition.WRITE_EMPTY)):
    """Copies all tables from source dataset to target.

    Args:
        bq_client (bigquery.Client): BigQuery client to use.
        source_project (str): Source project.
        source_dataset (str): Source dataset. Must exist in specified location.
        target_project (str): Target project.
        target_dataset (str): Target dataset. Must exist in specified location.
        location (str): BigQuery location.
        skip_existing_tables (bool): Skip tables that already exist.
                                    Defaults to False.
        write_disposition (bigquery.WriteDisposition): Write disposition,
                          Defaults to WRITE_EMPTY (skip if has data).
    """
    logging.info("Copying tables from `%s.%s` to `%s.%s`.",
                 source_project, source_dataset,
                 target_project, target_dataset)
    tables = get_table_list(bq_client, source_project, source_dataset)
    source_tables = [f"{source_project}.{source_dataset}.{t}" for t in tables]
    target_tables = [f"{target_project}.{target_dataset}.{t}" for t in tables]
    load_tables(bq_client,
                source_tables, target_tables, location,
                skip_existing_tables=skip_existing_tables,
                write_disposition=write_disposition)

def label_dataset(bq_client: bigquery.Client,
                  dataset: bigquery.Dataset) -> None:
    """
    Adds Cortex label to BigQuery Dataset.

    Only updates dataset if label does not exist already in dataset.

    Args:
        bq_client (bigquery.Client): BigQuery client to use.
        dataset (bigquery.Dataset): Dataset to label.

    Raises:
        NotFoundCError: If dataset does not exist.
    """
    try:
        labels = dataset.labels or {}
        update_labels = False

        for key, value in constants.BQ_DATASET_LABEL.items():
            if key not in labels:
                labels[key] = value
                update_labels = True
        if update_labels:
            dataset.labels = labels
            bq_client.update_dataset(dataset, ["labels"])

    except NotFound:
        error_msg = (
            f"Dataset {dataset.project}.{dataset.dataset_id} not found.")
        raise exc.NotFoundCError(error_msg) from NotFound

    except Forbidden:
        logger.info(
            "Permission to tag %s.%s denied. Skipping tagging dataset.",
                      dataset.project,
                      dataset.dataset_id)