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

import base64

import concurrent.futures

import csv

import datetime

import decimal

import json

import operator

import os

import time

import unittest

import uuid

import re

import six

import pytest

import pytz

try:

from google.cloud import bigquery_storage_v1beta1

except ImportError: # pragma: NO COVER

bigquery_storage_v1beta1 = None

try:

import pandas

except ImportError: # pragma: NO COVER

pandas = None

try:

import pyarrow

import pyarrow.types

except ImportError: # pragma: NO COVER

pyarrow = None

try:

import IPython

from IPython.utils import io

from IPython.testing import tools

from IPython.terminal import interactiveshell

except ImportError: # pragma: NO COVER

IPython = None

from google.api_core.exceptions import PreconditionFailed

from google.api_core.exceptions import BadRequest

from google.api_core.exceptions import Conflict

from google.api_core.exceptions import Forbidden

from google.api_core.exceptions import GoogleAPICallError

from google.api_core.exceptions import NotFound

from google.api_core.exceptions import InternalServerError

from google.api_core.exceptions import ServiceUnavailable

from google.api_core.exceptions import TooManyRequests

from google.cloud import bigquery

from google.cloud import bigquery_v2

from google.cloud.bigquery.dataset import Dataset

from google.cloud.bigquery.dataset import DatasetReference

from google.cloud.bigquery.table import Table

from google.cloud._helpers import UTC

from google.cloud.bigquery import dbapi

from google.cloud import storage

from test_utils.retry import RetryErrors

from test_utils.retry import RetryInstanceState

from test_utils.retry import RetryResult

from test_utils.system import unique_resource_id

JOB_TIMEOUT = 120 # 2 minutes

WHERE = os.path.abspath(os.path.dirname(__file__))

# Common table data used for many tests.

ROWS = [

("Phred Phlyntstone", 32),

("Bharney Rhubble", 33),

("Wylma Phlyntstone", 29),

("Bhettye Rhubble", 27),

]

HEADER_ROW = ("Full Name", "Age")

SCHEMA = [

bigquery.SchemaField("full_name", "STRING", mode="REQUIRED"),

bigquery.SchemaField("age", "INTEGER", mode="REQUIRED"),

]

TIME_PARTITIONING_CLUSTERING_FIELDS_SCHEMA = [

bigquery.SchemaField("transaction_time", "TIMESTAMP", mode="REQUIRED"),

bigquery.SchemaField("transaction_id", "INTEGER", mode="REQUIRED"),

bigquery.SchemaField("user_email", "STRING", mode="REQUIRED"),

bigquery.SchemaField("store_code", "STRING", mode="REQUIRED"),

bigquery.SchemaField(

"items",

"RECORD",

mode="REPEATED",

fields=[

bigquery.SchemaField("item_code", "STRING", mode="REQUIRED"),

bigquery.SchemaField("quantity", "INTEGER", mode="REQUIRED"),

bigquery.SchemaField("comments", "STRING", mode="NULLABLE"),

bigquery.SchemaField("expiration_date", "DATE", mode="REQUIRED"),

],

),

]

# The VPC-SC team maintains a mirror of the GCS bucket used for code

# samples. The public bucket crosses the configured security boundary.

# See: https://github.com/googleapis/google-cloud-python/issues/8550

SAMPLES_BUCKET = os.environ.get("GCLOUD_TEST_SAMPLES_BUCKET", "cloud-samples-data")

retry_storage_errors = RetryErrors(

(TooManyRequests, InternalServerError, ServiceUnavailable)

)

def _has_rows(result):

return len(result) > 0

def _make_dataset_id(prefix):

return "%s%s" % (prefix, unique_resource_id())

def _load_json_schema(filename="data/schema.json"):

from google.cloud.bigquery.table import _parse_schema_resource

json_filename = os.path.join(WHERE, filename)

with open(json_filename, "r") as schema_file:

return _parse_schema_resource(json.load(schema_file))

def _rate_limit_exceeded(forbidden):

"""Predicate: pass only exceptions with 'rateLimitExceeded' as reason."""

return any(error["reason"] == "rateLimitExceeded" for error in forbidden._errors)

# We need to wait to stay within the rate limits.

# The alternative outcome is a 403 Forbidden response from upstream, which

# they return instead of the more appropriate 429.

# See https://cloud.google.com/bigquery/quota-policy

retry_403 = RetryErrors(Forbidden, error_predicate=_rate_limit_exceeded)

class Config(object):

"""Run-time configuration to be modified at set-up.

This is a mutable stand-in to allow test set-up to modify

global state.

"""

CLIENT = None

