bigquery.py

# Copyright 2018 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
#
#     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.

"""IPython Magics

.. function:: ``%%bigquery`` or ``%%bqsql``

    IPython cell magic to run a query and display the result as a DataFrame

    .. code-block:: python

        %%bqsql [<destination_var>] [--project <project>] [--use_legacy_sql]
                   [--verbose] [--params <params>]
        <query>

    Parameters:

    * ``<destination_var>`` (Optional[line argument]):
        variable to store the query results. The results are not displayed if
        this parameter is used. If an error occurs during the query execution,
        the corresponding ``QueryJob`` instance (if available) is stored in
        the variable instead. Set ``bigquery_magics.context.default_variable``
        to set a destination variable without specifying this argument.
    * ``--destination_table`` (Optional[line argument]):
        A dataset and table to store the query results. If table does not exists,
        it will be created. If table already exists, its data will be overwritten.
        Variable should be in a format <dataset_id>.<table_id>.
    * ``--no_query_cache`` (Optional[line argument]):
        Do not use cached query results.
    * ``--project <project>`` (Optional[line argument]):
        Project to use for running the query. Defaults to the context
        :attr:`~google.cloud.bigquery.magics.Context.project`.
    * ``--use_bqstorage_api`` (Optional[line argument]):
        [Deprecated] Not used anymore, as BigQuery Storage API is used by default.
    * ``--use_rest_api`` (Optional[line argument]):
        Use the BigQuery REST API instead of the Storage API.
    * ``--use_legacy_sql`` (Optional[line argument]):
        Runs the query using Legacy SQL syntax. Defaults to Standard SQL if
        this argument not used.
    * ``--verbose`` (Optional[line argument]):
        If this flag is used, information including the query job ID and the
        amount of time for the query to complete will not be cleared after the
        query is finished. By default, this information will be displayed but
        will be cleared after the query is finished.
    * ``--params <params>`` (Optional[line argument]):
        If present, the argument following the ``--params`` flag must be
        either:

        * :class:`str` - A JSON string representation of a dictionary in the
          format ``{"param_name": "param_value"}`` (ex. ``{"num": 17}``). Use
          of the parameter in the query should be indicated with
          ``@param_name``. See ``In[5]`` in the Examples section below.

        * :class:`dict` reference - A reference to a ``dict`` in the format
          ``{"param_name": "param_value"}``, where the value types must be JSON
          serializable. The variable reference is indicated by a ``$`` before
          the variable name (ex. ``$my_dict_var``). See ``In[6]`` and ``In[7]``
          in the Examples section below.

    * ``<query>`` (required, cell argument):
        SQL query to run. If the query does not contain any whitespace (aside
        from leading and trailing whitespace), it is assumed to represent a
        fully-qualified table ID, and the latter's data will be fetched.

    Returns:
        A :class:`pandas.DataFrame` or :class:`bigframes.pandas.DataFrame`
        with the query results, depending on the ``engine`` chosen.

    .. note::
        All queries run using this magic will run using the context
        :attr:`~bigquery_magics.config.Context.credentials`.
"""

from __future__ import print_function

import ast
from concurrent import futures
import copy
import re
import sys
import time
from typing import Any, List, Tuple
import warnings

import IPython  # type: ignore
from IPython import display  # type: ignore
from IPython.core import magic_arguments  # type: ignore
from IPython.core.getipython import get_ipython
from google.api_core import client_info
from google.api_core.exceptions import NotFound
from google.cloud import bigquery
from google.cloud.bigquery import exceptions
from google.cloud.bigquery.dataset import DatasetReference
from google.cloud.bigquery.dbapi import _helpers
from google.cloud.bigquery.job import QueryJobConfig

from bigquery_magics import line_arg_parser as lap
import bigquery_magics._versions_helpers
import bigquery_magics.config
import bigquery_magics.line_arg_parser.exceptions
import bigquery_magics.version

try:
    from google.cloud import bigquery_storage  # type: ignore
except ImportError:
    bigquery_storage = None

try:
    import bigframes.pandas as bpd
except ImportError:
    bpd = None

USER_AGENT = f"ipython-{IPython.__version__} bigquery-magics/{bigquery_magics.version.__version__}"
context = bigquery_magics.config.context


def _handle_error(error, destination_var=None):
    """Process a query execution error.

    Args:
        error (Exception):
            An exception that occurred during the query execution.
        destination_var (Optional[str]):
            The name of the IPython session variable to store the query job.
    """
    if destination_var:
        query_job = getattr(error, "query_job", None)

        if query_job is not None:
            get_ipython().push({destination_var: query_job})
        else:
            # this is the case when previewing table rows by providing just
            # table ID to cell magic
            print(
                "Could not save output to variable '{}'.".format(destination_var),
                file=sys.stderr,
            )

    print("\nERROR:\n", str(error), file=sys.stderr)


