googleapis · mxmzdlv · Dec 12, 2017 · Dec 5, 2017 · Dec 9, 2017 · Dec 12, 2017
diff --git a/dev_requirements.txt b/dev_requirements.txt
@@ -1,5 +1,5 @@
 sqlalchemy>=1.1.9
-google-cloud-bigquery>=0.27.0
+google-cloud-bigquery>=0.28.0
 future==0.16.0
 
 pytest==3.2.2

diff --git a/pybigquery/sqlalchemy_bigquery.py b/pybigquery/sqlalchemy_bigquery.py
@@ -5,11 +5,12 @@
 
 from google.cloud.bigquery import dbapi
 from google.cloud import bigquery
-from google.api.core.exceptions import NotFound
+from google.api_core.exceptions import NotFound
 from sqlalchemy.exc import NoSuchTableError
 from sqlalchemy import types, util
 from sqlalchemy.sql.compiler import SQLCompiler, IdentifierPreparer
 from sqlalchemy.engine.default import DefaultDialect, DefaultExecutionContext
+from sqlalchemy.engine.base import Engine
 
 
 class UniversalSet(object):
@@ -118,24 +119,31 @@ def _split_table_name(self, full_table_name):
 
         return (project, dataset, table_name)
 
-    def _get_table(self, connection, table_name):
-        project, dataset, table_name = self._split_table_name(table_name)
-        table = connection.connection._client.dataset(dataset, project=project).table(table_name)
+    def _get_table(self, connection, table_name, schema=None):
+        if isinstance(connection, Engine):
+            connection = connection.connect()
+
+        project, dataset, table_name_prepared = self._split_table_name(table_name)
+        if dataset is None and schema is not None:
+            dataset = schema
+            table_name_prepared = table_name
+
+        table = connection.connection._client.dataset(dataset, project=project).table(table_name_prepared)
         try:
-            table.reload()
+            t = connection.connection._client.get_table(table)
         except NotFound as e:
             raise NoSuchTableError(table_name)
-        return table
+        return t
 
     def has_table(self, connection, table_name, schema=None):
         try:
-            self._get_table(connection, table_name)
+            self._get_table(connection, table_name, schema)
             return True
         except NoSuchTableError:
             return False
 
     def get_columns(self, connection, table_name, schema=None, **kw):
-        table = self._get_table(connection, table_name)
+        table = self._get_table(connection, table_name, schema)
         columns = table.schema
         result = []
         for col in columns:
@@ -165,13 +173,25 @@ def get_indexes(self, connection, table_name, schema=None, **kw):
         # BigQuery has no support for indexes.
         return []
 
+    def get_schema_names(self, connection, **kw):
+        if isinstance(connection, Engine):
+            connection = connection.connect()
+
+        datasets = connection.connection._client.list_datasets()
+        return [d.dataset_id for d in datasets]
+
     def get_table_names(self, connection, schema=None, **kw):
+        if isinstance(connection, Engine):
+            connection = connection.connect()
+
         datasets = connection.connection._client.list_datasets()
         result = []
         for d in datasets:
-            tables = d.list_tables()
+            if schema is not None and d.dataset_id != schema:
+                continue
+            tables = connection.connection._client.list_dataset_tables(d)
             for t in tables:
-                result.append(d.name + '.' + t.name)
+                result.append(d.dataset_id + '.' + t.table_id)
         return result
 
     def do_rollback(self, dbapi_connection):
@@ -185,4 +205,3 @@ def _check_unicode_returns(self, connection, additional_tests=None):
     def _check_unicode_description(self, connection):
         # requests gives back Unicode strings
         return True
