Added task for linting queries (#2630)

## Changes
Workspace sql queries are not currently linted.
This PR fixes that by introducing a QueryLinter which mimics for
dashboard queries what WorkflowLinter does for jobs

### Linked issues
Resolves #2350 

### Functionality
- [x] modified existing workflow: `experimental-workflow-linter`
- [x] added a new table `query_problems`

### Tests
- [x] added unit tests
- [x] added integration tests
- [x] manually tested schema upgrade:

![Screenshot 2024-09-13 at 16 20
42](https://github.com/user-attachments/assets/a64bf2fa-d8c6-4346-818d-1be5526df24b)

Please note that storage is simplified in PR
#2662

---------

Co-authored-by: Eric Vergnaud <[email protected]>
Co-authored-by: Serge Smertin <[email protected]>
Co-authored-by: Andrew Snare <[email protected]>
Co-authored-by: Cor <[email protected]>
Co-authored-by: Serge Smertin <[email protected]>

Author: 4 people

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

@@ -431,7 +431,11 @@ def workflow_linter(self):
 
     @cached_property
     def query_linter(self):
-        return QueryLinter(self.workspace_client, self.directfs_access_crawler_for_queries)
+        return QueryLinter(
+            self.workspace_client,
+            TableMigrationIndex([]),  # TODO: bring back self.tables_migrator.index()
+            self.directfs_access_crawler_for_queries,
+        )
 
     @cached_property
     def directfs_access_crawler_for_paths(self):

src/databricks/labs/ucx/framework/crawlers.py

@@ -136,8 +136,8 @@ def _snapshot(self, fetcher: ResultFn, loader: ResultFn, *, force_refresh: bool)
                 cached_results = list(fetcher())
                 if cached_results:
                     return cached_results
-            except NotFound:
-                pass
+            except NotFound as e:
+                logger.debug("Inventory table not found", exc_info=e)
         logger.debug(f"[{self.full_name}] crawling new set of snapshot data for {self._table}")
         loaded_records = list(loader())
         self._update_snapshot(loaded_records, mode="overwrite")

src/databricks/labs/ucx/install.py

@@ -75,8 +75,9 @@
 from databricks.labs.ucx.installer.workflows import WorkflowsDeployment
 from databricks.labs.ucx.recon.migration_recon import ReconResult
 from databricks.labs.ucx.runtime import Workflows
-from databricks.labs.ucx.source_code.directfs_access import DirectFsAccessInPath, DirectFsAccessInQuery
+from databricks.labs.ucx.source_code.directfs_access import DirectFsAccess
 from databricks.labs.ucx.source_code.jobs import JobProblem
+from databricks.labs.ucx.source_code.queries import QueryProblem
 from databricks.labs.ucx.workspace_access.base import Permissions
 from databricks.labs.ucx.workspace_access.generic import WorkspaceObjectInfo
 from databricks.labs.ucx.workspace_access.groups import ConfigureGroups, MigratedGroup
@@ -118,11 +119,12 @@ def deploy_schema(sql_backend: SqlBackend, inventory_schema: str):
             functools.partial(table, "policies", PolicyInfo),
             functools.partial(table, "migration_status", TableMigrationStatus),
             functools.partial(table, "workflow_problems", JobProblem),
+            functools.partial(table, "query_problems", QueryProblem),
             functools.partial(table, "udfs", Udf),
             functools.partial(table, "logs", LogRecord),
             functools.partial(table, "recon_results", ReconResult),
-            functools.partial(table, "directfs_in_paths", DirectFsAccessInPath),
-            functools.partial(table, "directfs_in_queries", DirectFsAccessInQuery),
+            functools.partial(table, "directfs_in_paths", DirectFsAccess),
+            functools.partial(table, "directfs_in_queries", DirectFsAccess),
         ],
     )
     deployer.deploy_view("grant_detail", "queries/views/grant_detail.sql")

src/databricks/labs/ucx/mixins/fixtures.py

@@ -6,10 +6,7 @@ SELECT
     source_timestamp,
     source_lineage,
     assessment_start_timestamp,
-    assessment_end_timestamp,
-    job_id,
-    job_name,
-    task_key
+    assessment_end_timestamp
 FROM $inventory.directfs_in_paths
 UNION ALL
 SELECT