CURSOR = None

def setUpModule():

Config.CLIENT = bigquery.Client()

Config.CURSOR = dbapi.connect(Config.CLIENT).cursor()

class TestBigQuery(unittest.TestCase):

def setUp(self):

self.to_delete = []

def tearDown(self):

def _still_in_use(bad_request):

return any(

error["reason"] == "resourceInUse" for error in bad_request._errors

)

retry_in_use = RetryErrors(BadRequest, error_predicate=_still_in_use)

retry_storage_errors_conflict = RetryErrors(

(Conflict, TooManyRequests, InternalServerError, ServiceUnavailable)

)

for doomed in self.to_delete:

if isinstance(doomed, storage.Bucket):

retry_storage_errors_conflict(doomed.delete)(force=True)

elif isinstance(doomed, (Dataset, bigquery.DatasetReference)):

retry_in_use(Config.CLIENT.delete_dataset)(doomed, delete_contents=True)

elif isinstance(doomed, (Table, bigquery.TableReference)):

retry_in_use(Config.CLIENT.delete_table)(doomed)

else:

doomed.delete()

def test_get_service_account_email(self):

client = Config.CLIENT

got = client.get_service_account_email()

self.assertIsInstance(got, six.text_type)

self.assertIn("@", got)

def _create_bucket(self, bucket_name, location=None):

storage_client = storage.Client()

bucket = storage_client.bucket(bucket_name)

retry_storage_errors(bucket.create)(location=location)

self.to_delete.append(bucket)

return bucket

def test_create_dataset(self):

DATASET_ID = _make_dataset_id("create_dataset")

dataset = self.temp_dataset(DATASET_ID)

self.assertTrue(_dataset_exists(dataset))

self.assertEqual(dataset.dataset_id, DATASET_ID)

self.assertEqual(dataset.project, Config.CLIENT.project)

def test_get_dataset(self):

dataset_id = _make_dataset_id("get_dataset")

client = Config.CLIENT

dataset_arg = Dataset(client.dataset(dataset_id))

dataset_arg.friendly_name = "Friendly"

dataset_arg.description = "Description"

dataset = retry_403(client.create_dataset)(dataset_arg)

self.to_delete.append(dataset)

dataset_ref = client.dataset(dataset_id)

# Get with a reference.

got = client.get_dataset(dataset_ref)

self.assertEqual(got.friendly_name, "Friendly")

self.assertEqual(got.description, "Description")

# Get with a string.

got = client.get_dataset(dataset_id)

self.assertEqual(got.friendly_name, "Friendly")

self.assertEqual(got.description, "Description")

# Get with a fully-qualified string.

got = client.get_dataset("{}.{}".format(client.project, dataset_id))

self.assertEqual(got.friendly_name, "Friendly")

self.assertEqual(got.description, "Description")

def test_update_dataset(self):

dataset = self.temp_dataset(_make_dataset_id("update_dataset"))

self.assertTrue(_dataset_exists(dataset))

self.assertIsNone(dataset.friendly_name)

self.assertIsNone(dataset.description)

self.assertEqual(dataset.labels, {})

dataset.friendly_name = "Friendly"

dataset.description = "Description"

dataset.labels = {"priority": "high", "color": "blue"}

ds2 = Config.CLIENT.update_dataset(

dataset, ("friendly_name", "description", "labels")

)

self.assertEqual(ds2.friendly_name, "Friendly")

self.assertEqual(ds2.description, "Description")

self.assertEqual(ds2.labels, {"priority": "high", "color": "blue"})

ds2.labels = {

"color": "green", # change

"shape": "circle", # add

"priority": None, # delete

}

ds3 = Config.CLIENT.update_dataset(ds2, ["labels"])

self.assertEqual(ds3.labels, {"color": "green", "shape": "circle"})

# If we try to update using d2 again, it will fail because the

# previous update changed the ETag.

ds2.description = "no good"

with self.assertRaises(PreconditionFailed):

Config.CLIENT.update_dataset(ds2, ["description"])

def test_list_datasets(self):

datasets_to_create = [

"new" + unique_resource_id(),

"newer" + unique_resource_id(),

"newest" + unique_resource_id(),

]

for dataset_id in datasets_to_create:

self.temp_dataset(dataset_id)

# Retrieve the datasets.

iterator = Config.CLIENT.list_datasets()

all_datasets = list(iterator)

self.assertIsNone(iterator.next_page_token)

created = [

dataset

for dataset in all_datasets

if dataset.dataset_id in datasets_to_create

and dataset.project == Config.CLIENT.project

]

self.assertEqual(len(created), len(datasets_to_create))

