BigQuery: option for job ID generation with user-supplied prefix (#4198)

adds job_id_prefix to _make_job_id()

Author: alixhami

bigquery/google/cloud/bigquery/client.py

@@ -599,8 +599,8 @@ def list_jobs(self, max_results=None, page_token=None, all_users=None,
             extra_params=extra_params)
 
     def load_table_from_storage(self, source_uris, destination,
-                                job_id=None, job_config=None,
-                                retry=DEFAULT_RETRY):
+                                job_id=None, job_id_prefix=None,
+                                job_config=None, retry=DEFAULT_RETRY):
         """Starts a job for loading data into a table from CloudStorage.
 
         See
@@ -618,6 +618,11 @@ def load_table_from_storage(self, source_uris, destination,
         :type job_id: str
         :param job_id: (Optional) Name of the job.
 
+        :type job_id_prefix: str or ``NoneType``
+        :param job_id_prefix: (Optional) the user-provided prefix for a
+                              randomly generated job ID. This parameter will be
+                              ignored if a ``job_id`` is also given.
+
         :type job_config: :class:`google.cloud.bigquery.job.LoadJobConfig`
         :param job_config: (Optional) Extra configuration options for the job.
 
@@ -627,7 +632,7 @@ def load_table_from_storage(self, source_uris, destination,
         :rtype: :class:`google.cloud.bigquery.job.LoadJob`
         :returns: a new ``LoadJob`` instance
         """
-        job_id = _make_job_id(job_id)
+        job_id = _make_job_id(job_id, job_id_prefix)
         if isinstance(source_uris, six.string_types):
             source_uris = [source_uris]
         job = LoadJob(job_id, source_uris, destination, self, job_config)
@@ -638,7 +643,7 @@ def load_table_from_file(self, file_obj, destination,
                              rewind=False,
                              size=None,
                              num_retries=_DEFAULT_NUM_RETRIES,
-                             job_id=None, job_config=None):
+                             job_id=None, job_id_prefix=None, job_config=None):
         """Upload the contents of this table from a file-like object.
 
         Like load_table_from_storage, this creates, starts and returns
@@ -665,6 +670,11 @@ def load_table_from_file(self, file_obj, destination,
         :type job_id: str
         :param job_id: (Optional) Name of the job.
 
+        :type job_id_prefix: str or ``NoneType``
+        :param job_id_prefix: (Optional) the user-provided prefix for a
+                              randomly generated job ID. This parameter will be
+                              ignored if a ``job_id`` is also given.
+
         :type job_config: :class:`google.cloud.bigquery.job.LoadJobConfig`
         :param job_config: (Optional) Extra configuration options for the job.
 
@@ -677,7 +687,7 @@ def load_table_from_file(self, file_obj, destination,
                  be determined, or if the ``file_obj`` can be detected to be
                  a file opened in text mode.
         """
-        job_id = _make_job_id(job_id)
+        job_id = _make_job_id(job_id, job_id_prefix)
         job = LoadJob(job_id, None, destination, self, job_config)
         job_resource = job._build_resource()
         if rewind:
@@ -801,8 +811,8 @@ def _do_multipart_upload(self, stream, metadata, size, num_retries):
 
         return response
 
-    def copy_table(self, sources, destination, job_id=None, job_config=None,
-                   retry=DEFAULT_RETRY):
+    def copy_table(self, sources, destination, job_id=None, job_id_prefix=None,
+                   job_config=None, retry=DEFAULT_RETRY):
         """Start a job for copying one or more tables into another table.
 
         See
@@ -821,6 +831,11 @@ def copy_table(self, sources, destination, job_id=None, job_config=None,
         :type job_id: str
         :param job_id: (Optional) The ID of the job.
 
+        :type job_id_prefix: str or ``NoneType``
+        :param job_id_prefix: (Optional) the user-provided prefix for a
+                              randomly generated job ID. This parameter will be
+                              ignored if a ``job_id`` is also given.
+
         :type job_config: :class:`google.cloud.bigquery.job.CopyJobConfig`
         :param job_config: (Optional) Extra configuration options for the job.
 
@@ -830,7 +845,7 @@ def copy_table(self, sources, destination, job_id=None, job_config=None,
         :rtype: :class:`google.cloud.bigquery.job.CopyJob`
         :returns: a new ``CopyJob`` instance
         """
-        job_id = _make_job_id(job_id)
+        job_id = _make_job_id(job_id, job_id_prefix)
 
         if not isinstance(sources, collections.Sequence):
             sources = [sources]
@@ -841,7 +856,7 @@ def copy_table(self, sources, destination, job_id=None, job_config=None,
 
     def extract_table(
             self, source, destination_uris, job_config=None, job_id=None,
-            retry=DEFAULT_RETRY):
+            job_id_prefix=None, retry=DEFAULT_RETRY):
         """Start a job to extract a table into Cloud Storage files.
 
         See
@@ -863,6 +878,11 @@ def extract_table(
         :type job_id: str
         :param job_id: (Optional) The ID of the job.
 
+        :type job_id_prefix: str or ``NoneType``
+        :param job_id_prefix: (Optional) the user-provided prefix for a
+                              randomly generated job ID. This parameter will be
+                              ignored if a ``job_id`` is also given.
+
         :type job_config: :class:`google.cloud.bigquery.job.ExtractJobConfig`
         :param job_config: (Optional) Extra configuration options for the job.
 
@@ -872,7 +892,7 @@ def extract_table(
         :rtype: :class:`google.cloud.bigquery.job.ExtractJob`
         :returns: a new ``ExtractJob`` instance
         """
