utils/gcp.py

import re

from airflow import models
from airflow.providers.cncf.kubernetes.secret import Secret
from airflow.providers.google.cloud.hooks.dataproc import DataProcJobBuilder
from airflow.providers.google.cloud.operators.dataproc import (
    ClusterGenerator,
    DataprocCreateClusterOperator,
    DataprocDeleteClusterOperator,
    DataprocSubmitJobOperator,
)
from airflow.providers.google.cloud.operators.gcs import GCSDeleteObjectsOperator
from airflow.providers.google.cloud.transfers.bigquery_to_gcs import (
    BigQueryToGCSOperator,
)

from operators.gcp_container_operator import GKEPodOperator
from utils.dataproc import get_dataproc_parameters

GCP_PROJECT_ID = "moz-fx-data-airflow-gke-prod"
DATAPROC_PROJECT_ID = "airflow-dataproc"
BIGQUERY_ETL_DOCKER_IMAGE = "gcr.io/moz-fx-data-airflow-prod-88e0/bigquery-etl:latest"


def export_to_parquet(
    table,
    destination_table=None,
    static_partitions=None,
    arguments=None,
    use_storage_api=False,
    dag_name="export_to_parquet",
    parent_dag_name=None,
    default_args=None,
    gcp_conn_id="google_cloud_airflow_dataproc",
    dataproc_storage_bucket="airflow-dataproc-bq-parquet-exports",
    num_workers=2,
    num_preemptible_workers=0,
    gcs_output_bucket="airflow-dataproc-bq-parquet-exports",
    region="us-west1",
):
    """
    Export a BigQuery table to Parquet.

    https://github.com/mozilla/bigquery-etl/blob/main/script/pyspark/export_to_parquet.py

    :param str table:                             [Required] BigQuery table name
    :param Optional[str] destination_table:       Output table name, defaults to table,
                                                  will have r'_v[0-9]+$' replaced with
                                                  r'/v[0-9]+'
    :param List[str] arguments:                   Additional pyspark arguments
    :param bool use_storage_api:                  Whether to read from the BigQuery
                                                  Storage API or an AVRO export
    :param str dag_name:                          Name of DAG
    :param Optional[str] parent_dag_name:         Parent DAG name
    :param Optional[Dict[str, Any]] default_args: DAG configuration
    :param str gcp_conn_id:                       Airflow connection id for GCP access
    :param str dataproc_storage_bucket:           Dataproc staging GCS bucket
    :param int num_preemptible_workers:           Number of Dataproc preemptible workers
    :param str region:                            Region where the dataproc cluster will
                                                  be located. Zone will be chosen
                                                  automatically

    :return: airflow.models.DAG
    """

    # remove the dataset prefix and partition suffix from table
    table_id = table.rsplit(".", 1)[-1]
    unqualified_table, _, partition_id = table_id.partition("$")
    # limit cluster name to 35 characters plus suffix of -export-YYYYMMDD (51 total)
    cluster_name = unqualified_table.replace("_", "-")
    if len(cluster_name) > 35:
        # preserve version when truncating cluster name to 42 characters
        prefix, version = re.match(r"(.*?)(-v[0-9]+)?$", cluster_name).groups("")
        cluster_name = prefix[: 35 - len(version)] + version
    cluster_name += "-export-{{ ds_nodash }}"

    dag_prefix = parent_dag_name + "." if parent_dag_name else ""
    project_id = DATAPROC_PROJECT_ID

    if destination_table is None:
        destination_table = unqualified_table
    # separate version using "/" instead of "_"
    export_prefix = re.sub(r"_(v[0-9]+)$", r"/\1", destination_table) + "/"
    if static_partitions:
        export_prefix += "/".join(static_partitions) + "/"
    avro_prefix = "avro/" + export_prefix
    if not static_partitions and partition_id:
        avro_prefix += "partition_id=" + partition_id + "/"
    avro_path = "gs://" + gcs_output_bucket + "/" + avro_prefix + "*.avro"

    params = get_dataproc_parameters("google_cloud_airflow_dataproc")

    if arguments is None:
        arguments = []

    with models.DAG(dag_id=dag_prefix + dag_name, default_args=default_args) as dag:
        create_dataproc_cluster = DataprocCreateClusterOperator(
            task_id="create_dataproc_cluster",
            cluster_name=cluster_name,
            project_id=project_id,
            gcp_conn_id=gcp_conn_id,
            region=region,
            cluster_config=ClusterGenerator(
                project_id=project_id,
                service_account=params.client_email,
                num_workers=num_workers,
                storage_bucket=dataproc_storage_bucket,
                init_actions_uris=[
                    # This is probably a google hosted bucket for
                    # https://github.com/GoogleCloudDataproc/initialization-actions/blob/master/python/pip-install.sh
                    "gs://dataproc-initialization-actions/python/pip-install.sh",
                ],
                metadata={"PIP_PACKAGES": "google-cloud-bigquery==1.20.0"},
                image_version="1.4-debian10",
                properties={},
                master_machine_type="n1-standard-8",
                worker_machine_type="n1-standard-8",
                num_preemptible_workers=num_preemptible_workers,
            ).make(),
        )

        dataproc_job_builder = DataProcJobBuilder(
            job_type="pyspark_job",
            task_id="run_dataproc_pyspark",
            cluster_name=cluster_name,
            project_id=project_id,
            properties={
                "spark.jars.packages": "org.apache.spark:spark-avro_2.11:2.4.4",
            },
        )
        dataproc_job_builder.add_jar_file_uris(
            ["gs://spark-lib/bigquery/spark-2.4-bigquery-latest.jar"]
        )
        dataproc_job_builder.set_python_main(
            "https://raw.githubusercontent.com/mozilla/bigquery-etl/main/script/legacy/export_to_parquet.py"
        )
        dataproc_job_builder.add_args(
            [table]
            + [
                "--" + key + "=" + value
                for key, value in {
                    "avro-path": (not use_storage_api) and avro_path,
                    "destination": "gs://" + gcs_output_bucket,
                    "destination-table": destination_table,
                }.items()
                if value
            ]
            + (["--static-partitions"] if static_partitions else [])
            + (static_partitions if static_partitions else [])
            + arguments
        )
        dataproc_job = dataproc_job_builder.build()

        run_dataproc_pyspark = DataprocSubmitJobOperator(
            task_id="run_dataproc_pyspark",
            job=dataproc_job["job"],
            gcp_conn_id=gcp_conn_id,
            project_id=project_id,
            region=region,
        )

        delete_dataproc_cluster = DataprocDeleteClusterOperator(
            task_id="delete_dataproc_cluster",
            cluster_name=cluster_name,
            gcp_conn_id=gcp_conn_id,
            project_id=project_id,
            trigger_rule="all_done",
            region=region,
        )

        if not use_storage_api:
            avro_export = BigQueryToGCSOperator(
                task_id="avro_export",
                source_project_dataset_table=table,
                destination_cloud_storage_uris=avro_path,
                compression=None,
                export_format="AVRO",
                gcp_conn_id=gcp_conn_id,
            )
            avro_delete = GCSDeleteObjectsOperator(
                task_id="avro_delete",
                bucket_name=gcs_output_bucket,
                prefix=avro_prefix,
                gcp_conn_id=gcp_conn_id,
                trigger_rule="all_done",
            )
            avro_export >> run_dataproc_pyspark >> avro_delete

        create_dataproc_cluster >> run_dataproc_pyspark >> delete_dataproc_cluster

        return dag