def _run_query(client, query, job_config=None):
    """Runs a query while printing status updates

    Args:
        client (google.cloud.bigquery.client.Client):
            Client to bundle configuration needed for API requests.
        query (str):
            SQL query to be executed. Defaults to the standard SQL dialect.
            Use the ``job_config`` parameter to change dialects.
        job_config (Optional[google.cloud.bigquery.job.QueryJobConfig]):
            Extra configuration options for the job.

    Returns:
        google.cloud.bigquery.job.QueryJob: the query job created

    Example:
        >>> client = bigquery.Client()
        >>> _run_query(client, "SELECT 17")
        Executing query with job ID: bf633912-af2c-4780-b568-5d868058632b
        Query executing: 1.66s
        Query complete after 2.07s
        'bf633912-af2c-4780-b568-5d868058632b'
    """
    start_time = time.perf_counter()
    query_job = client.query(query, job_config=job_config)

    if job_config and job_config.dry_run:
        return query_job

    print(f"Executing query with job ID: {query_job.job_id}")

    while True:
        print(
            f"\rQuery executing: {time.perf_counter() - start_time:.2f}s".format(),
            end="",
        )
        try:
            query_job.result(timeout=0.5)
            break
        except futures.TimeoutError:
            continue
    print(f"\nJob ID {query_job.job_id} successfully executed")
    return query_job


def _create_dataset_if_necessary(client, dataset_id):
    """Create a dataset in the current project if it doesn't exist.

    Args:
        client (google.cloud.bigquery.client.Client):
            Client to bundle configuration needed for API requests.
        dataset_id (str):
            Dataset id.
    """
    dataset_reference = DatasetReference(client.project, dataset_id)
    try:
        dataset = client.get_dataset(dataset_reference)
        return
    except NotFound:
        pass
    dataset = bigquery.Dataset(dataset_reference)
    dataset.location = client.location
    print(f"Creating dataset: {dataset_id}")
    dataset = client.create_dataset(dataset)


