Skip to content

Commit e41f360

Browse files
committed
Tidied up default_job_config creation
1 parent 2cd3161 commit e41f360

File tree

1 file changed

+16
-22
lines changed

1 file changed

+16
-22
lines changed

sqlalchemy_bigquery/base.py

Lines changed: 16 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1081,7 +1081,6 @@ def __init__(
10811081
self,
10821082
arraysize=5000,
10831083
credentials_path=None,
1084-
project_id=None,
10851084
billing_project_id=None,
10861085
location=None,
10871086
credentials_info=None,
@@ -1095,7 +1094,7 @@ def __init__(
10951094
self.credentials_path = credentials_path
10961095
self.credentials_info = credentials_info
10971096
self.credentials_base64 = credentials_base64
1098-
self.project_id = project_id
1097+
self.project_id = None
10991098
self.billing_project_id = billing_project_id
11001099
self.location = location
11011100
self.identifier_preparer = self.preparer(self)
@@ -1119,18 +1118,20 @@ def _build_formatted_table_id(table):
11191118
"""Build '<dataset_id>.<table_id>' string using given table."""
11201119
return "{}.{}".format(table.reference.dataset_id, table.table_id)
11211120

1122-
@staticmethod
1123-
def _add_default_dataset_to_job_config(job_config: QueryJobConfig, project_id, dataset_id):
1124-
# If dataset_id is set, then we know the job_config isn't None
1125-
if project_id:
1126-
dataset = ConnectionProperty(key='dataset_project_id', value=project_id)
1127-
job_config.connection_properties = [ dataset ]
1128-
if dataset_id:
1129-
# If project_id is missing, use default project_id for the current environment
1121+
1122+
def create_job_config(self, provided_config: QueryJobConfig):
1123+
project_id = self.project_id
1124+
if self.dataset_id is None and project_id == self.billing_project_id:
1125+
return provided_config
1126+
job_config = provided_config or QueryJobConfig()
1127+
if project_id != self.billing_project_id:
1128+
job_config.connection_properties = [ConnectionProperty(key="dataset_project_id", value=project_id)]
1129+
if self.dataset_id:
11301130
if not project_id:
11311131
_, project_id = auth.default()
1132+
job_config.default_dataset = "{}.{}".format(project_id, self.dataset_id)
1133+
return job_config
11321134

1133-
job_config.default_dataset = "{}.{}".format(project_id, dataset_id)
11341135

11351136
def do_execute(self, cursor, statement, parameters, context=None):
11361137
kwargs = {}
@@ -1140,13 +1141,13 @@ def do_execute(self, cursor, statement, parameters, context=None):
11401141

11411142
def create_connect_args(self, url):
11421143
(
1143-
project_id,
1144+
self.project_id,
11441145
location,
11451146
dataset_id,
11461147
arraysize,
11471148
credentials_path,
11481149
credentials_base64,
1149-
default_query_job_config,
1150+
provided_job_config,
11501151
list_tables_page_size,
11511152
user_supplied_client,
11521153
) = parse_url(url)
@@ -1157,17 +1158,10 @@ def create_connect_args(self, url):
11571158
self.credentials_path = credentials_path or self.credentials_path
11581159
self.credentials_base64 = credentials_base64 or self.credentials_base64
11591160
self.dataset_id = dataset_id
1160-
from google.cloud.bigquery import QueryJobConfig
1161-
if default_query_job_config is None:
1162-
default_query_job_config = QueryJobConfig()
11631161
if self.billing_project_id is None:
1164-
self.billing_project_id = project_id
1165-
if self.project_id is None:
1166-
self.project_id = project_id
1162+
self.billing_project_id = self.project_id
11671163

1168-
self._add_default_dataset_to_job_config(
1169-
default_query_job_config, project_id, dataset_id
1170-
)
1164+
default_query_job_config = self.create_job_config(provided_job_config)
11711165

11721166
if user_supplied_client:
11731167
# The user is expected to supply a client with

0 commit comments

Comments
 (0)