def test_list_datasets_w_project(self):

# Retrieve datasets from a different project.

iterator = Config.CLIENT.list_datasets(project="bigquery-public-data")

all_datasets = frozenset([dataset.dataset_id for dataset in iterator])

self.assertIn("usa_names", all_datasets)

def test_create_table(self):

dataset = self.temp_dataset(_make_dataset_id("create_table"))

table_id = "test_table"

table_arg = Table(dataset.table(table_id), schema=SCHEMA)

self.assertFalse(_table_exists(table_arg))

table = retry_403(Config.CLIENT.create_table)(table_arg)

self.to_delete.insert(0, table)

self.assertTrue(_table_exists(table))

self.assertEqual(table.table_id, table_id)

def test_create_table_w_time_partitioning_w_clustering_fields(self):

from google.cloud.bigquery.table import TimePartitioning

from google.cloud.bigquery.table import TimePartitioningType

dataset = self.temp_dataset(_make_dataset_id("create_table_tp_cf"))

table_id = "test_table"

table_arg = Table(

dataset.table(table_id), schema=TIME_PARTITIONING_CLUSTERING_FIELDS_SCHEMA

)

self.assertFalse(_table_exists(table_arg))

table_arg.time_partitioning = TimePartitioning(field="transaction_time")

table_arg.clustering_fields = ["user_email", "store_code"]

table = retry_403(Config.CLIENT.create_table)(table_arg)

self.to_delete.insert(0, table)

self.assertTrue(_table_exists(table))

self.assertEqual(table.table_id, table_id)

time_partitioning = table.time_partitioning

self.assertEqual(time_partitioning.type_, TimePartitioningType.DAY)

self.assertEqual(time_partitioning.field, "transaction_time")

self.assertEqual(table.clustering_fields, ["user_email", "store_code"])

def test_delete_dataset_with_string(self):

dataset_id = _make_dataset_id("delete_table_true")

dataset_ref = Config.CLIENT.dataset(dataset_id)

retry_403(Config.CLIENT.create_dataset)(Dataset(dataset_ref))

self.assertTrue(_dataset_exists(dataset_ref))

Config.CLIENT.delete_dataset(dataset_id)

self.assertFalse(_dataset_exists(dataset_ref))

def test_delete_dataset_delete_contents_true(self):

dataset_id = _make_dataset_id("delete_table_true")

dataset = retry_403(Config.CLIENT.create_dataset)(

Dataset(Config.CLIENT.dataset(dataset_id))

)

table_id = "test_table"

table_arg = Table(dataset.table(table_id), schema=SCHEMA)

table = retry_403(Config.CLIENT.create_table)(table_arg)

Config.CLIENT.delete_dataset(dataset, delete_contents=True)

self.assertFalse(_table_exists(table))

def test_delete_dataset_delete_contents_false(self):

from google.api_core import exceptions

dataset = self.temp_dataset(_make_dataset_id("delete_table_false"))

table_id = "test_table"

table_arg = Table(dataset.table(table_id), schema=SCHEMA)

retry_403(Config.CLIENT.create_table)(table_arg)

with self.assertRaises(exceptions.BadRequest):

Config.CLIENT.delete_dataset(dataset)

def test_get_table_w_public_dataset(self):

public = "bigquery-public-data"

dataset_id = "samples"

table_id = "shakespeare"

table_ref = DatasetReference(public, dataset_id).table(table_id)

# Get table with reference.

table = Config.CLIENT.get_table(table_ref)

self.assertEqual(table.table_id, table_id)

self.assertEqual(table.dataset_id, dataset_id)

self.assertEqual(table.project, public)

schema_names = [field.name for field in table.schema]

self.assertEqual(schema_names, ["word", "word_count", "corpus", "corpus_date"])

# Get table with string.

table = Config.CLIENT.get_table("{}.{}.{}".format(public, dataset_id, table_id))

self.assertEqual(table.table_id, table_id)

self.assertEqual(table.dataset_id, dataset_id)

self.assertEqual(table.project, public)

def test_list_partitions(self):

table_ref = DatasetReference(

"bigquery-public-data", "ethereum_blockchain"

).table("blocks")

all_rows = Config.CLIENT.list_partitions(table_ref)

self.assertIn("20180801", all_rows)

self.assertGreater(len(all_rows), 1000)

def test_list_tables(self):

dataset_id = _make_dataset_id("list_tables")

dataset = self.temp_dataset(dataset_id)

# Retrieve tables before any are created for the dataset.

iterator = Config.CLIENT.list_tables(dataset)

all_tables = list(iterator)

