Repository: googleapis/python-bigquery

Issue:
Allow enabling/disabling the storage API on a per-connection basis
In #16, support was added for the BigQuery storage API in the form of a `bqstorage_client` which is passed to or created by [`connect()`](https://github.com/googleapis/python-bigquery/blob/b739596f37b8c00b375cc811c316b618097d761a/google/cloud/bigquery/dbapi/connection.py#L98). Yet there is no way, other than installing or uninstalling the storage API itself, to toggle whether the storage API is enabled/disabled.

Since use of the storage API does affect [pricing](https://cloud.google.com/bigquery/pricing), it would be preferable to enable or disable on a per-connection basis, perhaps as a `use_storage_api: bool` parameter on `connect()`. This way, callers could choose at will whether to use the storage API without having to change the Python environment the code is running in.

As a hack to workaround this issue, I attempted calling `connect` passing a normal `Client` object rather than a `BigQueryReadClient` in the `bqstorage_client` parameter, but this does not create a functioning usable connection.

Discussion:
Linchin: > QUESTION: Would we consider this a breaking change since we are moving to having a default setting of `True` instead of `False` which is what I feel like the current setting is.
> 
> I am approving this, pending an answer to the above question. Would be happy to get in a chat OR a meeting to discuss, as need. be.

Thanks for the question! The default behavior in the past is `True`. As in https://github.com/googleapis/python-bigquery/issues/1920, the customer couldn't disable the storage client unless they uninstall the package altogether.

Patch:
diff --git a/google/cloud/bigquery/dbapi/connection.py b/google/cloud/bigquery/dbapi/connection.py
index 66dee7df..a1a69b8f 100644
--- a/google/cloud/bigquery/dbapi/connection.py
+++ b/google/cloud/bigquery/dbapi/connection.py
@@ -35,12 +35,18 @@ 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`` (provided that BigQuery Storage dependencies are installed).
-
-            If both clients are available, ``bqstorage_client`` is used for
-            fetching query results.
+        prefer_bqstorage_client (Optional[bool]):
+            Prefer the BigQuery Storage client over the REST client. If Storage
+            client isn't available, fall back to the REST client. Defaults to
+            ``True``.
     """
 
-    def __init__(self, client=None, bqstorage_client=None):
+    def __init__(
+        self,
+        client=None,
+        bqstorage_client=None,
+        prefer_bqstorage_client=True,
+    ):
         if client is None:
             client = bigquery.Client()
             self._owns_client = True
@@ -49,7 +55,10 @@ class Connection(object):
 
         # A warning is already raised by the BQ Storage client factory factory if
         # instantiation fails, or if the given BQ Storage client instance is outdated.
-        if bqstorage_client is None:
+        if not prefer_bqstorage_client:
+            bqstorage_client = None
+            self._owns_bqstorage_client = False
+        elif bqstorage_client is None:
             bqstorage_client = client._ensure_bqstorage_client()
             self._owns_bqstorage_client = bqstorage_client is not None
         else:
@@ -95,7 +104,7 @@ class Connection(object):
         return new_cursor
 
 
-def connect(client=None, bqstorage_client=None):
+def connect(client=None, bqstorage_client=None, prefer_bqstorage_client=True):
     """Construct a DB-API connection to Google BigQuery.
 
     Args:
@@ -108,11 +117,12 @@ 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`` (provided that BigQuery Storage dependencies are installed).
-
-            If both clients are available, ``bqstorage_client`` is used for
-            fetching query results.
+        prefer_bqstorage_client (Optional[bool]):
+            Prefer the BigQuery Storage client over the REST client. If Storage
+            client isn't available, fall back to the REST client. Defaults to
+            ``True``.
 
     Returns:
         google.cloud.bigquery.dbapi.Connection: A new DB-API connection to BigQuery.
     """
-    return Connection(client, bqstorage_client)
+    return Connection(client, bqstorage_client, prefer_bqstorage_client)


Test patch:
diff --git a/tests/unit/test_dbapi_connection.py b/tests/unit/test_dbapi_connection.py
index 4071e57e..f5c77c44 100644
--- a/tests/unit/test_dbapi_connection.py
+++ b/tests/unit/test_dbapi_connection.py
@@ -122,6 +122,26 @@ class TestConnection(unittest.TestCase):
         self.assertIs(connection._client, mock_client)
         self.assertIs(connection._bqstorage_client, mock_bqstorage_client)
 
+    def test_connect_prefer_bqstorage_client_false(self):
+        pytest.importorskip("google.cloud.bigquery_storage")
+        from google.cloud.bigquery.dbapi import connect
+        from google.cloud.bigquery.dbapi import Connection
+
+        mock_client = self._mock_client()
+        mock_bqstorage_client = self._mock_bqstorage_client()
+        mock_client._ensure_bqstorage_client.return_value = mock_bqstorage_client
+
+        connection = connect(
+            client=mock_client,
+            bqstorage_client=mock_bqstorage_client,
+            prefer_bqstorage_client=False,
+        )
+
+        mock_client._ensure_bqstorage_client.assert_not_called()
+        self.assertIsInstance(connection, Connection)
+        self.assertIs(connection._client, mock_client)
+        self.assertIs(connection._bqstorage_client, None)
+
     def test_raises_error_if_closed(self):
         from google.cloud.bigquery.dbapi.exceptions import ProgrammingError