tools/gcp/utils/big_query_utils.py

# Copyright 2015 gRPC authors.
#
# 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
#
#     http://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.

from __future__ import print_function

import argparse
import json
import uuid

from apiclient import discovery
from apiclient.errors import HttpError
import httplib2
from oauth2client.client import GoogleCredentials

# 30 days in milliseconds
_EXPIRATION_MS = 30 * 24 * 60 * 60 * 1000
NUM_RETRIES = 3


def create_big_query():
    """Authenticates with cloud platform and gets a BiqQuery service object"""
    creds = GoogleCredentials.get_application_default()
    return discovery.build(
        "bigquery", "v2", credentials=creds, cache_discovery=False
    )


def create_dataset(biq_query, project_id, dataset_id):
    is_success = True
    body = {
        "datasetReference": {"projectId": project_id, "datasetId": dataset_id}
    }

    try:
        dataset_req = biq_query.datasets().insert(
            projectId=project_id, body=body
        )
        dataset_req.execute(num_retries=NUM_RETRIES)
    except HttpError as http_error:
        if http_error.resp.status == 409:
            print("Warning: The dataset %s already exists" % dataset_id)
        else:
            # Note: For more debugging info, print "http_error.content"
            print(
                "Error in creating dataset: %s. Err: %s"
                % (dataset_id, http_error)
            )
            is_success = False
    return is_success


def create_table(
    big_query, project_id, dataset_id, table_id, table_schema, description
):
    fields = [
        {
            "name": field_name,
            "type": field_type,
            "description": field_description,
        }
        for (field_name, field_type, field_description) in table_schema
    ]
    return create_table2(
        big_query, project_id, dataset_id, table_id, fields, description
    )


def create_partitioned_table(
    big_query,
    project_id,
    dataset_id,
    table_id,
    table_schema,
    description,
    partition_type="DAY",
    expiration_ms=_EXPIRATION_MS,
):
    """Creates a partitioned table. By default, a date-paritioned table is created with
    each partition lasting 30 days after it was last modified.
    """
    fields = [
        {
            "name": field_name,
            "type": field_type,
            "description": field_description,
        }
        for (field_name, field_type, field_description) in table_schema
    ]
    return create_table2(
        big_query,
        project_id,
        dataset_id,
        table_id,
        fields,
        description,
        partition_type,
        expiration_ms,
    )


def create_table2(
    big_query,
    project_id,
    dataset_id,
    table_id,
    fields_schema,
    description,
    partition_type=None,
    expiration_ms=None,
):
    is_success = True

    body = {
        "description": description,
        "schema": {"fields": fields_schema},
        "tableReference": {
            "datasetId": dataset_id,
            "projectId": project_id,
            "tableId": table_id,
        },
    }

    if partition_type and expiration_ms:
        body["timePartitioning"] = {
            "type": partition_type,
            "expirationMs": expiration_ms,
        }

    try:
        table_req = big_query.tables().insert(
            projectId=project_id, datasetId=dataset_id, body=body
        )
        res = table_req.execute(num_retries=NUM_RETRIES)
        print('Successfully created %s "%s"' % (res["kind"], res["id"]))
    except HttpError as http_error:
        if http_error.resp.status == 409:
            print("Warning: Table %s already exists" % table_id)
        else:
            print(
                "Error in creating table: %s. Err: %s" % (table_id, http_error)
            )
            is_success = False
    return is_success


def patch_table(big_query, project_id, dataset_id, table_id, fields_schema):
    is_success = True

    body = {
        "schema": {"fields": fields_schema},
        "tableReference": {
            "datasetId": dataset_id,
            "projectId": project_id,
            "tableId": table_id,
        },
    }

    try:
        table_req = big_query.tables().patch(
            projectId=project_id,
            datasetId=dataset_id,
            tableId=table_id,
            body=body,
        )
        res = table_req.execute(num_retries=NUM_RETRIES)
        print('Successfully patched %s "%s"' % (res["kind"], res["id"]))
    except HttpError as http_error:
        print("Error in creating table: %s. Err: %s" % (table_id, http_error))
        is_success = False
    return is_success


def insert_rows(big_query, project_id, dataset_id, table_id, rows_list):
    is_success = True
    body = {"rows": rows_list}
    try:
        insert_req = big_query.tabledata().insertAll(
            projectId=project_id,
            datasetId=dataset_id,
            tableId=table_id,
            body=body,
        )
        res = insert_req.execute(num_retries=NUM_RETRIES)
        if res.get("insertErrors", None):
            print("Error inserting rows! Response: %s" % res)
            is_success = False
    except HttpError as http_error:
        print("Error inserting rows to the table %s" % table_id)
        print("Error message: %s" % http_error)
        is_success = False

    return is_success


def sync_query_job(big_query, project_id, query, timeout=5000):
    query_data = {"query": query, "timeoutMs": timeout}
    query_job = None
    try:
        query_job = (
            big_query.jobs()
            .query(projectId=project_id, body=query_data)
            .execute(num_retries=NUM_RETRIES)
        )
    except HttpError as http_error:
        print("Query execute job failed with error: %s" % http_error)
        print(http_error.content)
    return query_job

    # List of (column name, column type, description) tuples


def make_row(unique_row_id, row_values_dict):
    """row_values_dict is a dictionary of column name and column value."""
    return {"insertId": unique_row_id, "json": row_values_dict}