diff --git a/sqlalchemy_bigquery/base.py b/sqlalchemy_bigquery/base.py index 7f4102fa..0204bc92 100644 --- a/sqlalchemy_bigquery/base.py +++ b/sqlalchemy_bigquery/base.py @@ -1124,6 +1124,12 @@ def _add_default_dataset_to_job_config(job_config, project_id, dataset_id): job_config.default_dataset = "{}.{}".format(project_id, dataset_id) + def do_execute(self, cursor, statement, parameters, context=None): + kwargs = {} + if context is not None and context.execution_options.get("job_config"): + kwargs["job_config"] = context.execution_options.get("job_config") + cursor.execute(statement, parameters, **kwargs) + def create_connect_args(self, url): ( project_id, diff --git a/tests/unit/test_sqlalchemy_bigquery.py b/tests/unit/test_sqlalchemy_bigquery.py index db20e2f0..85408aef 100644 --- a/tests/unit/test_sqlalchemy_bigquery.py +++ b/tests/unit/test_sqlalchemy_bigquery.py @@ -240,3 +240,24 @@ def test_setting_user_supplied_client_skips_creating_client( ) assert result == ([], {}) assert not mock_create_bigquery_client.called + + +def test_do_execute(): + # Ensures the do_execute() method overrides that of the parent class. + import sqlalchemy_bigquery # noqa + from sqlalchemy_bigquery.base import BigQueryExecutionContext + + job_config_kwargs = {} + job_config_kwargs["use_query_cache"] = False + job_config = bigquery.QueryJobConfig(**job_config_kwargs) + execution_options = {"job_config": job_config} + context = mock.MagicMock(spec=BigQueryExecutionContext) + type(context).execution_options = mock.PropertyMock(return_value=execution_options) + + cursor = mock.MagicMock() + + sqlalchemy_bigquery.BigQueryDialect().do_execute( + cursor, sqlalchemy.text("SELECT 'a' AS `1`"), mock.MagicMock(), context=context + ) + + assert cursor.execute.call_args.kwargs["job_config"] is job_config