Allow job_id to be explicitly specified. (#3779)

Author: lukesneeringer

bigquery/google/cloud/bigquery/dbapi/cursor.py

@@ -100,7 +100,7 @@ def _set_rowcount(self, query_results):
             total_rows = num_dml_affected_rows
         self.rowcount = total_rows
 
-    def execute(self, operation, parameters=None):
+    def execute(self, operation, parameters=None, job_id=None):
         """Prepare and execute a database operation.
 
         .. note::
@@ -128,12 +128,17 @@ def execute(self, operation, parameters=None):
         :type parameters: Mapping[str, Any] or Sequence[Any]
         :param parameters:
             (Optional) dictionary or sequence of parameter values.
+
+        :type job_id: str
+        :param job_id: (Optional) The job_id to use. If not set, a job ID
+            is generated at random.
         """
         self._query_results = None
         self._page_token = None
         self._has_fetched_all_rows = False
         client = self.connection._client
-        job_id = str(uuid.uuid4())
+        if job_id is None:
+            job_id = str(uuid.uuid4())
 
         # The DB-API uses the pyformat formatting, since the way BigQuery does
         # query parameters was not one of the standard options. Convert both

bigquery/tests/unit/test_dbapi_cursor.py

@@ -170,6 +170,14 @@ def test_fetchall_w_row(self):
         self.assertEqual(len(rows), 1)
         self.assertEqual(rows[0], (1,))
 
+    def test_execute_custom_job_id(self):
+        from google.cloud.bigquery.dbapi import connect
+        client = self._mock_client(rows=[], num_dml_affected_rows=0)
+        connection = connect(client)
+        cursor = connection.cursor()
+        cursor.execute('SELECT 1;', job_id='foo')
+        self.assertEqual(client.run_async_query.mock_calls[0][1][0], 'foo')
+
     def test_execute_w_dml(self):
         from google.cloud.bigquery.dbapi import connect
         connection = connect(