@magic_arguments.magic_arguments()
@magic_arguments.argument(
    "destination_var",
    nargs="?",
    help=("If provided, save the output to this variable instead of displaying it."),
)
@magic_arguments.argument(
    "--destination_table",
    type=str,
    default=None,
    help=(
        "If provided, save the output of the query to a new BigQuery table. "
        "Variable should be in a format <dataset_id>.<table_id>. "
        "If table does not exists, it will be created. "
        "If table already exists, its data will be overwritten."
    ),
)
@magic_arguments.argument(
    "--project",
    type=str,
    default=None,
    help=("Project to use for executing this query. Defaults to the context project."),
)
@magic_arguments.argument(
    "--max_results",
    default=None,
    help=(
        "Maximum number of rows in dataframe returned from executing the query."
        "Defaults to returning all rows."
    ),
)
@magic_arguments.argument(
    "--maximum_bytes_billed",
    default=None,
    help=(
        "maximum_bytes_billed to use for executing this query. Defaults to "
        "the context default_query_job_config.maximum_bytes_billed."
    ),
)
@magic_arguments.argument(
    "--dry_run",
    action="store_true",
    default=False,
    help=(
        "Sets query to be a dry run to estimate costs. "
        "Defaults to executing the query instead of dry run if this argument is not used."
        "Does not work with engine 'bigframes'. "
    ),
)
@magic_arguments.argument(
    "--use_legacy_sql",
    action="store_true",
    default=False,
    help=(
        "Sets query to use Legacy SQL instead of Standard SQL. Defaults to "
        "Standard SQL if this argument is not used."
    ),
)
@magic_arguments.argument(
    "--bigquery_api_endpoint",
    type=str,
    default=None,
    help=(
        "The desired API endpoint, e.g., bigquery.googlepis.com. Defaults to this "
        "option's value in the context bigquery_client_options."
    ),
)
@magic_arguments.argument(
    "--bqstorage_api_endpoint",
    type=str,
    default=None,
    help=(
        "The desired API endpoint, e.g., bigquerystorage.googlepis.com. Defaults to "
        "this option's value in the context bqstorage_client_options."
    ),
)
@magic_arguments.argument(
    "--no_query_cache",
    action="store_true",
    default=False,
    help=("Do not use cached query results."),
)
@magic_arguments.argument(
    "--use_bqstorage_api",
    action="store_true",
    default=None,
    help=(
        "[Deprecated] The BigQuery Storage API is already used by default to "
        "download large query results, and this option has no effect. "
        "If you want to switch to the classic REST API instead, use the "
        "--use_rest_api option."
    ),
)
@magic_arguments.argument(
    "--use_rest_api",
    action="store_true",
    default=False,
    help=(
        "Use the classic REST API instead of the BigQuery Storage API to "
        "download query results."
    ),
)
@magic_arguments.argument(
    "--verbose",
    action="store_true",
    default=False,
    help=(
        "If set, print verbose output, including the query job ID and the "
        "amount of time for the query to finish. By default, this "
        "information will be displayed as the query runs, but will be "
        "cleared after the query is finished."
        "This flag is ignored when the engine is 'bigframes'."
    ),
)
@magic_arguments.argument(
    "--params",
    nargs="+",
    default=None,
    help=(
        "Parameters to format the query string. If present, the --params "
        "flag should be followed by a string representation of a dictionary "
        "in the format {'param_name': 'param_value'} (ex. {\"num\": 17}), "
        "or a reference to a dictionary in the same format. The dictionary "
        "reference can be made by including a '$' before the variable "
        "name (ex. $my_dict_var)."
    ),
)
@magic_arguments.argument(
    "--progress_bar_type",
    type=str,
    default=None,
    help=(
        "Sets progress bar type to display a progress bar while executing the query."
        "Defaults to use tqdm_notebook. Install the ``tqdm`` package to use this feature."
    ),
)
@magic_arguments.argument(
    "--location",
    type=str,
    default=None,
    help=(
        "Set the location to execute query."
        "Defaults to location set in query setting in console."
        "This flag is ignored when the engine is 'bigframes'."
    ),
)
@magic_arguments.argument(
    "--engine",
    type=str,
    default=None,
    help=(
        "Set the execution engine, either 'pandas' or 'bigframes'."
        "Defaults to engine set in the query setting in console."
    ),
)
def _cell_magic(line, query):
    """Underlying function for bigquery cell magic

    Note:
        This function contains the underlying logic for the 'bigquery' cell
        magic. This function is not meant to be called directly.

    Args:
        line (str): "%%bigquery" followed by arguments as required
        query (str): SQL query to run

    Returns:
        pandas.DataFrame: the query results.
    """

    params, args = _parse_magic_args(line)

    query = query.strip()
    if not query:
        error = ValueError("Query is missing.")
        _handle_error(error, args.destination_var)
        return
    query = _validate_and_resolve_query(query, args)

    engine = args.engine or context.engine

    if engine == "bigframes":
        return _query_with_bigframes(query, params, args)

    return _query_with_pandas(query, params, args)


def _parse_magic_args(line: str) -> Tuple[List[Any], Any]:
    # The built-in parser does not recognize Python structures such as dicts, thus
    # we extract the "--params" option and inteprpret it separately.
    try:
        params_option_value, rest_of_args = _split_args_line(line)

    except lap.QueryParamsParseError as exc:
        raise SyntaxError(
            "--params is not a correctly formatted JSON string or a JSON "
            "serializable dictionary"
        ) from exc

    except lap.DuplicateQueryParamsError as exc:
        raise ValueError("Duplicate --params option.") from exc

    except lap.ParseError as exc:
        raise ValueError(
            "Unrecognized input, are option values correct? "
            "Error details: {}".format(exc.args[0])
        ) from exc

    params = []
    if params_option_value:
        # A non-existing params variable is not expanded and ends up in the input
        # in its raw form, e.g. "$query_params".
        if params_option_value.startswith("$"):
            msg = 'Parameter expansion failed, undefined variable "{}".'.format(
                params_option_value[1:]
            )
            raise NameError(msg)

        params = _helpers.to_query_parameters(ast.literal_eval(params_option_value), {})

    args = magic_arguments.parse_argstring(_cell_magic, rest_of_args)

    if args.engine is not None and args.engine not in ("pandas", "bigframes"):
        raise ValueError(f"Invalid engine: {args.engine}")

    return params, args


def _split_args_line(line: str) -> Tuple[str, str]:
    """Split out the --params option value from the input line arguments.

    Args:
        line: The line arguments passed to the cell magic.

    Returns:
        A tuple of two strings. The first is param option value and
        the second is the rest of the arguments.
    """
    tree = lap.Parser(lap.Lexer(line)).input_line()

    extractor = lap.QueryParamsExtractor()
    params_option_value, rest_of_args = extractor.visit(tree)

    return params_option_value, rest_of_args


