From b43a733ff8af7c0fbf09b2bc87d6d1079c38ae43 Mon Sep 17 00:00:00 2001 From: steffnay Date: Mon, 11 Oct 2021 14:44:32 -0700 Subject: [PATCH 1/2] feat: allow queryJob.result() to be called on a dryRun --- google/cloud/bigquery/job/query.py | 4 +++- tests/unit/job/test_query.py | 34 ++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 1 deletion(-) diff --git a/google/cloud/bigquery/job/query.py b/google/cloud/bigquery/job/query.py index 0cb4798be..15b1ca6cd 100644 --- a/google/cloud/bigquery/job/query.py +++ b/google/cloud/bigquery/job/query.py @@ -1262,7 +1262,7 @@ def result( timeout: float = None, start_index: int = None, job_retry: "retries.Retry" = DEFAULT_JOB_RETRY, - ) -> Union["RowIterator", _EmptyRowIterator]: + ) -> Union["RowIterator", _EmptyRowIterator, None]: """Start the job and wait for it to complete and get the result. Args: @@ -1318,6 +1318,8 @@ def result( If Non-``None`` and non-default ``job_retry`` is provided and the job is not retryable. """ + if self.dry_run: + return try: retry_do_query = getattr(self, "_retry_do_query", None) if retry_do_query is not None: diff --git a/tests/unit/job/test_query.py b/tests/unit/job/test_query.py index 4c598d797..4e92f7541 100644 --- a/tests/unit/job/test_query.py +++ b/tests/unit/job/test_query.py @@ -989,6 +989,40 @@ def test_result(self): [query_results_call, query_results_call, reload_call, query_page_call] ) + def test_result_dry_run(self): + query_resource_done = { + "jobComplete": True, + "jobReference": { + "projectId": self.PROJECT, + "jobId": self.JOB_ID, + "location": "EU", + }, + "schema": {"fields": [{"name": "col1", "type": "STRING"}]}, + "totalRows": "2", + } + job_resource = self._make_resource(started=True, location="EU") + job_resource["configuration"]["dryRun"] = True + # job_resource_done = self._make_resource(started=True, ended=True, location="EU") + # job_resource_done["configuration"]["query"]["destinationTable"] = { + # "projectId": "dest-project", + # "datasetId": "dest_dataset", + # "tableId": "dest_table", + # } + results_page_resource = { + "totalRows": "1", + "pageToken": None, + "rows": [{"f": [{"v": "abc"}]}], + } + conn = make_connection(query_resource_done, results_page_resource) + client = _make_client(self.PROJECT, connection=conn) + job = self._get_target_class().from_api_repr(job_resource, client) + + result = job.result() + + calls = conn.api_request.mock_calls + self.assertEqual(result, None) + self.assertEqual(calls, []) + def test_result_with_done_job_calls_get_query_results(self): query_resource_done = { "jobComplete": True, From 39f75a939e6338b18aac9097cdeac3bcd7bd4c5d Mon Sep 17 00:00:00 2001 From: steffnay Date: Tue, 12 Oct 2021 14:01:52 -0700 Subject: [PATCH 2/2] update to return EmptyRowIterator --- google/cloud/bigquery/job/query.py | 4 ++-- tests/unit/job/test_query.py | 26 +++----------------------- 2 files changed, 5 insertions(+), 25 deletions(-) diff --git a/google/cloud/bigquery/job/query.py b/google/cloud/bigquery/job/query.py index 15b1ca6cd..d9c796cf7 100644 --- a/google/cloud/bigquery/job/query.py +++ b/google/cloud/bigquery/job/query.py @@ -1262,7 +1262,7 @@ def result( timeout: float = None, start_index: int = None, job_retry: "retries.Retry" = DEFAULT_JOB_RETRY, - ) -> Union["RowIterator", _EmptyRowIterator, None]: + ) -> Union["RowIterator", _EmptyRowIterator]: """Start the job and wait for it to complete and get the result. Args: @@ -1319,7 +1319,7 @@ def result( provided and the job is not retryable. """ if self.dry_run: - return + return _EmptyRowIterator() try: retry_do_query = getattr(self, "_retry_do_query", None) if retry_do_query is not None: diff --git a/tests/unit/job/test_query.py b/tests/unit/job/test_query.py index 4e92f7541..8c0b944b0 100644 --- a/tests/unit/job/test_query.py +++ b/tests/unit/job/test_query.py @@ -26,6 +26,7 @@ from google.cloud.bigquery.client import _LIST_ROWS_FROM_QUERY_RESULTS_FIELDS import google.cloud.bigquery.query +from google.cloud.bigquery.table import _EmptyRowIterator from ..helpers import make_connection @@ -990,37 +991,16 @@ def test_result(self): ) def test_result_dry_run(self): - query_resource_done = { - "jobComplete": True, - "jobReference": { - "projectId": self.PROJECT, - "jobId": self.JOB_ID, - "location": "EU", - }, - "schema": {"fields": [{"name": "col1", "type": "STRING"}]}, - "totalRows": "2", - } job_resource = self._make_resource(started=True, location="EU") job_resource["configuration"]["dryRun"] = True - # job_resource_done = self._make_resource(started=True, ended=True, location="EU") - # job_resource_done["configuration"]["query"]["destinationTable"] = { - # "projectId": "dest-project", - # "datasetId": "dest_dataset", - # "tableId": "dest_table", - # } - results_page_resource = { - "totalRows": "1", - "pageToken": None, - "rows": [{"f": [{"v": "abc"}]}], - } - conn = make_connection(query_resource_done, results_page_resource) + conn = make_connection() client = _make_client(self.PROJECT, connection=conn) job = self._get_target_class().from_api_repr(job_resource, client) result = job.result() calls = conn.api_request.mock_calls - self.assertEqual(result, None) + self.assertIsInstance(result, _EmptyRowIterator) self.assertEqual(calls, []) def test_result_with_done_job_calls_get_query_results(self):