fix: include external tables in 'get_table_names' (#363)

Also, include materialized views in 'get_view_names'. Closes #332.
googleapis · Oct 27, 2021 · 5e158fe · 5e158fe
1 parent d28cac5
commit 5e158fe
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 6 deletions.
diff --git 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
@@ -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")],