collect used tables as part of running report (#2835)

## Changes
Similarly to what is done for workflows, collect used tables as part of
running report

### Linked issues
None

### Functionality
- [ ] modified existing command: `databricks labs ucx install`, now
collects tables from dashboards as part of he assessment

### Tests
- [x] added integration tests

---------

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

Author: ericvergnaud

src/databricks/labs/ucx/contexts/application.py

@@ -449,6 +449,7 @@ def query_linter(self) -> QueryLinter:
             self.workspace_client,
             TableMigrationIndex([]),  # TODO: bring back self.tables_migrator.index()
             self.directfs_access_crawler_for_queries,
+            self.used_tables_crawler_for_queries,
             self.config.include_dashboard_ids,
         )
 

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

@@ -1,7 +1,7 @@
 import dataclasses
 import logging
 from collections.abc import Iterable
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, field
 from datetime import datetime, timezone
 
 from databricks.sdk import WorkspaceClient
@@ -12,11 +12,11 @@
 
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
 from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationIndex
-from databricks.labs.ucx.source_code.base import CurrentSessionState, LineageAtom
+from databricks.labs.ucx.source_code.base import CurrentSessionState, LineageAtom, UsedTable
 from databricks.labs.ucx.source_code.directfs_access import DirectFsAccessCrawler, DirectFsAccess
 from databricks.labs.ucx.source_code.linters.context import LinterContext
-from databricks.labs.ucx.source_code.linters.directfs import DirectFsAccessSqlLinter
 from databricks.labs.ucx.source_code.redash import Redash
+from databricks.labs.ucx.source_code.used_table import UsedTablesCrawler
 
 logger = logging.getLogger(__name__)
 
@@ -33,60 +33,92 @@ class QueryProblem:
     message: str
 
 
+@dataclass
+class _ReportingContext:
+    linted_queries: set[str] = field(default_factory=set)
+    all_problems: list[QueryProblem] = field(default_factory=list)
+    all_dfsas: list[DirectFsAccess] = field(default_factory=list)
+    all_tables: list[UsedTable] = field(default_factory=list)
+
+
 class QueryLinter:
 
     def __init__(
         self,
         ws: WorkspaceClient,
         migration_index: TableMigrationIndex,
         directfs_crawler: DirectFsAccessCrawler,
+        used_tables_crawler: UsedTablesCrawler,
         include_dashboard_ids: list[str] | None,
     ):
         self._ws = ws
         self._migration_index = migration_index
         self._directfs_crawler = directfs_crawler
+        self._used_tables_crawler = used_tables_crawler
         self._include_dashboard_ids = include_dashboard_ids
 
     def refresh_report(self, sql_backend: SqlBackend, inventory_database: str):
         assessment_start = datetime.now(timezone.utc)
-        dashboard_ids = self._dashboard_ids_in_scope()
-        logger.info(f"Running {len(dashboard_ids)} linting tasks...")
-        linted_queries: set[str] = set()
-        all_problems: list[QueryProblem] = []
-        all_dfsas: list[DirectFsAccess] = []
-        # first lint and collect queries from dashboards
-        for dashboard_id in dashboard_ids:
-            dashboard = self._ws.dashboards.get(dashboard_id=dashboard_id)
-            problems, dfsas = self._lint_and_collect_from_dashboard(dashboard, linted_queries)
-            all_problems.extend(problems)
-            all_dfsas.extend(dfsas)
-        for query in self._queries_in_scope():
-            if query.id in linted_queries:
-                continue
-            linted_queries.add(query.id)
-            problems = self.lint_query(query)
-            all_problems.extend(problems)
-            dfsas = self.collect_dfsas_from_query("no-dashboard-id", query)
-            all_dfsas.extend(dfsas)
-        # dump problems
-        logger.info(f"Saving {len(all_problems)} linting problems...")
+        context = _ReportingContext()
+        self._lint_dashboards(context)
+        self._lint_queries(context)
+        assessment_end = datetime.now(timezone.utc)
+        self._dump_problems(context, sql_backend, inventory_database)
+        self._dump_dfsas(context, assessment_start, assessment_end)
+        self._dump_used_tables(context, assessment_start, assessment_end)
+
+    def _dump_problems(self, context: _ReportingContext, sql_backend: SqlBackend, inventory_database: str):
+        logger.info(f"Saving {len(context.all_problems)} linting problems...")
         sql_backend.save_table(
             f'{escape_sql_identifier(inventory_database)}.query_problems',
-            all_problems,
+            context.all_problems,
             QueryProblem,
             mode='overwrite',
         )