def _query_with_bigframes(query: str, params: List[Any], args: Any):
    if args.dry_run:
        raise ValueError("Dry run is not supported by bigframes engine.")

    if bpd is None:
        raise ValueError("Bigframes package is not installed.")

    bpd.options.bigquery.project = context.project
    bpd.options.bigquery.credentials = context.credentials

    max_results = int(args.max_results) if args.max_results else None

    result = bpd.read_gbq_query(
        query,
        max_results=max_results,
        configuration=_create_job_config(args, params).to_api_repr(),
    )

    return _handle_result(result, args)


def _query_with_pandas(query: str, params: List[Any], args: Any):
    bq_client, bqstorage_client = _create_clients(args)

    try:
        return _make_bq_query(
            query,
            args=args,
            params=params,
            bq_client=bq_client,
            bqstorage_client=bqstorage_client,
        )
    finally:
        _close_transports(bq_client, bqstorage_client)


def _create_clients(args: Any) -> Tuple[bigquery.Client, Any]:
    bigquery_client_options = copy.deepcopy(context.bigquery_client_options)
    if args.bigquery_api_endpoint:
        if isinstance(bigquery_client_options, dict):
            bigquery_client_options["api_endpoint"] = args.bigquery_api_endpoint
        else:
            bigquery_client_options.api_endpoint = args.bigquery_api_endpoint

    bq_client = bigquery.Client(
        project=args.project or context.project,
        credentials=context.credentials,
        default_query_job_config=context.default_query_job_config,
        client_info=client_info.ClientInfo(user_agent=USER_AGENT),
        client_options=bigquery_client_options,
        location=args.location,
    )
    if context._connection:
        bq_client._connection = context._connection

    # Check and instantiate bq storage client
    if args.use_bqstorage_api is not None:
        warnings.warn(
            "Deprecated option --use_bqstorage_api, the BigQuery "
            "Storage API is already used by default.",
            category=DeprecationWarning,
        )
    use_bqstorage_api = not args.use_rest_api and (bigquery_storage is not None)

    if not use_bqstorage_api:
        return bq_client, None

    bqstorage_client_options = copy.deepcopy(context.bqstorage_client_options)
    if args.bqstorage_api_endpoint:
        if isinstance(bqstorage_client_options, dict):
            bqstorage_client_options["api_endpoint"] = args.bqstorage_api_endpoint
        else:
            bqstorage_client_options.api_endpoint = args.bqstorage_api_endpoint

    bqstorage_client = _make_bqstorage_client(
        bq_client,
        bqstorage_client_options,
    )

    return bq_client, bqstorage_client


def _handle_result(result, args):
    """Determine the output of the cell, depending on options set.

    If an explicit destination is set, that takes precedence. Write to that
    variable and skip showing any results.

    Otherwise, if there is a default variable set (such as if this module is
    initialized by bigframes), then set that but also show the output.

    Finally, there is no variable to save to, so just show the output.
    """
    if args.destination_var:
        get_ipython().push({args.destination_var: result})
        return None

    if context.default_variable:
        # If a default variable is set, save the result _and_ show the results.
        get_ipython().push({context.default_variable: result})

    return result


def _make_bq_query(
    query: str,
    args: Any,
    params: List[Any],
    bq_client: bigquery.Client,
    bqstorage_client: Any,
):
    max_results = int(args.max_results) if args.max_results else None

    # Any query that does not contain whitespace (aside from leading and trailing whitespace)
    # is assumed to be a table id
    if not re.search(r"\s", query):
        try:
            rows = bq_client.list_rows(query, max_results=max_results)
        except Exception as ex:
            _handle_error(ex, args.destination_var)
            return

        result = rows.to_dataframe(
            bqstorage_client=bqstorage_client,
            create_bqstorage_client=False,
        )
        return _handle_result(result, args)

    job_config = _create_job_config(args, params)
    if args.destination_table:
        split = args.destination_table.split(".")
        if len(split) != 2:
            raise ValueError(
                "--destination_table should be in a <dataset_id>.<table_id> format."
            )
        dataset_id, table_id = split
        job_config.allow_large_results = True
        dataset_ref = DatasetReference(bq_client.project, dataset_id)
        destination_table_ref = dataset_ref.table(table_id)
        job_config.destination = destination_table_ref
        job_config.create_disposition = "CREATE_IF_NEEDED"
        job_config.write_disposition = "WRITE_TRUNCATE"
        _create_dataset_if_necessary(bq_client, dataset_id)

    try:
        query_job = _run_query(bq_client, query, job_config=job_config)
    except Exception as ex:
        _handle_error(ex, args.destination_var)
        return

    if not args.verbose:
        display.clear_output()

    if args.dry_run:
        # TODO(tswast): Use _handle_result() here, too, but perhaps change the
        # format to match the dry run schema from bigframes and pandas-gbq.
        # See: https://github.com/googleapis/python-bigquery-pandas/issues/585
        if args.destination_var:
            get_ipython().push({args.destination_var: query_job})
            return
        else:
            print(
                "Query validated. This query will process {} bytes.".format(
                    query_job.total_bytes_processed
                )
            )
            return query_job

    progress_bar = context.progress_bar_type or args.progress_bar_type

    if max_results:
        result = query_job.result(max_results=max_results).to_dataframe(
            bqstorage_client=None,
            create_bqstorage_client=False,
            progress_bar_type=progress_bar,
        )
    else:
        result = query_job.to_dataframe(
            bqstorage_client=bqstorage_client,
            create_bqstorage_client=False,
            progress_bar_type=progress_bar,
        )

    return _handle_result(result, args)


