From 7c9c0cbd5c83c40e40fd86ea04907b7fe3c9c238 Mon Sep 17 00:00:00 2001
From: TobKed <tobiasz.kedzierski@polidea.com>
Date: Thu, 26 Sep 2019 09:36:44 +0200
Subject: [PATCH] feat(bigquery): add support for custom QueryJobConfig in
 BigQuery.cursor.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
---
 bigquery/google/cloud/bigquery/dbapi/cursor.py |  8 +++++---
 bigquery/tests/unit/test_dbapi_cursor.py       | 14 ++++++++++++++
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/bigquery/google/cloud/bigquery/dbapi/cursor.py b/bigquery/google/cloud/bigquery/dbapi/cursor.py
index 3fdc750951e2..9b7a895b367f 100644
--- a/bigquery/google/cloud/bigquery/dbapi/cursor.py
+++ b/bigquery/google/cloud/bigquery/dbapi/cursor.py
@@ -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::
@@ -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
@@ -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
         )
diff --git a/bigquery/tests/unit/test_dbapi_cursor.py b/bigquery/tests/unit/test_dbapi_cursor.py
index 4a675c73958d..4ccd5e71af72 100644
--- a/bigquery/tests/unit/test_dbapi_cursor.py
+++ b/bigquery/tests/unit/test_dbapi_cursor.py
@@ -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