From 5e158fe8bb2394369c020337092b5cfdb01880e0 Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Tue, 26 Oct 2021 21:32:30 -0400 Subject: [PATCH] fix: include external tables in 'get_table_names' (#363) Also, include materialized views in 'get_view_names'. Closes #332. --- sqlalchemy_bigquery/base.py | 10 ++++++---- tests/unit/test_sqlalchemy_bigquery.py | 8 ++++++-- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/sqlalchemy_bigquery/base.py b/sqlalchemy_bigquery/base.py index cbb0f8cc..f2da562d 100644 --- a/sqlalchemy_bigquery/base.py +++ b/sqlalchemy_bigquery/base.py @@ -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 @@ -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 @@ -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. diff --git a/tests/unit/test_sqlalchemy_bigquery.py b/tests/unit/test_sqlalchemy_bigquery.py index 75cbec42..e97b15ff 100644 --- a/tests/unit/test_sqlalchemy_bigquery.py +++ b/tests/unit/test_sqlalchemy_bigquery.py @@ -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")], @@ -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")],