@@ -20,8 +17,5 @@ SELECT
     source_timestamp,
     source_lineage,
     assessment_start_timestamp,
-    assessment_end_timestamp,
-    NULL as job_id,
-    NULL as job_name,
-    null as task_key
+    assessment_end_timestamp
 FROM $inventory.directfs_in_queries

src/databricks/labs/ucx/queries/views/directfs.sql

@@ -6,7 +6,7 @@
 from collections.abc import Sequence, Iterable
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any, TypeVar
+from typing import Any
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase
 from databricks.labs.lsql.backends import SqlBackend
@@ -39,7 +39,7 @@ class DirectFsAccess:
     def from_dict(cls, data: dict[str, Any]) -> Self:
         source_lineage = data.get("source_lineage", None)
         if isinstance(source_lineage, list) and len(source_lineage) > 0 and isinstance(source_lineage[0], dict):
-            lineage_atoms = [LineageAtom(*lineage) for lineage in source_lineage]
+            lineage_atoms = [LineageAtom(**lineage) for lineage in source_lineage]
             data["source_lineage"] = lineage_atoms
         return cls(**data)
 
@@ -77,63 +77,27 @@ def replace_assessment_infos(
         )
 
 
-@dataclass
-class DirectFsAccessInQuery(DirectFsAccess):
-
-    pass
-
-
-@dataclass
-class DirectFsAccessInPath(DirectFsAccess):
-    job_id: int = -1
-    job_name: str = DirectFsAccess.UNKNOWN
-    task_key: str = DirectFsAccess.UNKNOWN
-
-    def replace_job_infos(
-        self,
-        job_id: int | None = None,
-        job_name: str | None = None,
-        task_key: str | None = None,
-    ):
-        return dataclasses.replace(
-            self, job_id=job_id or self.job_id, job_name=job_name or self.job_name, task_key=task_key or self.task_key
-        )
-
-
-T = TypeVar("T", bound=DirectFsAccess)
-
-
-class DirectFsAccessCrawler(CrawlerBase[T]):
+class DirectFsAccessCrawler(CrawlerBase[DirectFsAccess]):
 
     @classmethod
     def for_paths(cls, backend: SqlBackend, schema) -> DirectFsAccessCrawler:
-        return DirectFsAccessCrawler[DirectFsAccessInPath](
-            backend,
-            schema,
-            "directfs_in_paths",
-            DirectFsAccessInPath,
-        )
+        return DirectFsAccessCrawler(backend, schema, "directfs_in_paths")
 
     @classmethod
     def for_queries(cls, backend: SqlBackend, schema) -> DirectFsAccessCrawler:
-        return DirectFsAccessCrawler[DirectFsAccessInQuery](
-            backend,
-            schema,
-            "directfs_in_queries",
-            DirectFsAccessInQuery,
-        )
+        return DirectFsAccessCrawler(backend, schema, "directfs_in_queries")
 
-    def __init__(self, backend: SqlBackend, schema: str, table: str, klass: type[T]):
+    def __init__(self, backend: SqlBackend, schema: str, table: str):
         """
         Initializes a DFSACrawler instance.
 
         Args:
             sql_backend (SqlBackend): The SQL Execution Backend abstraction (either REST API or Spark)
             schema: The schema name for the inventory persistence.
         """
-        super().__init__(backend, "hive_metastore", schema, table, klass)
+        super().__init__(backend=backend, catalog="hive_metastore", schema=schema, table=table, klass=DirectFsAccess)
 
