chore: remove workarounds for BQ Storage issue with small result sets (#133)

* chore: remove workarounds for BQ Storage issue with small result sets

* Fix two typos in docstrings

Author: plamut

google/cloud/bigquery/dbapi/connection.py

@@ -34,16 +34,10 @@ class Connection(object):
         ):
             A client that uses the faster BigQuery Storage API to fetch rows from
             BigQuery. If not passed, it is created using the same credentials
-            as ``client``.
+            as ``client`` (provided that BigQuery Storage dependencies are installed).
 
-            When fetching query results, ``bqstorage_client`` is used first, with
-            a fallback on ``client``, if necessary.
-
-            .. note::
-                There is a known issue with the BigQuery Storage API with small
-                anonymous result sets, which results in such fallback.
-
-                https://github.com/googleapis/python-bigquery-storage/issues/2
+            If both clients are available, ``bqstorage_client`` is used for
+            fetching query results.
     """
 
     def __init__(self, client=None, bqstorage_client=None):
@@ -110,16 +104,10 @@ def connect(client=None, bqstorage_client=None):
         ):
             A client that uses the faster BigQuery Storage API to fetch rows from
             BigQuery. If not passed, it is created using the same credentials
-            as ``client``.
-
-            When fetching query results, ``bqstorage_client`` is used first, with
-            a fallback on ``client``, if necessary.
-
-            .. note::
-                There is a known issue with the BigQuery Storage API with small
-                anonymous result sets, which results in such fallback.
+            as ``client`` (provided that BigQuery Storage dependencies are installed).
 
-                https://github.com/googleapis/python-bigquery-storage/issues/2
+            If both clients are available, ``bqstorage_client`` is used for
+            fetching query results.
 
     Returns:
         google.cloud.bigquery.dbapi.Connection: A new DB-API connection to BigQuery.

google/cloud/bigquery/dbapi/cursor.py

@@ -224,26 +224,9 @@ def _try_fetch(self, size=None):
             bqstorage_client = self.connection._bqstorage_client
 
             if bqstorage_client is not None:
-                try:
-                    rows_iterable = self._bqstorage_fetch(bqstorage_client)
-                    self._query_data = _helpers.to_bq_table_rows(rows_iterable)
-                    return
-                except google.api_core.exceptions.GoogleAPICallError as exc:
-                    # NOTE: Forbidden is a subclass of GoogleAPICallError
-                    if isinstance(exc, google.api_core.exceptions.Forbidden):
-                        # Don't hide errors such as insufficient permissions to create
-                        # a read session, or the API is not enabled. Both of those are
-                        # clearly problems if the developer has explicitly asked for
-                        # BigQuery Storage API support.
-                        raise
-
-                    # There is an issue with reading from small anonymous
-                    # query results tables. If such an error occurs, we silence
-                    # it in order to try again with the tabledata.list API.
-                    _LOGGER.debug(
-                        "Error fetching data with BigQuery Storage API, "
-                        "falling back to tabledata.list API."
-                    )
+                rows_iterable = self._bqstorage_fetch(bqstorage_client)
+                self._query_data = _helpers.to_bq_table_rows(rows_iterable)
+                return
 
             rows_iter = client.list_rows(
                 self._query_job.destination,

google/cloud/bigquery/job.py

@@ -3335,9 +3335,6 @@ def to_dataframe(
                 Reading from a specific partition or snapshot is not
                 currently supported by this method.
 
-                **Caution**: There is a known issue reading small anonymous
-                query result tables with the BQ Storage API. Write your query
-                results to a destination table to work around this issue.
             dtypes (Map[str, Union[str, pandas.Series.dtype]]):
                 Optional. A dictionary of column names pandas ``dtype``s. The
                 provided ``dtype`` is used when constructing the series for

google/cloud/bigquery/table.py

@@ -1431,30 +1431,10 @@ def _to_page_iterable(
         self, bqstorage_download, tabledata_list_download, bqstorage_client=None
     ):
         if bqstorage_client is not None:
-            try:
-                # Iterate over the stream so that read errors are raised (and
-                # the method can then fallback to tabledata.list).
-                for item in bqstorage_download():
-                    yield item
-                return
-            except google.api_core.exceptions.Forbidden:
-                # Don't hide errors such as insufficient permissions to create
-                # a read session, or the API is not enabled. Both of those are
-                # clearly problems if the developer has explicitly asked for
-                # BigQuery Storage API support.
-                raise
-            except google.api_core.exceptions.GoogleAPICallError:
-                # There is a known issue with reading from small anonymous
-                # query results tables, so some errors are expected. Rather
-                # than throw those errors, try reading the DataFrame again, but
-                # with the tabledata.list API.
-                pass
-
-        _LOGGER.debug(
-            "Started reading table '{}.{}.{}' with tabledata.list.".format(
-                self._table.project, self._table.dataset_id, self._table.table_id
-            )
-        )
+            for item in bqstorage_download():
+                yield item
+            return
+
         for item in tabledata_list_download():
             yield item
 
@@ -1599,14 +1579,10 @@ def to_dataframe_iterable(self, bqstorage_client=None, dtypes=None):
                 This method requires the ``pyarrow`` and
                 ``google-cloud-bigquery-storage`` libraries.
 
-                This method only  exposes a subset of the capabilities of the
-                BigQuery Storage API.  For full access to all features
+                This method only exposes a subset of the capabilities of the
+                BigQuery Storage API. For full access to all features
                 (projections, filters, snapshots) use the Storage API directly.
 
-                **Caution**: There is a known issue reading small anonymous
-                query result tables with the BQ Storage API. When a problem
-                is encountered reading a table, the tabledata.list method
-                from the BigQuery API is used, instead.
             dtypes (Map[str, Union[str, pandas.Series.dtype]]):
                 Optional. A dictionary of column names pandas ``dtype``s. The
                 provided ``dtype`` is used when constructing the series for
@@ -1668,14 +1644,10 @@ def to_dataframe(
                 This method requires the ``pyarrow`` and
                 ``google-cloud-bigquery-storage`` libraries.
 
-                This method only  exposes a subset of the capabilities of the
-                BigQuery Storage API.  For full access to all features
+                This method only exposes a subset of the capabilities of the
+                BigQuery Storage API. For full access to all features
                 (projections, filters, snapshots) use the Storage API directly.
 
-                **Caution**: There is a known issue reading small anonymous
-                query result tables with the BQ Storage API. When a problem
-                is encountered reading a table, the tabledata.list method
-                from the BigQuery API is used, instead.
             dtypes (Map[str, Union[str, pandas.Series.dtype]]):
                 Optional. A dictionary of column names pandas ``dtype``s. The
                 provided ``dtype`` is used when constructing the series for

tests/system.py

@@ -1690,50 +1690,6 @@ def test_dbapi_fetchall(self):
             row_tuples = [r.values() for r in rows]
             self.assertEqual(row_tuples, [(1, 2), (3, 4), (5, 6)])
 
-    @unittest.skipIf(
-        bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`"
-    )
-    def test_dbapi_fetch_w_bqstorage_client_small_result_set(self):
-        bqstorage_client = bigquery_storage_v1.BigQueryReadClient(
-            credentials=Config.CLIENT._credentials
-        )
-        cursor = dbapi.connect(Config.CLIENT, bqstorage_client).cursor()
-
-        # Reading small result sets causes an issue with BQ storage client,
-        # and the DB API should transparently fall back to the default client.
-        cursor.execute(
-            """
-            SELECT id, `by`, time_ts
-            FROM `bigquery-public-data.hacker_news.comments`
-            ORDER BY `id` ASC
-            LIMIT 10
-        """
-        )
-
-        result_rows = [cursor.fetchone(), cursor.fetchone(), cursor.fetchone()]
-
-        field_name = operator.itemgetter(0)
-        fetched_data = [sorted(row.items(), key=field_name) for row in result_rows]
-
-        expected_data = [
-            [
-                ("by", "sama"),
-                ("id", 15),
-                ("time_ts", datetime.datetime(2006, 10, 9, 19, 51, 1, tzinfo=UTC)),
-            ],
-            [
-                ("by", "pg"),
-                ("id", 17),
-                ("time_ts", datetime.datetime(2006, 10, 9, 19, 52, 45, tzinfo=UTC)),
-            ],
-            [
-                ("by", "pg"),
-                ("id", 22),
-                ("time_ts", datetime.datetime(2006, 10, 10, 2, 18, 22, tzinfo=UTC)),
-            ],
-        ]
-        self.assertEqual(fetched_data, expected_data)
-
     @unittest.skipIf(
         bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`"
     )
@@ -1744,10 +1700,6 @@ def test_dbapi_fetch_w_bqstorage_client_large_result_set(self):
         )
         cursor = dbapi.connect(Config.CLIENT, bqstorage_client).cursor()
 
-        # Pick a large enough LIMIT value to assure that the fallback to the
-        # default client is not needed due to the result set being too small
-        # (a known issue that causes problems when reading such result sets with
-        # BQ storage client).
         cursor.execute(
             """
             SELECT id, `by`, time_ts