self.assertEqual(all_tables, [])

self.assertIsNone(iterator.next_page_token)

# Insert some tables to be listed.

tables_to_create = [

"new" + unique_resource_id(),

"newer" + unique_resource_id(),

"newest" + unique_resource_id(),

]

for table_name in tables_to_create:

table = Table(dataset.table(table_name), schema=SCHEMA)

created_table = retry_403(Config.CLIENT.create_table)(table)

self.to_delete.insert(0, created_table)

# Retrieve the tables.

iterator = Config.CLIENT.list_tables(dataset)

all_tables = list(iterator)

self.assertIsNone(iterator.next_page_token)

created = [

table

for table in all_tables

if (table.table_id in tables_to_create and table.dataset_id == dataset_id)

]

self.assertEqual(len(created), len(tables_to_create))

# List tables with a string ID.

iterator = Config.CLIENT.list_tables(dataset_id)

self.assertGreater(len(list(iterator)), 0)

# List tables with a fully-qualified string ID.

iterator = Config.CLIENT.list_tables(

"{}.{}".format(Config.CLIENT.project, dataset_id)

)

self.assertGreater(len(list(iterator)), 0)

def test_update_table(self):

dataset = self.temp_dataset(_make_dataset_id("update_table"))

TABLE_NAME = "test_table"

table_arg = Table(dataset.table(TABLE_NAME), schema=SCHEMA)

self.assertFalse(_table_exists(table_arg))

table = retry_403(Config.CLIENT.create_table)(table_arg)

self.to_delete.insert(0, table)

self.assertTrue(_table_exists(table))

self.assertIsNone(table.friendly_name)

self.assertIsNone(table.description)

self.assertEqual(table.labels, {})

table.friendly_name = "Friendly"

table.description = "Description"

table.labels = {"priority": "high", "color": "blue"}

table2 = Config.CLIENT.update_table(

table, ["friendly_name", "description", "labels"]

)

self.assertEqual(table2.friendly_name, "Friendly")

self.assertEqual(table2.description, "Description")

self.assertEqual(table2.labels, {"priority": "high", "color": "blue"})

table2.description = None

table2.labels = {

"color": "green", # change

"shape": "circle", # add

"priority": None, # delete

}

table3 = Config.CLIENT.update_table(table2, ["description", "labels"])

self.assertIsNone(table3.description)

self.assertEqual(table3.labels, {"color": "green", "shape": "circle"})

# If we try to update using table2 again, it will fail because the

# previous update changed the ETag.

table2.description = "no good"

with self.assertRaises(PreconditionFailed):

Config.CLIENT.update_table(table2, ["description"])

def test_update_table_schema(self):

dataset = self.temp_dataset(_make_dataset_id("update_table"))

TABLE_NAME = "test_table"

table_arg = Table(dataset.table(TABLE_NAME), schema=SCHEMA)

self.assertFalse(_table_exists(table_arg))

table = retry_403(Config.CLIENT.create_table)(table_arg)

self.to_delete.insert(0, table)

self.assertTrue(_table_exists(table))

voter = bigquery.SchemaField("voter", "BOOLEAN", mode="NULLABLE")

schema = table.schema

schema.append(voter)

table.schema = schema

updated_table = Config.CLIENT.update_table(table, ["schema"])

self.assertEqual(len(updated_table.schema), len(schema))

for found, expected in zip(updated_table.schema, schema):

self.assertEqual(found.name, expected.name)

self.assertEqual(found.field_type, expected.field_type)

self.assertEqual(found.mode, expected.mode)

@staticmethod

def _fetch_single_page(table, selected_fields=None):

iterator = Config.CLIENT.list_rows(table, selected_fields=selected_fields)

page = six.next(iterator.pages)

return list(page)

def _create_table_many_columns(self, rowcount):

# Generate a table of maximum width via CREATE TABLE AS SELECT.

# first column is named 'rowval', and has a value from 1..rowcount

# Subsequent column is named col_<N> and contains the value N*rowval,

# where N is between 1 and 9999 inclusive.

dsname = _make_dataset_id("wide_schema")

dataset = self.temp_dataset(dsname)

table_id = "many_columns"

table_ref = dataset.table(table_id)

self.to_delete.insert(0, table_ref)

colprojections = ",".join(

["r * {} as col_{}".format(n, n) for n in range(1, 10000)]

)

sql = """

CREATE TABLE {}.{}

AS

SELECT

r as rowval,

{}

FROM

UNNEST(GENERATE_ARRAY(1,{},1)) as r

""".format(

dsname, table_id, colprojections, rowcount

)

query_job = Config.CLIENT.query(sql)