def bigquery_etl_query(
    destination_table,
    dataset_id,
    parameters=(),
    arguments=(),
    project_id=None,
    sql_file_path=None,
    reattach_on_restart=True,
    gcp_conn_id="google_cloud_airflow_gke",
    gke_project_id=GCP_PROJECT_ID,
    gke_location="us-west1",
    gke_cluster_name="workloads-prod-v1",
    gke_namespace="default",
    docker_image=BIGQUERY_ETL_DOCKER_IMAGE,
    date_partition_parameter="submission_date",
    multipart=False,
    is_delete_operator_pod=False,
    table_partition_template="${{ds_nodash}}",
    **kwargs,
):
    """
    Generate.

    :param str destination_table:                  [Required] BigQuery destination table
    :param str dataset_id:                         [Required] BigQuery default dataset id
    :param Tuple[str] parameters:                  Parameters passed to bq query
    :param Tuple[str] arguments:                   Additional bq query arguments
    :param Optional[str] project_id:               BigQuery default project id
    :param Optional[str] sql_file_path:            Optional override for path to the
                                                   SQL query file to run
    :param str gcp_conn_id:                        Airflow connection id for GCP access
    :param str gke_project_id:                     GKE cluster project id
    :param str gke_location:                       GKE cluster location
    :param str gke_cluster_name:                   GKE cluster name
    :param str gke_namespace:                      GKE cluster namespace
    :param str docker_image:                       docker image to use
    :param Optional[str] date_partition_parameter: Parameter for indicating destination
                                                   partition to generate, if None
                                                   destination should be whole table
                                                   rather than partition
    :param Dict[str, Any] kwargs:                  Additional keyword arguments for
                                                   GKEPodOperator
    :param is_delete_operator_pod                  Optional, What to do when the pod reaches its final
                                                   state, or the execution is interrupted.
                                                   If False (default): do nothing, If True: delete the pod
    :param str table_partition_template:           Template for the datestring that's used for
                                                   the partition table id. Defaults to "{{ds_nodash}}".
    :return: GKEPodOperator
    """
    kwargs["task_id"] = kwargs.get("task_id", destination_table)
    kwargs["name"] = kwargs.get("name", kwargs["task_id"].replace("_", "-"))
    if not project_id:
        project_id = "moz-fx-data-shared-prod"
    destination_table_no_partition = (
        destination_table.split("$")[0] if destination_table is not None else None
    )
    sql_file_path = (
        sql_file_path
        or f"sql/{project_id}/{dataset_id}/{destination_table_no_partition}/query.sql"
    )
    if destination_table is not None and date_partition_parameter is not None:
        destination_table = destination_table + table_partition_template
        parameters += (date_partition_parameter + ":DATE:{{ds}}",)
    args = ["script/bqetl", "query", "run-multipart"] if multipart else ["query"]
    return GKEPodOperator(
        reattach_on_restart=reattach_on_restart,
        gcp_conn_id=gcp_conn_id,
        project_id=gke_project_id,
        location=gke_location,
        cluster_name=gke_cluster_name,
        namespace=gke_namespace,
        image=docker_image,
        arguments=args
        + (["--destination_table=" + destination_table] if destination_table else [])
        + ["--dataset_id=" + dataset_id]
        + (["--project_id=" + project_id] if project_id else [])
        + ["--parameter=" + parameter for parameter in parameters]
        + list(arguments)
        + [sql_file_path],
        is_delete_operator_pod=is_delete_operator_pod,
        **kwargs,
    )


