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)