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

import datetime

import decimal

import email

import gzip

import http.client

import io

import itertools

import json

import operator

import unittest

import warnings

import mock

import packaging

import requests

import pytest

import pkg_resources

try:

import pandas

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

pandas = None

try:

import opentelemetry

from opentelemetry import trace

from opentelemetry.sdk.trace import TracerProvider

from opentelemetry.sdk.trace.export import SimpleExportSpanProcessor

from opentelemetry.sdk.trace.export.in_memory_span_exporter import (

InMemorySpanExporter,

)

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

opentelemetry = None

try:

import pyarrow

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

pyarrow = None

import google.api_core.exceptions

from google.api_core import client_info

import google.cloud._helpers

from google.cloud import bigquery_v2

from google.cloud.bigquery.dataset import DatasetReference

from google.cloud.bigquery.retry import DEFAULT_TIMEOUT

try:

from google.cloud import bigquery_storage

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

bigquery_storage = None

from test_utils.imports import maybe_fail_import

from tests.unit.helpers import make_connection

PANDAS_MINIUM_VERSION = pkg_resources.parse_version("1.0.0")

if pandas is not None:

PANDAS_INSTALLED_VERSION = pkg_resources.get_distribution("pandas").parsed_version

else:

# Set to less than MIN version.

PANDAS_INSTALLED_VERSION = pkg_resources.parse_version("0.0.0")

def _make_credentials():

import google.auth.credentials

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

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/us/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)

self.assertEqual(

client._connection.API_BASE_URL, Connection.DEFAULT_API_ENDPOINT

)

def test_ctor_w_empty_client_options(self):

from google.api_core.client_options import ClientOptions

creds = _make_credentials()

http = object()

client_options = ClientOptions()

client = self._make_one(

project=self.PROJECT,

credentials=creds,

_http=http,

client_options=client_options,

)

self.assertEqual(

client._connection.API_BASE_URL, client._connection.DEFAULT_API_ENDPOINT

)

def test_ctor_w_client_options_dict(self):

creds = _make_credentials()

http = object()

client_options = {"api_endpoint": "https://www.foo-googleapis.com"}

client = self._make_one(

project=self.PROJECT,

credentials=creds,

_http=http,

client_options=client_options,

)

self.assertEqual(

client._connection.API_BASE_URL, "https://www.foo-googleapis.com"

)

def test_ctor_w_client_options_object(self):

from google.api_core.client_options import ClientOptions

creds = _make_credentials()

http = object()

client_options = ClientOptions(api_endpoint="https://www.foo-googleapis.com")

client = self._make_one(

project=self.PROJECT,

credentials=creds,

_http=http,

client_options=client_options,

)

self.assertEqual(

client._connection.API_BASE_URL, "https://www.foo-googleapis.com"

)

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__call_api_applying_custom_retry_on_timeout(self):

from concurrent.futures import TimeoutError

from google.cloud.bigquery.retry import DEFAULT_RETRY

creds = _make_credentials()

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

api_request_patcher = mock.patch.object(

client._connection, "api_request", side_effect=[TimeoutError, "result"],

)

retry = DEFAULT_RETRY.with_deadline(1).with_predicate(

lambda exc: isinstance(exc, TimeoutError)

)

with api_request_patcher as fake_api_request:

result = client._call_api(retry, foo="bar")

self.assertEqual(result, "result")

self.assertEqual(

fake_api_request.call_args_list,

[mock.call(foo="bar"), mock.call(foo="bar")], # was retried once

)

def test__call_api_span_creator_not_called(self):

from concurrent.futures import TimeoutError

from google.cloud.bigquery.retry import DEFAULT_RETRY

creds = _make_credentials()

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

api_request_patcher = mock.patch.object(

client._connection, "api_request", side_effect=[TimeoutError, "result"],

)

retry = DEFAULT_RETRY.with_deadline(1).with_predicate(

lambda exc: isinstance(exc, TimeoutError)

)

with api_request_patcher:

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

client._call_api(retry)

final_attributes.assert_not_called()

def test__call_api_span_creator_called(self):

from concurrent.futures import TimeoutError

from google.cloud.bigquery.retry import DEFAULT_RETRY

creds = _make_credentials()

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

api_request_patcher = mock.patch.object(

client._connection, "api_request", side_effect=[TimeoutError, "result"],

)