def bigquery_etl_copy_deduplicate(
    task_id,
    target_project_id,
    billing_projects=(),
    only_tables=None,
    except_tables=None,
    parallelism=4,
    priority="INTERACTIVE",
    hourly=False,
    slices=None,
    gcp_conn_id="google_cloud_airflow_gke",
    gke_project_id=GCP_PROJECT_ID,
    gke_location="us-west1",
    gke_cluster_name="workloads-prod-v1",
    gke_namespace="default",
    docker_image=BIGQUERY_ETL_DOCKER_IMAGE,
    **kwargs,
):
    """
    Copy a day's data from live ping tables to stable ping tables, deduplicating on document_id.

    :param str task_id:              [Required] ID for the task
    :param str target_project_id:    [Required] ID of project where target tables live
    :param Tuple[str] billing_projects: ID of projects where queries will be executed,
                                     defaults to gcp_conn_id project
    :param Tuple[str] only_tables:   Only process tables matching the given globs of form 'telemetry_live.main_v*'
    :param Tuple[str] except_tables: Process all tables except those matching the given globs
    :param int parallelism:          Maximum number of queries to execute concurrently
    :param str priority:             BigQuery query priority to use, must be BATCH or INTERACTIVE
    :param bool hourly:              Alias for --slices=24
    :param int slices:               Number of time-based slices to deduplicate in, rather than for whole days at once
    :param str gcp_conn_id:          Airflow connection id for GCP access
    :param str gke_project_id:       GKE cluster project id
    :param str gke_location:         GKE cluster location
    :param str gke_cluster_name:     GKE cluster name
    :param str gke_namespace:        GKE cluster namespace
    :param str docker_image:         docker image to use
    :param Dict[str, Any] kwargs:    Additional keyword arguments for
                                     GKEPodOperator

    :return: GKEPodOperator
    """
    kwargs["name"] = kwargs.get("name", task_id.replace("_", "-"))
    table_qualifiers = []
    if only_tables:
        # bqetl expects multiple args as --only ... --only ...
        for only_table in only_tables:
            table_qualifiers.append("--only")
            table_qualifiers.append(only_table)
    if except_tables:
        for except_table in except_tables:
            table_qualifiers.append("--except")
            table_qualifiers.append(except_table)
    return GKEPodOperator(
        task_id=task_id,
        reattach_on_restart=True,
        gcp_conn_id=gcp_conn_id,
        project_id=gke_project_id,
        location=gke_location,
        cluster_name=gke_cluster_name,
        namespace=gke_namespace,
        image=docker_image,
        arguments=["script/bqetl", "copy_deduplicate"]
        + ["--project-id=" + target_project_id]
        + (["--billing-projects", *list(billing_projects)] if billing_projects else [])
        + ["--date={{ds}}"]
        + [f"--parallelism={parallelism}"]
        + [f"--priority={priority}"]
        + (["--hourly"] if hourly else [])
        + ([f"--slices={slices}"] if slices is not None else [])
        + table_qualifiers,
        **kwargs,
    )