def _validate_and_resolve_query(query: str, args: Any) -> str:
    # Check if query is given as a reference to a variable.
    if not query.startswith("$"):
        return query

    query_var_name = query[1:]

    if not query_var_name:
        missing_msg = 'Missing query variable name, empty "$" is not allowed.'
        raise NameError(missing_msg)

    if query_var_name.isidentifier():
        ip = get_ipython()
        query = ip.user_ns.get(query_var_name, ip)  # ip serves as a sentinel

        if query is ip:
            raise NameError(f"Unknown query, variable {query_var_name} does not exist.")
        elif not isinstance(query, (str, bytes)):
            raise TypeError(
                f"Query variable {query_var_name} must be a string "
                "or a bytes-like value."
            )
    return query


def _create_job_config(args: Any, params: List[Any]) -> QueryJobConfig:
    job_config = QueryJobConfig()
    job_config.query_parameters = params
    job_config.use_legacy_sql = args.use_legacy_sql
    job_config.dry_run = args.dry_run

    # Don't override context job config unless --no_query_cache is explicitly set.
    if args.no_query_cache:
        job_config.use_query_cache = False

    if args.maximum_bytes_billed == "None":
        job_config.maximum_bytes_billed = 0
    elif args.maximum_bytes_billed is not None:
        value = int(args.maximum_bytes_billed)
        job_config.maximum_bytes_billed = value

    return job_config


def _make_bqstorage_client(client, client_options):
    """Creates a BigQuery Storage client.

    Args:
        client (:class:`~google.cloud.bigquery.client.Client`): BigQuery client.
        client_options (:class:`google.api_core.client_options.ClientOptions`):
            Custom options used with a new BigQuery Storage client instance
            if one is created.

    Raises:
        ImportError: if google-cloud-bigquery-storage is not installed, or
            grpcio package is not installed.


    Returns:
        None: if ``use_bqstorage_api == False``, or google-cloud-bigquery-storage
            is outdated.
        BigQuery Storage Client:
    """
    try:
        bigquery_magics._versions_helpers.BQ_STORAGE_VERSIONS.try_import(
            raise_if_error=True
        )
    except exceptions.BigQueryStorageNotFoundError as err:
        customized_error = ImportError(
            "The default BigQuery Storage API client cannot be used, install "
            "the missing google-cloud-bigquery-storage and pyarrow packages "
            "to use it. Alternatively, use the classic REST API by specifying "
            "the --use_rest_api magic option."
        )
        raise customized_error from err

    try:
        from google.api_core.gapic_v1 import client_info as gapic_client_info
    except ImportError as err:
        customized_error = ImportError(
            "Install the grpcio package to use the BigQuery Storage API."
        )
        raise customized_error from err

    return client._ensure_bqstorage_client(
        client_options=client_options,
        client_info=gapic_client_info.ClientInfo(user_agent=USER_AGENT),
    )


def _close_transports(client, bqstorage_client):
    """Close the given clients' underlying transport channels.

    Closing the transport is needed to release system resources, namely open
    sockets.

    Args:
        client (:class:`~google.cloud.bigquery.client.Client`):
        bqstorage_client
            (Optional[:class:`~google.cloud.bigquery_storage.BigQueryReadClient`]):
            A client for the BigQuery Storage API.

    """
    client.close()
    if bqstorage_client is not None:
        bqstorage_client._transport.grpc_channel.close()