Merge pull request #53 from bluelabsio/default_dataset_in_has_table

Prefer explicitly provided dataset over default dataset in lookup

Author: tswast

packages/sqlalchemy-bigquery/pybigquery/sqlalchemy_bigquery.py

@@ -357,7 +357,9 @@ def _split_table_name(self, full_table_name):
         project = None
 
         table_name_split = full_table_name.split('.')
-        if len(table_name_split) == 2:
+        if len(table_name_split) == 1:
+            table_name = full_table_name
+        elif len(table_name_split) == 2:
             dataset, table_name = table_name_split
         elif len(table_name_split) == 3:
             project, dataset, table_name = table_name_split
@@ -368,17 +370,12 @@ def _get_table(self, connection, table_name, schema=None):
         if isinstance(connection, Engine):
             connection = connection.connect()
 
-        table_name_prepared = dataset = project = None
-        if self.dataset_id:
-            table_name_prepared = table_name
-            dataset = self.dataset_id
-            project = None
-
-        else:
-            project, dataset, table_name_prepared = self._split_table_name(table_name)
-            if dataset is None and schema is not None:
-                table_name_prepared = table_name
+        project, dataset, table_name_prepared = self._split_table_name(table_name)
+        if dataset is None:
+            if schema is not None:
                 dataset = schema
+            elif self.dataset_id:
+                dataset = self.dataset_id
 
         table = connection.connection._client.dataset(dataset, project=project).table(table_name_prepared)
         try:

packages/sqlalchemy-bigquery/scripts/load_test_data.sh

@@ -1,12 +1,15 @@
 bq mk test_pybigquery
 bq mk --data_location=asia-northeast1 test_pybigquery_location
+bq mk test_pybigquery_alt
 
 bq rm -f -t test_pybigquery.sample
+bq rm -f -t test_pybigquery_alt.sample_alt
 bq rm -f -t test_pybigquery.sample_one_row
 bq rm -f -t test_pybigquery.sample_dml
 bq rm -f -t test_pybigquery_location.sample_one_row
 
 bq mk --table --schema=$(dirname $0)/schema.json --time_partitioning_field timestamp --clustering_fields integer,string test_pybigquery.sample
+bq mk --table --schema=$(dirname $0)/schema.json --time_partitioning_field timestamp --clustering_fields integer,string test_pybigquery_alt.sample_alt
 bq load --source_format=NEWLINE_DELIMITED_JSON --schema=$(dirname $0)/schema.json test_pybigquery.sample $(dirname $0)/sample.json
 
 bq mk --table --schema=$(dirname $0)/schema.json --time_partitioning_field timestamp --clustering_fields integer,string test_pybigquery.sample_one_row

packages/sqlalchemy-bigquery/test/test_sqlalchemy_bigquery.py

@@ -199,9 +199,12 @@ def test_engine_with_dataset(engine_using_test_dataset):
     rows = table_one_row.select().execute().fetchall()
     assert list(rows[0]) == ONE_ROW_CONTENTS_EXPANDED
 
-    # Table name shouldn't include dataset
-    with pytest.raises(Exception):
-        table_one_row = Table('test_pybigquery.sample_one_row', MetaData(bind=engine_using_test_dataset), autoload=True)
+    table_one_row = Table('test_pybigquery.sample_one_row', MetaData(bind=engine_using_test_dataset), autoload=True)
+    rows = table_one_row.select().execute().fetchall()
+    # verify that we are pulling from the specifically-named dataset,
+    # instead of pulling from the default dataset of the engine (which
+    # does not have this table at all)
+    assert list(rows[0]) == ONE_ROW_CONTENTS_EXPANDED
 
 
 def test_dataset_location(engine_with_location):
@@ -478,6 +481,14 @@ def test_has_table(engine, engine_using_test_dataset):
     assert engine.has_table('sample', 'test_pybigquery') is True
     assert engine.has_table('test_pybigquery.sample') is True
 
+    assert engine.has_table('sample_alt', 'test_pybigquery_alt') is True
+    assert engine.has_table('test_pybigquery_alt.sample_alt') is True
+
     assert engine_using_test_dataset.has_table('sample') is True
-    with pytest.raises(Exception):
-        assert engine_using_test_dataset.has_table('test_pybigquery.sample') is True
+    assert engine_using_test_dataset.has_table('sample', 'test_pybigquery') is True
+    assert engine_using_test_dataset.has_table('test_pybigquery.sample') is True
+
+    assert engine_using_test_dataset.has_table('sample_alt') is False
+
+    assert engine_using_test_dataset.has_table('sample_alt', 'test_pybigquery_alt') is True
+    assert engine_using_test_dataset.has_table('test_pybigquery_alt.sample_alt') is True