query_job.result()

self.assertEqual(query_job.statement_type, "CREATE_TABLE_AS_SELECT")

self.assertEqual(query_job.ddl_operation_performed, "CREATE")

self.assertEqual(query_job.ddl_target_table, table_ref)

return table_ref

def test_query_many_columns(self):

# Test working with the widest schema BigQuery supports, 10k columns.

row_count = 2

table_ref = self._create_table_many_columns(row_count)

rows = list(

Config.CLIENT.query(

"SELECT * FROM `{}.{}`".format(table_ref.dataset_id, table_ref.table_id)

)

)

self.assertEqual(len(rows), row_count)

# check field representations adhere to expected values.

correctwidth = 0

badvals = 0

for r in rows:

vals = r._xxx_values

rowval = vals[0]

if len(vals) == 10000:

correctwidth = correctwidth + 1

for n in range(1, 10000):

if vals[n] != rowval * (n):

badvals = badvals + 1

self.assertEqual(correctwidth, row_count)

self.assertEqual(badvals, 0)

def test_insert_rows_then_dump_table(self):

NOW_SECONDS = 1448911495.484366

NOW = datetime.datetime.utcfromtimestamp(NOW_SECONDS).replace(tzinfo=UTC)

ROWS = [

("Phred Phlyntstone", 32, NOW),

("Bharney Rhubble", 33, NOW + datetime.timedelta(seconds=10)),

("Wylma Phlyntstone", 29, NOW + datetime.timedelta(seconds=20)),

("Bhettye Rhubble", 27, None),

]

ROW_IDS = range(len(ROWS))

dataset = self.temp_dataset(_make_dataset_id("insert_rows_then_dump"))

TABLE_ID = "test_table"

schema = [

bigquery.SchemaField("full_name", "STRING", mode="REQUIRED"),

bigquery.SchemaField("age", "INTEGER", mode="REQUIRED"),

bigquery.SchemaField("now", "TIMESTAMP"),

]

table_arg = Table(dataset.table(TABLE_ID), schema=schema)

self.assertFalse(_table_exists(table_arg))

table = retry_403(Config.CLIENT.create_table)(table_arg)

self.to_delete.insert(0, table)

self.assertTrue(_table_exists(table))

errors = Config.CLIENT.insert_rows(table, ROWS, row_ids=ROW_IDS)

self.assertEqual(len(errors), 0)

rows = ()

# Allow for "warm up" before rows visible. See

# https://cloud.google.com/bigquery/streaming-data-into-bigquery#dataavailability

# 8 tries -> 1 + 2 + 4 + 8 + 16 + 32 + 64 = 127 seconds

retry = RetryResult(_has_rows, max_tries=8)

rows = retry(self._fetch_single_page)(table)

row_tuples = [r.values() for r in rows]

by_age = operator.itemgetter(1)

self.assertEqual(sorted(row_tuples, key=by_age), sorted(ROWS, key=by_age))

def test_load_table_from_local_avro_file_then_dump_table(self):

from google.cloud.bigquery.job import SourceFormat

from google.cloud.bigquery.job import WriteDisposition

TABLE_NAME = "test_table_avro"

ROWS = [

("violet", 400),

("indigo", 445),

("blue", 475),

("green", 510),

("yellow", 570),

("orange", 590),

("red", 650),

]

dataset = self.temp_dataset(_make_dataset_id("load_local_then_dump"))

table_ref = dataset.table(TABLE_NAME)

table = Table(table_ref)

self.to_delete.insert(0, table)

with open(os.path.join(WHERE, "data", "colors.avro"), "rb") as avrof:

config = bigquery.LoadJobConfig()

config.source_format = SourceFormat.AVRO

config.write_disposition = WriteDisposition.WRITE_TRUNCATE

job = Config.CLIENT.load_table_from_file(

avrof, table_ref, job_config=config

)

# Retry until done.

job.result(timeout=JOB_TIMEOUT)

self.assertEqual(job.output_rows, len(ROWS))

table = Config.CLIENT.get_table(table)

rows = self._fetch_single_page(table)

row_tuples = [r.values() for r in rows]

by_wavelength = operator.itemgetter(1)

self.assertEqual(

sorted(row_tuples, key=by_wavelength), sorted(ROWS, key=by_wavelength)

)

@unittest.skipIf(pandas is None, "Requires `pandas`")

@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")

def test_load_table_from_dataframe_w_nulls(self):

"""Test that a DataFrame with null columns can be uploaded if a

BigQuery schema is specified.

See: https://github.com/googleapis/google-cloud-python/issues/7370

"""

