SNOW-733235: Returning correct exception when table does not exist (#375)

Author: davidov541

DESCRIPTION.md

@@ -12,6 +12,7 @@ Source code is also available at:
 - v1.4.7(Unreleased)
 
   - Re-applied the application name of driver connection `SnowflakeConnection` to `SnowflakeSQLAlchemy`.
+  - `SnowflakeDialect.get_columns` now throws a `NoSuchTableError` exception when the specified table doesn't exist, instead of the more vague `KeyError`.
 
 - v1.4.6(Feb 8, 2023)
 

src/snowflake/sqlalchemy/snowdialect.py

@@ -659,6 +659,10 @@ def _get_table_columns(self, connection, table_name, schema=None, **kw):
                     else False,
                 }
             )
+
+        # If we didn't find any columns for the table, the table doesn't exist.
+        if len(ans) == 0:
+            raise sa_exc.NoSuchTableError()
         return ans
 
     def get_columns(self, connection, table_name, schema=None, **kw):
@@ -673,7 +677,10 @@ def get_columns(self, connection, table_name, schema=None, **kw):
         if schema_columns is None:
             # Too many results, fall back to only query about single table
             return self._get_table_columns(connection, table_name, schema, **kw)
-        return schema_columns[self.normalize_name(table_name)]
+        normalized_table_name = self.normalize_name(table_name)
+        if normalized_table_name not in schema_columns:
+            raise sa_exc.NoSuchTableError()
+        return schema_columns[normalized_table_name]
 
     @reflection.cache
     def get_table_names(self, connection, schema=None, **kw):

tests/test_core.py

@@ -32,7 +32,7 @@
     inspect,
     text,
 )
-from sqlalchemy.exc import DBAPIError
+from sqlalchemy.exc import DBAPIError, NoSuchTableError
 from sqlalchemy.pool import NullPool
 from sqlalchemy.sql import and_, not_, or_, select
 
@@ -406,6 +406,15 @@ def test_insert_tables(engine_testaccount):
             users.drop(engine_testaccount)
 
 
+def test_table_does_not_exist(engine_testaccount):
+    """
+    Tests Correct Exception Thrown When Table Does Not Exist
+    """
+    meta = MetaData()
+    with pytest.raises(NoSuchTableError):
+        Table("does_not_exist", meta, autoload=True, autoload_with=engine_testaccount)
+
+
 @pytest.mark.skip(
     """
 Reflection is not implemented yet.