feat(bigquery): add support for custom QueryJobConfig in BigQuery.cursor.execute method (googleapis#9278)

TobKed · plamut · commit 7c9c0cbd5c83 · 2019-09-26T09:36:44.000+02:00
* 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
diff --git 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
@@ -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
 
