From 9ca87b2a77331f228199cf2ce410209b370f6360 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Tue, 27 Aug 2019 23:21:01 -0700 Subject: [PATCH 01/13] added max_results flag and property to QueryJobConfig --- bigquery/google/cloud/bigquery/client.py | 14 +++++++- bigquery/google/cloud/bigquery/job.py | 43 ++++++++++++++++++++++-- bigquery/google/cloud/bigquery/magics.py | 15 +++++++++ 3 files changed, 69 insertions(+), 3 deletions(-) diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py index 047633642e1e..6a2f86ee16ec 100644 --- a/bigquery/google/cloud/bigquery/client.py +++ b/bigquery/google/cloud/bigquery/client.py @@ -1055,7 +1055,13 @@ def delete_table(self, table, retry=DEFAULT_RETRY, not_found_ok=False): raise def _get_query_results( - self, job_id, retry, project=None, timeout_ms=None, location=None + self, + job_id, + retry, + project=None, + timeout_ms=None, + location=None, + max_results=None ): """Get the query results object for a query job. @@ -1070,6 +1076,8 @@ def _get_query_results( (Optional) number of milliseconds the the API call should wait for the query to complete before the request times out. location (str): Location of the query job. + max_results (int): + (Optional) Maximum number of results to read. Returns: google.cloud.bigquery.query._QueryResults: @@ -1089,6 +1097,9 @@ def _get_query_results( if location is not None: extra_params["location"] = location + + if max_results is not None: + extra_params['maxResults'] = max_results path = "/projects/{}/queries/{}".format(project, job_id) @@ -1986,6 +1997,7 @@ def query( to the client's project. retry (google.api_core.retry.Retry): (Optional) How to retry the RPC. + Returns: google.cloud.bigquery.job.QueryJob: A new query job instance. diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index 6e1eb81648f5..c25a689013bb 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -2151,6 +2151,19 @@ def flatten_results(self): def flatten_results(self, value): self._set_sub_prop("flattenResults", value) + @property + def max_results(self): + """int: Maximum number of results to read. + + See + https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query + """ + return self._properties.get("maxResults") + + @max_results.setter + def max_results(self, value): + self._properties["maxResults"] = value + @property def maximum_billing_tier(self): """int: Deprecated. Changes the billing tier to allow high-compute @@ -2385,6 +2398,7 @@ def __init__(self, job_id, query, client, job_config=None): if job_config is None: job_config = QueryJobConfig() + if job_config.use_legacy_sql is None: job_config.use_legacy_sql = False @@ -2451,6 +2465,13 @@ def flatten_results(self): """ return self._configuration.flatten_results + @property + def max_results(self): + """See + :attr:`google.cloud.bigquery.job.QueryJobConfig.max_results`. + """ + return self._configuration.max_results + @property def priority(self): """See @@ -2837,6 +2858,7 @@ def done(self, retry=DEFAULT_RETRY): project=self.project, timeout_ms=timeout_ms, location=self.location, + max_results=self.max_results ) # Only reload the job once we know the query is complete. @@ -2933,13 +2955,25 @@ def result(self, timeout=None, page_size=None, retry=DEFAULT_RETRY): concurrent.futures.TimeoutError: If the job did not complete in the given timeout. """ + + if not re.search(r"[\s]", self.query): + table_id = self.query + rows = self._client.list_rows( + table_id, page_size=page_size, retry=retry, max_results=self.max_results + ) + return rows + try: super(QueryJob, self).result(timeout=timeout) # Return an iterator instead of returning the job. if not self._query_results: self._query_results = self._client._get_query_results( - self.job_id, retry, project=self.project, location=self.location + self.job_id, + retry, + project=self.project, + location=self.location, + max_results=self.max_results ) except exceptions.GoogleCloudError as exc: exc.message += self._format_for_exception(self.query, self.job_id) @@ -2956,7 +2990,12 @@ def result(self, timeout=None, page_size=None, retry=DEFAULT_RETRY): dest_table_ref = self.destination dest_table = Table(dest_table_ref, schema=schema) dest_table._properties["numRows"] = self._query_results.total_rows - rows = self._client.list_rows(dest_table, page_size=page_size, retry=retry) + rows = self._client.list_rows( + dest_table, + page_size=page_size, + retry=retry, + max_results=self.max_results + ) rows._preserve_order = _contains_order_by(self.query) return rows diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index b2dae2511ec8..fc4fefc30b5b 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -278,6 +278,7 @@ def _run_query(client, query, job_config=None): job_config (google.cloud.bigquery.job.QueryJobConfig, optional): Extra configuration options for the job. + Returns: google.cloud.bigquery.job.QueryJob: the query job created @@ -320,6 +321,14 @@ def _run_query(client, query, job_config=None): default=None, help=("Project to use for executing this query. Defaults to the context project."), ) +@magic_arguments.argument( + "--max_results", + default=None, + help=( + "Maximum number of rows in dataframe returned from executing the query." + "Defaults to returning all rows." + ), +) @magic_arguments.argument( "--maximum_bytes_billed", default=None, @@ -409,6 +418,7 @@ def _cell_magic(line, query): ) project = args.project or context.project + client = bigquery.Client( project=project, credentials=context.credentials, @@ -425,6 +435,11 @@ def _cell_magic(line, query): job_config.use_legacy_sql = args.use_legacy_sql job_config.dry_run = args.dry_run + if args.max_results: + job_config.max_results = int(args.max_results) + else: + job_config.max_results = None + if args.maximum_bytes_billed == "None": job_config.maximum_bytes_billed = 0 elif args.maximum_bytes_billed is not None: From 718bb2e676449eb91545bd4521cdcc004129e60d Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Wed, 28 Aug 2019 20:10:14 -0700 Subject: [PATCH 02/13] tests for setting max_results and using table_id instead of query passing --- bigquery/google/cloud/bigquery/client.py | 9 +-- bigquery/google/cloud/bigquery/job.py | 26 ++----- bigquery/google/cloud/bigquery/magics.py | 18 +++-- bigquery/tests/unit/test_magics.py | 95 ++++++++++++++++++++++++ 4 files changed, 120 insertions(+), 28 deletions(-) diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py index 6a2f86ee16ec..3c086a3dea08 100644 --- a/bigquery/google/cloud/bigquery/client.py +++ b/bigquery/google/cloud/bigquery/client.py @@ -1061,7 +1061,7 @@ def _get_query_results( project=None, timeout_ms=None, location=None, - max_results=None + max_results=None, ): """Get the query results object for a query job. @@ -1076,7 +1076,7 @@ def _get_query_results( (Optional) number of milliseconds the the API call should wait for the query to complete before the request times out. location (str): Location of the query job. - max_results (int): + max_results (int): (Optional) Maximum number of results to read. Returns: @@ -1097,9 +1097,9 @@ def _get_query_results( if location is not None: extra_params["location"] = location - + if max_results is not None: - extra_params['maxResults'] = max_results + extra_params["maxResults"] = max_results path = "/projects/{}/queries/{}".format(project, job_id) @@ -1997,7 +1997,6 @@ def query( to the client's project. retry (google.api_core.retry.Retry): (Optional) How to retry the RPC. - Returns: google.cloud.bigquery.job.QueryJob: A new query job instance. diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index c25a689013bb..aeb8dada6329 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -2154,7 +2154,7 @@ def flatten_results(self, value): @property def max_results(self): """int: Maximum number of results to read. - + See https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query """ @@ -2858,7 +2858,7 @@ def done(self, retry=DEFAULT_RETRY): project=self.project, timeout_ms=timeout_ms, location=self.location, - max_results=self.max_results + max_results=self.max_results, ) # Only reload the job once we know the query is complete. @@ -2956,24 +2956,17 @@ def result(self, timeout=None, page_size=None, retry=DEFAULT_RETRY): If the job did not complete in the given timeout. """ - if not re.search(r"[\s]", self.query): - table_id = self.query - rows = self._client.list_rows( - table_id, page_size=page_size, retry=retry, max_results=self.max_results - ) - return rows - try: super(QueryJob, self).result(timeout=timeout) # Return an iterator instead of returning the job. if not self._query_results: self._query_results = self._client._get_query_results( - self.job_id, - retry, - project=self.project, - location=self.location, - max_results=self.max_results + self.job_id, + retry, + project=self.project, + location=self.location, + max_results=self.max_results, ) except exceptions.GoogleCloudError as exc: exc.message += self._format_for_exception(self.query, self.job_id) @@ -2991,10 +2984,7 @@ def result(self, timeout=None, page_size=None, retry=DEFAULT_RETRY): dest_table = Table(dest_table_ref, schema=schema) dest_table._properties["numRows"] = self._query_results.total_rows rows = self._client.list_rows( - dest_table, - page_size=page_size, - retry=retry, - max_results=self.max_results + dest_table, page_size=page_size, retry=retry, max_results=self.max_results ) rows._preserve_order = _contains_order_by(self.query) return rows diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index fc4fefc30b5b..17283436f8be 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -129,6 +129,7 @@ from __future__ import print_function +import re import ast import sys import time @@ -430,15 +431,22 @@ def _cell_magic(line, query): bqstorage_client = _make_bqstorage_client( args.use_bqstorage_api or context.use_bqstorage_api, context.credentials ) + + if args.max_results: + max_results = int(args.max_results) + else: + max_results = None + + if not re.search(r"[\s]", query): + table_id = query + rows = client.list_rows(table_id, max_results=max_results) + return rows.to_dataframe(bqstorage_client=bqstorage_client) + job_config = bigquery.job.QueryJobConfig() job_config.query_parameters = params job_config.use_legacy_sql = args.use_legacy_sql job_config.dry_run = args.dry_run - - if args.max_results: - job_config.max_results = int(args.max_results) - else: - job_config.max_results = None + job_config.max_results = max_results if args.maximum_bytes_billed == "None": job_config.maximum_bytes_billed = 0 diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index fbea9bdd9050..119aace2042d 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -743,6 +743,101 @@ def test_bigquery_magic_dryrun_option_saves_query_job_to_variable(): assert isinstance(q_job, job.QueryJob) +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_w_max_results_invalid(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + default_patch = mock.patch( + "google.auth.default", return_value=(credentials_mock, "general-project") + ) + client_query_patch = mock.patch( + "google.cloud.bigquery.client.Client.query", autospec=True + ) + query_job_mock = mock.create_autospec( + google.cloud.bigquery.job.QueryJob, instance=True + ) + + sql = "SELECT 17 AS num" + + with pytest.raises(ValueError), default_patch, client_query_patch: + ip.run_cell_magic("bigquery", "--max_results=abc", sql) + + +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_w_max_results_valid_sets_job_config(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + default_patch = mock.patch( + "google.auth.default", return_value=(credentials_mock, "general-project") + ) + run_query_patch = mock.patch( + "google.cloud.bigquery.magics._run_query", autospec=True + ) + + sql = "SELECT 17 AS num" + + with run_query_patch as run_query_mock: + ip.run_cell_magic("bigquery", "--max_results=5", sql) + + job_config_used = run_query_mock.call_args_list[0][0][-1] + assert job_config_used.max_results == 5 + + +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_w_table_id_instead_of_sql(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + credentials_mock = mock.create_autospec( + google.auth.credentials.Credentials, instance=True + ) + default_patch = mock.patch( + "google.auth.default", return_value=(credentials_mock, "general-project") + ) + + row_iterator_mock = mock.create_autospec( + google.cloud.bigquery.table.RowIterator, instance=True + ) + + query_job_mock = mock.create_autospec( + google.cloud.bigquery.job.QueryJob, instance=True + ) + + client_patch = mock.patch( + "google.cloud.bigquery.magics.bigquery.Client", autospec=True + ) + + run_query_patch = mock.patch( + "google.cloud.bigquery.magics._run_query", autospec=True + ) + + table_id = "12345" + result = pandas.DataFrame([17], columns=["num"]) + + with client_patch as client_mock, run_query_patch as run_query_mock: + client_mock().list_rows.return_value = row_iterator_mock + row_iterator_mock.to_dataframe.return_value = result + + return_value = ip.run_cell_magic("bigquery", "--max_results=5", table_id) + client_mock.return_value.list_rows.assert_called_once_with( + table_id, max_results=5 + ) + run_query_mock.assert_not_called() + + assert isinstance(return_value, pandas.DataFrame) + + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_w_maximum_bytes_billed_invalid(): ip = IPython.get_ipython() From 5eb56ce1df86d180ee80ef338e926326fce92bd4 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Thu, 29 Aug 2019 15:38:46 -0700 Subject: [PATCH 03/13] adjusted regex whitespace check to account for trailing newline added by notebook cell --- bigquery/google/cloud/bigquery/job.py | 4 ++-- bigquery/google/cloud/bigquery/magics.py | 12 +++++++++--- bigquery/tests/unit/test_job.py | 5 +++++ bigquery/tests/unit/test_magics.py | 22 +--------------------- 4 files changed, 17 insertions(+), 26 deletions(-) diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index aeb8dada6329..fadd766c2c2d 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -2158,11 +2158,11 @@ def max_results(self): See https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query """ - return self._properties.get("maxResults") + return _helpers._int_or_none(self._get_sub_prop("maxResults")) @max_results.setter def max_results(self, value): - self._properties["maxResults"] = value + self._set_sub_prop("maxResults", str(value)) @property def maximum_billing_tier(self): diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 17283436f8be..4140ee11f340 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -404,6 +404,7 @@ def _cell_magic(line, query): Returns: pandas.DataFrame: the query results. """ + args = magic_arguments.parse_argstring(_cell_magic, line) params = [] @@ -437,10 +438,15 @@ def _cell_magic(line, query): else: max_results = None - if not re.search(r"[\s]", query): - table_id = query + if not re.search(r"\b\s+\b", query): + table_id = query.rstrip() rows = client.list_rows(table_id, max_results=max_results) - return rows.to_dataframe(bqstorage_client=bqstorage_client) + result = rows.to_dataframe(bqstorage_client=bqstorage_client) + if args.destination_var: + IPython.get_ipython().push({args.destination_var: result}) + return + else: + return result job_config = bigquery.job.QueryJobConfig() job_config.query_parameters = params diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py index 19409d8d43c3..341a3968a7fe 100644 --- a/bigquery/tests/unit/test_job.py +++ b/bigquery/tests/unit/test_job.py @@ -3455,6 +3455,11 @@ def _verifyIntegerResourceProperties(self, job, config): self.assertIsInstance(job.maximum_bytes_billed, int) else: self.assertIsNone(job.maximum_bytes_billed) + if "maxResults" in config: + self.assertEqual(str(job.max_results), config["maxResults"]) + self.assertIsInstance(job.max_results, int) + else: + self.assertIsNone(job.max_results) def _verify_udf_resources(self, job, config): udf_resources = config.get("userDefinedFunctionResources", ()) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 119aace2042d..eeec44a9f649 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -758,9 +758,6 @@ def test_bigquery_magic_w_max_results_invalid(): client_query_patch = mock.patch( "google.cloud.bigquery.client.Client.query", autospec=True ) - query_job_mock = mock.create_autospec( - google.cloud.bigquery.job.QueryJob, instance=True - ) sql = "SELECT 17 AS num" @@ -774,12 +771,6 @@ def test_bigquery_magic_w_max_results_valid_sets_job_config(): ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._project = None - credentials_mock = mock.create_autospec( - google.auth.credentials.Credentials, instance=True - ) - default_patch = mock.patch( - "google.auth.default", return_value=(credentials_mock, "general-project") - ) run_query_patch = mock.patch( "google.cloud.bigquery.magics._run_query", autospec=True ) @@ -799,21 +790,10 @@ def test_bigquery_magic_w_table_id_instead_of_sql(): ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._project = None - credentials_mock = mock.create_autospec( - google.auth.credentials.Credentials, instance=True - ) - default_patch = mock.patch( - "google.auth.default", return_value=(credentials_mock, "general-project") - ) - row_iterator_mock = mock.create_autospec( google.cloud.bigquery.table.RowIterator, instance=True ) - query_job_mock = mock.create_autospec( - google.cloud.bigquery.job.QueryJob, instance=True - ) - client_patch = mock.patch( "google.cloud.bigquery.magics.bigquery.Client", autospec=True ) @@ -822,7 +802,7 @@ def test_bigquery_magic_w_table_id_instead_of_sql(): "google.cloud.bigquery.magics._run_query", autospec=True ) - table_id = "12345" + table_id = "bigquery-public-data.samples.shakespeare" result = pandas.DataFrame([17], columns=["num"]) with client_patch as client_mock, run_query_patch as run_query_mock: From 904c31b456709529ad4047114d479fab8729a2a4 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Thu, 29 Aug 2019 23:44:21 -0700 Subject: [PATCH 04/13] preserve value of max_results after QueryJob._set_properties called --- bigquery/google/cloud/bigquery/job.py | 21 ++++++++++++++------- bigquery/google/cloud/bigquery/magics.py | 2 +- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index fadd766c2c2d..2c00c8dd84c0 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -2162,7 +2162,7 @@ def max_results(self): @max_results.setter def max_results(self, value): - self._set_sub_prop("maxResults", str(value)) + self._set_sub_prop("maxResults", value) @property def maximum_billing_tier(self): @@ -2579,6 +2579,18 @@ def to_api_repr(self): return resource + def _set_properties(self, api_response): + """Update properties from resource in body of ``api_response`` + + :type api_response: dict + :param api_response: response returned from an API call + """ + + # prevent max_results value from being erased + max_results = self.max_results + super(QueryJob, self)._set_properties(api_response) + self._configuration.max_results = max_results + def _copy_configuration_properties(self, configuration): """Helper: assign subclass configuration properties in cleaned.""" self._configuration._properties = copy.deepcopy(configuration) @@ -2858,7 +2870,6 @@ def done(self, retry=DEFAULT_RETRY): project=self.project, timeout_ms=timeout_ms, location=self.location, - max_results=self.max_results, ) # Only reload the job once we know the query is complete. @@ -2962,11 +2973,7 @@ def result(self, timeout=None, page_size=None, retry=DEFAULT_RETRY): # Return an iterator instead of returning the job. if not self._query_results: self._query_results = self._client._get_query_results( - self.job_id, - retry, - project=self.project, - location=self.location, - max_results=self.max_results, + self.job_id, retry, project=self.project, location=self.location ) except exceptions.GoogleCloudError as exc: exc.message += self._format_for_exception(self.query, self.job_id) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 4140ee11f340..9e62a28268dd 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -438,7 +438,7 @@ def _cell_magic(line, query): else: max_results = None - if not re.search(r"\b\s+\b", query): + if not re.search(r"\s", query.rstrip()): table_id = query.rstrip() rows = client.list_rows(table_id, max_results=max_results) result = rows.to_dataframe(bqstorage_client=bqstorage_client) From 53d0b0b13b726cb3c75892078f2be0bfb86de599 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Fri, 30 Aug 2019 01:19:19 -0700 Subject: [PATCH 05/13] added tests for using --max_results with destination_var and bq_storage_api --- bigquery/google/cloud/bigquery/client.py | 3 -- bigquery/tests/unit/test_job.py | 5 -- bigquery/tests/unit/test_magics.py | 65 +++++++++++++++++++++++- 3 files changed, 64 insertions(+), 9 deletions(-) diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py index 3c086a3dea08..8dc445b63b4c 100644 --- a/bigquery/google/cloud/bigquery/client.py +++ b/bigquery/google/cloud/bigquery/client.py @@ -1098,9 +1098,6 @@ def _get_query_results( if location is not None: extra_params["location"] = location - if max_results is not None: - extra_params["maxResults"] = max_results - path = "/projects/{}/queries/{}".format(project, job_id) # This call is typically made in a polling loop that checks whether the diff --git a/bigquery/tests/unit/test_job.py b/bigquery/tests/unit/test_job.py index 341a3968a7fe..19409d8d43c3 100644 --- a/bigquery/tests/unit/test_job.py +++ b/bigquery/tests/unit/test_job.py @@ -3455,11 +3455,6 @@ def _verifyIntegerResourceProperties(self, job, config): self.assertIsInstance(job.maximum_bytes_billed, int) else: self.assertIsNone(job.maximum_bytes_billed) - if "maxResults" in config: - self.assertEqual(str(job.max_results), config["maxResults"]) - self.assertIsInstance(job.max_results, int) - else: - self.assertIsNone(job.max_results) def _verify_udf_resources(self, job, config): udf_resources = config.get("userDefinedFunctionResources", ()) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index eeec44a9f649..84a6b34e0b54 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -785,7 +785,7 @@ def test_bigquery_magic_w_max_results_valid_sets_job_config(): @pytest.mark.usefixtures("ipython_interactive") -def test_bigquery_magic_w_table_id_instead_of_sql(): +def test_bigquery_magic_w_table_id_and_max_results(): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._project = None @@ -817,6 +817,69 @@ def test_bigquery_magic_w_table_id_instead_of_sql(): assert isinstance(return_value, pandas.DataFrame) +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_w_table_id_and_destination_var(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + row_iterator_mock = mock.create_autospec( + google.cloud.bigquery.table.RowIterator, instance=True + ) + + client_patch = mock.patch( + "google.cloud.bigquery.magics.bigquery.Client", autospec=True + ) + + table_id = "bigquery-public-data.samples.shakespeare" + result = pandas.DataFrame([17], columns=["num"]) + + with client_patch as client_mock: + client_mock().list_rows.return_value = row_iterator_mock + row_iterator_mock.to_dataframe.return_value = result + + return_value = ip.run_cell_magic("bigquery", "df --max_results=5", table_id) + + assert "df" in ip.user_ns + df = ip.user_ns['df'] + + assert isinstance(df, pandas.DataFrame) + +@pytest.mark.usefixtures("ipython_interactive") +def test_bigquery_magic_w_table_id_and_bqstorage_client(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + row_iterator_mock = mock.create_autospec( + google.cloud.bigquery.table.RowIterator, instance=True + ) + + client_patch = mock.patch( + "google.cloud.bigquery.magics.bigquery.Client", autospec=True + ) + + + bqstorage_mock = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient + ) + bqstorage_instance_mock = mock.create_autospec( + bigquery_storage_v1beta1.BigQueryStorageClient, instance=True + ) + bqstorage_mock.return_value = bqstorage_instance_mock + bqstorage_client_patch = mock.patch( + "google.cloud.bigquery_storage_v1beta1.BigQueryStorageClient", bqstorage_mock + ) + + table_id = "bigquery-public-data.samples.shakespeare" + + with client_patch as client_mock, bqstorage_client_patch: + client_mock().list_rows.return_value = row_iterator_mock + + return_value = ip.run_cell_magic("bigquery", "--use_bqstorage_api --max_results=5", table_id) + row_iterator_mock.to_dataframe.assert_called_once_with( + bqstorage_client=bqstorage_instance_mock) + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_w_maximum_bytes_billed_invalid(): From 3706a0a477379a00851f10db5ae5f9487df5b388 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Fri, 30 Aug 2019 01:31:24 -0700 Subject: [PATCH 06/13] blacken and lint --- bigquery/tests/unit/test_magics.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 84a6b34e0b54..027f2f9d288b 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -817,6 +817,7 @@ def test_bigquery_magic_w_table_id_and_max_results(): assert isinstance(return_value, pandas.DataFrame) + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_w_table_id_and_destination_var(): ip = IPython.get_ipython() @@ -838,13 +839,14 @@ def test_bigquery_magic_w_table_id_and_destination_var(): client_mock().list_rows.return_value = row_iterator_mock row_iterator_mock.to_dataframe.return_value = result - return_value = ip.run_cell_magic("bigquery", "df --max_results=5", table_id) + ip.run_cell_magic("bigquery", "df --max_results=5", table_id) assert "df" in ip.user_ns - df = ip.user_ns['df'] + df = ip.user_ns["df"] assert isinstance(df, pandas.DataFrame) + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_w_table_id_and_bqstorage_client(): ip = IPython.get_ipython() @@ -859,7 +861,6 @@ def test_bigquery_magic_w_table_id_and_bqstorage_client(): "google.cloud.bigquery.magics.bigquery.Client", autospec=True ) - bqstorage_mock = mock.create_autospec( bigquery_storage_v1beta1.BigQueryStorageClient ) @@ -875,10 +876,11 @@ def test_bigquery_magic_w_table_id_and_bqstorage_client(): with client_patch as client_mock, bqstorage_client_patch: client_mock().list_rows.return_value = row_iterator_mock - - return_value = ip.run_cell_magic("bigquery", "--use_bqstorage_api --max_results=5", table_id) + + ip.run_cell_magic("bigquery", "--use_bqstorage_api --max_results=5", table_id) row_iterator_mock.to_dataframe.assert_called_once_with( - bqstorage_client=bqstorage_instance_mock) + bqstorage_client=bqstorage_instance_mock + ) @pytest.mark.usefixtures("ipython_interactive") From 0e9e03214c24a0399b85a62728c47e9b75321e67 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Fri, 30 Aug 2019 01:49:07 -0700 Subject: [PATCH 07/13] removed unused max_results parameter from client._get_query_results --- bigquery/google/cloud/bigquery/client.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py index 8dc445b63b4c..c8b69138ed49 100644 --- a/bigquery/google/cloud/bigquery/client.py +++ b/bigquery/google/cloud/bigquery/client.py @@ -1055,13 +1055,7 @@ def delete_table(self, table, retry=DEFAULT_RETRY, not_found_ok=False): raise def _get_query_results( - self, - job_id, - retry, - project=None, - timeout_ms=None, - location=None, - max_results=None, + self, job_id, retry, project=None, timeout_ms=None, location=None ): """Get the query results object for a query job. @@ -1076,9 +1070,6 @@ def _get_query_results( (Optional) number of milliseconds the the API call should wait for the query to complete before the request times out. location (str): Location of the query job. - max_results (int): - (Optional) Maximum number of results to read. - Returns: google.cloud.bigquery.query._QueryResults: A new ``_QueryResults`` instance. From 4756ec34e7cf8749af6801863a8e9c7be89a84d8 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Fri, 30 Aug 2019 13:43:36 -0700 Subject: [PATCH 08/13] added error messaging and tests for failure case --- bigquery/google/cloud/bigquery/magics.py | 16 +++++++++++++++- bigquery/tests/unit/test_magics.py | 24 ++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index 9e62a28268dd..e9134b52c1b6 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -440,7 +440,21 @@ def _cell_magic(line, query): if not re.search(r"\s", query.rstrip()): table_id = query.rstrip() - rows = client.list_rows(table_id, max_results=max_results) + error = None + try: + rows = client.list_rows(table_id, max_results=max_results) + except Exception as ex: + error = str(ex) + if error: + if args.destination_var: + print( + "Could not save output to variable '{}'.".format( + args.destination_var), + file=sys.stderr, + ) + print("\nERROR:\n", error, file=sys.stderr) + return + result = rows.to_dataframe(bqstorage_client=bqstorage_client) if args.destination_var: IPython.get_ipython().push({args.destination_var: result}) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 027f2f9d288b..cdc710c2ec99 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -817,6 +817,30 @@ def test_bigquery_magic_w_table_id_and_max_results(): assert isinstance(return_value, pandas.DataFrame) +def test_bigquery_magic_w_table_id_invalid(): + ip = IPython.get_ipython() + ip.extension_manager.load_extension("google.cloud.bigquery") + magics.context._project = None + + + + list_rows_patch = mock.patch( + "google.cloud.bigquery.magics.bigquery.Client.list_rows", + autospec=True, + side_effect=exceptions.BadRequest("Not a valid table ID") + ) + + table_id = "not-a-real-table" + + + with list_rows_patch, io.capture_output() as captured_io: + + ip.run_cell_magic("bigquery", "", table_id) + + output = captured_io.stderr + assert "400 Not a valid table ID" in output + assert "Traceback (most recent call last)" not in output + @pytest.mark.usefixtures("ipython_interactive") def test_bigquery_magic_w_table_id_and_destination_var(): From 304c76162ee2e621b2e0a3eed7189f7e0e74dc9b Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Fri, 30 Aug 2019 13:51:17 -0700 Subject: [PATCH 09/13] reformatted docstrings --- bigquery/google/cloud/bigquery/job.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index 2c00c8dd84c0..9421439f2ba6 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -499,8 +499,8 @@ def _copy_configuration_properties(self, configuration): def _set_properties(self, api_response): """Update properties from resource in body of ``api_response`` - :type api_response: dict - :param api_response: response returned from an API call + Args: + api_response (dict): response returned from an API call. """ cleaned = api_response.copy() self._scrub_local_properties(cleaned) @@ -2582,8 +2582,8 @@ def to_api_repr(self): def _set_properties(self, api_response): """Update properties from resource in body of ``api_response`` - :type api_response: dict - :param api_response: response returned from an API call + Args: + api_response (dict): response returned from an API call. """ # prevent max_results value from being erased From 8d0781d751fc0cb373bd4cd797f4887d6ffaaf21 Mon Sep 17 00:00:00 2001 From: Shubha Rajan Date: Fri, 30 Aug 2019 14:18:00 -0700 Subject: [PATCH 10/13] Fix docstring formatting --- bigquery/google/cloud/bigquery/job.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/bigquery/google/cloud/bigquery/job.py b/bigquery/google/cloud/bigquery/job.py index 9421439f2ba6..8aa504d0874b 100644 --- a/bigquery/google/cloud/bigquery/job.py +++ b/bigquery/google/cloud/bigquery/job.py @@ -500,7 +500,7 @@ def _set_properties(self, api_response): """Update properties from resource in body of ``api_response`` Args: - api_response (dict): response returned from an API call. + api_response (Dict): response returned from an API call. """ cleaned = api_response.copy() self._scrub_local_properties(cleaned) @@ -2583,7 +2583,7 @@ def _set_properties(self, api_response): """Update properties from resource in body of ``api_response`` Args: - api_response (dict): response returned from an API call. + api_response (Dict): response returned from an API call. """ # prevent max_results value from being erased From b93b01d5014cc89c27ac789b83a9ab46797a9218 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Fri, 30 Aug 2019 14:30:13 -0700 Subject: [PATCH 11/13] blacken/lint --- bigquery/google/cloud/bigquery/magics.py | 5 +++-- bigquery/tests/unit/test_magics.py | 10 ++++------ 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index e9134b52c1b6..b982330f8ce4 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -449,12 +449,13 @@ def _cell_magic(line, query): if args.destination_var: print( "Could not save output to variable '{}'.".format( - args.destination_var), + args.destination_var + ), file=sys.stderr, ) print("\nERROR:\n", error, file=sys.stderr) return - + result = rows.to_dataframe(bqstorage_client=bqstorage_client) if args.destination_var: IPython.get_ipython().push({args.destination_var: result}) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index cdc710c2ec99..c3861a9190b4 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -817,24 +817,22 @@ def test_bigquery_magic_w_table_id_and_max_results(): assert isinstance(return_value, pandas.DataFrame) + def test_bigquery_magic_w_table_id_invalid(): ip = IPython.get_ipython() ip.extension_manager.load_extension("google.cloud.bigquery") magics.context._project = None - - list_rows_patch = mock.patch( - "google.cloud.bigquery.magics.bigquery.Client.list_rows", + "google.cloud.bigquery.magics.bigquery.Client.list_rows", autospec=True, - side_effect=exceptions.BadRequest("Not a valid table ID") + side_effect=exceptions.BadRequest("Not a valid table ID"), ) table_id = "not-a-real-table" - with list_rows_patch, io.capture_output() as captured_io: - + ip.run_cell_magic("bigquery", "", table_id) output = captured_io.stderr From 32a1ebcae6bd7104a4beea13fa214c85fc4c3f49 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Fri, 30 Aug 2019 19:43:25 -0700 Subject: [PATCH 12/13] update error messaging test --- bigquery/tests/unit/test_magics.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index c3861a9190b4..8a51dccbeab6 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -833,9 +833,10 @@ def test_bigquery_magic_w_table_id_invalid(): with list_rows_patch, io.capture_output() as captured_io: - ip.run_cell_magic("bigquery", "", table_id) + ip.run_cell_magic("bigquery", "df", table_id) output = captured_io.stderr + assert "Could not save output to variable" in output assert "400 Not a valid table ID" in output assert "Traceback (most recent call last)" not in output From aceee811a35cf9b39a8e6a8d425b44df2fba7de4 Mon Sep 17 00:00:00 2001 From: shubha-rajan Date: Fri, 30 Aug 2019 23:18:57 -0700 Subject: [PATCH 13/13] refactored error message display into its own method.fixed coverage failure --- bigquery/google/cloud/bigquery/magics.py | 30 +++++++++++------------- bigquery/tests/unit/test_magics.py | 1 - 2 files changed, 14 insertions(+), 17 deletions(-) diff --git a/bigquery/google/cloud/bigquery/magics.py b/bigquery/google/cloud/bigquery/magics.py index b982330f8ce4..33083a31f4bc 100644 --- a/bigquery/google/cloud/bigquery/magics.py +++ b/bigquery/google/cloud/bigquery/magics.py @@ -267,6 +267,15 @@ def default_query_job_config(self, value): context = Context() +def _print_error(error, destination_var=None): + if destination_var: + print( + "Could not save output to variable '{}'.".format(destination_var), + file=sys.stderr, + ) + print("\nERROR:\n", error, file=sys.stderr) + + def _run_query(client, query, job_config=None): """Runs a query while printing status updates @@ -438,22 +447,17 @@ def _cell_magic(line, query): else: max_results = None + error = None + if not re.search(r"\s", query.rstrip()): table_id = query.rstrip() - error = None + try: rows = client.list_rows(table_id, max_results=max_results) except Exception as ex: error = str(ex) if error: - if args.destination_var: - print( - "Could not save output to variable '{}'.".format( - args.destination_var - ), - file=sys.stderr, - ) - print("\nERROR:\n", error, file=sys.stderr) + _print_error(error, args.destination_var) return result = rows.to_dataframe(bqstorage_client=bqstorage_client) @@ -475,7 +479,6 @@ def _cell_magic(line, query): value = int(args.maximum_bytes_billed) job_config.maximum_bytes_billed = value - error = None try: query_job = _run_query(client, query, job_config) except Exception as ex: @@ -485,12 +488,7 @@ def _cell_magic(line, query): display.clear_output() if error: - if args.destination_var: - print( - "Could not save output to variable '{}'.".format(args.destination_var), - file=sys.stderr, - ) - print("\nERROR:\n", error, file=sys.stderr) + _print_error(error, args.destination_var) return if args.dry_run and args.destination_var: diff --git a/bigquery/tests/unit/test_magics.py b/bigquery/tests/unit/test_magics.py index 8a51dccbeab6..da85b1aac53c 100644 --- a/bigquery/tests/unit/test_magics.py +++ b/bigquery/tests/unit/test_magics.py @@ -832,7 +832,6 @@ def test_bigquery_magic_w_table_id_invalid(): table_id = "not-a-real-table" with list_rows_patch, io.capture_output() as captured_io: - ip.run_cell_magic("bigquery", "df", table_id) output = captured_io.stderr