feat: add dry run to the read_gbq function

Author: antoineeripret

pandas_gbq/gbq.py

@@ -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.
 
@@ -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
@@ -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(

pandas_gbq/gbq_connector.py

@@ -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 = {
@@ -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(
@@ -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
+
         return self._download_results(
             rows_iter,
             max_results=max_results,

tests/system/test_gbq.py

@@ -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"
+        job = gbq.read_gbq(
+            query,
+            project_id=project_id,
+            credentials=self.credentials,
+            dialect="standard",
+            dry_run=True,
+        )
+        assert job.dry_run
+        assert job.state == "DONE"
+        assert job.total_bytes_processed > 0
+
 
 class TestToGBQIntegration(object):
     @pytest.fixture(autouse=True, scope="function")

tests/unit/test_gbq.py

@@ -937,3 +937,10 @@ 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):
+    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