Skip to content

Commit 155bacc

Browse files
fix: create_job method accepts dictionary arguments (#300)
* fix: broken create_job method * fix: changes in unit tests * fix: fix sourceTable thing * fix: handle sourceTable passed in job resource * fix: remove delete destination table from query * fix: revert destination table for query
1 parent 82290c3 commit 155bacc

File tree

2 files changed

+36
-33
lines changed

2 files changed

+36
-33
lines changed

google/cloud/bigquery/client.py

+8-4
Original file line numberDiff line numberDiff line change
@@ -48,11 +48,11 @@
4848
from google.cloud import exceptions
4949
from google.cloud.client import ClientWithProject
5050

51+
from google.cloud.bigquery._helpers import _del_sub_prop
5152
from google.cloud.bigquery._helpers import _get_sub_prop
5253
from google.cloud.bigquery._helpers import _record_field_to_json
5354
from google.cloud.bigquery._helpers import _str_or_none
5455
from google.cloud.bigquery._helpers import _verify_job_config_type
55-
from google.cloud.bigquery._helpers import _del_sub_prop
5656
from google.cloud.bigquery._http import Connection
5757
from google.cloud.bigquery import _pandas_helpers
5858
from google.cloud.bigquery.dataset import Dataset
@@ -1619,6 +1619,7 @@ def create_job(self, job_config, retry=DEFAULT_RETRY, timeout=None):
16191619
)
16201620
destination = _get_sub_prop(job_config, ["load", "destinationTable"])
16211621
source_uris = _get_sub_prop(job_config, ["load", "sourceUris"])
1622+
destination = TableReference.from_api_repr(destination)
16221623
return self.load_table_from_uri(
16231624
source_uris,
16241625
destination,
@@ -1631,9 +1632,9 @@ def create_job(self, job_config, retry=DEFAULT_RETRY, timeout=None):
16311632
job_config
16321633
)
16331634
destination = _get_sub_prop(job_config, ["copy", "destinationTable"])
1635+
destination = TableReference.from_api_repr(destination)
16341636
sources = []
16351637
source_configs = _get_sub_prop(job_config, ["copy", "sourceTables"])
1636-
16371638
if source_configs is None:
16381639
source_configs = [_get_sub_prop(job_config, ["copy", "sourceTable"])]
16391640
for source_config in source_configs:
@@ -1651,10 +1652,13 @@ def create_job(self, job_config, retry=DEFAULT_RETRY, timeout=None):
16511652
job_config
16521653
)
16531654
source = _get_sub_prop(job_config, ["extract", "sourceTable"])
1654-
source_type = "Table"
1655-
if not source:
1655+
if source:
1656+
source_type = "Table"
1657+
source = TableReference.from_api_repr(source)
1658+
else:
16561659
source = _get_sub_prop(job_config, ["extract", "sourceModel"])
16571660
source_type = "Model"
1661+
source = ModelReference.from_api_repr(source)
16581662
destination_uris = _get_sub_prop(job_config, ["extract", "destinationUris"])
16591663
return self.extract_table(
16601664
source,

tests/unit/test_client.py

+28-29
Original file line numberDiff line numberDiff line change
@@ -3573,21 +3573,28 @@ def test_delete_table_w_not_found_ok_true(self):
35733573

35743574
conn.api_request.assert_called_with(method="DELETE", path=path, timeout=None)
35753575

3576-
def _create_job_helper(self, job_config, client_method):
3576+
def _create_job_helper(self, job_config):
3577+
from google.cloud.bigquery import _helpers
3578+
35773579
creds = _make_credentials()
35783580
http = object()
35793581
client = self._make_one(project=self.PROJECT, credentials=creds, _http=http)
35803582

3581-
client._connection = make_connection()
3582-
rf1 = mock.Mock()
3583-
get_config_patch = mock.patch(
3584-
"google.cloud.bigquery.job._JobConfig.from_api_repr", return_value=rf1,
3585-
)
3586-
load_patch = mock.patch(client_method, autospec=True)
3583+
RESOURCE = {
3584+
"jobReference": {"projectId": self.PROJECT, "jobId": mock.ANY},
3585+
"configuration": job_config,
3586+
}
3587+
conn = client._connection = make_connection(RESOURCE)
3588+
client.create_job(job_config=job_config)
3589+
if "query" in job_config:
3590+
_helpers._del_sub_prop(job_config, ["query", "destinationTable"])
35873591

3588-
with load_patch as client_method, get_config_patch:
3589-
client.create_job(job_config=job_config)
3590-
client_method.assert_called_once()
3592+
conn.api_request.assert_called_once_with(
3593+
method="POST",
3594+
path="/projects/%s/jobs" % self.PROJECT,
3595+
data=RESOURCE,
3596+
timeout=None,
3597+
)
35913598

35923599
def test_create_job_load_config(self):
35933600
configuration = {
@@ -3601,9 +3608,7 @@ def test_create_job_load_config(self):
36013608
}
36023609
}
36033610

3604-
self._create_job_helper(
3605-
configuration, "google.cloud.bigquery.client.Client.load_table_from_uri"
3606-
)
3611+
self._create_job_helper(configuration)
36073612

36083613
def test_create_job_copy_config(self):
36093614
configuration = {
@@ -3623,9 +3628,7 @@ def test_create_job_copy_config(self):
36233628
}
36243629
}
36253630

3626-
self._create_job_helper(
3627-
configuration, "google.cloud.bigquery.client.Client.copy_table",
3628-
)
3631+
self._create_job_helper(configuration)
36293632

36303633
def test_create_job_copy_config_w_single_source(self):
36313634
configuration = {
@@ -3643,9 +3646,7 @@ def test_create_job_copy_config_w_single_source(self):
36433646
}
36443647
}
36453648

3646-
self._create_job_helper(
3647-
configuration, "google.cloud.bigquery.client.Client.copy_table",
3648-
)
3649+
self._create_job_helper(configuration)
36493650

36503651
def test_create_job_extract_config(self):
36513652
configuration = {
@@ -3658,9 +3659,7 @@ def test_create_job_extract_config(self):
36583659
"destinationUris": ["gs://test_bucket/dst_object*"],
36593660
}
36603661
}
3661-
self._create_job_helper(
3662-
configuration, "google.cloud.bigquery.client.Client.extract_table",
3663-
)
3662+
self._create_job_helper(configuration)
36643663

36653664
def test_create_job_extract_config_for_model(self):
36663665
configuration = {
@@ -3673,17 +3672,17 @@ def test_create_job_extract_config_for_model(self):
36733672
"destinationUris": ["gs://test_bucket/dst_object*"],
36743673
}
36753674
}
3676-
self._create_job_helper(
3677-
configuration, "google.cloud.bigquery.client.Client.extract_table",
3678-
)
3675+
self._create_job_helper(configuration)
36793676

36803677
def test_create_job_query_config(self):
36813678
configuration = {
3682-
"query": {"query": "query", "destinationTable": {"tableId": "table_id"}}
3679+
"query": {
3680+
"query": "query",
3681+
"destinationTable": {"tableId": "table_id"},
3682+
"useLegacySql": False,
3683+
}
36833684
}
3684-
self._create_job_helper(
3685-
configuration, "google.cloud.bigquery.client.Client.query",
3686-
)
3685+
self._create_job_helper(configuration)
36873686

36883687
def test_create_job_query_config_w_rateLimitExceeded_error(self):
36893688
from google.cloud.exceptions import Forbidden

0 commit comments

Comments
 (0)