Skip to content

Commit

Permalink
feat(bigquery): add support for custom QueryJobConfig in BigQuery.cur…
Browse files Browse the repository at this point in the history
…sor.execute method (#9278)

* feat(big_query): add support for custom QueryJobConfig in execute cursor method

* fixup! feat(big_query): add support for custom QueryJobConfig in execute cursor method
  • Loading branch information
TobKed authored and plamut committed Sep 26, 2019
1 parent d8ce06e commit 7c9c0cb
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 3 deletions.
8 changes: 5 additions & 3 deletions bigquery/google/cloud/bigquery/dbapi/cursor.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ def _set_rowcount(self, query_results):
total_rows = num_dml_affected_rows
self.rowcount = total_rows

def execute(self, operation, parameters=None, job_id=None):
def execute(self, operation, parameters=None, job_id=None, job_config=None):
"""Prepare and execute a database operation.
.. note::
Expand Down Expand Up @@ -148,6 +148,9 @@ def execute(self, operation, parameters=None, job_id=None):
:type job_id: str
:param job_id: (Optional) The job_id to use. If not set, a job ID
is generated at random.
:type job_config: :class:`~google.cloud.bigquery.job.QueryJobConfig`
:param job_config: (Optional) Extra configuration options for the query job.
"""
self._query_data = None
self._query_job = None
Expand All @@ -160,9 +163,8 @@ def execute(self, operation, parameters=None, job_id=None):
formatted_operation = _format_operation(operation, parameters=parameters)
query_parameters = _helpers.to_query_parameters(parameters)

config = job.QueryJobConfig()
config = job_config or job.QueryJobConfig(use_legacy_sql=False)
config.query_parameters = query_parameters
config.use_legacy_sql = False
self._query_job = client.query(
formatted_operation, job_config=config, job_id=job_id
)
Expand Down
14 changes: 14 additions & 0 deletions bigquery/tests/unit/test_dbapi_cursor.py
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,20 @@ def test_execute_custom_job_id(self):
self.assertEqual(args[0], "SELECT 1;")
self.assertEqual(kwargs["job_id"], "foo")

def test_execute_custom_job_config(self):
from google.cloud.bigquery.dbapi import connect
from google.cloud.bigquery import job

config = job.QueryJobConfig(use_legacy_sql=True)
client = self._mock_client(rows=[], num_dml_affected_rows=0)
connection = connect(client)
cursor = connection.cursor()
cursor.execute("SELECT 1;", job_id="foo", job_config=config)
args, kwargs = client.query.call_args
self.assertEqual(args[0], "SELECT 1;")
self.assertEqual(kwargs["job_id"], "foo")
self.assertEqual(kwargs["job_config"], config)

def test_execute_w_dml(self):
from google.cloud.bigquery.dbapi import connect

Expand Down

0 comments on commit 7c9c0cb

Please sign in to comment.