@@ -1794,10 +1746,6 @@ def test_dbapi_fetch_w_bqstorage_client_v1beta1_large_result_set(self):
         )
         cursor = dbapi.connect(Config.CLIENT, bqstorage_client).cursor()
 
-        # Pick a large enouhg LIMIT value to assure that the fallback to the
-        # default client is not needed due to the result set being too small
-        # (a known issue that causes problems when reding such result sets with
-        # BQ storage client).
         cursor.execute(
             """
             SELECT id, `by`, time_ts
@@ -1845,10 +1793,6 @@ def test_dbapi_connection_does_not_leak_sockets(self):
         connection = dbapi.connect()
         cursor = connection.cursor()
 
-        # Pick a large enough LIMIT value to assure that the fallback to the
-        # default client is not needed due to the result set being too small
-        # (a known issue that causes problems when reding such result sets with
-        # BQ storage client).
         cursor.execute(
             """
             SELECT id, `by`, time_ts
@@ -2272,9 +2216,6 @@ def test_query_results_to_dataframe(self):
         bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`"
     )
     def test_query_results_to_dataframe_w_bqstorage(self):
-        dest_dataset = self.temp_dataset(_make_dataset_id("bqstorage_to_dataframe_"))
-        dest_ref = dest_dataset.table("query_results")
-
         query = """
             SELECT id, author, time_ts, dead
             FROM `bigquery-public-data.hacker_news.comments`
@@ -2285,50 +2226,29 @@ def test_query_results_to_dataframe_w_bqstorage(self):
             credentials=Config.CLIENT._credentials
         )
 
-        job_configs = (
-            # There is a known issue reading small anonymous query result
-            # tables with the BQ Storage API. Writing to a destination
-            # table works around this issue.
-            bigquery.QueryJobConfig(
-                destination=dest_ref, write_disposition="WRITE_TRUNCATE"
-            ),
-            # Check that the client is able to work around the issue with
-            # reading small anonymous query result tables by falling back to
-            # the tabledata.list API.
-            None,
-        )
-
-        for job_config in job_configs:
-            df = (
-                Config.CLIENT.query(query, job_config=job_config)
-                .result()
-                .to_dataframe(bqstorage_client)
-            )
+        df = Config.CLIENT.query(query).result().to_dataframe(bqstorage_client)
 
-            self.assertIsInstance(df, pandas.DataFrame)
-            self.assertEqual(len(df), 10)  # verify the number of rows
-            column_names = ["id", "author", "time_ts", "dead"]
-            self.assertEqual(list(df), column_names)
-            exp_datatypes = {
-                "id": int,
-                "author": six.text_type,
-                "time_ts": pandas.Timestamp,
-                "dead": bool,
-            }
-            for index, row in df.iterrows():
-                for col in column_names:
-                    # all the schema fields are nullable, so None is acceptable
-                    if not row[col] is None:
-                        self.assertIsInstance(row[col], exp_datatypes[col])
+        self.assertIsInstance(df, pandas.DataFrame)
+        self.assertEqual(len(df), 10)  # verify the number of rows
+        column_names = ["id", "author", "time_ts", "dead"]
+        self.assertEqual(list(df), column_names)
+        exp_datatypes = {
+            "id": int,
+            "author": six.text_type,
+            "time_ts": pandas.Timestamp,
+            "dead": bool,
+        }
+        for index, row in df.iterrows():
+            for col in column_names:
+                # all the schema fields are nullable, so None is acceptable
+                if not row[col] is None:
+                    self.assertIsInstance(row[col], exp_datatypes[col])
 
     @unittest.skipIf(pandas is None, "Requires `pandas`")
     @unittest.skipIf(
         bigquery_storage_v1beta1 is None, "Requires `google-cloud-bigquery-storage`"
     )
     def test_query_results_to_dataframe_w_bqstorage_v1beta1(self):
-        dest_dataset = self.temp_dataset(_make_dataset_id("bqstorage_to_dataframe_"))
-        dest_ref = dest_dataset.table("query_results")
-
         query = """
             SELECT id, author, time_ts, dead
             FROM `bigquery-public-data.hacker_news.comments`
@@ -2339,41 +2259,23 @@ def test_query_results_to_dataframe_w_bqstorage_v1beta1(self):
             credentials=Config.CLIENT._credentials
         )
 
-        job_configs = (
-            # There is a known issue reading small anonymous query result
-            # tables with the BQ Storage API. Writing to a destination
-            # table works around this issue.
-            bigquery.QueryJobConfig(
-                destination=dest_ref, write_disposition="WRITE_TRUNCATE"
-            ),
-            # Check that the client is able to work around the issue with
-            # reading small anonymous query result tables by falling back to
-            # the tabledata.list API.
-            None,
-        )
-
-        for job_config in job_configs:
-            df = (
-                Config.CLIENT.query(query, job_config=job_config)
-                .result()
-                .to_dataframe(bqstorage_client)
-            )
+        df = Config.CLIENT.query(query).result().to_dataframe(bqstorage_client)
 
-            self.assertIsInstance(df, pandas.DataFrame)
-            self.assertEqual(len(df), 10)  # verify the number of rows
-            column_names = ["id", "author", "time_ts", "dead"]
-            self.assertEqual(list(df), column_names)
-            exp_datatypes = {
-                "id": int,
-                "author": six.text_type,
-                "time_ts": pandas.Timestamp,
-                "dead": bool,
-            }
-            for index, row in df.iterrows():
-                for col in column_names:
-                    # all the schema fields are nullable, so None is acceptable
-                    if not row[col] is None:
-                        self.assertIsInstance(row[col], exp_datatypes[col])
+        self.assertIsInstance(df, pandas.DataFrame)
+        self.assertEqual(len(df), 10)  # verify the number of rows
+        column_names = ["id", "author", "time_ts", "dead"]
+        self.assertEqual(list(df), column_names)
+        exp_datatypes = {
+            "id": int,
+            "author": six.text_type,
+            "time_ts": pandas.Timestamp,
+            "dead": bool,
+        }
+        for index, row in df.iterrows():
+            for col in column_names:
+                # all the schema fields are nullable, so None is acceptable
+                if not row[col] is None:
+                    self.assertIsInstance(row[col], exp_datatypes[col])
 
     @unittest.skipIf(pandas is None, "Requires `pandas`")
     def test_insert_rows_from_dataframe(self):