-
diff --git a/setup.py b/setup.py
@@ -18,7 +18,7 @@
     ],
     install_requires=[
         'sqlalchemy>=1.1.9',
-        'google-cloud-bigquery>=0.27.0',
+        'google-cloud-bigquery>=0.28.0',
         'future',
     ],
     tests_require=[

diff --git a/test/test_sqlalchemy_bigquery.py b/test/test_sqlalchemy_bigquery.py
@@ -3,7 +3,7 @@
 
 from sqlalchemy.engine import create_engine
 from sqlalchemy.schema import Table, MetaData, Column
-from sqlalchemy import types, func, case
+from sqlalchemy import types, func, case, inspect
 from sqlalchemy.sql import expression, select, literal_column
 from sqlalchemy.exc import NoSuchTableError
 from sqlalchemy.orm import sessionmaker
@@ -38,6 +38,18 @@
     'test_bytes'
 ]
 
+SAMPLE_COLUMNS = [
+    {'name': 'integer', 'type': types.Integer(), 'nullable': True, 'default': None},
+    {'name': 'timestamp', 'type': types.TIMESTAMP(), 'nullable': True, 'default': None},
+    {'name': 'string', 'type': types.String(), 'nullable': True, 'default': None},
+    {'name': 'float', 'type': types.Float(), 'nullable': True, 'default': None},
+    {'name': 'boolean', 'type': types.Boolean(), 'nullable': True, 'default': None},
+    {'name': 'date', 'type': types.DATE(), 'nullable': True, 'default': None},
+    {'name': 'datetime', 'type': types.DATETIME(), 'nullable': True, 'default': None},
+    {'name': 'time', 'type': types.TIME(), 'nullable': True, 'default': None},
+    {'name': 'bytes', 'type': types.BINARY(), 'nullable': True, 'default': None}
+]
+
 
 @pytest.fixture(scope='session')
 def engine():
@@ -67,6 +79,11 @@ def session(engine):
     return session
 
 
+@pytest.fixture(scope='session')
+def inspector(engine):
+    return inspect(engine)
+
+
 @pytest.fixture(scope='session')
 def query(table):
     col1 = literal_column("TIMESTAMP_TRUNC(timestamp, DAY)").label("timestamp_label")
@@ -138,8 +155,10 @@ def test_reflect_dataset_does_not_exist(engine):
 
 
 def test_tables_list(engine):
-    assert 'test_pybigquery.sample' in engine.table_names()
-    assert 'test_pybigquery.sample_one_row' in engine.table_names()
+    tables = engine.table_names()
+    assert 'test_pybigquery.sample' in tables
+    assert 'test_pybigquery.sample_one_row' in tables
+    assert 'test_pybigquery.sample_dml' in tables
 
 
 def test_group_by(session, table):
@@ -208,3 +227,34 @@ def test_dml(engine, session, table_dml):
     session.query(table_dml).filter(table_dml.c.string == 'updated_row').delete(synchronize_session=False)
     result = table_dml.select().execute().fetchall()
     assert len(result) == 0
+
+
+def test_schemas_names(inspector):
+    datasets = inspector.get_schema_names()
+    assert 'test_pybigquery' in datasets
+
+
+def test_table_names_in_schema(inspector):
+    tables = inspector.get_table_names('test_pybigquery')
+    assert 'test_pybigquery.sample' in tables
+    assert 'test_pybigquery.sample_one_row' in tables
+    assert 'test_pybigquery.sample_dml' in tables
+    assert len(tables) == 3
+
+
+def test_get_columns(inspector):
+    columns_without_schema = inspector.get_columns('test_pybigquery.sample')
+    columns_schema = inspector.get_columns('sample', 'test_pybigquery')
+    columns_queries = [columns_without_schema, columns_schema]
+    for columns in columns_queries:
+        for i, col in enumerate(columns):
+            sample_col = SAMPLE_COLUMNS[i]
+            assert col['name'] == sample_col['name']
+            assert col['nullable'] == sample_col['nullable']
+            assert col['default'] == sample_col['default']
+            assert col['type'].__class__.__name__ == sample_col['type'].__class__.__name__
+
+
+def test_has_table(engine):
+    assert engine.has_table('sample', 'test_pybigquery') is True
+    assert engine.has_table('test_pybigquery.sample') is True