-        # dump dfsas
-        assessment_end = datetime.now(timezone.utc)
-        processed = []
-        for dfsa in all_dfsas:
+
+    def _dump_dfsas(self, context: _ReportingContext, assessment_start: datetime, assessment_end: datetime):
+        processed_dfsas = []
+        for dfsa in context.all_dfsas:
             dfsa = dataclasses.replace(
                 dfsa,
                 assessment_start_timestamp=assessment_start,
                 assessment_end_timestamp=assessment_end,
             )
-            processed.append(dfsa)
-        self._directfs_crawler.dump_all(processed)
+            processed_dfsas.append(dfsa)
+        self._directfs_crawler.dump_all(processed_dfsas)
+
+    def _dump_used_tables(self, context: _ReportingContext, assessment_start: datetime, assessment_end: datetime):
+        processed_tables = []
+        for table in context.all_tables:
+            table = dataclasses.replace(
+                table,
+                assessment_start_timestamp=assessment_start,
+                assessment_end_timestamp=assessment_end,
+            )
+            processed_tables.append(table)
+        self._used_tables_crawler.dump_all(processed_tables)
+
+    def _lint_dashboards(self, context: _ReportingContext):
+        dashboard_ids = self._dashboard_ids_in_scope()
+        logger.info(f"Running {len(dashboard_ids)} linting tasks...")
+        for dashboard_id in dashboard_ids:
+            dashboard = self._ws.dashboards.get(dashboard_id=dashboard_id)
+            problems, dfsas, tables = self._lint_and_collect_from_dashboard(dashboard, context.linted_queries)
+            context.all_problems.extend(problems)
+            context.all_dfsas.extend(dfsas)
+            context.all_tables.extend(tables)
+
+    def _lint_queries(self, context: _ReportingContext):
+        for query in self._queries_in_scope():
+            if query.id in context.linted_queries:
+                continue
+            context.linted_queries.add(query.id)
+            problems = self.lint_query(query)
+            context.all_problems.extend(problems)
+            dfsas = self.collect_dfsas_from_query("no-dashboard-id", query)
+            context.all_dfsas.extend(dfsas)
+            tables = self.collect_used_tables_from_query("no-dashboard-id", query)
+            context.all_tables.extend(tables)
 
     def _dashboard_ids_in_scope(self) -> list[str]:
         if self._include_dashboard_ids is not None:  # an empty list is accepted
@@ -102,10 +134,11 @@ def _queries_in_scope(self):
 
     def _lint_and_collect_from_dashboard(
         self, dashboard: Dashboard, linted_queries: set[str]
-    ) -> tuple[Iterable[QueryProblem], Iterable[DirectFsAccess]]:
+    ) -> tuple[Iterable[QueryProblem], Iterable[DirectFsAccess], Iterable[UsedTable]]:
         dashboard_queries = Redash.get_queries_from_dashboard(dashboard)
         query_problems: list[QueryProblem] = []
         query_dfsas: list[DirectFsAccess] = []
+        query_tables: list[UsedTable] = []
         dashboard_id = dashboard.id or "<no-id>"
         dashboard_parent = dashboard.parent or "<orphan>"
         dashboard_name = dashboard.name or "<anonymous>"
@@ -134,7 +167,16 @@ def _lint_and_collect_from_dashboard(
                 )
                 source_lineage = [atom] + dfsa.source_lineage
                 query_dfsas.append(dataclasses.replace(dfsa, source_lineage=source_lineage))
-        return query_problems, query_dfsas
+            tables = self.collect_used_tables_from_query(dashboard_id, query)
+            for table in tables:
+                atom = LineageAtom(
+                    object_type="DASHBOARD",
+                    object_id=dashboard_id,
+                    other={"parent": dashboard_parent, "name": dashboard_name},
+                )
+                source_lineage = [atom] + table.source_lineage
+                query_tables.append(dataclasses.replace(table, source_lineage=source_lineage))
+        return query_problems, query_dfsas, query_tables
 
     def lint_query(self, query: LegacyQuery) -> Iterable[QueryProblem]:
         if not query.query:
@@ -156,20 +198,34 @@ def lint_query(self, query: LegacyQuery) -> Iterable[QueryProblem]:
                 message=advice.message,
             )
 
-    @classmethod
-    def collect_dfsas_from_query(cls, dashboard_id: str, query: LegacyQuery) -> Iterable[DirectFsAccess]:
+    def collect_dfsas_from_query(self, dashboard_id: str, query: LegacyQuery) -> Iterable[DirectFsAccess]:
         if query.query is None:
             return
-        linter = DirectFsAccessSqlLinter()
+        ctx = LinterContext(self._migration_index, CurrentSessionState())
+        collector = ctx.dfsa_collector(Language.SQL)
         source_id = f"{dashboard_id}/{query.id}"
         source_name = query.name or "<anonymous>"