def bigquery_dq_check(
    source_table,
    dataset_id,
    task_id,
    parameters=(),
    arguments=(),
    project_id="moz-fx-data-shared-prod",
    gcp_conn_id="google_cloud_airflow_gke",
    gke_project_id=GCP_PROJECT_ID,
    gke_location="us-west1",
    gke_cluster_name="workloads-prod-v1",
    gke_namespace="default",
    docker_image=BIGQUERY_ETL_DOCKER_IMAGE,
    date_partition_parameter="submission_date",
    is_dq_check_fail=True,
    **kwargs,
):
    """
    Run `bqetl check run` to run data quality checks against BigQuery table.

    :param str destination_table:                  [Required] BigQuery destination table the DQ checks are run on
    :param str dataset_id:                         [Required] BigQuery default dataset id

    :param str task_id:              [             [Required] ID for the task
    :param Optional[str] sql_file_path:            Optional override for path to the
                                                   SQL query file to run
    :param Tuple[str] parameters:                  Parameters passed to bq query
    :param Optional[str] project_id:               BigQuery default project id
    :param str gcp_conn_id:                        Airflow connection id for GCP access
    :param str gke_project_id:                     GKE cluster project id
    :param str gke_location:                       GKE cluster location
    :param str gke_cluster_name:                   GKE cluster name
    :param str gke_namespace:                      GKE cluster namespace
    :param str docker_image:                       docker image to use
    :param Optional[str] date_partition_parameter: Parameter for indicating destination
                                                   partition to generate, if None
                                                   destination should be whole table
                                                   rather than partition
    :param Dict[str, Any] kwargs:                  Additional keyword arguments for
                                                   GKEPodOperator
    :return: GKEPodOperator
    """
    kwargs["task_id"] = kwargs.get("task_id", task_id)
    kwargs["name"] = kwargs.get("name", task_id.replace("_", "-"))
    destination_table_no_partition = (
        source_table.split("$")[0] if source_table is not None else None
    )

    sql_file_path = f"sql/{project_id}/{dataset_id}/{destination_table_no_partition}"

    marker = ["--marker"]
    if is_dq_check_fail:
        marker += ["fail"]
    else:
        marker += ["warn"]

    args = ["script/bqetl", "check", "run", *marker, sql_file_path]
    return GKEPodOperator(
        gcp_conn_id=gcp_conn_id,
        project_id=gke_project_id,
        location=gke_location,
        cluster_name=gke_cluster_name,
        namespace=gke_namespace,
        image=docker_image,
        arguments=args
        + ["--parameter=" + parameter for parameter in parameters]
        + list(arguments),
        **kwargs,
    )


