Refactor linters (#2506)

## Changes
Naming of current linters is inconsistent, making the existing code very
hard to read.
Also, SQL linters do not derive from a specialized base linter, leading
to code duplication and suboptimal performance
This PR:
 - introduces specialized SqlLinter and SqlSequentialLinter
 - renames existing linters such that their names are self-explanatory

### Linked issues
None

### Functionality
None

### Tests
- [x] ran existing unit tests

---------

Co-authored-by: Eric Vergnaud <[email protected]>

Author: ericvergnaud

src/databricks/labs/ucx/account/aggregate.py

@@ -16,7 +16,7 @@
 from databricks.labs.ucx.hive_metastore.locations import LocationTrie
 from databricks.labs.ucx.hive_metastore.tables import Table
 from databricks.labs.ucx.source_code.base import CurrentSessionState
-from databricks.labs.ucx.source_code.queries import FromTable
+from databricks.labs.ucx.source_code.queries import FromTableSqlLinter
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -81,7 +81,7 @@ def _federated_ucx_query(self, query: str, table_name='objects') -> Iterable[tup
                         )
                     ]
                 )
-                from_table = FromTable(stub_index, CurrentSessionState(schema=inventory_database))
+                from_table = FromTableSqlLinter(stub_index, CurrentSessionState(schema=inventory_database))
                 logger.info(f"Querying Schema {inventory_database}")
 
                 workspace_specific_query = from_table.apply(query)

src/databricks/labs/ucx/hive_metastore/view_migrate.py

@@ -9,7 +9,7 @@
 from databricks.labs.ucx.hive_metastore.migration_status import MigrationIndex, TableView
 from databricks.labs.ucx.hive_metastore.mapping import TableToMigrate
 from databricks.labs.ucx.source_code.base import CurrentSessionState
-from databricks.labs.ucx.source_code.queries import FromTable
+from databricks.labs.ucx.source_code.queries import FromTableSqlLinter
 
 logger = logging.getLogger(__name__)
 
@@ -53,7 +53,7 @@ def _read_aliases(self, statement: expressions.Expression):
         return aliases
 
     def sql_migrate_view(self, index: MigrationIndex) -> str:
-        from_table = FromTable(index, CurrentSessionState(self.src.database))
+        from_table = FromTableSqlLinter(index, CurrentSessionState(self.src.database))
         assert self.src.view_text is not None, 'Expected a view text'
         migrated_select = from_table.apply(self.src.view_text)
         statements = sqlglot.parse(migrated_select, read='databricks')

src/databricks/labs/ucx/source_code/base.py

@@ -9,11 +9,12 @@
 from pathlib import Path
 
 from astroid import AstroidSyntaxError, NodeNG  # type: ignore
-from databricks.labs.blueprint.paths import WorkspacePath
+from sqlglot import Expression, parse as parse_sql, ParseError as SqlParseError
 
 from databricks.sdk.service import compute
 from databricks.sdk.service.workspace import Language
 
+from databricks.labs.blueprint.paths import WorkspacePath
 from databricks.labs.ucx.source_code.linters.python_ast import Tree
 
 # Code mapping between LSP, PyLint, and our own diagnostics:
@@ -136,6 +137,35 @@ class Linter:
     def lint(self, code: str) -> Iterable[Advice]: ...
 
 
+class SqlLinter(Linter):
+
+    def lint(self, code: str) -> Iterable[Advice]:
+        try:
+            expressions = parse_sql(code, read='databricks')
+            for expression in expressions:
+                if not expression:
+                    continue
+                yield from self.lint_expression(expression)
+        except SqlParseError as e:
+            logger.debug(f"Failed to parse SQL: {code}", exc_info=e)
+            yield self.sql_parse_failure(code)
+
+    @staticmethod
+    def sql_parse_failure(code: str):
+        return Failure(
+            code='sql-parse-error',
+            message=f"SQL expression is not supported yet: {code}",
+            # SQLGlot does not propagate tokens yet. See https://github.com/tobymao/sqlglot/issues/3159
+            start_line=0,
+            start_col=0,
+            end_line=0,
+            end_col=1024,
+        )
+
+    @abstractmethod
+    def lint_expression(self, expression: Expression) -> Iterable[Advice]: ...
+
+
 class PythonLinter(Linter):
 
     def lint(self, code: str) -> Iterable[Advice]:
@@ -201,13 +231,14 @@ def parse_security_mode(mode_str: str | None) -> compute.DataSecurityMode | None
             return None
 
 
-class SequentialLinter(Linter):
-    def __init__(self, linters: list[Linter]):
+class SqlSequentialLinter(SqlLinter):
+
+    def __init__(self, linters: list[SqlLinter]):
         self._linters = linters
 
