feat(bigquery): add create job method (#32)

* feat(bigquery): add create job method * feat(bigquery): Addressed comments and add unit test * feat(bigquery): make copy of job config for query job Co-authored-by: Peter Lamut <plamut@users.noreply.github.com>

Author: HemangChothani

google/cloud/bigquery/client.py

@@ -51,9 +51,11 @@
 from google.cloud import exceptions
 from google.cloud.client import ClientWithProject
 
+from google.cloud.bigquery._helpers import _get_sub_prop
 from google.cloud.bigquery._helpers import _record_field_to_json
 from google.cloud.bigquery._helpers import _str_or_none
 from google.cloud.bigquery._helpers import _verify_job_config_type
+from google.cloud.bigquery._helpers import _del_sub_prop
 from google.cloud.bigquery._http import Connection
 from google.cloud.bigquery import _pandas_helpers
 from google.cloud.bigquery.dataset import Dataset
@@ -1313,6 +1315,70 @@ def job_from_resource(self, resource):
  return job.QueryJob.from_api_repr(resource, self)
  return job.UnknownJob.from_api_repr(resource, self)
 
+ def create_job(self, job_config, retry=DEFAULT_RETRY):
+ """Create a new job.
+ Arguments:
+ job_config (dict): configuration job representation returned from the API.
+
+ Keyword Arguments:
+ retry (google.api_core.retry.Retry):
+ (Optional) How to retry the RPC.
+
+ Returns:
+ Union[ \
+ google.cloud.bigquery.job.LoadJob, \
+ google.cloud.bigquery.job.CopyJob, \
+ google.cloud.bigquery.job.ExtractJob, \
+ google.cloud.bigquery.job.QueryJob \
+ ]:
+ A new job instance.
+ """
+
+ if "load" in job_config:
+ load_job_config = google.cloud.bigquery.job.LoadJobConfig.from_api_repr(
+ job_config
+ )
+ destination = _get_sub_prop(job_config, ["load", "destinationTable"])
+ source_uris = _get_sub_prop(job_config, ["load", "sourceUris"])
+ return self.load_table_from_uri(
+ source_uris, destination, job_config=load_job_config, retry=retry
+ )
+ elif "copy" in job_config:
+ copy_job_config = google.cloud.bigquery.job.CopyJobConfig.from_api_repr(
+ job_config
+ )
+ destination = _get_sub_prop(job_config, ["copy", "destinationTable"])
+ sources = []
+ source_configs = _get_sub_prop(job_config, ["copy", "sourceTables"])
+
+ if source_configs is None:
+ source_configs = [_get_sub_prop(job_config, ["copy", "sourceTable"])]
+ for source_config in source_configs:
+ table_ref = TableReference.from_api_repr(source_config)
+ sources.append(table_ref)
+ return self.copy_table(
+ sources, destination, job_config=copy_job_config, retry=retry
+ )
+ elif "extract" in job_config:
+ extract_job_config = google.cloud.bigquery.job.ExtractJobConfig.from_api_repr(
+ job_config
+ )
+ source = _get_sub_prop(job_config, ["extract", "sourceTable"])
+ destination_uris = _get_sub_prop(job_config, ["extract", "destinationUris"])
+ return self.extract_table(
+ source, destination_uris, job_config=extract_job_config, retry=retry
+ )
+ elif "query" in job_config:
+ copy_config = copy.deepcopy(job_config)
+ _del_sub_prop(copy_config, ["query", "destinationTable"])
+ query_job_config = google.cloud.bigquery.job.QueryJobConfig.from_api_repr(
+ copy_config
+ )
+ query = _get_sub_prop(copy_config, ["query", "query"])
+ return self.query(query, job_config=query_job_config, retry=retry)
+ else:
+ raise TypeError("Invalid job configuration received.")
+
  def get_job(
  self, job_id, project=None, location=None, retry=DEFAULT_RETRY, timeout=None
  ):

tests/unit/test_client.py

@@ -2795,6 +2795,176 @@ def test_delete_table_w_not_found_ok_true(self):
 
  conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None)
 
