diff --git a/bigframes/session/clients.py b/bigframes/session/clients.py index a8e1ab71f1..86312eb9ba 100644 --- a/bigframes/session/clients.py +++ b/bigframes/session/clients.py @@ -181,6 +181,13 @@ def _create_bigquery_client(self): location=self._location, ) + # If a new enough client library is available, we opt-in to the faster + # backend behavior. This only affects code paths where query_and_wait is + # used, which doesn't expose a query job directly. See internal issue + # b/417985981. + if hasattr(bq_client, "default_job_creation_mode"): + bq_client.default_job_creation_mode = "JOB_CREATION_OPTIONAL" + if self._bq_kms_key_name: # Note: Key configuration only applies automatically to load and query jobs, not copy jobs. encryption_config = bigquery.EncryptionConfiguration( diff --git a/tests/unit/session/test_clients.py b/tests/unit/session/test_clients.py index c9a12be584..5d577a52ed 100644 --- a/tests/unit/session/test_clients.py +++ b/tests/unit/session/test_clients.py @@ -46,6 +46,8 @@ def create_clients_provider(application_name: Optional[str] = None): def monkeypatch_client_constructors(monkeypatch): bqclient = mock.create_autospec(google.cloud.bigquery.Client) bqclient.return_value = bqclient + # Assume we have a new client library in the unit tests. + bqclient.default_job_creation_mode = None # type: ignore monkeypatch.setattr(google.cloud.bigquery, "Client", bqclient) bqconnectionclient = mock.create_autospec( @@ -83,6 +85,11 @@ def monkeypatch_client_constructors(monkeypatch): ) +def assert_bqclient_sets_default_job_creation_mode(provider: clients.ClientsProvider): + bqclient = provider.bqclient + assert bqclient.default_job_creation_mode == "JOB_CREATION_OPTIONAL" + + def assert_constructed_w_user_agent(mock_client: mock.Mock, expected_user_agent: str): assert ( expected_user_agent