# 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 copy

import datetime

import decimal

import email

import gzip

import io

import json

import unittest

import mock

import requests

import six

from six.moves import http_client

import pytest

import pytz

try:

import pandas

except (ImportError, AttributeError): # pragma: NO COVER

pandas = None

try:

import pyarrow

except (ImportError, AttributeError): # pragma: NO COVER

pyarrow = None

import google.api_core.exceptions

from google.api_core.gapic_v1 import client_info

import google.cloud._helpers

from google.cloud.bigquery.dataset import DatasetReference

def _make_credentials():

import google.auth.credentials

return mock.Mock(spec=google.auth.credentials.Credentials)

def _make_connection(*responses):

import google.cloud.bigquery._http

from google.cloud.exceptions import NotFound

mock_conn = mock.create_autospec(google.cloud.bigquery._http.Connection)

mock_conn.USER_AGENT = "testing 1.2.3"

mock_conn.api_request.side_effect = list(responses) + [NotFound("miss")]

return mock_conn

def _make_list_partitons_meta_info(project, dataset_id, table_id, num_rows=0):

return {

"tableReference": {

"projectId": project,

"datasetId": dataset_id,

"tableId": "{}$__PARTITIONS_SUMMARY__".format(table_id),

},

"schema": {

"fields": [

{"name": "project_id", "type": "STRING", "mode": "NULLABLE"},

{"name": "dataset_id", "type": "STRING", "mode": "NULLABLE"},

{"name": "table_id", "type": "STRING", "mode": "NULLABLE"},

{"name": "partition_id", "type": "STRING", "mode": "NULLABLE"},

]

},

"etag": "ETAG",

"numRows": num_rows,

}

class TestClient(unittest.TestCase):

PROJECT = "PROJECT"

DS_ID = "DATASET_ID"

TABLE_ID = "TABLE_ID"

MODEL_ID = "MODEL_ID"

TABLE_REF = DatasetReference(PROJECT, DS_ID).table(TABLE_ID)

KMS_KEY_NAME = "projects/1/locations/global/keyRings/1/cryptoKeys/1"

LOCATION = "us-central"

@staticmethod

def _get_target_class():

from google.cloud.bigquery.client import Client

return Client

def _make_one(self, *args, **kw):

return self._get_target_class()(*args, **kw)

def _make_table_resource(self):

return {

"id": "%s:%s:%s" % (self.PROJECT, self.DS_ID, self.TABLE_ID),

"tableReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

"tableId": self.TABLE_ID,

},

}

def test_ctor_defaults(self):

from google.cloud.bigquery._http import Connection

creds = _make_credentials()

http = object()

client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)

self.assertIsInstance(client._connection, Connection)

self.assertIs(client._connection.credentials, creds)

self.assertIs(client._connection.http, http)

self.assertIsNone(client.location)

def test_ctor_w_location(self):

from google.cloud.bigquery._http import Connection

creds = _make_credentials()

http = object()

location = "us-central"

client = self._make_one(

project=self.PROJECT, credentials=creds, _http=http, location=location

)

self.assertIsInstance(client._connection, Connection)

self.assertIs(client._connection.credentials, creds)

self.assertIs(client._connection.http, http)

self.assertEqual(client.location, location)

def test_ctor_w_query_job_config(self):

from google.cloud.bigquery._http import Connection

from google.cloud.bigquery import QueryJobConfig

creds = _make_credentials()

http = object()

location = "us-central"

job_config = QueryJobConfig()

job_config.dry_run = True

client = self._make_one(

project=self.PROJECT,

credentials=creds,

_http=http,

location=location,

default_query_job_config=job_config,

)

self.assertIsInstance(client._connection, Connection)

self.assertIs(client._connection.credentials, creds)

self.assertIs(client._connection.http, http)

self.assertEqual(client.location, location)

self.assertIsInstance(client._default_query_job_config, QueryJobConfig)

self.assertTrue(client._default_query_job_config.dry_run)

def test__get_query_results_miss_w_explicit_project_and_timeout(self):

from google.cloud.exceptions import NotFound

creds = _make_credentials()

client = self._make_one(self.PROJECT, creds)

conn = client._connection = _make_connection()

