From 9e4823855de45f2853e44b30c336473e8105064d Mon Sep 17 00:00:00 2001 From: Linchin Date: Thu, 27 Feb 2025 15:04:56 -0800 Subject: [PATCH 1/3] feat: support job_cofig at query level --- sqlalchemy_bigquery/base.py | 6 ++++++ 1 file changed, 6 insertions(+) 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, From 81ba45c97a9b07a3977b5c6295972743705e4b85 Mon Sep 17 00:00:00 2001 From: Linchin Date: Fri, 7 Mar 2025 14:27:08 -0800 Subject: [PATCH 2/3] add unit test --- sqlalchemy_bigquery/base.py | 1 + tests/unit/test_sqlalchemy_bigquery.py | 21 +++++++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/sqlalchemy_bigquery/base.py b/sqlalchemy_bigquery/base.py index 0204bc92..5485ccca 100644 --- a/sqlalchemy_bigquery/base.py +++ b/sqlalchemy_bigquery/base.py @@ -1128,6 +1128,7 @@ 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") + # breakpoint() cursor.execute(statement, parameters, **kwargs) def create_connect_args(self, url): 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 From de7405ef26e663dc89cb6e9c5e94be695b5a4e5d Mon Sep 17 00:00:00 2001 From: Linchin Date: Mon, 10 Mar 2025 10:42:27 -0700 Subject: [PATCH 3/3] remove breakpoint --- sqlalchemy_bigquery/base.py | 1 - 1 file changed, 1 deletion(-) diff --git a/sqlalchemy_bigquery/base.py b/sqlalchemy_bigquery/base.py index 5485ccca..0204bc92 100644 --- a/sqlalchemy_bigquery/base.py +++ b/sqlalchemy_bigquery/base.py @@ -1128,7 +1128,6 @@ 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") - # breakpoint() cursor.execute(statement, parameters, **kwargs) def create_connect_args(self, url):