Skip to content

Commit

Permalink
fix: include external tables in 'get_table_names' (#363)
Browse files Browse the repository at this point in the history
Also, include materialized views in 'get_view_names'.

Closes #332.
  • Loading branch information
tseaver committed Oct 27, 2021
1 parent d28cac5 commit 5e158fe
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 6 deletions.
10 changes: 6 additions & 4 deletions sqlalchemy_bigquery/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -802,7 +802,7 @@ def create_connect_args(self, url):
)
return ([client], {})

def _get_table_or_view_names(self, connection, table_type, schema=None):
def _get_table_or_view_names(self, connection, item_types, schema=None):
current_schema = schema or self.dataset_id
get_table_name = (
self._build_formatted_table_id
Expand All @@ -823,7 +823,7 @@ def _get_table_or_view_names(self, connection, table_type, schema=None):
dataset.reference, page_size=self.list_tables_page_size
)
for table in tables:
if table_type == table.table_type:
if table.table_type in item_types:
result.append(get_table_name(table))
except google.api_core.exceptions.NotFound:
# It's possible that the dataset was deleted between when we
Expand Down Expand Up @@ -976,13 +976,15 @@ def get_table_names(self, connection, schema=None, **kw):
if isinstance(connection, Engine):
connection = connection.connect()

return self._get_table_or_view_names(connection, "TABLE", schema)
item_types = ["TABLE", "EXTERNAL"]
return self._get_table_or_view_names(connection, item_types, schema)

def get_view_names(self, connection, schema=None, **kw):
if isinstance(connection, Engine):
connection = connection.connect()

return self._get_table_or_view_names(connection, "VIEW", schema)
item_types = ["VIEW", "MATERIALIZED_VIEW"]
return self._get_table_or_view_names(connection, item_types, schema)

def do_rollback(self, dbapi_connection):
# BigQuery has no support for transactions.
Expand Down
8 changes: 6 additions & 2 deletions tests/unit/test_sqlalchemy_bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,9 +77,11 @@ def table_item(dataset_id, table_id, type_="TABLE"):
[
table_item("dataset_2", "d2t1"),
table_item("dataset_2", "d2view", type_="VIEW"),
table_item("dataset_2", "d2ext", type_="EXTERNAL"),
table_item("dataset_2", "d2mv", type_="MATERIALIZED_VIEW"),
],
],
["dataset_1.d1t1", "dataset_1.d1t2", "dataset_2.d2t1"],
["dataset_1.d1t1", "dataset_1.d1t2", "dataset_2.d2t1", "dataset_2.d2ext"],
),
(
[dataset_item("dataset_1"), dataset_item("dataset_deleted")],
Expand Down Expand Up @@ -117,9 +119,11 @@ def test_get_table_names(
[
table_item("dataset_2", "d2t1"),
table_item("dataset_2", "d2view", type_="VIEW"),
table_item("dataset_2", "d2ext", type_="EXTERNAL"),
table_item("dataset_2", "d2mv", type_="MATERIALIZED_VIEW"),
],
],
["dataset_1.d1view", "dataset_2.d2view"],
["dataset_1.d1view", "dataset_2.d2view", "dataset_2.d2mv"],
),
(
[dataset_item("dataset_1"), dataset_item("dataset_deleted")],
Expand Down

0 comments on commit 5e158fe

Please sign in to comment.