with self.assertRaises(NotFound):

client._get_query_results(

"nothere",

None,

project="other-project",

location=self.LOCATION,

timeout_ms=500,

)

conn.api_request.assert_called_once_with(

method="GET",

path="/projects/other-project/queries/nothere",

query_params={"maxResults": 0, "timeoutMs": 500, "location": self.LOCATION},

)

def test__get_query_results_miss_w_client_location(self):

from google.cloud.exceptions import NotFound

creds = _make_credentials()

client = self._make_one(self.PROJECT, creds, location=self.LOCATION)

conn = client._connection = _make_connection()

with self.assertRaises(NotFound):

client._get_query_results("nothere", None)

conn.api_request.assert_called_once_with(

method="GET",

path="/projects/PROJECT/queries/nothere",

query_params={"maxResults": 0, "location": self.LOCATION},

)

def test__get_query_results_hit(self):

job_id = "query_job"

data = {

"kind": "bigquery#getQueryResultsResponse",

"etag": "some-tag",

"schema": {

"fields": [

{"name": "title", "type": "STRING", "mode": "NULLABLE"},

{"name": "unique_words", "type": "INTEGER", "mode": "NULLABLE"},

]

},

"jobReference": {"projectId": self.PROJECT, "jobId": job_id},

"totalRows": "10",

"totalBytesProcessed": "2464625",

"jobComplete": True,

"cacheHit": False,

}

creds = _make_credentials()

client = self._make_one(self.PROJECT, creds)

client._connection = _make_connection(data)

query_results = client._get_query_results(job_id, None)

self.assertEqual(query_results.total_rows, 10)

self.assertTrue(query_results.complete)

def test_get_service_account_email(self):

path = "/projects/%s/serviceAccount" % (self.PROJECT,)

creds = _make_credentials()

http = object()

client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)

email = "bq-123@bigquery-encryption.iam.gserviceaccount.com"

resource = {"kind": "bigquery#getServiceAccountResponse", "email": email}

conn = client._connection = _make_connection(resource)

service_account_email = client.get_service_account_email()

conn.api_request.assert_called_once_with(method="GET", path=path)

self.assertEqual(service_account_email, email)

def test_get_service_account_email_w_alternate_project(self):

project = "my-alternate-project"

path = "/projects/%s/serviceAccount" % (project,)

creds = _make_credentials()

http = object()

client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)

email = "bq-123@bigquery-encryption.iam.gserviceaccount.com"

resource = {"kind": "bigquery#getServiceAccountResponse", "email": email}

conn = client._connection = _make_connection(resource)

service_account_email = client.get_service_account_email(project=project)

conn.api_request.assert_called_once_with(method="GET", path=path)

self.assertEqual(service_account_email, email)

def test_list_projects_defaults(self):

from google.cloud.bigquery.client import Project

PROJECT_1 = "PROJECT_ONE"

PROJECT_2 = "PROJECT_TWO"

TOKEN = "TOKEN"

DATA = {

"nextPageToken": TOKEN,

"projects": [

{

"kind": "bigquery#project",

"id": PROJECT_1,

"numericId": 1,

"projectReference": {"projectId": PROJECT_1},

"friendlyName": "One",

},

{

"kind": "bigquery#project",

"id": PROJECT_2,

"numericId": 2,

"projectReference": {"projectId": PROJECT_2},

"friendlyName": "Two",

},

],

}

creds = _make_credentials()

client = self._make_one(PROJECT_1, creds)

conn = client._connection = _make_connection(DATA)

iterator = client.list_projects()

page = six.next(iterator.pages)

projects = list(page)

token = iterator.next_page_token

self.assertEqual(len(projects), len(DATA["projects"]))

for found, expected in zip(projects, DATA["projects"]):

self.assertIsInstance(found, Project)

self.assertEqual(found.project_id, expected["id"])

self.assertEqual(found.numeric_id, expected["numericId"])

self.assertEqual(found.friendly_name, expected["friendlyName"])

self.assertEqual(token, TOKEN)

conn.api_request.assert_called_once_with(

method="GET", path="/projects", query_params={}

)

def test_list_projects_explicit_response_missing_projects_key(self):

TOKEN = "TOKEN"

DATA = {}