retry = DEFAULT_RETRY.with_deadline(1).with_predicate(

lambda exc: isinstance(exc, TimeoutError)

)

with api_request_patcher:

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

client._call_api(

retry,

span_name="test_name",

span_attributes={"test_attribute": "test_attribute-value"},

)

final_attributes.assert_called_once()

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()

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

with self.assertRaises(NotFound):

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

client._get_query_results(

"nothere",

None,

project="other-project",

location=self.LOCATION,

timeout_ms=500,

timeout=420,

)

final_attributes.assert_called_once_with({"path": path}, client, None)

conn.api_request.assert_called_once_with(

method="GET",

path=path,

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

timeout=420,

)

def test__get_query_results_miss_w_short_timeout(self):

import google.cloud.bigquery.client

from google.cloud.exceptions import NotFound

creds = _make_credentials()

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

conn = client._connection = make_connection()

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

with self.assertRaises(NotFound):

client._get_query_results(

"nothere",

None,

project="other-project",

location=self.LOCATION,

timeout_ms=500,

timeout=1,

)

conn.api_request.assert_called_once_with(

method="GET",

path=path,

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

timeout=google.cloud.bigquery.client._MIN_GET_QUERY_RESULTS_TIMEOUT,

)

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},

timeout=DEFAULT_TIMEOUT,

)

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)

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

service_account_email = client.get_service_account_email(timeout=7.5)

final_attributes.assert_called_once_with({"path": path}, client, None)

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

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)

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

service_account_email = client.get_service_account_email(project=project)

final_attributes.assert_called_once_with({"path": path}, client, None)

conn.api_request.assert_called_once_with(

method="GET", path=path, timeout=DEFAULT_TIMEOUT

)

self.assertEqual(service_account_email, email)

def test_get_service_account_email_w_custom_retry(self):

from google.cloud.bigquery.retry import DEFAULT_RETRY

api_path = "/projects/{}/serviceAccount".format(self.PROJECT)

creds = _make_credentials()

http = object()

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

resource = {

"kind": "bigquery#getServiceAccountResponse",

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

}

api_request_patcher = mock.patch.object(

client._connection, "api_request", side_effect=[ValueError, resource],

)

retry = DEFAULT_RETRY.with_deadline(1).with_predicate(

lambda exc: isinstance(exc, ValueError)

)

with api_request_patcher as fake_api_request:

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

service_account_email = client.get_service_account_email(

retry=retry, timeout=7.5

)

final_attributes.assert_called_once_with({"path": api_path}, client, None)

self.assertEqual(

service_account_email, "bq-123@bigquery-encryption.iam.gserviceaccount.com"

)

self.assertEqual(

fake_api_request.call_args_list,

[

mock.call(method="GET", path=api_path, timeout=7.5),

mock.call(method="GET", path=api_path, timeout=7.5), # was retried once

],

)

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)

catch_warnings = warnings.catch_warnings(record=True)

with catch_warnings as warned:

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

matches = [

warning

for warning in warned

if warning.category in (DeprecationWarning, PendingDeprecationWarning)

and "Client.dataset" in str(warning)

and "my_project.my_dataset" in str(warning)

]

assert matches, "A Client.dataset deprecation warning was not raised."

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)

catch_warnings = warnings.catch_warnings(record=True)

with catch_warnings as warned:

dataset = client.dataset(self.DS_ID)

matches = [

warning

for warning in warned

if warning.category in (DeprecationWarning, PendingDeprecationWarning)

and "Client.dataset" in str(warning)

and "my_project.my_dataset" in str(warning)

]

assert matches, "A Client.dataset deprecation warning was not raised."

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 = DatasetReference(self.PROJECT, self.DS_ID)

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

dataset = client.get_dataset(dataset_ref, timeout=7.5)

final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None)

conn.api_request.assert_called_once_with(

method="GET", path="/%s" % path, timeout=7.5

)

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):

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

client.get_dataset(dataset_ref)

final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None)

# Zero-length errors field.

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

with self.assertRaises(ServerError):

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

client.get_dataset(dataset_ref)

final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None)

# Non-retryable reason.

client._connection = make_connection(

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

)

with self.assertRaises(ServerError):

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

client.get_dataset(dataset_ref)

final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None)

# Retryable reason, but retry is disabled.

client._connection = make_connection(

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

)