-        source_timestamp = cls._read_timestamp(query.updated_at)
+        source_timestamp = self._read_timestamp(query.updated_at)
         source_lineage = [LineageAtom(object_type="QUERY", object_id=source_id, other={"name": source_name})]
-        for dfsa in linter.collect_dfsas(query.query):
+        for dfsa in collector.collect_dfsas(query.query):
             yield DirectFsAccess(**asdict(dfsa)).replace_source(
                 source_id=source_id, source_timestamp=source_timestamp, source_lineage=source_lineage
             )
 
+    def collect_used_tables_from_query(self, dashboard_id: str, query: LegacyQuery) -> Iterable[UsedTable]:
+        if query.query is None:
+            return
+        ctx = LinterContext(self._migration_index, CurrentSessionState())
+        collector = ctx.tables_collector(Language.SQL)
+        source_id = f"{dashboard_id}/{query.id}"
+        source_name = query.name or "<anonymous>"
+        source_timestamp = self._read_timestamp(query.updated_at)
+        source_lineage = [LineageAtom(object_type="QUERY", object_id=source_id, other={"name": source_name})]
+        for table in collector.collect_tables(query.query):
+            yield UsedTable(**asdict(table)).replace_source(
+                source_id=source_id, source_timestamp=source_timestamp, source_lineage=source_lineage
+            )
+
     @classmethod
     def _read_timestamp(cls, timestamp: str | None) -> datetime:
         if timestamp is not None:

tests/integration/source_code/test_queries.py

@@ -1,31 +1,58 @@
 from databricks.labs.ucx.hive_metastore.table_migration_status import TableMigrationIndex
 from databricks.labs.ucx.source_code.directfs_access import DirectFsAccessCrawler
 from databricks.labs.ucx.source_code.queries import QueryLinter
+from databricks.labs.ucx.source_code.used_table import UsedTablesCrawler
 
 
-def test_query_linter_lints_queries_and_stores_dfsas(simple_ctx, ws, sql_backend, make_query, make_dashboard):
-    query = make_query(sql_query="SELECT * from csv.`dbfs://some_folder/some_file.csv`")
-    _dashboard = make_dashboard(query=query)
-    linter = QueryLinter(ws, TableMigrationIndex([]), simple_ctx.directfs_access_crawler_for_queries, None)
+def test_query_linter_lints_queries_and_stores_dfsas_and_tables(
+    simple_ctx, ws, sql_backend, make_query, make_dashboard
+):
+    queries = [make_query(sql_query="SELECT * from csv.`dbfs://some_folder/some_file.csv`")]
+    dashboards = [make_dashboard(query=queries[0])]
+    queries.append(make_query(sql_query="SELECT * from some_schema.some_table"))
+    dashboards.append(make_dashboard(query=queries[1]))
+    linter = QueryLinter(
+        ws,
+        TableMigrationIndex([]),
+        simple_ctx.directfs_access_crawler_for_queries,
+        simple_ctx.used_tables_crawler_for_queries,
+        None,
+    )
     linter.refresh_report(sql_backend, simple_ctx.inventory_database)
     all_problems = sql_backend.fetch("SELECT * FROM query_problems", schema=simple_ctx.inventory_database)
-    problems = [row for row in all_problems if row["query_name"] == query.name]
+    problems = [row for row in all_problems if row["query_name"] == queries[0].name]
     assert len(problems) == 1
-    crawler = DirectFsAccessCrawler.for_queries(sql_backend, simple_ctx.inventory_database)
-    all_dfsas = crawler.snapshot()
-    source_id = f"{_dashboard.id}/{query.id}"
+    dfsa_crawler = DirectFsAccessCrawler.for_queries(sql_backend, simple_ctx.inventory_database)
+    all_dfsas = dfsa_crawler.snapshot()
+    source_id = f"{dashboards[0].id}/{queries[0].id}"
     dfsas = [dfsa for dfsa in all_dfsas if dfsa.source_id == source_id]
     assert len(dfsas) == 1
-    dfsa = dfsas[0]
-    assert len(dfsa.source_lineage) == 2
-    lineage = dfsa.source_lineage[0]
+    assert len(dfsas[0].source_lineage) == 2
+    lineage = dfsas[0].source_lineage[0]
     assert lineage.object_type == "DASHBOARD"
-    assert lineage.object_id == _dashboard.id
+    assert lineage.object_id == dashboards[0].id
     assert lineage.other
-    assert lineage.other.get("parent", None) == _dashboard.parent
-    assert lineage.other.get("name", None) == _dashboard.name
-    lineage = dfsa.source_lineage[1]
+    assert lineage.other.get("parent", None) == dashboards[0].parent
+    assert lineage.other.get("name", None) == dashboards[0].name
+    lineage = dfsas[0].source_lineage[1]
     assert lineage.object_type == "QUERY"
     assert lineage.object_id == source_id
     assert lineage.other