creds = _make_credentials()

client = self._make_one(self.PROJECT, creds)

conn = client._connection = _make_connection(DATA)

iterator = client.list_projects(max_results=3, page_token=TOKEN)

page = six.next(iterator.pages)

projects = list(page)

token = iterator.next_page_token

self.assertEqual(len(projects), 0)

self.assertIsNone(token)

conn.api_request.assert_called_once_with(

method="GET",

path="/projects",

query_params={"maxResults": 3, "pageToken": TOKEN},

)

def test_list_datasets_defaults(self):

from google.cloud.bigquery.dataset import DatasetListItem

DATASET_1 = "dataset_one"

DATASET_2 = "dataset_two"

PATH = "projects/%s/datasets" % self.PROJECT

TOKEN = "TOKEN"

DATA = {

"nextPageToken": TOKEN,

"datasets": [

{

"kind": "bigquery#dataset",

"id": "%s:%s" % (self.PROJECT, DATASET_1),

"datasetReference": {

"datasetId": DATASET_1,

"projectId": self.PROJECT,

},

"friendlyName": None,

},

{

"kind": "bigquery#dataset",

"id": "%s:%s" % (self.PROJECT, DATASET_2),

"datasetReference": {

"datasetId": DATASET_2,

"projectId": self.PROJECT,

},

"friendlyName": "Two",

},

],

}

creds = _make_credentials()

client = self._make_one(self.PROJECT, creds)

conn = client._connection = _make_connection(DATA)

iterator = client.list_datasets()

page = six.next(iterator.pages)

datasets = list(page)

token = iterator.next_page_token

self.assertEqual(len(datasets), len(DATA["datasets"]))

for found, expected in zip(datasets, DATA["datasets"]):

self.assertIsInstance(found, DatasetListItem)

self.assertEqual(found.full_dataset_id, expected["id"])

self.assertEqual(found.friendly_name, expected["friendlyName"])

self.assertEqual(token, TOKEN)

conn.api_request.assert_called_once_with(

method="GET", path="/%s" % PATH, query_params={}

)

def test_list_datasets_w_project(self):

creds = _make_credentials()

client = self._make_one(self.PROJECT, creds)

conn = client._connection = _make_connection({})

list(client.list_datasets(project="other-project"))

conn.api_request.assert_called_once_with(

method="GET", path="/projects/other-project/datasets", query_params={}

)

def test_list_datasets_explicit_response_missing_datasets_key(self):

PATH = "projects/%s/datasets" % self.PROJECT

TOKEN = "TOKEN"

FILTER = "FILTER"

DATA = {}

creds = _make_credentials()

client = self._make_one(self.PROJECT, creds)

conn = client._connection = _make_connection(DATA)

iterator = client.list_datasets(

include_all=True, filter=FILTER, max_results=3, page_token=TOKEN

)

page = six.next(iterator.pages)

datasets = list(page)

token = iterator.next_page_token

self.assertEqual(len(datasets), 0)

self.assertIsNone(token)

conn.api_request.assert_called_once_with(

method="GET",

path="/%s" % PATH,

query_params={

"all": True,

"filter": FILTER,

"maxResults": 3,

"pageToken": TOKEN,

},

)

def test_dataset_with_specified_project(self):

from google.cloud.bigquery.dataset import DatasetReference

creds = _make_credentials()

http = object()

client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)

dataset = client.dataset(self.DS_ID, self.PROJECT)

self.assertIsInstance(dataset, DatasetReference)

self.assertEqual(dataset.dataset_id, self.DS_ID)

self.assertEqual(dataset.project, self.PROJECT)

def test_dataset_with_default_project(self):

from google.cloud.bigquery.dataset import DatasetReference

creds = _make_credentials()

http = object()

client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)

dataset = client.dataset(self.DS_ID)

self.assertIsInstance(dataset, DatasetReference)

self.assertEqual(dataset.dataset_id, self.DS_ID)

self.assertEqual(dataset.project, self.PROJECT)

def test_get_dataset(self):

from google.cloud.exceptions import ServerError

path = "projects/%s/datasets/%s" % (self.PROJECT, self.DS_ID)

creds = _make_credentials()

http = object()

client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)