-    def dump_all(self, dfsas: Sequence[T]):
+    def dump_all(self, dfsas: Sequence[DirectFsAccess]):
         """This crawler doesn't follow the pull model because the fetcher fetches data for 2 crawlers, not just one
         It's not **bad** because all records are pushed at once.
         Providing a multi-entity crawler is out-of-scope of this PR
@@ -144,9 +108,11 @@ def dump_all(self, dfsas: Sequence[T]):
         except DatabricksError as e:
             logger.error("Failed to store DFSAs", exc_info=e)
 
-    def _try_fetch(self) -> Iterable[T]:
+    def _try_fetch(self) -> Iterable[DirectFsAccess]:
         sql = f"SELECT * FROM {escape_sql_identifier(self.full_name)}"
-        yield from self._backend.fetch(sql)
+        for row in self._backend.fetch(sql):
+            yield self._klass.from_dict(row.as_dict())
 
-    def _crawl(self) -> Iterable[T]:
-        raise NotImplementedError()
+    def _crawl(self) -> Iterable[DirectFsAccess]:
+        return []
+        # TODO raise NotImplementedError() once CrawlerBase supports empty snapshots

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

@@ -337,7 +337,7 @@ def __repr__(self):
 
     @property
     def lineage(self) -> list[LineageAtom]:
-        return [LineageAtom("path", str(self.path))]
+        return [LineageAtom(object_type="PATH", object_id=str(self.path))]
 
 
 class SourceContainer(abc.ABC):

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

@@ -1,10 +1,11 @@
+import dataclasses
 import functools
 import logging
 import shutil
 import tempfile
 from collections.abc import Generator, Iterable
 from contextlib import contextmanager
-from dataclasses import dataclass, asdict
+from dataclasses import dataclass
 from datetime import datetime, timezone
 from importlib import metadata
 from pathlib import Path
@@ -28,10 +29,9 @@
     guess_encoding,
 )
 from databricks.labs.ucx.source_code.directfs_access import (
-    DirectFsAccess,
     LineageAtom,
     DirectFsAccessCrawler,
-    DirectFsAccessInPath,
+    DirectFsAccess,
 )
 from databricks.labs.ucx.source_code.graph import (
     Dependency,
@@ -86,8 +86,8 @@ def __repr__(self):
     @property
     def lineage(self) -> list[LineageAtom]:
         job_name = (None if self._job.settings is None else self._job.settings.name) or "unknown job"
-        job_lineage = LineageAtom("job", str(self._job.job_id), {"name": job_name})
-        task_lineage = LineageAtom("task", self._task.task_key)
+        job_lineage = LineageAtom("JOB", str(self._job.job_id), {"name": job_name})
+        task_lineage = LineageAtom("TASK", self._task.task_key)
         return [job_lineage, task_lineage]
 
 
@@ -363,7 +363,7 @@ def refresh_report(self, sql_backend: SqlBackend, inventory_database: str):
         logger.info(f"Running {tasks} linting tasks in parallel...")
         job_results, errors = Threads.gather('linting workflows', tasks)
         job_problems: list[JobProblem] = []
-        job_dfsas: list[DirectFsAccessInPath] = []
+        job_dfsas: list[DirectFsAccess] = []
         for problems, dfsas in job_results:
             job_problems.extend(problems)
             job_dfsas.extend(dfsas)
@@ -378,7 +378,7 @@ def refresh_report(self, sql_backend: SqlBackend, inventory_database: str):
         if len(errors) > 0:
             raise ManyError(errors)
 
-    def lint_job(self, job_id: int) -> tuple[list[JobProblem], list[DirectFsAccessInPath]]:
+    def lint_job(self, job_id: int) -> tuple[list[JobProblem], list[DirectFsAccess]]:
         try:
             job = self._ws.jobs.get(job_id)
         except NotFound:
@@ -393,9 +393,9 @@ def lint_job(self, job_id: int) -> tuple[list[JobProblem], list[DirectFsAccessIn
 
     _UNKNOWN = Path('<UNKNOWN>')
 
-    def _lint_job(self, job: jobs.Job) -> tuple[list[JobProblem], list[DirectFsAccessInPath]]:
+    def _lint_job(self, job: jobs.Job) -> tuple[list[JobProblem], list[DirectFsAccess]]:
         problems: list[JobProblem] = []
-        dfsas: list[DirectFsAccessInPath] = []
+        dfsas: list[DirectFsAccess] = []
         assert job.job_id is not None
         assert job.settings is not None
         assert job.settings.name is not None
@@ -420,9 +420,9 @@ def _lint_job(self, job: jobs.Job) -> tuple[list[JobProblem], list[DirectFsAcces
                     end_col=advice.advice.end_col,
                 )
                 problems.append(job_problem)
-            assessment_start = datetime.now()
-            task_dfsas = self._collect_task_dfsas(task, job, graph, session_state)
-            assessment_end = datetime.now()
+            assessment_start = datetime.now(timezone.utc)
+            task_dfsas = self._collect_task_dfsas(job, task, graph, session_state)
+            assessment_end = datetime.now(timezone.utc)
             for dfsa in task_dfsas:
                 dfsa = dfsa.replace_assessment_infos(assessment_start=assessment_start, assessment_end=assessment_end)
                 dfsas.append(dfsa)
@@ -462,15 +462,17 @@ def _lint_task(
         yield from walker
 
     def _collect_task_dfsas(
-        self, task: jobs.Task, job: jobs.Job, graph: DependencyGraph, session_state: CurrentSessionState
-    ) -> Iterable[DirectFsAccessInPath]:
-        collector = DfsaCollectorWalker(graph, set(), self._path_lookup, session_state)
-        assert job.settings is not None  # as already done in _lint_job
-        job_name = job.settings.name
-        for dfsa in collector:
-            yield DirectFsAccessInPath(**asdict(dfsa)).replace_job_infos(
-                job_id=job.job_id, job_name=job_name, task_key=task.task_key
-            )
+        self, job: jobs.Job, task: jobs.Task, graph: DependencyGraph, session_state: CurrentSessionState
+    ) -> Iterable[DirectFsAccess]:
+        # walker doesn't register lineage for job/task
+        job_id = str(job.job_id)
+        job_name = job.settings.name if job.settings and job.settings.name else "<anonymous>"
+        for dfsa in DfsaCollectorWalker(graph, set(), self._path_lookup, session_state):
+            atoms = [
+                LineageAtom(object_type="JOB", object_id=job_id, other={"name": job_name}),
+                LineageAtom(object_type="TASK", object_id=task.task_key),
+            ]
+            yield dataclasses.replace(dfsa, source_lineage=atoms + dfsa.source_lineage)
 
 
 class LintingWalker(DependencyGraphWalker[LocatedAdvice]):

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

@@ -1823,6 +1823,89 @@
   "databricks-labs-pylint": {
     "databricks.labs.pylint": []
   },
+  "databricks-labs-pytester": {
+    "databricks": [],
+    "databricks.labs": [],
+    "databricks.labs.pytester": [],
+    "databricks.labs.pytester.fixtures": [],
+    "databricks.labs.pytester.fixtures.baseline": [
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /"
+      }
+    ],
+    "databricks.labs.pytester.fixtures.catalog": [
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /"
+      },
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: dbfs:/tmp/ucx_test_"
+      },
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: dbfs:/user/hive/warehouse/"
+      },
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: s3a://databricks-datasets-oregon/delta-sharing/share/open-datasets.share"
+      }
+    ],
+    "databricks.labs.pytester.fixtures.compute": [],
+    "databricks.labs.pytester.fixtures.connect": [],
+    "databricks.labs.pytester.fixtures.environment": [],
+    "databricks.labs.pytester.fixtures.iam": [
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /members"
+      },
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /users/groups/"
+      }
+    ],
+    "databricks.labs.pytester.fixtures.ml": [
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /"
+      },
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /api/2.0/feature-store/feature-tables/create"
+      },
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /api/2.0/feature-store/feature-tables/delete"
+      }
+    ],
+    "databricks.labs.pytester.fixtures.notebooks": [
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /Repos/"
+      },
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /Users/"
+      },
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /dummy-"
+      },
+      {
+        "code": "direct-filesystem-access",
+        "message": "The use of direct filesystem references is deprecated: /sdk-"
+      }
+    ],
+    "databricks.labs.pytester.fixtures.permissions": [],
+    "databricks.labs.pytester.fixtures.plugin": [],
+    "databricks.labs.pytester.fixtures.redash": [],
+    "databricks.labs.pytester.fixtures.secrets": [],
+    "databricks.labs.pytester.fixtures.sql": [],
+    "databricks.labs.pytester.fixtures.unwrap": [],
+    "databricks.labs.pytester.fixtures.watchdog": [],
+    "tests.resources.hatchling-whl.src.hatchling_whl": []
+  },
   "databricks-labs-ucx": {
     "databricks.labs.ucx": []
   },