# Schema with all scalar types.

scalars_schema = (

bigquery.SchemaField("bool_col", "BOOLEAN"),

bigquery.SchemaField("bytes_col", "BYTES"),

bigquery.SchemaField("date_col", "DATE"),

bigquery.SchemaField("dt_col", "DATETIME"),

bigquery.SchemaField("float_col", "FLOAT"),

bigquery.SchemaField("geo_col", "GEOGRAPHY"),

bigquery.SchemaField("int_col", "INTEGER"),

bigquery.SchemaField("num_col", "NUMERIC"),

bigquery.SchemaField("str_col", "STRING"),

bigquery.SchemaField("time_col", "TIME"),

bigquery.SchemaField("ts_col", "TIMESTAMP"),

)

table_schema = scalars_schema + (

# TODO: Array columns can't be read due to NULLABLE versus REPEATED

# mode mismatch. See:

# https://issuetracker.google.com/133415569#comment3

# bigquery.SchemaField("array_col", "INTEGER", mode="REPEATED"),

# TODO: Support writing StructArrays to Parquet. See:

# https://jira.apache.org/jira/browse/ARROW-2587

# bigquery.SchemaField("struct_col", "RECORD", fields=scalars_schema),

)

num_rows = 100

nulls = [None] * num_rows

dataframe = pandas.DataFrame(

{

"bool_col": nulls,

"bytes_col": nulls,

"date_col": nulls,

"dt_col": nulls,

"float_col": nulls,

"geo_col": nulls,

"int_col": nulls,

"num_col": nulls,

"str_col": nulls,

"time_col": nulls,

"ts_col": nulls,

}

)

dataset_id = _make_dataset_id("bq_load_test")

self.temp_dataset(dataset_id)

table_id = "{}.{}.load_table_from_dataframe_w_nulls".format(

Config.CLIENT.project, dataset_id

)

# Create the table before loading so that schema mismatch errors are

# identified.

table = retry_403(Config.CLIENT.create_table)(

Table(table_id, schema=table_schema)

)

self.to_delete.insert(0, table)

job_config = bigquery.LoadJobConfig(schema=table_schema)

load_job = Config.CLIENT.load_table_from_dataframe(

dataframe, table_id, job_config=job_config

)

load_job.result()

table = Config.CLIENT.get_table(table)

self.assertEqual(tuple(table.schema), table_schema)

self.assertEqual(table.num_rows, num_rows)

@unittest.skipIf(pandas is None, "Requires `pandas`")

@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")

def test_load_table_from_dataframe_w_required(self):

"""Test that a DataFrame with required columns can be uploaded if a

BigQuery schema is specified.

See: https://github.com/googleapis/google-cloud-python/issues/8093

"""

table_schema = (

bigquery.SchemaField("name", "STRING", mode="REQUIRED"),

bigquery.SchemaField("age", "INTEGER", mode="REQUIRED"),

)

records = [{"name": "Chip", "age": 2}, {"name": "Dale", "age": 3}]

dataframe = pandas.DataFrame(records)

job_config = bigquery.LoadJobConfig(schema=table_schema)

dataset_id = _make_dataset_id("bq_load_test")

self.temp_dataset(dataset_id)

table_id = "{}.{}.load_table_from_dataframe_w_required".format(

Config.CLIENT.project, dataset_id

)

# Create the table before loading so that schema mismatch errors are

# identified.

table = retry_403(Config.CLIENT.create_table)(

Table(table_id, schema=table_schema)

)

self.to_delete.insert(0, table)

job_config = bigquery.LoadJobConfig(schema=table_schema)

load_job = Config.CLIENT.load_table_from_dataframe(

dataframe, table_id, job_config=job_config

)

load_job.result()

table = Config.CLIENT.get_table(table)

self.assertEqual(tuple(table.schema), table_schema)

self.assertEqual(table.num_rows, 2)

@unittest.skipIf(pandas is None, "Requires `pandas`")

@unittest.skipIf(pyarrow is None, "Requires `pyarrow`")

def test_load_table_from_dataframe_w_explicit_schema(self):

# Schema with all scalar types.

scalars_schema = (

bigquery.SchemaField("bool_col", "BOOLEAN"),

bigquery.SchemaField("bytes_col", "BYTES"),

bigquery.SchemaField("date_col", "DATE"),

bigquery.SchemaField("dt_col", "DATETIME"),

bigquery.SchemaField("float_col", "FLOAT"),

bigquery.SchemaField("geo_col", "GEOGRAPHY"),

bigquery.SchemaField("int_col", "INTEGER"),

bigquery.SchemaField("num_col", "NUMERIC"),

bigquery.SchemaField("str_col", "STRING"),

bigquery.SchemaField("time_col", "TIME"),

bigquery.SchemaField("ts_col", "TIMESTAMP"),

)