resource = {

"id": "%s:%s" % (self.PROJECT, self.DS_ID),

"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},

}

conn = client._connection = _make_connection(resource)

dataset_ref = client.dataset(self.DS_ID)

dataset = client.get_dataset(dataset_ref)

conn.api_request.assert_called_once_with(method="GET", path="/%s" % path)

self.assertEqual(dataset.dataset_id, self.DS_ID)

# Test retry.

# Not a cloud API exception (missing 'errors' field).

client._connection = _make_connection(Exception(""), resource)

with self.assertRaises(Exception):

client.get_dataset(dataset_ref)

# Zero-length errors field.

client._connection = _make_connection(ServerError(""), resource)

with self.assertRaises(ServerError):

client.get_dataset(dataset_ref)

# Non-retryable reason.

client._connection = _make_connection(

ServerError("", errors=[{"reason": "serious"}]), resource

)

with self.assertRaises(ServerError):

client.get_dataset(dataset_ref)

# Retryable reason, but retry is disabled.

client._connection = _make_connection(

ServerError("", errors=[{"reason": "backendError"}]), resource

)

with self.assertRaises(ServerError):

client.get_dataset(dataset_ref, retry=None)

# Retryable reason, default retry: success.

client._connection = _make_connection(

ServerError("", errors=[{"reason": "backendError"}]), resource

)

dataset = client.get_dataset(

# Test with a string for dataset ID.

dataset_ref.dataset_id

)

self.assertEqual(dataset.dataset_id, self.DS_ID)

def test_create_dataset_minimal(self):

from google.cloud.bigquery.dataset import Dataset

PATH = "projects/%s/datasets" % self.PROJECT

RESOURCE = {

"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},

"etag": "etag",

"id": "%s:%s" % (self.PROJECT, self.DS_ID),

}

creds = _make_credentials()

client = self._make_one(project=self.PROJECT, credentials=creds)

conn = client._connection = _make_connection(RESOURCE)

ds_ref = client.dataset(self.DS_ID)

before = Dataset(ds_ref)

after = client.create_dataset(before)

self.assertEqual(after.dataset_id, self.DS_ID)

self.assertEqual(after.project, self.PROJECT)

self.assertEqual(after.etag, RESOURCE["etag"])

self.assertEqual(after.full_dataset_id, RESOURCE["id"])

conn.api_request.assert_called_once_with(

method="POST",

path="/%s" % PATH,

data={

"datasetReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

},

"labels": {},

},

)

def test_create_dataset_w_attrs(self):

from google.cloud.bigquery.dataset import Dataset, AccessEntry

PATH = "projects/%s/datasets" % self.PROJECT

DESCRIPTION = "DESC"

FRIENDLY_NAME = "FN"

LOCATION = "US"

USER_EMAIL = "phred@example.com"

LABELS = {"color": "red"}

VIEW = {

"projectId": "my-proj",

"datasetId": "starry-skies",

"tableId": "northern-hemisphere",

}

RESOURCE = {

"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},

"etag": "etag",

"id": "%s:%s" % (self.PROJECT, self.DS_ID),

"description": DESCRIPTION,

"friendlyName": FRIENDLY_NAME,

"location": LOCATION,

"defaultTableExpirationMs": "3600",

"labels": LABELS,

"access": [{"role": "OWNER", "userByEmail": USER_EMAIL}, {"view": VIEW}],

}

creds = _make_credentials()

client = self._make_one(project=self.PROJECT, credentials=creds)

conn = client._connection = _make_connection(RESOURCE)

entries = [

AccessEntry("OWNER", "userByEmail", USER_EMAIL),

AccessEntry(None, "view", VIEW),

]

ds_ref = client.dataset(self.DS_ID)

before = Dataset(ds_ref)

before.access_entries = entries

before.description = DESCRIPTION

before.friendly_name = FRIENDLY_NAME

before.default_table_expiration_ms = 3600

before.location = LOCATION

before.labels = LABELS

after = client.create_dataset(before)

self.assertEqual(after.dataset_id, self.DS_ID)

self.assertEqual(after.project, self.PROJECT)

self.assertEqual(after.etag, RESOURCE["etag"])

self.assertEqual(after.full_dataset_id, RESOURCE["id"])

