Improve robustness of sqlglot failure handling (#2952)

This PR fixes inadequate parsing quality from SQLglot

Author: nfx

src/databricks/labs/ucx/source_code/linters/from_table.py

@@ -53,7 +53,7 @@ def schema(self) -> str:
         return self._session_state.schema
 
     def lint_expression(self, expression: Expression) -> Iterable[Deprecation]:
-        for info in SqlExpression(expression).collect_table_infos("hive_metastore", self._session_state):
+        for info in SqlExpression(expression).collect_used_tables("hive_metastore", self._session_state):
             dst = self._index.get(info.schema_name, info.table_name)
             if not dst:
                 return
@@ -70,7 +70,7 @@ def lint_expression(self, expression: Expression) -> Iterable[Deprecation]:
     def collect_tables(self, source_code: str) -> Iterable[UsedTable]:
         try:
             for info in SqlParser.walk_expressions(
-                source_code, lambda e: e.collect_table_infos("hive_metastore", self._session_state)
+                source_code, lambda e: e.collect_used_tables("hive_metastore", self._session_state)
             ):
                 if any(pattern.matches(info.table_name) for pattern in DIRECT_FS_ACCESS_PATTERNS):
                     continue

src/databricks/labs/ucx/source_code/sql/sql_parser.py

@@ -1,5 +1,5 @@
 import logging
-from collections.abc import Callable, Iterable, Iterator
+from collections.abc import Callable, Iterable
 from typing import TypeVar
 
 from sqlglot import parse
@@ -19,8 +19,8 @@ class SqlExpression:
     def __init__(self, expression: Expression):
         self._expression = expression
 
-    def collect_table_infos(self, required_catalog: str, session_state: CurrentSessionState) -> Iterable[UsedTable]:
-        for table in self._expression.find_all(Table):
+    def collect_used_tables(self, required_catalog: str, session_state: CurrentSessionState) -> Iterable[UsedTable]:
+        for table in self.find_all(Table):
             info = self._collect_table_info(table, required_catalog, session_state)
             if info:
                 yield info
@@ -62,8 +62,12 @@ def _collect_table_info(
             is_write=isinstance(self._expression, (Create, Update, Delete)),
         )
 
-    def find_all(self, klass: type[E]) -> Iterator[E]:
-        return self._expression.find_all(klass)
+    def find_all(self, klass: type[E]) -> Iterable[E]:
+        try:
+            return self._expression.find_all(klass)
+        except SqlglotError as e:
+            logger.warning(f"Failed to find all {klass} in expression: {self._expression}", exc_info=e)
+            return []
 
 
 class SqlParser:
@@ -77,5 +81,4 @@ def walk_expressions(cls, sql_code: str, callback: Callable[[SqlExpression], Ite
                     continue
                 yield from callback(SqlExpression(expression))
         except SqlglotError as e:
-            logger.debug(f"Failed to parse SQL: {sql_code}", exc_info=e)
-            raise e
+            logger.warning(f"Failed to parse SQL: {sql_code}", exc_info=e)

tests/unit/source_code/sql/test_sql_parser.py

@@ -1,9 +1,5 @@
-import pytest
-from sqlglot import ParseError
-
 from databricks.labs.ucx.source_code.sql.sql_parser import SqlParser
 
 
-def test_raises_exception_with_unsupported_sql() -> None:
-    with pytest.raises(ParseError):
-        list(SqlParser.walk_expressions("XSELECT * from nowhere", lambda _: []))
+def test_does_not_raise_exception_with_unsupported_sql() -> None:
+    assert len(list(SqlParser.walk_expressions("XSELECT * from nowhere", lambda _: []))) == 0