-    assert lineage.other.get("name", None) == query.name
+    assert lineage.other.get("name", None) == queries[0].name
+    used_tables_crawler = UsedTablesCrawler.for_queries(sql_backend, simple_ctx.inventory_database)
+    all_tables = used_tables_crawler.snapshot()
+    source_id = f"{dashboards[1].id}/{queries[1].id}"
+    tables = [table for table in all_tables if table.source_id == source_id]
+    assert len(tables) == 1
+    assert len(tables[0].source_lineage) == 2
+    lineage = tables[0].source_lineage[0]
+    assert lineage.object_type == "DASHBOARD"
+    assert lineage.object_id == dashboards[1].id
+    assert lineage.other
+    assert lineage.other.get("parent", None) == dashboards[1].parent
+    assert lineage.other.get("name", None) == dashboards[1].name
+    lineage = tables[0].source_lineage[1]
+    assert lineage.object_type == "QUERY"
+    assert lineage.object_id == source_id
+    assert lineage.other
+    assert lineage.other.get("name", None) == queries[1].name

tests/unit/source_code/test_queries.py

@@ -8,6 +8,7 @@
 
 from databricks.labs.ucx.source_code.directfs_access import DirectFsAccessCrawler
 from databricks.labs.ucx.source_code.queries import QueryLinter
+from databricks.labs.ucx.source_code.used_table import UsedTablesCrawler
 
 
 @pytest.mark.parametrize(
@@ -25,38 +26,44 @@
 )
 def test_query_linter_collects_dfsas_from_queries(name, query, dfsa_paths, is_read, is_write, migration_index):
     ws = create_autospec(WorkspaceClient)
-    crawlers = create_autospec(DirectFsAccessCrawler)
+    dfsa_crawler = create_autospec(DirectFsAccessCrawler)
+    used_tables_crawler = create_autospec(UsedTablesCrawler)
     query = LegacyQuery.from_dict({"parent": "workspace", "name": name, "query": query})
-    linter = QueryLinter(ws, migration_index, crawlers, None)
+    linter = QueryLinter(ws, migration_index, dfsa_crawler, used_tables_crawler, None)
     dfsas = linter.collect_dfsas_from_query("no-dashboard-id", query)
     ws.assert_not_called()
-    crawlers.assert_not_called()
+    dfsa_crawler.assert_not_called()
+    used_tables_crawler.assert_not_called()
     assert set(dfsa.path for dfsa in dfsas) == set(dfsa_paths)
     assert all(dfsa.is_read == is_read for dfsa in dfsas)
     assert all(dfsa.is_write == is_write for dfsa in dfsas)
 
 
 def test_query_liner_refresh_report_writes_query_problems(migration_index, mock_backend) -> None:
     ws = create_autospec(WorkspaceClient)
-    crawlers = create_autospec(DirectFsAccessCrawler)
-    linter = QueryLinter(ws, migration_index, crawlers, None)
+    dfsa_crawler = create_autospec(DirectFsAccessCrawler)
+    used_tables_crawler = create_autospec(UsedTablesCrawler)
+    linter = QueryLinter(ws, migration_index, dfsa_crawler, used_tables_crawler, None)
 
     linter.refresh_report(mock_backend, inventory_database="test")
 
     assert mock_backend.has_rows_written_for("`test`.query_problems")
     ws.dashboards.list.assert_called_once()
-    crawlers.assert_not_called()
+    dfsa_crawler.assert_not_called()
+    used_tables_crawler.assert_not_called()
 
 
 def test_lints_queries(migration_index, mock_backend) -> None:
     with mock.patch("databricks.labs.ucx.source_code.queries.Redash") as mocked_redash:
         query = LegacyQuery(id="123", query="SELECT * from nowhere")
         mocked_redash.get_queries_from_dashboard.return_value = [query]
         ws = create_autospec(WorkspaceClient)
-        crawlers = create_autospec(DirectFsAccessCrawler)
-        linter = QueryLinter(ws, migration_index, crawlers, ["1"])
+        dfsa_crawler = create_autospec(DirectFsAccessCrawler)
+        used_tables_crawler = create_autospec(UsedTablesCrawler)
+        linter = QueryLinter(ws, migration_index, dfsa_crawler, used_tables_crawler, ["1"])
         linter.refresh_report(mock_backend, inventory_database="test")
 
         assert mock_backend.has_rows_written_for("`test`.query_problems")
         ws.dashboards.list.assert_not_called()
-        crawlers.assert_not_called()
+        dfsa_crawler.assert_not_called()
+        used_tables_crawler.assert_not_called()