Skip to content

Commit c34ab1c

Browse files
Prefer explicitly provided dataset over default dataset in lookup
Fixes #51
1 parent c3baf5f commit c34ab1c

File tree

3 files changed

+22
-16
lines changed

3 files changed

+22
-16
lines changed

pybigquery/sqlalchemy_bigquery.py

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -356,7 +356,9 @@ def _split_table_name(self, full_table_name):
356356
project = None
357357

358358
table_name_split = full_table_name.split('.')
359-
if len(table_name_split) == 2:
359+
if len(table_name_split) == 1:
360+
table_name = full_table_name
361+
elif len(table_name_split) == 2:
360362
dataset, table_name = table_name_split
361363
elif len(table_name_split) == 3:
362364
project, dataset, table_name = table_name_split
@@ -367,17 +369,12 @@ def _get_table(self, connection, table_name, schema=None):
367369
if isinstance(connection, Engine):
368370
connection = connection.connect()
369371

370-
table_name_prepared = dataset = project = None
371-
if self.dataset_id:
372-
table_name_prepared = table_name
373-
dataset = self.dataset_id
374-
project = None
375-
376-
else:
377-
project, dataset, table_name_prepared = self._split_table_name(table_name)
378-
if dataset is None and schema is not None:
379-
table_name_prepared = table_name
372+
project, dataset, table_name_prepared = self._split_table_name(table_name)
373+
if dataset is None:
374+
if schema is not None:
380375
dataset = schema
376+
elif self.dataset_id:
377+
dataset = self.dataset_id
381378

382379
table = connection.connection._client.dataset(dataset, project=project).table(table_name_prepared)
383380
try:

scripts/load_test_data.sh

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
11
bq mk test_pybigquery
22
bq mk --data_location=asia-northeast1 test_pybigquery_location
3+
bq mk test_pybigquery_alt
34

45
bq rm -f -t test_pybigquery.sample
6+
bq rm -f -t test_pybigquery_alt.sample_alt
57
bq rm -f -t test_pybigquery.sample_one_row
68
bq rm -f -t test_pybigquery.sample_dml
79
bq rm -f -t test_pybigquery_location.sample_one_row
810

911
bq mk --table --schema=$(dirname $0)/schema.json --time_partitioning_field timestamp --clustering_fields integer,string test_pybigquery.sample
12+
bq mk --table --schema=$(dirname $0)/schema.json --time_partitioning_field timestamp --clustering_fields integer,string test_pybigquery_alt.sample_alt
1013
bq load --source_format=NEWLINE_DELIMITED_JSON --schema=$(dirname $0)/schema.json test_pybigquery.sample $(dirname $0)/sample.json
1114

1215
bq mk --table --schema=$(dirname $0)/schema.json --time_partitioning_field timestamp --clustering_fields integer,string test_pybigquery.sample_one_row

test/test_sqlalchemy_bigquery.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -199,9 +199,7 @@ def test_engine_with_dataset(engine_using_test_dataset):
199199
rows = table_one_row.select().execute().fetchall()
200200
assert list(rows[0]) == ONE_ROW_CONTENTS_EXPANDED
201201

202-
# Table name shouldn't include dataset
203-
with pytest.raises(Exception):
204-
table_one_row = Table('test_pybigquery.sample_one_row', MetaData(bind=engine_using_test_dataset), autoload=True)
202+
table_one_row = Table('test_pybigquery.sample_one_row', MetaData(bind=engine_using_test_dataset), autoload=True)
205203

206204

207205
def test_dataset_location(engine_with_location):
@@ -478,6 +476,14 @@ def test_has_table(engine, engine_using_test_dataset):
478476
assert engine.has_table('sample', 'test_pybigquery') is True
479477
assert engine.has_table('test_pybigquery.sample') is True
480478

479+
assert engine.has_table('sample_alt', 'test_pybigquery_alt') is True
480+
assert engine.has_table('test_pybigquery_alt.sample_alt') is True
481+
481482
assert engine_using_test_dataset.has_table('sample') is True
482-
with pytest.raises(Exception):
483-
assert engine_using_test_dataset.has_table('test_pybigquery.sample') is True
483+
assert engine_using_test_dataset.has_table('sample', 'test_pybigquery') is True
484+
assert engine_using_test_dataset.has_table('test_pybigquery.sample') is True
485+
486+
assert engine_using_test_dataset.has_table('sample_alt') is False
487+
488+
assert engine_using_test_dataset.has_table('sample_alt', 'test_pybigquery_alt') is True
489+
assert engine_using_test_dataset.has_table('test_pybigquery_alt.sample_alt') is True

0 commit comments

Comments
 (0)