fix: query_and_wait now retains unknown query configuration _properties (#1793)

* fix: `query_and_wait` now retains unknown query configuration `_properties`

fix: raise `ValueError` in `query_and_wait` with wrong `job_config` type

Author: tswast

google/cloud/bigquery/_job_helpers.py

@@ -166,6 +166,14 @@ def do_query():
     return future
 
 
+def _validate_job_config(request_body: Dict[str, Any], invalid_key: str):
+    """Catch common mistakes, such as passing in a *JobConfig object of the
+    wrong type.
+    """
+    if invalid_key in request_body:
+        raise ValueError(f"got unexpected key {repr(invalid_key)} in job_config")
+
+
 def _to_query_request(
     job_config: Optional[job.QueryJobConfig] = None,
     *,
@@ -179,17 +187,15 @@ def _to_query_request(
     QueryRequest. If any configuration property is set that is not available in
     jobs.query, it will result in a server-side error.
     """
-    request_body = {}
-    job_config_resource = job_config.to_api_repr() if job_config else {}
-    query_config_resource = job_config_resource.get("query", {})
+    request_body = copy.copy(job_config.to_api_repr()) if job_config else {}
 
-    request_body.update(query_config_resource)
+    _validate_job_config(request_body, job.CopyJob._JOB_TYPE)
+    _validate_job_config(request_body, job.ExtractJob._JOB_TYPE)
+    _validate_job_config(request_body, job.LoadJob._JOB_TYPE)
 
-    # These keys are top level in job resource and query resource.
-    if "labels" in job_config_resource:
-        request_body["labels"] = job_config_resource["labels"]
-    if "dryRun" in job_config_resource:
-        request_body["dryRun"] = job_config_resource["dryRun"]
+    # Move query.* properties to top-level.
+    query_config_resource = request_body.pop("query", {})
+    request_body.update(query_config_resource)
 
     # Default to standard SQL.
     request_body.setdefault("useLegacySql", False)

tests/unit/test__job_helpers.py

@@ -23,6 +23,9 @@
 
 from google.cloud.bigquery.client import Client
 from google.cloud.bigquery import _job_helpers
+from google.cloud.bigquery.job import copy_ as job_copy
+from google.cloud.bigquery.job import extract as job_extract
+from google.cloud.bigquery.job import load as job_load
 from google.cloud.bigquery.job import query as job_query
 from google.cloud.bigquery.query import ConnectionProperty, ScalarQueryParameter
 
@@ -57,9 +60,34 @@ def make_query_response(
 @pytest.mark.parametrize(
     ("job_config", "expected"),
     (
-        (None, make_query_request()),
-        (job_query.QueryJobConfig(), make_query_request()),
-        (
+        pytest.param(
+            None,
+            make_query_request(),
+            id="job_config=None-default-request",
+        ),
+        pytest.param(
+            job_query.QueryJobConfig(),
+            make_query_request(),
+            id="job_config=QueryJobConfig()-default-request",
+        ),
+        pytest.param(
+            job_query.QueryJobConfig.from_api_repr(
+                {
+                    "unknownTopLevelProperty": "some-test-value",
+                    "query": {
+                        "unknownQueryProperty": "some-other-value",
+                    },
+                },
+            ),
+            make_query_request(
+                {
+                    "unknownTopLevelProperty": "some-test-value",
+                    "unknownQueryProperty": "some-other-value",
+                }
+            ),
+            id="job_config-with-unknown-properties-includes-all-properties-in-request",
+        ),
+        pytest.param(
             job_query.QueryJobConfig(default_dataset="my-project.my_dataset"),
             make_query_request(
                 {
@@ -69,17 +97,24 @@ def make_query_response(
                     }
                 }
             ),
+            id="job_config-with-default_dataset",
         ),
-        (job_query.QueryJobConfig(dry_run=True), make_query_request({"dryRun": True})),
-        (
+        pytest.param(
+            job_query.QueryJobConfig(dry_run=True),
+            make_query_request({"dryRun": True}),
+            id="job_config-with-dry_run",
+        ),
+        pytest.param(
             job_query.QueryJobConfig(use_query_cache=False),
             make_query_request({"useQueryCache": False}),
+            id="job_config-with-use_query_cache",
         ),
-        (
+        pytest.param(
             job_query.QueryJobConfig(use_legacy_sql=True),
             make_query_request({"useLegacySql": True}),
+            id="job_config-with-use_legacy_sql",
         ),
-        (
+        pytest.param(
             job_query.QueryJobConfig(
                 query_parameters=[
                     ScalarQueryParameter("named_param1", "STRING", "param-value"),
@@ -103,8 +138,9 @@ def make_query_response(
                     ],
                 }
             ),
+            id="job_config-with-query_parameters-named",
         ),
-        (
+        pytest.param(
             job_query.QueryJobConfig(
                 query_parameters=[
                     ScalarQueryParameter(None, "STRING", "param-value"),
@@ -126,8 +162,9 @@ def make_query_response(
                     ],
                 }
             ),
+            id="job_config-with-query_parameters-positional",
         ),
-        (
+        pytest.param(
             job_query.QueryJobConfig(
                 connection_properties=[
                     ConnectionProperty(key="time_zone", value="America/Chicago"),
@@ -142,14 +179,17 @@ def make_query_response(
                     ]
                 }
             ),
+            id="job_config-with-connection_properties",
         ),
-        (
+        pytest.param(
             job_query.QueryJobConfig(labels={"abc": "def"}),
             make_query_request({"labels": {"abc": "def"}}),
+            id="job_config-with-labels",
         ),
-        (
+        pytest.param(
             job_query.QueryJobConfig(maximum_bytes_billed=987654),
             make_query_request({"maximumBytesBilled": "987654"}),
+            id="job_config-with-maximum_bytes_billed",
         ),
     ),
 )
@@ -159,6 +199,19 @@ def test__to_query_request(job_config, expected):
     assert result == expected
 
 
+@pytest.mark.parametrize(
+    ("job_config", "invalid_key"),
+    (
+        pytest.param(job_copy.CopyJobConfig(), "copy", id="copy"),
+        pytest.param(job_extract.ExtractJobConfig(), "extract", id="extract"),
+        pytest.param(job_load.LoadJobConfig(), "load", id="load"),
+    ),
+)
+def test__to_query_request_raises_for_invalid_config(job_config, invalid_key):
+    with pytest.raises(ValueError, match=f"{repr(invalid_key)} in job_config"):
+        _job_helpers._to_query_request(job_config, query="SELECT 1")
+
+
 def test__to_query_job_defaults():
     mock_client = mock.create_autospec(Client)
     response = make_query_response(