-    def lint(self, code: str) -> Iterable[Advice]:
+    def lint_expression(self, expression: Expression) -> Iterable[Advice]:
         for linter in self._linters:
-            yield from linter.lint(code)
+            yield from linter.lint_expression(expression)
 
 
 class PythonSequentialLinter(Linter):

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

@@ -6,18 +6,19 @@
 from databricks.labs.ucx.source_code.base import (
     Fixer,
     Linter,
-    SequentialLinter,
+    SqlSequentialLinter,
     CurrentSessionState,
     PythonSequentialLinter,
     PythonLinter,
+    SqlLinter,
 )
-from databricks.labs.ucx.source_code.linters.dbfs import FromDbfsFolder, DBFSUsageLinter
-from databricks.labs.ucx.source_code.linters.imports import DbutilsLinter
+from databricks.labs.ucx.source_code.linters.dbfs import DbfsUsageSqlLinter, DBFSUsagePyLinter
+from databricks.labs.ucx.source_code.linters.imports import DbutilsPyLinter
 
-from databricks.labs.ucx.source_code.linters.pyspark import SparkSql
-from databricks.labs.ucx.source_code.linters.spark_connect import SparkConnectLinter
-from databricks.labs.ucx.source_code.linters.table_creation import DBRv8d0Linter
-from databricks.labs.ucx.source_code.queries import FromTable
+from databricks.labs.ucx.source_code.linters.pyspark import SparkSqlPyLinter
+from databricks.labs.ucx.source_code.linters.spark_connect import SparkConnectPyLinter
+from databricks.labs.ucx.source_code.linters.table_creation import DBRv8d0PyLinter
+from databricks.labs.ucx.source_code.queries import FromTableSqlLinter
 
 
 class LinterContext:
