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 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_non_empty_iterable(self):
rows_iterable = [
dict(
one=_to_pyarrow(1.1),
four=_to_pyarrow(1.4),
two=_to_pyarrow(1.2),
three=_to_pyarrow(1.3),
),
dict(
one=_to_pyarrow(2.1),
four=_to_pyarrow(2.4),
two=_to_pyarrow(2.2),
three=_to_pyarrow(2.3),
),
]
result = _helpers.to_bq_table_rows(rows_iterable)
rows = list(result)
self.assertEqual(len(rows), 2)
row_1, row_2 = rows
self.assertIsInstance(row_1, table.Row)
self.assertIsInstance(row_2, table.Row)
field_value = op.itemgetter(1)
items = sorted(row_1.items(), key=field_value)
expected_items = [("one", 1.1), ("two", 1.2), ("three", 1.3),
("four", 1.4)]
self.assertEqual(items, expected_items)
items = sorted(row_2.items(), key=field_value)
expected_items = [("one", 2.1), ("two", 2.2), ("three", 2.3),
("four", 2.4)]
self.assertEqual(items, expected_items)