with self.assertRaises(ServerError):

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

client.get_dataset(dataset_ref, retry=None)

final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None)

# Retryable reason, default retry: success.

client._connection = make_connection(

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

)

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

dataset = client.get_dataset(

# Test with a string for dataset ID.

dataset_ref.dataset_id

)

final_attributes.assert_called_once_with({"path": "/%s" % path}, client, None)

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

@unittest.skipIf(

bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"

)

def test_ensure_bqstorage_client_creating_new_instance(self):

mock_client = mock.create_autospec(bigquery_storage.BigQueryReadClient)

mock_client_instance = object()

mock_client.return_value = mock_client_instance

creds = _make_credentials()

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

with mock.patch(

"google.cloud.bigquery_storage.BigQueryReadClient", mock_client

):

bqstorage_client = client._ensure_bqstorage_client(

client_options=mock.sentinel.client_options,

client_info=mock.sentinel.client_info,

)

self.assertIs(bqstorage_client, mock_client_instance)

mock_client.assert_called_once_with(

credentials=creds,

client_options=mock.sentinel.client_options,

client_info=mock.sentinel.client_info,

)

def test_ensure_bqstorage_client_missing_dependency(self):

creds = _make_credentials()

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

def fail_bqstorage_import(name, globals, locals, fromlist, level):

# NOTE: *very* simplified, assuming a straightforward absolute import

return "bigquery_storage" in name or (

fromlist is not None and "bigquery_storage" in fromlist

)

no_bqstorage = maybe_fail_import(predicate=fail_bqstorage_import)

with no_bqstorage, warnings.catch_warnings(record=True) as warned:

bqstorage_client = client._ensure_bqstorage_client()

self.assertIsNone(bqstorage_client)

matching_warnings = [

warning

for warning in warned

if "not installed" in str(warning)

and "google-cloud-bigquery-storage" in str(warning)

]

assert matching_warnings, "Missing dependency warning not raised."

@unittest.skipIf(

bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"

)

def test_ensure_bqstorage_client_obsolete_dependency(self):

from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError

creds = _make_credentials()

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

patcher = mock.patch(

"google.cloud.bigquery.client.BQ_STORAGE_VERSIONS.verify_version",

side_effect=LegacyBigQueryStorageError("BQ Storage too old"),

)

with patcher, warnings.catch_warnings(record=True) as warned:

bqstorage_client = client._ensure_bqstorage_client()

self.assertIsNone(bqstorage_client)

matching_warnings = [

warning for warning in warned if "BQ Storage too old" in str(warning)

]

assert matching_warnings, "Obsolete dependency warning not raised."

@unittest.skipIf(

bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"

)

def test_ensure_bqstorage_client_existing_client_check_passes(self):

creds = _make_credentials()

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

mock_storage_client = mock.sentinel.mock_storage_client

bqstorage_client = client._ensure_bqstorage_client(

bqstorage_client=mock_storage_client

)

self.assertIs(bqstorage_client, mock_storage_client)

@unittest.skipIf(

bigquery_storage is None, "Requires `google-cloud-bigquery-storage`"

)

def test_ensure_bqstorage_client_existing_client_check_fails(self):

from google.cloud.bigquery.exceptions import LegacyBigQueryStorageError

creds = _make_credentials()

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

mock_storage_client = mock.sentinel.mock_storage_client

patcher = mock.patch(

"google.cloud.bigquery.client.BQ_STORAGE_VERSIONS.verify_version",

side_effect=LegacyBigQueryStorageError("BQ Storage too old"),

)

with patcher, warnings.catch_warnings(record=True) as warned:

bqstorage_client = client._ensure_bqstorage_client(mock_storage_client)

self.assertIsNone(bqstorage_client)

matching_warnings = [

warning for warning in warned if "BQ Storage too old" in str(warning)

]

assert matching_warnings, "Obsolete dependency warning not raised."

def test_create_routine_w_minimal_resource(self):

from google.cloud.bigquery.routine import Routine

from google.cloud.bigquery.routine import RoutineReference

creds = _make_credentials()

path = "/projects/test-routine-project/datasets/test_routines/routines"

resource = {

"routineReference": {

"projectId": "test-routine-project",

"datasetId": "test_routines",

"routineId": "minimal_routine",

}

}

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

conn = client._connection = make_connection(resource)