self.assertEqual(after.description, DESCRIPTION)

self.assertEqual(after.friendly_name, FRIENDLY_NAME)

self.assertEqual(after.location, LOCATION)

self.assertEqual(after.default_table_expiration_ms, 3600)

self.assertEqual(after.labels, LABELS)

conn.api_request.assert_called_once_with(

method="POST",

path="/%s" % PATH,

data={

"datasetReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

},

"description": DESCRIPTION,

"friendlyName": FRIENDLY_NAME,

"location": LOCATION,

"defaultTableExpirationMs": "3600",

"access": [

{"role": "OWNER", "userByEmail": USER_EMAIL},

{"view": VIEW},

],

"labels": LABELS,

},

)

def test_create_dataset_w_custom_property(self):

# The library should handle sending properties to the API that are not

# yet part of the library

from google.cloud.bigquery.dataset import Dataset

path = "/projects/%s/datasets" % self.PROJECT

resource = {

"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},

"newAlphaProperty": "unreleased property",

}

creds = _make_credentials()

client = self._make_one(project=self.PROJECT, credentials=creds)

conn = client._connection = _make_connection(resource)

ds_ref = client.dataset(self.DS_ID)

before = Dataset(ds_ref)

before._properties["newAlphaProperty"] = "unreleased property"

after = client.create_dataset(before)

self.assertEqual(after.dataset_id, self.DS_ID)

self.assertEqual(after.project, self.PROJECT)

self.assertEqual(after._properties["newAlphaProperty"], "unreleased property")

conn.api_request.assert_called_once_with(

method="POST",

path=path,

data={

"datasetReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

},

"newAlphaProperty": "unreleased property",

"labels": {},

},

)

def test_create_dataset_w_client_location_wo_dataset_location(self):

from google.cloud.bigquery.dataset import Dataset

PATH = "projects/%s/datasets" % self.PROJECT

RESOURCE = {

"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},

"etag": "etag",

"id": "%s:%s" % (self.PROJECT, self.DS_ID),

"location": self.LOCATION,

}

creds = _make_credentials()

client = self._make_one(

project=self.PROJECT, credentials=creds, location=self.LOCATION

)

conn = client._connection = _make_connection(RESOURCE)

ds_ref = client.dataset(self.DS_ID)

before = Dataset(ds_ref)

after = client.create_dataset(before)

self.assertEqual(after.dataset_id, self.DS_ID)

self.assertEqual(after.project, self.PROJECT)

self.assertEqual(after.etag, RESOURCE["etag"])

self.assertEqual(after.full_dataset_id, RESOURCE["id"])

self.assertEqual(after.location, self.LOCATION)

conn.api_request.assert_called_once_with(

method="POST",

path="/%s" % PATH,

data={

"datasetReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

},

"labels": {},

"location": self.LOCATION,

},

)

def test_create_dataset_w_client_location_w_dataset_location(self):

from google.cloud.bigquery.dataset import Dataset

PATH = "projects/%s/datasets" % self.PROJECT

OTHER_LOCATION = "EU"

RESOURCE = {

"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},

"etag": "etag",

"id": "%s:%s" % (self.PROJECT, self.DS_ID),

"location": OTHER_LOCATION,

}

creds = _make_credentials()

client = self._make_one(

project=self.PROJECT, credentials=creds, location=self.LOCATION

)

conn = client._connection = _make_connection(RESOURCE)

ds_ref = client.dataset(self.DS_ID)

before = Dataset(ds_ref)

before.location = OTHER_LOCATION

after = client.create_dataset(before)

self.assertEqual(after.dataset_id, self.DS_ID)

self.assertEqual(after.project, self.PROJECT)

self.assertEqual(after.etag, RESOURCE["etag"])

self.assertEqual(after.full_dataset_id, RESOURCE["id"])

self.assertEqual(after.location, OTHER_LOCATION)

conn.api_request.assert_called_once_with(

method="POST",

path="/%s" % PATH,

data={

"datasetReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

},

"labels": {},

"location": OTHER_LOCATION,

},

)

def test_create_dataset_w_reference(self):

path = "/projects/%s/datasets" % self.PROJECT