table_schema = scalars_schema + (

# TODO: Array columns can't be read due to NULLABLE versus REPEATED

# mode mismatch. See:

# https://issuetracker.google.com/133415569#comment3

# bigquery.SchemaField("array_col", "INTEGER", mode="REPEATED"),

# TODO: Support writing StructArrays to Parquet. See:

# https://jira.apache.org/jira/browse/ARROW-2587

# bigquery.SchemaField("struct_col", "RECORD", fields=scalars_schema),

)

dataframe = pandas.DataFrame(

{

"bool_col": [True, None, False],

"bytes_col": [b"abc", None, b"def"],

"date_col": [datetime.date(1, 1, 1), None, datetime.date(9999, 12, 31)],

"dt_col": [

datetime.datetime(1, 1, 1, 0, 0, 0),

None,

datetime.datetime(9999, 12, 31, 23, 59, 59, 999999),

],

"float_col": [float("-inf"), float("nan"), float("inf")],

"geo_col": [

"POINT(30 10)",

None,

"POLYGON ((30 10, 40 40, 20 40, 10 20, 30 10))",

],

"int_col": [-9223372036854775808, None, 9223372036854775807],

"num_col": [

decimal.Decimal("-99999999999999999999999999999.999999999"),

None,

decimal.Decimal("99999999999999999999999999999.999999999"),

],

"str_col": ["abc", None, "def"],

"time_col": [

datetime.time(0, 0, 0),

None,

datetime.time(23, 59, 59, 999999),

],

"ts_col": [

datetime.datetime(1, 1, 1, 0, 0, 0, tzinfo=pytz.utc),

None,

datetime.datetime(

9999, 12, 31, 23, 59, 59, 999999, tzinfo=pytz.utc

),

],

},

dtype="object",

)

dataset_id = _make_dataset_id("bq_load_test")

self.temp_dataset(dataset_id)

table_id = "{}.{}.load_table_from_dataframe_w_explicit_schema".format(

Config.CLIENT.project, dataset_id

)

job_config = bigquery.LoadJobConfig(schema=table_schema)

load_job = Config.CLIENT.load_table_from_dataframe(

dataframe, table_id, job_config=job_config

)

load_job.result()

table = Config.CLIENT.get_table(table_id)

self.assertEqual(tuple(table.schema), table_schema)

self.assertEqual(table.num_rows, 3)

def test_load_avro_from_uri_then_dump_table(self):

from google.cloud.bigquery.job import CreateDisposition

from google.cloud.bigquery.job import SourceFormat

from google.cloud.bigquery.job import WriteDisposition

table_name = "test_table"

rows = [

("violet", 400),

("indigo", 445),

("blue", 475),

("green", 510),

("yellow", 570),

("orange", 590),

("red", 650),

]

with open(os.path.join(WHERE, "data", "colors.avro"), "rb") as f:

GS_URL = self._write_avro_to_storage(

"bq_load_test" + unique_resource_id(), "colors.avro", f

)

dataset = self.temp_dataset(_make_dataset_id("bq_load_test"))

table_arg = dataset.table(table_name)

table = retry_403(Config.CLIENT.create_table)(Table(table_arg))

self.to_delete.insert(0, table)

config = bigquery.LoadJobConfig()

config.create_disposition = CreateDisposition.CREATE_NEVER

config.source_format = SourceFormat.AVRO

config.write_disposition = WriteDisposition.WRITE_EMPTY

job = Config.CLIENT.load_table_from_uri(GS_URL, table_arg, job_config=config)

job.result(timeout=JOB_TIMEOUT)

self.assertEqual(job.output_rows, len(rows))

table = Config.CLIENT.get_table(table)

fetched = self._fetch_single_page(table)

row_tuples = [r.values() for r in fetched]

self.assertEqual(

sorted(row_tuples, key=lambda x: x[1]), sorted(rows, key=lambda x: x[1])

)

def test_load_table_from_uri_then_dump_table(self):

from google.cloud.bigquery.job import CreateDisposition

from google.cloud.bigquery.job import SourceFormat

from google.cloud.bigquery.job import WriteDisposition

TABLE_ID = "test_table"

GS_URL = self._write_csv_to_storage(

"bq_load_test" + unique_resource_id(), "person_ages.csv", HEADER_ROW, ROWS

)

dataset = self.temp_dataset(_make_dataset_id("load_gcs_then_dump"))