full_routine_id = "test-routine-project.test_routines.minimal_routine"

routine = Routine(full_routine_id)

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

actual_routine = client.create_routine(routine, timeout=7.5)

final_attributes.assert_called_once_with({"path": path}, client, None)

conn.api_request.assert_called_once_with(

method="POST", path=path, data=resource, timeout=7.5,

)

self.assertEqual(

actual_routine.reference, RoutineReference.from_string(full_routine_id)

)

def test_create_routine_w_conflict(self):

from google.cloud.bigquery.routine import Routine

creds = _make_credentials()

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

conn = client._connection = make_connection(

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

)

path = "/projects/test-routine-project/datasets/test_routines/routines"

full_routine_id = "test-routine-project.test_routines.minimal_routine"

routine = Routine(full_routine_id)

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

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

client.create_routine(routine)

final_attributes.assert_called_once_with({"path": path}, client, None)

resource = {

"routineReference": {

"projectId": "test-routine-project",

"datasetId": "test_routines",

"routineId": "minimal_routine",

}

}

conn.api_request.assert_called_once_with(

method="POST", path=path, data=resource, timeout=DEFAULT_TIMEOUT,

)

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

def test_span_status_is_set(self):

from google.cloud.bigquery.routine import Routine

tracer_provider = TracerProvider()

memory_exporter = InMemorySpanExporter()

span_processor = SimpleExportSpanProcessor(memory_exporter)

tracer_provider.add_span_processor(span_processor)

trace.set_tracer_provider(tracer_provider)

creds = _make_credentials()

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

conn = client._connection = make_connection(

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

)

path = "/projects/test-routine-project/datasets/test_routines/routines"

full_routine_id = "test-routine-project.test_routines.minimal_routine"

routine = Routine(full_routine_id)

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

client.create_routine(routine)

span_list = memory_exporter.get_finished_spans()

self.assertTrue(span_list[0].status is not None)

resource = {

"routineReference": {

"projectId": "test-routine-project",

"datasetId": "test_routines",

"routineId": "minimal_routine",

}

}

conn.api_request.assert_called_once_with(

method="POST", path=path, data=resource, timeout=DEFAULT_TIMEOUT,

)

def test_create_routine_w_conflict_exists_ok(self):

from google.cloud.bigquery.routine import Routine

creds = _make_credentials()

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

resource = {

"routineReference": {

"projectId": "test-routine-project",

"datasetId": "test_routines",

"routineId": "minimal_routine",

}

}

path = "/projects/test-routine-project/datasets/test_routines/routines"

conn = client._connection = make_connection(

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

)

full_routine_id = "test-routine-project.test_routines.minimal_routine"

routine = Routine(full_routine_id)

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

actual_routine = client.create_routine(routine, exists_ok=True)

final_attributes.assert_called_with(

{"path": "%s/minimal_routine" % path}, client, None

)

self.assertEqual(actual_routine.project, "test-routine-project")

self.assertEqual(actual_routine.dataset_id, "test_routines")

self.assertEqual(actual_routine.routine_id, "minimal_routine")

conn.api_request.assert_has_calls(

[

mock.call(

method="POST", path=path, data=resource, timeout=DEFAULT_TIMEOUT,

),

mock.call(

method="GET",

path="/projects/test-routine-project/datasets/test_routines/routines/minimal_routine",

timeout=DEFAULT_TIMEOUT,

),

]

)

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()

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

got = client.create_table(table, timeout=7.5)

final_attributes.assert_called_once_with(

{"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None

)

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": {},

},

timeout=7.5,

)

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"

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

got = client.create_table(table)

final_attributes.assert_called_once_with(

{"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None

)

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": {},

},

timeout=DEFAULT_TIMEOUT,

)

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.encryption_configuration 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

)

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

got = client.create_table(table)

final_attributes.assert_called_once_with(

{"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None

)

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},

},

timeout=DEFAULT_TIMEOUT,

)

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)

with mock.patch(

"google.cloud.bigquery.opentelemetry_tracing._get_final_span_attributes"

) as final_attributes:

got = client.create_table(table)

final_attributes.assert_called_once_with(

{"path": "/%s" % path, "dataset_id": table.dataset_id}, client, None

)

conn.api_request.assert_called_once_with(

method="POST",

path="/%s" % path,

data={

"tableReference": {