-        job_id = _make_job_id(job_id)
+        job_id = _make_job_id(job_id, job_id_prefix)
 
         if isinstance(destination_uris, six.string_types):
             destination_uris = [destination_uris]
@@ -883,7 +903,8 @@ def extract_table(
         job.begin(retry=retry)
         return job
 
-    def query(self, query, job_config=None, job_id=None, retry=DEFAULT_RETRY):
+    def query(self, query, job_config=None, job_id=None, job_id_prefix=None,
+              retry=DEFAULT_RETRY):
         """Start a job that runs a SQL query.
 
         See
@@ -900,13 +921,18 @@ def query(self, query, job_config=None, job_id=None, retry=DEFAULT_RETRY):
         :type job_id: str
         :param job_id: (Optional) ID to use for the query job.
 
+        :type job_id_prefix: str or ``NoneType``
+        :param job_id_prefix: (Optional) the user-provided prefix for a
+                              randomly generated job ID. This parameter will be
+                              ignored if a ``job_id`` is also given.
+
         :type retry: :class:`google.api.core.retry.Retry`
         :param retry: (Optional) How to retry the RPC.
 
         :rtype: :class:`google.cloud.bigquery.job.QueryJob`
         :returns: a new ``QueryJob`` instance
         """
-        job_id = _make_job_id(job_id)
+        job_id = _make_job_id(job_id, job_id_prefix)
         job = QueryJob(job_id, query, client=self, job_config=job_config)
         job.begin(retry=retry)
         return job
@@ -1269,18 +1295,24 @@ def _item_to_table(iterator, resource):
     return Table.from_api_repr(resource)
 
 
-def _make_job_id(job_id):
+def _make_job_id(job_id, prefix=None):
     """Construct an ID for a new job.
 
     :type job_id: str or ``NoneType``
     :param job_id: the user-provided job ID
 
+    :type prefix: str or ``NoneType``
+    :param prefix: (Optional) the user-provided prefix for a job ID
+
     :rtype: str
     :returns: A job ID
     """
-    if job_id is None:
+    if job_id is not None:
+        return job_id
+    elif prefix is not None:
+        return str(prefix) + str(uuid.uuid4())
+    else:
         return str(uuid.uuid4())
-    return job_id
 
 
 def _check_mode(stream):

bigquery/tests/system.py

@@ -633,7 +633,7 @@ def test_copy_table(self):
 
     def test_job_cancel(self):
         DATASET_ID = _make_dataset_id('job_cancel')
-        JOB_ID = 'fetch_' + DATASET_ID + str(uuid.uuid4())
+        JOB_ID_PREFIX = 'fetch_' + DATASET_ID
         TABLE_NAME = 'test_table'
         QUERY = 'SELECT * FROM %s.%s' % (DATASET_ID, TABLE_NAME)
 
@@ -643,7 +643,7 @@ def test_job_cancel(self):
         table = retry_403(Config.CLIENT.create_table)(table_arg)
         self.to_delete.insert(0, table)
 
-        job = Config.CLIENT.query(QUERY, job_id=JOB_ID)
+        job = Config.CLIENT.query(QUERY, job_id_prefix=JOB_ID_PREFIX)
         job.cancel()
 
         retry = RetryInstanceState(_job_done, max_tries=8)
@@ -866,7 +866,7 @@ def test_query_w_dml(self):
 
         query_job = Config.CLIENT.query(
             query_template.format(dataset_name, table_name),
-            job_id='test_query_w_dml_{}'.format(str(uuid.uuid4())))
+            job_id_prefix='test_query_w_dml_')
         query_job.result()
 
         self.assertEqual(query_job.num_dml_affected_rows, 1)
@@ -1053,8 +1053,7 @@ def test_query_w_query_params(self):
             query_job = Config.CLIENT.query(
                 example['sql'],
                 job_config=jconfig,
-                job_id='test_query_w_query_params{}'.format(
-                    str(uuid.uuid4())))
+                job_id_prefix='test_query_w_query_params')
             rows = list(query_job.result())
             self.assertEqual(len(rows), 1)
             self.assertEqual(len(rows[0]), 1)

bigquery/tests/unit/test_client.py

@@ -2859,6 +2859,35 @@ def test_list_partitions(self):
                          [20160804, 20160805])
 
 
+class Test_make_job_id(unittest.TestCase):
+    def _call_fut(self, job_id, prefix=None):
+        from google.cloud.bigquery.client import _make_job_id
+
+        return _make_job_id(job_id, prefix=prefix)
+
+    def test__make_job_id_wo_suffix(self):
+        job_id = self._call_fut('job_id')
+
+        self.assertEqual(job_id, 'job_id')
+
+    def test__make_job_id_w_suffix(self):
+        with mock.patch('uuid.uuid4', side_effect=['212345']):
+            job_id = self._call_fut(None, prefix='job_id')
+
+        self.assertEqual(job_id, 'job_id212345')
+
+    def test__make_random_job_id(self):
+        with mock.patch('uuid.uuid4', side_effect=['212345']):
+            job_id = self._call_fut(None)
+
+        self.assertEqual(job_id, '212345')
+
+    def test__make_job_id_w_job_id_overrides_prefix(self):
+        job_id = self._call_fut('job_id', prefix='unused_prefix')
+
+        self.assertEqual(job_id, 'job_id')
+
+
 class TestClientUpload(object):
     # NOTE: This is a "partner" to `TestClient` meant to test some of the
     #       "load_table_from_file" portions of `Client`. It also uses