resource = {

"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},

"etag": "etag",

"id": "%s:%s" % (self.PROJECT, self.DS_ID),

"location": self.LOCATION,

}

creds = _make_credentials()

client = self._make_one(

project=self.PROJECT, credentials=creds, location=self.LOCATION

)

conn = client._connection = _make_connection(resource)

dataset = client.create_dataset(client.dataset(self.DS_ID))

self.assertEqual(dataset.dataset_id, self.DS_ID)

self.assertEqual(dataset.project, self.PROJECT)

self.assertEqual(dataset.etag, resource["etag"])

self.assertEqual(dataset.full_dataset_id, resource["id"])

self.assertEqual(dataset.location, self.LOCATION)

conn.api_request.assert_called_once_with(

method="POST",

path=path,

data={

"datasetReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

},

"labels": {},

"location": self.LOCATION,

},

)

def test_create_dataset_w_fully_qualified_string(self):

path = "/projects/%s/datasets" % self.PROJECT

resource = {

"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},

"etag": "etag",

"id": "%s:%s" % (self.PROJECT, self.DS_ID),

"location": self.LOCATION,

}

creds = _make_credentials()

client = self._make_one(

project=self.PROJECT, credentials=creds, location=self.LOCATION

)

conn = client._connection = _make_connection(resource)

dataset = client.create_dataset("{}.{}".format(self.PROJECT, self.DS_ID))

self.assertEqual(dataset.dataset_id, self.DS_ID)

self.assertEqual(dataset.project, self.PROJECT)

self.assertEqual(dataset.etag, resource["etag"])

self.assertEqual(dataset.full_dataset_id, resource["id"])

self.assertEqual(dataset.location, self.LOCATION)

conn.api_request.assert_called_once_with(

method="POST",

path=path,

data={

"datasetReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

},

"labels": {},

"location": self.LOCATION,

},

)

def test_create_dataset_w_string(self):

path = "/projects/%s/datasets" % self.PROJECT

resource = {

"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},

"etag": "etag",

"id": "%s:%s" % (self.PROJECT, self.DS_ID),

"location": self.LOCATION,

}

creds = _make_credentials()

client = self._make_one(

project=self.PROJECT, credentials=creds, location=self.LOCATION

)

conn = client._connection = _make_connection(resource)

dataset = client.create_dataset(self.DS_ID)

self.assertEqual(dataset.dataset_id, self.DS_ID)

self.assertEqual(dataset.project, self.PROJECT)

self.assertEqual(dataset.etag, resource["etag"])

self.assertEqual(dataset.full_dataset_id, resource["id"])

self.assertEqual(dataset.location, self.LOCATION)

conn.api_request.assert_called_once_with(

method="POST",

path=path,

data={

"datasetReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

},

"labels": {},

"location": self.LOCATION,

},

)

def test_create_dataset_alreadyexists_w_exists_ok_false(self):

creds = _make_credentials()

client = self._make_one(

project=self.PROJECT, credentials=creds, location=self.LOCATION

)

client._connection = _make_connection(

google.api_core.exceptions.AlreadyExists("dataset already exists")

)

with pytest.raises(google.api_core.exceptions.AlreadyExists):

client.create_dataset(self.DS_ID)

def test_create_dataset_alreadyexists_w_exists_ok_true(self):

post_path = "/projects/{}/datasets".format(self.PROJECT)

get_path = "/projects/{}/datasets/{}".format(self.PROJECT, self.DS_ID)

resource = {

"datasetReference": {"projectId": self.PROJECT, "datasetId": self.DS_ID},

"etag": "etag",

"id": "{}:{}".format(self.PROJECT, self.DS_ID),

"location": self.LOCATION,

}

creds = _make_credentials()

client = self._make_one(

project=self.PROJECT, credentials=creds, location=self.LOCATION

)

conn = client._connection = _make_connection(

google.api_core.exceptions.AlreadyExists("dataset already exists"), resource

)

dataset = client.create_dataset(self.DS_ID, exists_ok=True)

self.assertEqual(dataset.dataset_id, self.DS_ID)

self.assertEqual(dataset.project, self.PROJECT)

self.assertEqual(dataset.etag, resource["etag"])

self.assertEqual(dataset.full_dataset_id, resource["id"])

