Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions sqlalchemy_bigquery/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -1124,6 +1124,13 @@ 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")
# breakpoint()
cursor.execute(statement, parameters, **kwargs)

def create_connect_args(self, url):
(
project_id,
Expand Down
21 changes: 21 additions & 0 deletions tests/unit/test_sqlalchemy_bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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