Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion pandas_gbq/gbq.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,7 @@ def read_gbq(
*,
col_order=None,
bigquery_client=None,
dry_run: bool = False,
):
r"""Read data from Google BigQuery to a pandas DataFrame.

Expand Down Expand Up @@ -269,7 +270,8 @@ def read_gbq(
bigquery_client : google.cloud.bigquery.Client, optional
A Google Cloud BigQuery Python Client instance. If provided, it will be used for reading
data, while the project and credentials parameters will be ignored.

dry_run : bool, default False
If True, run a dry run query.
Returns
-------
df: DataFrame
Expand Down Expand Up @@ -328,6 +330,7 @@ def read_gbq(
max_results=max_results,
progress_bar_type=progress_bar_type,
dtypes=dtypes,
dry_run=dry_run,
)
else:
final_df = connector.download_table(
Expand Down
14 changes: 13 additions & 1 deletion pandas_gbq/gbq_connector.py
Original file line number Diff line number Diff line change
Expand Up @@ -199,7 +199,14 @@ def download_table(
user_dtypes=dtypes,
)

def run_query(self, query, max_results=None, progress_bar_type=None, **kwargs):
def run_query(
self,
query,
max_results=None,
progress_bar_type=None,
dry_run: bool = False,
**kwargs,
):
from google.cloud import bigquery

job_config_dict = {
Expand Down Expand Up @@ -235,6 +242,7 @@ def run_query(self, query, max_results=None, progress_bar_type=None, **kwargs):

self._start_timer()
job_config = bigquery.QueryJobConfig.from_api_repr(job_config_dict)
job_config.dry_run = dry_run

if FEATURES.bigquery_has_query_and_wait:
rows_iter = pandas_gbq.query.query_and_wait_via_client_library(
Expand All @@ -260,6 +268,10 @@ def run_query(self, query, max_results=None, progress_bar_type=None, **kwargs):
)

dtypes = kwargs.get("dtypes")

if dry_run:
return rows_iter.total_bytes_processed / 1024**3

return self._download_results(
rows_iter,
max_results=max_results,
Expand Down
13 changes: 13 additions & 0 deletions tests/system/test_gbq.py
Original file line number Diff line number Diff line change
Expand Up @@ -656,6 +656,19 @@ def test_columns_and_col_order_raises_error(self, project_id):
dialect="standard",
)

def test_read_gbq_with_dry_run(self, project_id):
query = "SELECT 1"
cost = gbq.read_gbq(
query,
project_id=project_id,
credentials=self.credentials,
dialect="standard",
dry_run=True,
)
assert isinstance(cost, float)
assert cost > 0



class TestToGBQIntegration(object):
@pytest.fixture(autouse=True, scope="function")
Expand Down
9 changes: 9 additions & 0 deletions tests/unit/test_gbq.py
Original file line number Diff line number Diff line change
Expand Up @@ -937,3 +937,12 @@ def test_run_query_with_dml_query(mock_bigquery_client, mock_query_job):
type(mock_query_job).destination = mock.PropertyMock(return_value=None)
connector.run_query("UPDATE tablename SET value = '';")
mock_bigquery_client.list_rows.assert_not_called()


def test_read_gbq_with_dry_run(mock_bigquery_client, mock_query_job):
type(mock_query_job).total_bytes_processed = mock.PropertyMock(return_value=12345)
cost = gbq.read_gbq("SELECT 1", project_id="my-project", dry_run=True)
_, kwargs = mock_bigquery_client.query.call_args
job_config = kwargs["job_config"]
assert job_config.dry_run is True
assert cost == 12345 / 1024**3
Loading