def bigquery_bigeye_check(
    task_id,
    table_id,
    warehouse_id,
    project_id="moz-fx-data-shared-prod",
    gcp_conn_id="google_cloud_airflow_gke",
    gke_project_id=GCP_PROJECT_ID,
    gke_location="us-west1",
    gke_cluster_name="workloads-prod-v1",
    gke_namespace="default",
    docker_image=BIGQUERY_ETL_DOCKER_IMAGE,
    **kwargs,
):
    """
    Run `bqetl monitoring run` to run Bigeye checks against BigQuery table.

    :param str table_id:                           [Required] BigQuery table the Bigeye checks are run against
    :param str warehouse_id:                       [Required] Bigeye warehouse ID where checks located
    :param str task_id:                            [Required] ID for the task
    :param Optional[str] project_id:               BigQuery default project id
    :param str gcp_conn_id:                        Airflow connection id for GCP access
    :param str gke_project_id:                     GKE cluster project id
    :param str gke_location:                       GKE cluster location
    :param str gke_cluster_name:                   GKE cluster name
    :param str gke_namespace:                      GKE cluster namespace
    :param str docker_image:                       docker image to use
    :param Dict[str, Any] kwargs:                  Additional keyword arguments for
                                                   GKEPodOperator
    :return: GKEPodOperator
    """
    kwargs["task_id"] = kwargs.get("task_id", task_id)
    kwargs["name"] = kwargs.get("name", task_id.replace("_", "-"))

    args = [
        "script/bqetl",
        "monitoring",
        "run",
        table_id,
        "--warehouse_id",
        warehouse_id,
        "--project_id",
        project_id,
    ]

    bigeye_api_key = Secret(
        deploy_type="env",
        deploy_target="BIGEYE_API_KEY",
        secret="airflow-gke-secrets",
        key="bqetl_artifact_deployment__bigeye_api_key",
    )

    return GKEPodOperator(
        gcp_conn_id=gcp_conn_id,
        project_id=gke_project_id,
        location=gke_location,
        cluster_name=gke_cluster_name,
        namespace=gke_namespace,
        image=docker_image,
        arguments=list(args),
        secrets=[bigeye_api_key],
        **kwargs,
    )


def bigquery_xcom_query(
    destination_table,
    dataset_id,
    xcom_task_id,
    parameters=(),
    arguments=(),
    project_id=None,
    gcp_conn_id="google_cloud_airflow_gke",
    gke_project_id=GCP_PROJECT_ID,
    gke_location="us-west1",
    gke_cluster_name="workloads-prod-v1",
    gke_namespace="default",
    docker_image=BIGQUERY_ETL_DOCKER_IMAGE,
    date_partition_parameter="submission_date",
    table_partition_template="${{ds_nodash}}",
    **kwargs,
):
    """
    Generate a GKEPodOperator which runs an xcom result as a bigquery query.

    :param str destination_table:                  [Required] BigQuery destination table
    :param str dataset_id:                         [Required] BigQuery default dataset id
    :param str xcom_task_id:                       [Required] task_id which generated the xcom to pull
    :param Tuple[str] parameters:                  Parameters passed to bq query
    :param Tuple[str] arguments:                   Additional bq query arguments
    :param Optional[str] project_id:               BigQuery default project id
    :param str gcp_conn_id:                        Airflow connection id for GCP access
    :param str gke_project_id:                     GKE cluster project id
    :param str gke_location:                       GKE cluster location
    :param str gke_cluster_name:                   GKE cluster name
    :param str gke_namespace:                      GKE cluster namespace
    :param str docker_image:                       docker image to use
    :param Optional[str] date_partition_parameter: Parameter for indicating destination
                                                   partition to generate, if None
                                                   destination should be whole table
                                                   rather than partition
    :param str table_partition_template:           Template for the datestring that's used for
                                                   the partition table id. Defaults to "{{ds_nodash}}".
    :param Dict[str, Any] kwargs:                  Additional keyword arguments for
                                                   GKEPodOperator

    :return: GKEPodOperator
    """
    kwargs["task_id"] = kwargs.get("task_id", destination_table)
    kwargs["name"] = kwargs.get("name", kwargs["task_id"].replace("_", "-"))
    if destination_table is not None and date_partition_parameter is not None:
        destination_table = destination_table + table_partition_template
        parameters += (date_partition_parameter + ":DATE:{{ds}}",)
    query = "{{ " + f"task_instance.xcom_pull({xcom_task_id!r})" + " }}"

    return GKEPodOperator(
        gcp_conn_id=gcp_conn_id,
        project_id=gke_project_id,
        location=gke_location,
        cluster_name=gke_cluster_name,
        namespace=gke_namespace,
        image=docker_image,
        arguments=["bq"]
        + ["query"]
        + (["--destination_table=" + destination_table] if destination_table else [])
        + ["--dataset_id=" + dataset_id]
        + (["--project_id=" + project_id] if project_id else [])
        + ["--parameter=" + parameter for parameter in parameters]
        + list(arguments)
        + [query],
        **kwargs,
    )


def normalize_table_id(table_name):
    """
    Normalize table name for use with BigQuery.

    * Contain up to 1,024 characters
    * Contain letters (upper or lower case), numbers, and underscores
    We intentionally lower case the table_name.
    https://cloud.google.com/bigquery/docs/tables.
    """
    if len(table_name) > 1024:
        raise ValueError("table_name cannot contain more than 1024 characters")
    else:
        return re.sub("\\W+", "_", table_name).lower()