Fixed AttributeError: 'UsedTable' has no attribute 'table' by adding more type checks (#2895)

Fix #2887

Author: nfx

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

@@ -270,9 +270,16 @@ class TableInfoNode:
 class TablePyCollector(TableCollector, ABC):
 
     def collect_tables(self, source_code: str):
-        tree = Tree.normalize_and_parse(source_code)
-        for table_node in self.collect_tables_from_tree(tree):
-            yield table_node.table
+        try:
+            tree = Tree.normalize_and_parse(source_code)
+            for table_node in self.collect_tables_from_tree(tree):
+                # see https://github.com/databrickslabs/ucx/issues/2887
+                if isinstance(table_node, UsedTable):
+                    yield table_node
+                else:
+                    yield table_node.table
+        except AstroidSyntaxError as e:
+            logger.warning('syntax-error', exc_info=e)
 
     @abstractmethod
     def collect_tables_from_tree(self, tree: Tree) -> Iterable[TableInfoNode]: ...
@@ -451,7 +458,12 @@ def collect_tables(self, source_code: str) -> Iterable[UsedTable]:
         try:
             tree = self._parse_and_append(source_code)
             for table_node in self.collect_tables_from_tree(tree):
-                yield table_node.table
+                # there's a bug in the code that causes this to be necessary
+                # see https://github.com/databrickslabs/ucx/issues/2887
+                if isinstance(table_node, UsedTable):
+                    yield table_node
+                else:
+                    yield table_node.table
         except AstroidSyntaxError as e:
             logger.warning('syntax-error', exc_info=e)
 

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

@@ -412,7 +412,7 @@ def lint_job(self, job_id: int) -> tuple[list[JobProblem], list[DirectFsAccess],
     def _lint_job(self, job: jobs.Job) -> tuple[list[JobProblem], list[DirectFsAccess], list[UsedTable]]:
         problems: list[JobProblem] = []
         dfsas: list[DirectFsAccess] = []
-        table_infos: list[UsedTable] = []
+        used_tables: list[UsedTable] = []
 
         assert job.job_id is not None
         assert job.settings is not None
@@ -447,13 +447,14 @@ def _lint_job(self, job: jobs.Job) -> tuple[list[JobProblem], list[DirectFsAcces
             assessment_start = datetime.now(timezone.utc)
             task_tables = self._collect_task_tables(job, task, graph, session_state)
             assessment_end = datetime.now(timezone.utc)
-            for table_info in task_tables:
-                table_info = table_info.replace_assessment_infos(
-                    assessment_start=assessment_start, assessment_end=assessment_end
+            for used_table in task_tables:
+                used_table = used_table.replace_assessment_infos(
+                    assessment_start=assessment_start,
+                    assessment_end=assessment_end,
                 )
-                table_infos.append(table_info)
+                used_tables.append(used_table)
 
-        return problems, dfsas, table_infos
+        return problems, dfsas, used_tables
 
     def _build_task_dependency_graph(
         self, task: jobs.Task, job: jobs.Job
@@ -502,17 +503,21 @@ def _collect_task_dfsas(
             yield dataclasses.replace(dfsa, source_lineage=atoms + dfsa.source_lineage)
 
     def _collect_task_tables(
-        self, job: jobs.Job, task: jobs.Task, graph: DependencyGraph, session_state: CurrentSessionState
+        self,
+        job: jobs.Job,
+        task: jobs.Task,
+        graph: DependencyGraph,
+        session_state: CurrentSessionState,
     ) -> Iterable[UsedTable]:
         # need to add lineage for job/task because walker doesn't register them
         job_id = str(job.job_id)
         job_name = job.settings.name if job.settings and job.settings.name else "<anonymous>"
-        for dfsa in TablesCollectorWalker(graph, set(), self._path_lookup, session_state, self._migration_index):
+        for used_table in TablesCollectorWalker(graph, set(), self._path_lookup, session_state, self._migration_index):
             atoms = [
                 LineageAtom(object_type="WORKFLOW", object_id=job_id, other={"name": job_name}),
                 LineageAtom(object_type="TASK", object_id=f"{job_id}/{task.task_key}"),
             ]
-            yield dataclasses.replace(dfsa, source_lineage=atoms + dfsa.source_lineage)
+            yield dataclasses.replace(used_table, source_lineage=atoms + used_table.source_lineage)
 
 
 class LintingWalker(DependencyGraphWalker[LocatedAdvice]):

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

@@ -395,7 +395,7 @@ def collect_tables_from_tree(self, tree: Tree) -> Iterable[TableInfoNode]:
                 continue
             assert isinstance(node, Call)
             for used_table in matcher.collect_tables(self._from_table, self._index, self._session_state, node):
-                yield TableInfoNode(used_table, node)
+                yield TableInfoNode(used_table, node)  # B
 
 
 class _SparkSqlAnalyzer:
@@ -475,4 +475,4 @@ def collect_tables_from_tree(self, tree: Tree) -> Iterable[TableInfoNode]:
                 if not value.is_inferred():
                     continue  # TODO error handling strategy
                 for table in self._sql_collector.collect_tables(value.as_string()):
-                    yield TableInfoNode(table, call_node)
+                    yield TableInfoNode(table, call_node)  # A

tests/unit/source_code/test_context.py

@@ -0,0 +1,25 @@
+import pytest
+from databricks.sdk.service.workspace import Language
+
+from databricks.labs.ucx.source_code.linters.context import LinterContext
+
+
+@pytest.mark.parametrize(
+    'code,expected',
+    [
+        ('spark.table("a.b").count()', {'r:a.b'}),
+        ('spark.getTable("a.b")', {'r:a.b'}),
+        ('spark.cacheTable("a.b")', {'r:a.b'}),
+        ('spark.range(10).saveAsTable("a.b")', {'r:a.b'}),  # TODO: bug: has to be w:a.b
+        ('spark.sql("SELECT * FROM b.c LEFT JOIN c.d USING (e)")', {'r:b.c', 'r:c.d'}),
+        ('spark.sql("SELECT * FROM delta.`/foo/bar`")', set()),
+    ],
+)
+def test_collector_walker_from_python(code, expected, migration_index) -> None:
+    used = set()
+    ctx = LinterContext(migration_index)
+    collector = ctx.tables_collector(Language.PYTHON)
+    for used_table in collector.collect_tables(code):
+        prefix = 'r' if used_table.is_read else 'w'
+        used.add(f'{prefix}:{used_table.schema_name}.{used_table.table_name}')
+    assert used == expected