feat: add prefer_bqstorage_client option for Connection

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 @@ def __init__(self, client=None, bqstorage_client=None):
 
  # 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 @@ def cursor(self):
  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)

tests/unit/test_dbapi_connection.py

@@ -122,6 +122,26 @@ def test_connect_w_both_clients(self):
  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