Skip to content

Commit

Permalink
Return no rows on dry run instead of processed bytes count
Browse files Browse the repository at this point in the history
  • Loading branch information
plamut committed Jun 12, 2020
1 parent a3b5d2e commit b8f393c
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 43 deletions.
25 changes: 6 additions & 19 deletions google/cloud/bigquery/dbapi/cursor.py
Expand Up @@ -28,8 +28,6 @@
import six

from google.cloud.bigquery import job
from google.cloud.bigquery import schema
from google.cloud.bigquery import table
from google.cloud.bigquery.dbapi import _helpers
from google.cloud.bigquery.dbapi import exceptions
import google.cloud.exceptions
Expand Down Expand Up @@ -188,11 +186,8 @@ def execute(self, operation, parameters=None, job_id=None, job_config=None):
)

if self._query_job.dry_run:
schema_field = schema.SchemaField(
name="estimated_bytes", field_type="INTEGER", mode="REQUIRED",
)
self._set_description(schema=[schema_field])
self.rowcount = 1
self._set_description(schema=None)
self.rowcount = 0
return

# Wait for the query to finish.
Expand Down Expand Up @@ -228,9 +223,7 @@ def _try_fetch(self, size=None):
)

if self._query_job.dry_run:
estimated_bytes = self._query_job.total_bytes_processed
row = table.Row((estimated_bytes,), {"estimated_bytes": 0})
self._query_data = iter([row])
self._query_data = iter([])
return

is_dml = (
Expand Down Expand Up @@ -347,9 +340,7 @@ def fetchone(self):
"""Fetch a single row from the results of the last ``execute*()`` call.
.. note::
If a dry run query was executed, a row with a single value is
returned representing the estimated number of bytes that would be
processed by the query.
If a dry run query was executed, no rows are returned.
Returns:
Tuple:
Expand All @@ -369,9 +360,7 @@ def fetchmany(self, size=None):
"""Fetch multiple results from the last ``execute*()`` call.
.. note::
If a dry run query was executed, a row with a single value is
returned representing the estimated number of bytes that would be
processed by the query.
If a dry run query was executed, no rows are returned.
.. note::
The size parameter is not used for the request/response size.
Expand Down Expand Up @@ -410,9 +399,7 @@ def fetchall(self):
"""Fetch all remaining results from the last ``execute*()`` call.
.. note::
If a dry run query was executed, a row with a single value is
returned representing the estimated number of bytes that would be
processed by the query.
If a dry run query was executed, no rows are returned.
Returns:
List[Tuple]: A list of all the rows in the results.
Expand Down
6 changes: 2 additions & 4 deletions tests/system.py
Expand Up @@ -1844,13 +1844,11 @@ def test_dbapi_dry_run_query(self):
"""

Config.CURSOR.execute(query, job_config=QueryJobConfig(dry_run=True))
self.assertEqual(Config.CURSOR.rowcount, 1, "expected a single row")
self.assertEqual(Config.CURSOR.rowcount, 0, "expected no rows")

rows = Config.CURSOR.fetchall()

row_tuples = [r.values() for r in rows]
expected = [(3473,)]
self.assertEqual(row_tuples, expected)
self.assertEqual(list(rows), [])

@unittest.skipIf(
bigquery_storage_v1 is None, "Requires `google-cloud-bigquery-storage`"
Expand Down
22 changes: 2 additions & 20 deletions tests/unit/test_dbapi_cursor.py
Expand Up @@ -662,27 +662,9 @@ def test_execute_w_query_dry_run(self):
job_config=QueryJobConfig(dry_run=True),
)

expected_description = (
dbapi.cursor.Column(
name="estimated_bytes",
type_code="INTEGER",
display_size=None,
internal_size=None,
precision=None,
scale=None,
null_ok=False,
),
)
self.assertEqual(cursor.description, expected_description)
self.assertEqual(cursor.rowcount, 1)

self.assertEqual(cursor.rowcount, 0)
rows = cursor.fetchall()

# We expect a single row with one column - the estimated numbe of bytes
# that will be processed by the query.
self.assertEqual(len(rows), 1)
self.assertEqual(len(rows[0]), 1)
self.assertEqual(rows[0][0], 12345)
self.assertEqual(list(rows), [])

def test_execute_raises_if_result_raises(self):
import google.cloud.exceptions
Expand Down

0 comments on commit b8f393c

Please sign in to comment.