Fix BigQuery \list, \describe, and \introspect (#92)

* Harden for BigQuery dialect foibles per ENG-5350

Author: James Robinson

noteable_magics/sql/meta_commands.py

@@ -15,6 +15,7 @@
 from pandas import DataFrame
 from sqlalchemy import inspect
 from sqlalchemy.engine.reflection import Inspector
+from sqlalchemy.exc import NoSuchTableError
 
 from noteable_magics.sql.connection import Connection
 from noteable_magics.sql.gate_messaging_types import (
@@ -399,21 +400,17 @@ def run(self, invoked_as: str, args: List[str]) -> Tuple[DataFrame, bool]:
 
         inspector = self.get_inspector()
 
-        is_view = relation_name in inspector.get_view_names(schema)
-
-        if not is_view:
-            # Ensure is a table
-            if relation_name not in inspector.get_table_names(schema):
-                if schema:
-                    msg = f'Relation {schema}.{relation_name} does not exist'
-                else:
-                    msg = f'Relation {relation_name} does not exist'
-                raise MetaCommandException(msg)
-            rtype = 'Table'
-        else:
-            rtype = 'View'
+        try:
+            # In some dialects (BigQuery), this will raise NoSuchTableError if
+            # the table doesn't exist. Yay, sane.
 
-        column_dicts = inspector.get_columns(relation_name, schema=schema)
+            # On some dialects (sigh, CockroachDB, what are you doing??),
+            # this call may succeed returning empty list even if
+            # the named relation does not exist. But the call to get_pk_constraint()
+            # down below will then raise NoSuchTableError.
+            column_dicts = inspector.get_columns(relation_name, schema=schema)
+        except NoSuchTableError:
+            self._raise_from_no_such_table(schema, relation_name)
 
         # 'Pivot' the dicts from get_columns()
         names = []
@@ -454,11 +451,19 @@ def run(self, invoked_as: str, args: List[str]) -> Tuple[DataFrame, bool]:
 
         displayable_rname = displayable_relation_name(schema, relation_name)
 
+        if is_view := relation_name in inspector.get_view_names(schema):
+            rtype = 'View'
+        else:
+            rtype = 'Table'
+
         main_relation_df = set_dataframe_metadata(
             DataFrame(data=data), title=f'{rtype} "{displayable_rname}" Structure'
         )
 
-        display(main_relation_df)
+        # Keep a list of things to call display() on. Only do so at the very
+        # end if we don't hit any exceptions.
+
+        displayables = [main_relation_df]
 
         if is_view:
             view_definition = inspector.get_view_definition(relation_name, schema)
@@ -472,7 +477,7 @@ def run(self, invoked_as: str, args: List[str]) -> Tuple[DataFrame, bool]:
                 html_buf.append('<br />')
                 html_buf.append(f'<pre>{view_definition}</pre>')
 
-                display(HTML('\n'.join(html_buf)))
+                displayables.append(HTML('\n'.join(html_buf)))
         else:
             # Is a table. Let's go get indices, foreign keys, other table constraints.
             # If meaningful dataframe returned for any of these, transform to
@@ -484,7 +489,11 @@ def run(self, invoked_as: str, args: List[str]) -> Tuple[DataFrame, bool]:
             ):
                 secondary_df = secondary_function(inspector, relation_name, schema)
                 if len(secondary_df):
-                    display(secondary_dataframe_to_html(secondary_df))
+                    displayables.append(secondary_dataframe_to_html(secondary_df))
+
+        # Make it this far? Display what we should display.
+        for displayable in displayables:
+            display(displayable)
 
         return main_relation_df, False
 
@@ -598,7 +607,7 @@ def all_table_and_views(self, inspector) -> List[Tuple[str, str, str]]:
         default_schema = inspector.default_schema_name
         all_schemas = set(inspector.get_schema_names())
         all_schemas.difference_update(self.AVOID_SCHEMAS)
-        if default_schema not in all_schemas:
+        if default_schema and default_schema not in all_schemas:
             all_schemas.add(default_schema)
 
         for schema_name in sorted(all_schemas):
@@ -971,7 +980,10 @@ def index_dataframe(
     uniques: List[bool] = []
 
     # Primary key index is ... treated special by SQLA for some reason. Sigh.
-    primary_index_dict = inspector.get_pk_constraint(table_name, schema)
+    try:
+        primary_index_dict = inspector.get_pk_constraint(table_name, schema)
+    except NoSuchTableError:
+        _raise_from_no_such_table(schema, table_name)
 
     # If it returned something truthy with nonempty constrained_columns, then
     # we assume it described a real primary key constraint here.
@@ -1164,7 +1176,8 @@ def __init__(self, underlying_inspector: Inspector):
 
     # Direct passthrough attributes / methods
     @property
-    def default_schema_name(self) -> str:
+    def default_schema_name(self) -> Optional[str]:
+        # BigQuery, Trino dialects may end up returning None.
         return self.underlying_inspector.default_schema_name
 
     def get_schema_names(self) -> List[str]:
@@ -1183,16 +1196,19 @@ def get_foreign_keys(self, table_name: str, schema: Optional[str] = None) -> Lis
         return self.underlying_inspector.get_foreign_keys(table_name, schema=schema)
 
     def get_check_constraints(self, table_name: str, schema: Optional[str] = None) -> List[dict]:
-        return self.underlying_inspector.get_check_constraints(table_name, schema=schema)
+        try:
+            return self.underlying_inspector.get_check_constraints(table_name, schema=schema)
+        except NotImplementedError:
+            return []
 
     def get_indexes(self, table_name: str, schema: Optional[str] = None) -> List[dict]:
         return self.underlying_inspector.get_indexes(table_name, schema=schema)
 
     def get_unique_constraints(self, table_name: str, schema: Optional[str] = None) -> List[dict]:
-        return self.underlying_inspector.get_unique_constraints(table_name, schema=schema)
-
-    def get_check_constraints(self, table_name: str, schema: Optional[str] = None) -> List[dict]:
-        return self.underlying_inspector.get_check_constraints(table_name, schema=schema)
+        try:
+            return self.underlying_inspector.get_unique_constraints(table_name, schema=schema)
+        except NotImplementedError:
+            return []
 
     # Now the value-adding filtering methods.
     def get_table_names(self, schema: Optional[str] = None) -> List[str]:
@@ -1211,3 +1227,12 @@ def _strip_schema(self, names: List[str], schema: Optional[str] = None) -> List[
         # Remove "schema." from the start of each name if starts with.
         # (name[False:] is equiv to name[0:], 'cause python bools are subclasses of ints)
         return [name[name.startswith(prefix) and len(prefix) :] for name in names]
+
+
+def _raise_from_no_such_table(schema: str, relation_name: str):
+    """Raise a MetaCommandException when eaten a NoSuchTableException"""
+    if schema:
+        msg = f'Relation {schema}.{relation_name} does not exist'
+    else:
+        msg = f'Relation {relation_name} does not exist'
+    raise MetaCommandException(msg)