@@ -28,26 +29,25 @@ def __init__(self, index: MigrationIndex | None = None, session_state: CurrentSe
         python_linters: list[PythonLinter] = []
         python_fixers: list[Fixer] = []
 
-        sql_linters: list[Linter] = []
+        sql_linters: list[SqlLinter] = []
         sql_fixers: list[Fixer] = []
 
         if index is not None:
-            from_table = FromTable(index, session_state=session_state)
-            python_linters.append(SparkSql(from_table, index, session_state))
-            python_fixers.append(SparkSql(from_table, index, session_state))
-
+            from_table = FromTableSqlLinter(index, session_state=session_state)
             sql_linters.append(from_table)
             sql_fixers.append(from_table)
+            python_linters.append(SparkSqlPyLinter(from_table, index, session_state))
+            python_fixers.append(SparkSqlPyLinter(from_table, index, session_state))
 
         python_linters += [
-            DBFSUsageLinter(session_state),
-            DBRv8d0Linter(dbr_version=session_state.dbr_version),
-            SparkConnectLinter(session_state),
-            DbutilsLinter(session_state),
+            DBFSUsagePyLinter(session_state),
+            DBRv8d0PyLinter(dbr_version=session_state.dbr_version),
+            SparkConnectPyLinter(session_state),
+            DbutilsPyLinter(session_state),
         ]
-        sql_linters.append(FromDbfsFolder())
+        sql_linters.append(DbfsUsageSqlLinter())
 
-        self._linters: dict[Language, list[Linter] | list[PythonLinter]] = {
+        self._linters: dict[Language, list[SqlLinter] | list[PythonLinter]] = {
             Language.PYTHON: python_linters,
             Language.SQL: sql_linters,
         }
@@ -64,7 +64,9 @@ def linter(self, language: Language) -> Linter:
             raise ValueError(f"Unsupported language: {language}")
         if language is Language.PYTHON:
             return PythonSequentialLinter(cast(list[PythonLinter], self._linters[language]))
-        return SequentialLinter(cast(list[Linter], self._linters[language]))
+        if language is Language.SQL:
+            return SqlSequentialLinter(cast(list[SqlLinter], self._linters[language]))
+        raise ValueError(f"Unsupported language: {language}")
 
     def fixer(self, language: Language, diagnostic_code: str) -> Fixer | None:
         if language not in self._fixers:

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

@@ -2,13 +2,18 @@
 from collections.abc import Iterable
 
 from astroid import Call, Const, InferenceError, NodeNG  # type: ignore
-from sqlglot import Expression, parse as parse_sql, ParseError as SqlParseError
+from sqlglot import Expression
 from sqlglot.expressions import Table
 
-from databricks.labs.ucx.source_code.base import Advice, Linter, Deprecation, CurrentSessionState, PythonLinter
+from databricks.labs.ucx.source_code.base import (
+    Advice,
+    Deprecation,
+    CurrentSessionState,
+    PythonLinter,
+    SqlLinter,
+)
 from databricks.labs.ucx.source_code.linters.python_ast import Tree, TreeVisitor
 from databricks.labs.ucx.source_code.linters.python_infer import InferredValue
-from databricks.labs.ucx.source_code.queries import FromTable
 
 logger = logging.getLogger(__name__)
 
@@ -68,7 +73,7 @@ def get_advices(self) -> Iterable[Advice]:
         yield from self._advices
 
 
-class DBFSUsageLinter(PythonLinter):
+class DBFSUsagePyLinter(PythonLinter):
 
     def __init__(self, session_state: CurrentSessionState):
         self._session_state = session_state
@@ -89,27 +94,16 @@ def lint_tree(self, tree: Tree) -> Iterable[Advice]:
         yield from visitor.get_advices()
 
 
-class FromDbfsFolder(Linter):
+class DbfsUsageSqlLinter(SqlLinter):
     def __init__(self):
         self._dbfs_prefixes = ["/dbfs/mnt", "dbfs:/", "/mnt/", "/dbfs/", "/"]
 
     @staticmethod
     def name() -> str:
         return 'dbfs-query'
 
-    def lint(self, code: str) -> Iterable[Advice]:
-        try:
-            queries = parse_sql(code, read='databricks')
-            for query in queries:
-                if not query:
-                    continue
-                yield from self._lint_query(query)
-        except SqlParseError as e:
-            logger.debug(f"Failed to parse SQL: {code}", exc_info=e)
-            yield FromTable.sql_parse_failure(code)
-
-    def _lint_query(self, query: Expression):
-        for table in query.find_all(Table):
+    def lint_expression(self, expression: Expression):
+        for table in expression.find_all(Table):
             # Check table names for deprecated DBFS table names
             yield from self._check_dbfs_folder(table)
 

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

@@ -91,7 +91,7 @@ def get_notebook_paths(self, session_state: CurrentSessionState) -> tuple[bool,
         for path in paths:
             dbutils.notebook.run(path)
         """
-        arg = DbutilsLinter.get_dbutils_notebook_run_path_arg(self.node)
+        arg = DbutilsPyLinter.get_dbutils_notebook_run_path_arg(self.node)
         try:
             all_inferred = InferredValue.infer_from_node(arg, session_state)
             return self._get_notebook_paths(all_inferred)
@@ -113,7 +113,7 @@ def _get_notebook_paths(cls, all_inferred: Iterable[InferredValue]) -> tuple[boo
         return has_unresolved, paths
 
 
-class DbutilsLinter(PythonLinter):
+class DbutilsPyLinter(PythonLinter):
 
     def __init__(self, session_state: CurrentSessionState):
         self._session_state = session_state

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

@@ -13,7 +13,7 @@
     PythonLinter,
 )
 from databricks.labs.ucx.source_code.linters.python_infer import InferredValue
-from databricks.labs.ucx.source_code.queries import FromTable
+from databricks.labs.ucx.source_code.queries import FromTableSqlLinter
 from databricks.labs.ucx.source_code.linters.python_ast import Tree, TreeHelper
 
 
@@ -37,12 +37,12 @@ def matches(self, node: NodeNG):
 
     @abstractmethod
     def lint(
-        self, from_table: FromTable, index: MigrationIndex, session_state: CurrentSessionState, node: Call
+        self, from_table: FromTableSqlLinter, index: MigrationIndex, session_state: CurrentSessionState, node: Call
     ) -> Iterator[Advice]:
         """raises Advices by linting the code"""
 
     @abstractmethod
-    def apply(self, from_table: FromTable, index: MigrationIndex, node: Call) -> None:
+    def apply(self, from_table: FromTableSqlLinter, index: MigrationIndex, node: Call) -> None:
         """applies recommendations"""
 
     def _get_table_arg(self, node: Call):
@@ -78,7 +78,7 @@ def _check_call_context(self, node: Call) -> bool:
 class QueryMatcher(Matcher):
 
     def lint(
-        self, from_table: FromTable, index: MigrationIndex, session_state: CurrentSessionState, node: Call
+        self, from_table: FromTableSqlLinter, index: MigrationIndex, session_state: CurrentSessionState, node: Call
     ) -> Iterator[Advice]:
         table_arg = self._get_table_arg(node)
         if table_arg:
@@ -93,7 +93,7 @@ def lint(
                 )
 
     @classmethod
-    def _lint_table_arg(cls, from_table: FromTable, call_node: NodeNG, inferred: InferredValue):
+    def _lint_table_arg(cls, from_table: FromTableSqlLinter, call_node: NodeNG, inferred: InferredValue):
         if inferred.is_inferred():
             for advice in from_table.lint(inferred.as_string()):
                 yield advice.replace_from_node(call_node)
@@ -104,7 +104,7 @@ def _lint_table_arg(cls, from_table: FromTable, call_node: NodeNG, inferred: Inf
                 node=call_node,
             )
 
-    def apply(self, from_table: FromTable, index: MigrationIndex, node: Call) -> None:
+    def apply(self, from_table: FromTableSqlLinter, index: MigrationIndex, node: Call) -> None:
         table_arg = self._get_table_arg(node)
         assert isinstance(table_arg, Const)
         new_query = from_table.apply(table_arg.value)
@@ -115,7 +115,7 @@ def apply(self, from_table: FromTable, index: MigrationIndex, node: Call) -> Non
 class TableNameMatcher(Matcher):
 
     def lint(
-        self, from_table: FromTable, index: MigrationIndex, session_state: CurrentSessionState, node: Call
+        self, from_table: FromTableSqlLinter, index: MigrationIndex, session_state: CurrentSessionState, node: Call
     ) -> Iterator[Advice]:
         table_arg = self._get_table_arg(node)
         table_name = table_arg.as_string().strip("'").strip('"')
@@ -137,7 +137,7 @@ def lint(
                 node=node,
             )
 
-    def apply(self, from_table: FromTable, index: MigrationIndex, node: Call) -> None:
+    def apply(self, from_table: FromTableSqlLinter, index: MigrationIndex, node: Call) -> None:
         table_arg = self._get_table_arg(node)
         assert isinstance(table_arg, Const)
         dst = self._find_dest(index, table_arg.value, from_table.schema)
@@ -162,7 +162,7 @@ def matches(self, node: NodeNG):
         )
 
     def lint(
-        self, from_table: FromTable, index: MigrationIndex, session_state: CurrentSessionState, node: Call
+        self, from_table: FromTableSqlLinter, index: MigrationIndex, session_state: CurrentSessionState, node: Call
     ) -> Iterator[Advice]:
         assert isinstance(node.func, Attribute)  # always true, avoids a pylint warning
         yield Advisory.from_node(
@@ -171,7 +171,7 @@ def lint(
             node=node,
         )
 
-    def apply(self, from_table: FromTable, index: MigrationIndex, node: Call) -> None:
+    def apply(self, from_table: FromTableSqlLinter, index: MigrationIndex, node: Call) -> None:
         # No transformations to apply
         return
 
@@ -200,7 +200,7 @@ def matches(self, node: NodeNG):
         )
 
     def lint(
-        self, from_table: FromTable, index: MigrationIndex, session_state: CurrentSessionState, node: NodeNG
+        self, from_table: FromTableSqlLinter, index: MigrationIndex, session_state: CurrentSessionState, node: NodeNG
     ) -> Iterator[Advice]:
         table_arg = self._get_table_arg(node)
         if not isinstance(table_arg, Const):
@@ -223,7 +223,7 @@ def lint(
                 node=node,
             )
 
-    def apply(self, from_table: FromTable, index: MigrationIndex, node: Call) -> None:
+    def apply(self, from_table: FromTableSqlLinter, index: MigrationIndex, node: Call) -> None:
         # No transformations to apply
         return
 
@@ -327,11 +327,11 @@ def matchers(self):
         return self._matchers
 
 
-class SparkSql(PythonLinter, Fixer):
+class SparkSqlPyLinter(PythonLinter, Fixer):
 
     _spark_matchers = SparkMatchers()
 
-    def __init__(self, from_table: FromTable, index: MigrationIndex, session_state):
+    def __init__(self, from_table: FromTableSqlLinter, index: MigrationIndex, session_state):
         self._from_table = from_table
         self._index = index
         self._session_state = session_state

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

@@ -242,7 +242,7 @@ def lint(self, node: NodeNG) -> Iterator[Advice]:
             )
 
 
-class SparkConnectLinter(PythonLinter):
+class SparkConnectPyLinter(PythonLinter):
     def __init__(self, session_state: CurrentSessionState):
         if session_state.data_security_mode != DataSecurityMode.USER_ISOLATION:
             self._matchers = []

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

@@ -96,7 +96,7 @@ def lint(self, node: NodeNG) -> Iterator[Advice]:
                 )
 
 
-class DBRv8d0Linter(PythonLinter):
+class DBRv8d0PyLinter(PythonLinter):
     """Performs Python linting for backwards incompatible changes in DBR version 8.0.
     Specifically, it yields advice for table-creation with implicit format.
     """