diff --git a/bigquery/google/cloud/bigquery/dbapi/cursor.py b/bigquery/google/cloud/bigquery/dbapi/cursor.py index 7519c762ae1e..167afb45e285 100644 --- a/bigquery/google/cloud/bigquery/dbapi/cursor.py +++ b/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 diff --git a/bigquery/tests/unit/test_dbapi_cursor.py b/bigquery/tests/unit/test_dbapi_cursor.py index 2a2ccfd989a6..49a332999f7e 100644 --- a/bigquery/tests/unit/test_dbapi_cursor.py +++ b/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(