+ def _create_job_helper(self, job_config, client_method):
+ creds = _make_credentials()
+ http = object()
+ client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+
+ client._connection = make_connection()
+ rf1 = mock.Mock()
+ get_config_patch = mock.patch(
+ "google.cloud.bigquery.job._JobConfig.from_api_repr", return_value=rf1,
+ )
+ load_patch = mock.patch(client_method, autospec=True)
+
+ with load_patch as client_method, get_config_patch:
+ client.create_job(job_config=job_config)
+ client_method.assert_called_once()
+
+ def test_create_job_load_config(self):
+ configuration = {
+ "load": {
+ "destinationTable": {
+ "projectId": self.PROJECT,
+ "datasetId": self.DS_ID,
+ "tableId": "source_table",
+ },
+ "sourceUris": ["gs://test_bucket/src_object*"],
+ }
+ }
+
+ self._create_job_helper(
+ configuration, "google.cloud.bigquery.client.Client.load_table_from_uri"
+ )
+
+ def test_create_job_copy_config(self):
+ configuration = {
+ "copy": {
+ "sourceTables": [
+ {
+ "projectId": self.PROJECT,
+ "datasetId": self.DS_ID,
+ "tableId": "source_table",
+ }
+ ],
+ "destinationTable": {
+ "projectId": self.PROJECT,
+ "datasetId": self.DS_ID,
+ "tableId": "destination_table",
+ },
+ }
+ }
+
+ self._create_job_helper(
+ configuration, "google.cloud.bigquery.client.Client.copy_table",
+ )
+
+ def test_create_job_copy_config_w_single_source(self):
+ configuration = {
+ "copy": {
+ "sourceTable": {
+ "projectId": self.PROJECT,
+ "datasetId": self.DS_ID,
+ "tableId": "source_table",
+ },
+ "destinationTable": {
+ "projectId": self.PROJECT,
+ "datasetId": self.DS_ID,
+ "tableId": "destination_table",
+ },
+ }
+ }
+
+ self._create_job_helper(
+ configuration, "google.cloud.bigquery.client.Client.copy_table",
+ )
+
+ def test_create_job_extract_config(self):
+ configuration = {
+ "extract": {
+ "sourceTable": {
+ "projectId": self.PROJECT,
+ "datasetId": self.DS_ID,
+ "tableId": "source_table",
+ },
+ "destinationUris": ["gs://test_bucket/dst_object*"],
+ }
+ }
+ self._create_job_helper(
+ configuration, "google.cloud.bigquery.client.Client.extract_table",
+ )
+
+ def test_create_job_query_config(self):
+ configuration = {
+ "query": {"query": "query", "destinationTable": {"tableId": "table_id"}}
+ }
+ self._create_job_helper(
+ configuration, "google.cloud.bigquery.client.Client.query",
+ )
+
+ def test_create_job_query_config_w_rateLimitExceeded_error(self):
+ from google.cloud.exceptions import Forbidden
+ from google.cloud.bigquery.retry import DEFAULT_RETRY
+
+ query = "select count(*) from persons"
+ configuration = {
+ "query": {
+ "query": query,
+ "useLegacySql": False,
+ "destinationTable": {"tableId": "table_id"},
+ }
+ }
+ resource = {
+ "jobReference": {"projectId": self.PROJECT, "jobId": mock.ANY},
+ "configuration": {
+ "query": {
+ "query": query,
+ "useLegacySql": False,
+ "destinationTable": {
+ "projectId": self.PROJECT,
+ "datasetId": self.DS_ID,
+ "tableId": "query_destination_table",
+ },
+ }
+ },
+ }
+ data_without_destination = {
+ "jobReference": {"projectId": self.PROJECT, "jobId": mock.ANY},
+ "configuration": {"query": {"query": query, "useLegacySql": False}},
+ }
+
+ creds = _make_credentials()
+ http = object()
+ retry = DEFAULT_RETRY.with_deadline(1).with_predicate(
+ lambda exc: isinstance(exc, Forbidden)
+ )
+ client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+
+ api_request_patcher = mock.patch.object(
+ client._connection,
+ "api_request",
+ side_effect=[
+ Forbidden("", errors=[{"reason": "rateLimitExceeded"}]),
+ resource,
+ ],
+ )
+
+ with api_request_patcher as fake_api_request:
+ job = client.create_job(job_config=configuration, retry=retry)
+
+ self.assertEqual(job.destination.table_id, "query_destination_table")
+ self.assertEqual(len(fake_api_request.call_args_list), 2) # was retried once
+ self.assertEqual(
+ fake_api_request.call_args_list[1],
+ mock.call(
+ method="POST",
+ path="/projects/PROJECT/jobs",
+ data=data_without_destination,
+ timeout=None,
+ ),
+ )
+
+ def test_create_job_w_invalid_job_config(self):
+ configuration = {"unknown": {}}
+ creds = _make_credentials()
+ http = object()
+ client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
+
+ with self.assertRaises(TypeError) as exc:
+ client.create_job(job_config=configuration)
+
+ self.assertIn("Invalid job configuration", exc.exception.args[0])
+
  def test_job_from_resource_unknown_type(self):
  from google.cloud.bigquery.job import UnknownJob