From 7afa3d70f8564dcdacda2b9acbbd7207b50b186e Mon Sep 17 00:00:00 2001
From: Peter Lamut <plamut@users.noreply.github.com>
Date: Thu, 23 Jul 2020 17:23:45 +0200
Subject: [PATCH] fix: RowIterator.to_arrow() error when BQ Storage client
 cannot be created (#181)

* fix: to_arrow() when can't create BQ Storage client

* Clarify using BQ Storage client by default
---
 CHANGELOG.md                   |  2 +-
 google/cloud/bigquery/table.py |  2 +-
 tests/unit/test_table.py       | 27 +++++++++++++++++++++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6da5b01a0..b1adaa8cf 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,7 +9,7 @@
 
 ### Features
 
-* use BigQuery Storage client by default ([#55](https://www.github.com/googleapis/python-bigquery/issues/55)) ([e75ff82](https://www.github.com/googleapis/python-bigquery/commit/e75ff8297c65981545b097f75a17cf9e78ac6772)), closes [#91](https://www.github.com/googleapis/python-bigquery/issues/91)
+* use BigQuery Storage client by default (if dependencies available) ([#55](https://www.github.com/googleapis/python-bigquery/issues/55)) ([e75ff82](https://www.github.com/googleapis/python-bigquery/commit/e75ff8297c65981545b097f75a17cf9e78ac6772)), closes [#91](https://www.github.com/googleapis/python-bigquery/issues/91)
 * **bigquery:** add __eq__ method for class PartitionRange and RangePartitioning ([#162](https://www.github.com/googleapis/python-bigquery/issues/162)) ([0d2a88d](https://www.github.com/googleapis/python-bigquery/commit/0d2a88d8072154cfc9152afd6d26a60ddcdfbc73))
 * **bigquery:** expose date_as_object parameter to users ([#150](https://www.github.com/googleapis/python-bigquery/issues/150)) ([a2d5ce9](https://www.github.com/googleapis/python-bigquery/commit/a2d5ce9e97992318d7dc85c51c053cab74e25a11))
 * **bigquery:** expose date_as_object parameter to users ([#150](https://www.github.com/googleapis/python-bigquery/issues/150)) ([cbd831e](https://www.github.com/googleapis/python-bigquery/commit/cbd831e08024a67148723afd49e1db085e0a862c))
diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
index f1575ffb2..10b4198d3 100644
--- a/google/cloud/bigquery/table.py
+++ b/google/cloud/bigquery/table.py
@@ -1534,8 +1534,8 @@ def to_arrow(
 
         owns_bqstorage_client = False
         if not bqstorage_client and create_bqstorage_client:
-            owns_bqstorage_client = True
             bqstorage_client = self.client._create_bqstorage_client()
+            owns_bqstorage_client = bqstorage_client is not None
 
         try:
             progress_bar = self._get_progress_bar(progress_bar_type)
diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py
index 3aabebb77..2c9d0f64e 100644
--- a/tests/unit/test_table.py
+++ b/tests/unit/test_table.py
@@ -1990,6 +1990,33 @@ def test_to_arrow_w_bqstorage_creates_client(self):
         mock_client._create_bqstorage_client.assert_called_once()
         bqstorage_client.transport.channel.close.assert_called_once()
 
+    @unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
+    def test_to_arrow_create_bqstorage_client_wo_bqstorage(self):
+        from google.cloud.bigquery.schema import SchemaField
+
+        schema = [
+            SchemaField("name", "STRING", mode="REQUIRED"),
+            SchemaField("age", "INTEGER", mode="REQUIRED"),
+        ]
+        rows = [
+            {"f": [{"v": "Alice"}, {"v": "98"}]},
+            {"f": [{"v": "Bob"}, {"v": "99"}]},
+        ]
+        path = "/foo"
+        api_request = mock.Mock(return_value={"rows": rows})
+
+        mock_client = _mock_client()
+        mock_client._create_bqstorage_client.return_value = None
+        row_iterator = self._make_one(mock_client, api_request, path, schema)
+
+        tbl = row_iterator.to_arrow(create_bqstorage_client=True)
+
+        # The client attempted to create a BQ Storage client, and even though
+        # that was not possible, results were still returned without errors.
+        mock_client._create_bqstorage_client.assert_called_once()
+        self.assertIsInstance(tbl, pyarrow.Table)
+        self.assertEqual(tbl.num_rows, 2)
+
     @unittest.skipIf(pyarrow is None, "Requires `pyarrow`")
     @unittest.skipIf(
         bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`"