def test_fetchall_w_bqstorage_client_v1beta1_fetch_success(self):
from google.cloud.bigquery import dbapi
from google.cloud.bigquery import table
# use unordered data to also test any non-determenistic key order in dicts
row_data = [
table.Row([1.4, 1.1, 1.3, 1.2], {"bar": 3, "baz": 2, "foo": 1, "quux": 0}),
table.Row([2.4, 2.1, 2.3, 2.2], {"bar": 3, "baz": 2, "foo": 1, "quux": 0}),
]
bqstorage_streamed_rows = [
{
"bar": _to_pyarrow(1.2),
"foo": _to_pyarrow(1.1),
"quux": _to_pyarrow(1.4),
"baz": _to_pyarrow(1.3),
},
{
"bar": _to_pyarrow(2.2),
"foo": _to_pyarrow(2.1),
"quux": _to_pyarrow(2.4),
"baz": _to_pyarrow(2.3),
},
]
mock_client = self._mock_client(rows=row_data)
mock_bqstorage_client = self._mock_bqstorage_client(
stream_count=1, rows=bqstorage_streamed_rows, v1beta1=True
)
connection = dbapi.connect(
client=mock_client, bqstorage_client=mock_bqstorage_client,
)
cursor = connection.cursor()
cursor.execute("SELECT foo, bar FROM some_table")
with warnings.catch_warnings(record=True) as warned:
rows = cursor.fetchall()
# a deprecation warning should have been emitted
expected_warnings = [
warning
for warning in warned
if issubclass(warning.category, DeprecationWarning)
and "v1beta1" in str(warning)
]
self.assertEqual(len(expected_warnings), 1, "Deprecation warning not raised.")
# the default client was not used
mock_client.list_rows.assert_not_called()
# check the data returned
field_value = op.itemgetter(1)
sorted_row_data = [sorted(row.items(), key=field_value) for row in rows]
expected_row_data = [
[("foo", 1.1), ("bar", 1.2), ("baz", 1.3), ("quux", 1.4)],
[("foo", 2.1), ("bar", 2.2), ("baz", 2.3), ("quux", 2.4)],
]
self.assertEqual(sorted_row_data, expected_row_data)
def test_fetchall_w_bqstorage_client_fetch_success(self):
from google.cloud.bigquery import dbapi
from google.cloud.bigquery import table
# use unordered data to also test any non-determenistic key order in dicts
row_data = [
table.Row([1.4, 1.1, 1.3, 1.2], {"bar": 3, "baz": 2, "foo": 1, "quux": 0}),
table.Row([2.4, 2.1, 2.3, 2.2], {"bar": 3, "baz": 2, "foo": 1, "quux": 0}),
]
bqstorage_streamed_rows = [
{
"bar": _to_pyarrow(1.2),
"foo": _to_pyarrow(1.1),
"quux": _to_pyarrow(1.4),
"baz": _to_pyarrow(1.3),
},
{
"bar": _to_pyarrow(2.2),
"foo": _to_pyarrow(2.1),
"quux": _to_pyarrow(2.4),
"baz": _to_pyarrow(2.3),
},
]
mock_client = self._mock_client(rows=row_data)
mock_bqstorage_client = self._mock_bqstorage_client(
stream_count=1, rows=bqstorage_streamed_rows,
)
mock_client._ensure_bqstorage_client.return_value = mock_bqstorage_client
connection = dbapi.connect(
client=mock_client, bqstorage_client=mock_bqstorage_client,
)
cursor = connection.cursor()
cursor.execute("SELECT foo, bar FROM some_table")
rows = cursor.fetchall()
# the default client was not used
mock_client.list_rows.assert_not_called()
# check the data returned
field_value = op.itemgetter(1)
sorted_row_data = [sorted(row.items(), key=field_value) for row in rows]
expected_row_data = [
[("foo", 1.1), ("bar", 1.2), ("baz", 1.3), ("quux", 1.4)],
[("foo", 2.1), ("bar", 2.2), ("baz", 2.3), ("quux", 2.4)],
]
self.assertEqual(sorted_row_data, expected_row_data)
def test__run_query():
magics.context._credentials = None
job_id = "job_1234"
sql = "SELECT 17"
responses = [
futures.TimeoutError,
futures.TimeoutError,
[table.Row((17, ), {"num": 0})],
]
client_patch = mock.patch("google.cloud.bigquery.magics.bigquery.Client",
autospec=True)
with client_patch as client_mock, io.capture_output() as captured:
client_mock().query(sql).result.side_effect = responses
client_mock().query(sql).job_id = job_id
query_job = magics._run_query(client_mock(), sql)
lines = re.split("\n|\r", captured.stdout)
# Removes blanks & terminal code (result of display clearing)
updates = list(filter(lambda x: bool(x) and x != "\x1b[2K", lines))
assert query_job.job_id == job_id
expected_first_line = "Executing query with job ID: {}".format(job_id)
assert updates[0] == expected_first_line
execution_updates = updates[1:-1]
assert len(execution_updates) == 3 # one update per API response
assert all(
re.match("Query executing: .*s", line) for line in execution_updates)
assert re.match("Query complete after .*s", updates[-1])
def test_fetchall_w_bqstorage_client_fetch_error_no_fallback(self):
from google.cloud.bigquery import dbapi
from google.cloud.bigquery import table
row_data = [table.Row([1.1, 1.2], {"foo": 0, "bar": 1})]
def fake_ensure_bqstorage_client(bqstorage_client=None, **kwargs):
return bqstorage_client
mock_client = self._mock_client(rows=row_data)
mock_client._ensure_bqstorage_client.side_effect = fake_ensure_bqstorage_client
mock_bqstorage_client = self._mock_bqstorage_client(
stream_count=1, rows=row_data,
)
no_access_error = exceptions.Forbidden("invalid credentials")
mock_bqstorage_client.create_read_session.side_effect = no_access_error
connection = dbapi.connect(
client=mock_client, bqstorage_client=mock_bqstorage_client,
)
cursor = connection.cursor()
cursor.execute("SELECT foo, bar FROM some_table")
with self.assertRaisesRegex(exceptions.Forbidden, "invalid credentials"):
cursor.fetchall()
# the default client was not used
mock_client.list_rows.assert_not_called()
def test_fetchall_w_bqstorage_client_no_arrow_compression(self):
from google.cloud.bigquery import dbapi
from google.cloud.bigquery import table
# Use unordered data to also test any non-determenistic key order in dicts.
row_data = [table.Row([1.2, 1.1], {"bar": 1, "foo": 0})]
bqstorage_streamed_rows = [{
"bar": _to_pyarrow(1.2),
"foo": _to_pyarrow(1.1)
}]
mock_client = self._mock_client(rows=row_data)
mock_bqstorage_client = self._mock_bqstorage_client(
stream_count=1,
rows=bqstorage_streamed_rows,
)
connection = dbapi.connect(
client=mock_client,
bqstorage_client=mock_bqstorage_client,
)
cursor = connection.cursor()
cursor.execute("SELECT foo, bar FROM some_table")
with mock.patch(
"google.cloud.bigquery.dbapi.cursor._ARROW_COMPRESSION_SUPPORT",
new=False):
rows = cursor.fetchall()
mock_client.list_rows.assert_not_called(
) # The default client was not used.
# Check the BQ Storage session config.
expected_session = bigquery_storage.ReadSession(
table="projects/P/datasets/DS/tables/T",
data_format=bigquery_storage.DataFormat.ARROW,
)
mock_bqstorage_client.create_read_session.assert_called_once_with(
parent="projects/P",
read_session=expected_session,
max_stream_count=1)
# Check the data returned.
field_value = op.itemgetter(1)
sorted_row_data = [
sorted(row.items(), key=field_value) for row in rows
]
expected_row_data = [[("foo", 1.1), ("bar", 1.2)]]
self.assertEqual(sorted_row_data, expected_row_data)
def to_table_row(row):
# NOTE: We fetch ARROW values, thus we need to convert them to Python
# objects with as_py().
values = tuple(value.as_py() for value in row.values())
keys_to_index = {key: i for i, key in enumerate(row.keys())}
return table.Row(values, keys_to_index)
def test_fetchall_w_bqstorage_client_fetch_error_fallback_on_client(self):
from google.cloud.bigquery import dbapi
from google.cloud.bigquery import table
# use unordered data to also test any non-determenistic key order in dicts
row_data = [
table.Row([1.4, 1.1, 1.3, 1.2], {
"bar": 3,
"baz": 2,
"foo": 1,
"quux": 0
}),
table.Row([2.4, 2.1, 2.3, 2.2], {
"bar": 3,
"baz": 2,
"foo": 1,
"quux": 0
}),
]
bqstorage_streamed_rows = [
{
"bar": 1.2,
"foo": 1.1,
"quux": 1.4,
"baz": 1.3
},
{
"bar": 2.2,
"foo": 2.1,
"quux": 2.4,
"baz": 2.3
},
]
mock_client = self._mock_client(rows=row_data)
mock_bqstorage_client = self._mock_bqstorage_client(
stream_count=1,
rows=bqstorage_streamed_rows,
)
request_error = exceptions.BadRequest("BQ storage what??")
mock_bqstorage_client.create_read_session.side_effect = request_error
connection = dbapi.connect(
client=mock_client,
bqstorage_client=mock_bqstorage_client,
)
cursor = connection.cursor()
cursor.execute("SELECT foo, bar FROM some_table")
logger_patcher = mock.patch(
"google.cloud.bigquery.dbapi.cursor._LOGGER")
with logger_patcher as mock_logger:
rows = cursor.fetchall()
# both client were used
mock_bqstorage_client.create_read_session.assert_called()
mock_client.list_rows.assert_called()
# fallback to default API should have been logged
relevant_calls = [
call for call in mock_logger.debug.call_args_list
if call.args and "tabledata.list API" in call.args[0]
]
self.assertTrue(relevant_calls)
# check the data returned
field_value = op.itemgetter(1)
sorted_row_data = [
sorted(row.items(), key=field_value) for row in rows
]
expected_row_data = [
[("foo", 1.1), ("bar", 1.2), ("baz", 1.3), ("quux", 1.4)],
[("foo", 2.1), ("bar", 2.2), ("baz", 2.3), ("quux", 2.4)],
]
self.assertEqual(sorted_row_data, expected_row_data)
