Skip to content

Commit

Permalink
running the bigquery magic in a notebook with the dryrun option print…
Browse files Browse the repository at this point in the history
…s total processed bytes to console
  • Loading branch information
shubha-rajan committed Aug 22, 2019
1 parent 3336697 commit 6efa408
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 4 deletions.
7 changes: 5 additions & 2 deletions bigquery/google/cloud/bigquery/magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -293,8 +293,6 @@ def _run_query(client, query, job_config=None):
query_job = client.query(query, job_config=job_config)

if job_config and job_config.dry_run:
assert query_job.state == "DONE"
print("Query successfully validated. This query will process {} bytes.".format(query_job.total_bytes_processed))
return query_job

print("Executing query with job ID: {}".format(query_job.job_id))
Expand Down Expand Up @@ -450,6 +448,11 @@ def _cell_magic(line, query):
IPython.get_ipython().push({args.destination_var: query_job})
return
elif args.dry_run:
print(
"Query validated. This query will process {} bytes.".format(
query_job.total_bytes_processed
)
)
return query_job

result = query_job.to_dataframe(bqstorage_client=bqstorage_client)
Expand Down
25 changes: 23 additions & 2 deletions bigquery/tests/unit/test_magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,26 @@ def test__run_query():
assert re.match("Query complete after .*s", updates[-1])


def test__run_query_dry_run_is_silent():
magics.context._credentials = None

sql = "SELECT 17"

client_patch = mock.patch(
"google.cloud.bigquery.magics.bigquery.Client", autospec=True
)

job_config = job.QueryJobConfig()
job_config.dry_run = True
with client_patch as client_mock, io.capture_output() as captured:
client_mock().query(sql).job_id = None

magics._run_query(client_mock(), sql, job_config=job_config)

assert len(captured.stderr) == 0
assert len(captured.stdout) == 0


def test__make_bqstorage_client_false():
credentials_mock = mock.create_autospec(
google.auth.credentials.Credentials, instance=True
Expand Down Expand Up @@ -662,12 +682,13 @@ def test_bigquery_magic_dryrun_option_returns_query_job():
sql = "SELECT 17 AS num"
result = pandas.DataFrame([17], columns=["num"])

with run_query_patch as run_query_mock:
with run_query_patch as run_query_mock, io.capture_output() as captured_io:
run_query_mock.return_value = query_job_mock
query_job_mock.to_dataframe.return_value = result
return_value = ip.run_cell_magic("bigquery", "--dry_run", sql)

assert isinstance(return_value, job.QueryJob)
assert "Query validated. This query will process" in captured_io.stdout
assert isinstance(return_value, job.QueryJob)


def test_bigquery_magic_dryrun_option_saves_query_job_to_variable():
Expand Down

0 comments on commit 6efa408

Please sign in to comment.