self.assertEqual(dataset.location, self.LOCATION)

conn.api_request.assert_has_calls(

[

mock.call(

method="POST",

path=post_path,

data={

"datasetReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

},

"labels": {},

"location": self.LOCATION,

},

),

mock.call(method="GET", path=get_path),

]

)

def test_create_table_w_day_partition(self):

from google.cloud.bigquery.table import Table

from google.cloud.bigquery.table import TimePartitioning

path = "projects/%s/datasets/%s/tables" % (self.PROJECT, self.DS_ID)

creds = _make_credentials()

client = self._make_one(project=self.PROJECT, credentials=creds)

resource = self._make_table_resource()

conn = client._connection = _make_connection(resource)

table = Table(self.TABLE_REF)

table.time_partitioning = TimePartitioning()

got = client.create_table(table)

conn.api_request.assert_called_once_with(

method="POST",

path="/%s" % path,

data={

"tableReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

"tableId": self.TABLE_ID,

},

"timePartitioning": {"type": "DAY"},

"labels": {},

},

)

self.assertEqual(table.time_partitioning.type_, "DAY")

self.assertEqual(got.table_id, self.TABLE_ID)

def test_create_table_w_custom_property(self):

# The library should handle sending properties to the API that are not

# yet part of the library

from google.cloud.bigquery.table import Table

path = "projects/%s/datasets/%s/tables" % (self.PROJECT, self.DS_ID)

creds = _make_credentials()

client = self._make_one(project=self.PROJECT, credentials=creds)

resource = self._make_table_resource()

resource["newAlphaProperty"] = "unreleased property"

conn = client._connection = _make_connection(resource)

table = Table(self.TABLE_REF)

table._properties["newAlphaProperty"] = "unreleased property"

got = client.create_table(table)

conn.api_request.assert_called_once_with(

method="POST",

path="/%s" % path,

data={

"tableReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

"tableId": self.TABLE_ID,

},

"newAlphaProperty": "unreleased property",

"labels": {},

},

)

self.assertEqual(got._properties["newAlphaProperty"], "unreleased property")

self.assertEqual(got.table_id, self.TABLE_ID)

def test_create_table_w_encryption_configuration(self):

from google.cloud.bigquery.table import EncryptionConfiguration

from google.cloud.bigquery.table import Table

path = "projects/%s/datasets/%s/tables" % (self.PROJECT, self.DS_ID)

creds = _make_credentials()

client = self._make_one(project=self.PROJECT, credentials=creds)

resource = self._make_table_resource()

conn = client._connection = _make_connection(resource)

table = Table(self.TABLE_REF)

table.encryption_configuration = EncryptionConfiguration(

kms_key_name=self.KMS_KEY_NAME

)

got = client.create_table(table)

conn.api_request.assert_called_once_with(

method="POST",

path="/%s" % path,

data={

"tableReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

"tableId": self.TABLE_ID,

},

"labels": {},

"encryptionConfiguration": {"kmsKeyName": self.KMS_KEY_NAME},

},

)

self.assertEqual(got.table_id, self.TABLE_ID)

def test_create_table_w_day_partition_and_expire(self):

from google.cloud.bigquery.table import Table

from google.cloud.bigquery.table import TimePartitioning

path = "projects/%s/datasets/%s/tables" % (self.PROJECT, self.DS_ID)

creds = _make_credentials()

client = self._make_one(project=self.PROJECT, credentials=creds)

resource = self._make_table_resource()

conn = client._connection = _make_connection(resource)

table = Table(self.TABLE_REF)

table.time_partitioning = TimePartitioning(expiration_ms=100)

got = client.create_table(table)

conn.api_request.assert_called_once_with(

method="POST",

path="/%s" % path,

data={

"tableReference": {

"projectId": self.PROJECT,

"datasetId": self.DS_ID,

"tableId": self.TABLE_ID,

},

"timePartitioning": {"type": "DAY", "expirationMs": "100"},

"labels": {},

},

)

self.assertEqual(table.time_partitioning.type_, "DAY")

self.assertEqual(table.time_partitioning.expiration_ms, 100)

self.assertEqual(got.table_id, self.TABLE_ID)