table_arg = Table(dataset.table(TABLE_ID), schema=SCHEMA)

table = retry_403(Config.CLIENT.create_table)(table_arg)

self.to_delete.insert(0, table)

config = bigquery.LoadJobConfig()

config.create_disposition = CreateDisposition.CREATE_NEVER

config.skip_leading_rows = 1

config.source_format = SourceFormat.CSV

config.write_disposition = WriteDisposition.WRITE_EMPTY

job = Config.CLIENT.load_table_from_uri(

GS_URL, dataset.table(TABLE_ID), job_config=config

)

# Allow for 90 seconds of "warm up" before rows visible. See

# https://cloud.google.com/bigquery/streaming-data-into-bigquery#dataavailability

# 8 tries -> 1 + 2 + 4 + 8 + 16 + 32 + 64 = 127 seconds

retry = RetryInstanceState(_job_done, max_tries=8)

retry(job.reload)()

rows = self._fetch_single_page(table)

row_tuples = [r.values() for r in rows]

by_age = operator.itemgetter(1)

self.assertEqual(sorted(row_tuples, key=by_age), sorted(ROWS, key=by_age))

def test_load_table_from_file_w_explicit_location(self):

# Create a temporary bucket for extract files.

bucket_name = "bq_load_table_eu_extract_test" + unique_resource_id()

self._create_bucket(bucket_name, location="eu")

# Create a temporary dataset & table in the EU.

table_bytes = six.BytesIO(b"a,3\nb,2\nc,1\n")

client = Config.CLIENT

dataset = self.temp_dataset(_make_dataset_id("eu_load_file"), location="EU")

table_ref = dataset.table("letters")

job_config = bigquery.LoadJobConfig()

job_config.skip_leading_rows = 0

job_config.schema = [

bigquery.SchemaField("letter", "STRING"),

bigquery.SchemaField("value", "INTEGER"),

]

# Load the file to an EU dataset with an EU load job.

load_job = client.load_table_from_file(

table_bytes, table_ref, location="EU", job_config=job_config

)

load_job.result()

job_id = load_job.job_id

# Can get the job from the EU.

load_job = client.get_job(job_id, location="EU")

self.assertEqual(job_id, load_job.job_id)

self.assertEqual("EU", load_job.location)

self.assertTrue(load_job.exists())

# Cannot get the job from the US.

with self.assertRaises(NotFound):

client.get_job(job_id, location="US")

load_job_us = client.get_job(job_id)

load_job_us._properties["jobReference"]["location"] = "US"

self.assertFalse(load_job_us.exists())

with self.assertRaises(NotFound):

load_job_us.reload()

# Can cancel the job from the EU.

self.assertTrue(load_job.cancel())

load_job = client.cancel_job(job_id, location="EU")

self.assertEqual(job_id, load_job.job_id)

self.assertEqual("EU", load_job.location)

# Cannot cancel the job from the US.

with self.assertRaises(NotFound):

client.cancel_job(job_id, location="US")

with self.assertRaises(NotFound):

load_job_us.cancel()

# Can list the table rows.

table = client.get_table(table_ref)

self.assertEqual(table.num_rows, 3)

rows = [(row.letter, row.value) for row in client.list_rows(table)]

self.assertEqual(list(sorted(rows)), [("a", 3), ("b", 2), ("c", 1)])

# Verify location behavior with queries

query_config = bigquery.QueryJobConfig()

query_config.dry_run = True

query_string = "SELECT * FROM `{}.letters` LIMIT 1".format(dataset.dataset_id)

eu_query = client.query(query_string, location="EU", job_config=query_config)

self.assertTrue(eu_query.done)

# Cannot query from US.

with self.assertRaises(GoogleAPICallError):

list(client.query(query_string, location="US", job_config=query_config))

# Cannot copy from US.

with self.assertRaises(GoogleAPICallError):

client.copy_table(

table_ref, dataset.table("letters2_us"), location="US"

).result()

# Cannot extract from US.

with self.assertRaises(GoogleAPICallError):

client.extract_table(

table_ref, "gs://{}/letters-us.csv".format(bucket_name), location="US"

).result()

def _write_csv_to_storage(self, bucket_name, blob_name, header_row, data_rows):

from google.cloud._testing import _NamedTemporaryFile

bucket = self._create_bucket(bucket_name)

blob = bucket.blob(blob_name)

with _NamedTemporaryFile() as temp:

with open(temp.name, "w") as csv_write:

writer = csv.writer(csv_write)

writer.writerow(header_row)

writer.writerows(data_rows)

with open(temp.name, "rb") as csv_read: