Crawlers: append snapshots to history journal, if available

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

@@ -1,3 +1,5 @@
+from __future__ import annotations
+import datetime as dt
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Iterable, Sequence
@@ -6,6 +8,7 @@
 from databricks.labs.lsql.backends import SqlBackend
 from databricks.sdk.errors import NotFound
 
+from databricks.labs.ucx.framework.history import HistoryLog
 from databricks.labs.ucx.framework.utils import escape_sql_identifier
 
 logger = logging.getLogger(__name__)
@@ -21,13 +24,22 @@ class DataclassInstance(Protocol):
 
 
 class CrawlerBase(ABC, Generic[Result]):
-    def __init__(self, backend: SqlBackend, catalog: str, schema: str, table: str, klass: type[Result]):
+    def __init__(
+        self,
+        backend: SqlBackend,
+        catalog: str,
+        schema: str,
+        table: str,
+        klass: type[Result],
+        history_log: HistoryLog | None = None,
+    ):
         """
         Initializes a CrawlerBase instance.
 
         Args:
             backend (SqlBackend): The backend that executes SQL queries:
                 Statement Execution API or Databricks Runtime.
+            history_log: The (optional) history log where (new) snapshots should be saved.
             catalog (str): The catalog name for the inventory persistence.
             schema: The schema name for the inventory persistence.
             table: The table name for the inventory persistence.
@@ -36,6 +48,7 @@ def __init__(self, backend: SqlBackend, catalog: str, schema: str, table: str, k
         self._schema = self._valid(schema)
         self._table = self._valid(table)
         self._backend = backend
+        self._history_appender = history_log.appender(klass) if history_log is not None else None
         self._fetch = backend.fetch
         self._exec = backend.execute
         self._klass = klass
@@ -155,10 +168,15 @@ def _snapshot(self, fetcher: ResultFn, loader: ResultFn, *, force_refresh: bool)
             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}")
+        crawl_start_time = dt.datetime.now(tz=dt.timezone.utc)
         loaded_records = list(loader())
-        self._update_snapshot(loaded_records, mode="overwrite")
+        self._update_snapshot(loaded_records, crawl_start_time=crawl_start_time, mode="overwrite")
         return loaded_records
 
-    def _update_snapshot(self, items: Sequence[Result], mode: Literal["append", "overwrite"] = "append") -> None:
+    def _update_snapshot(
+        self, items: Sequence[Result], *, crawl_start_time: dt.datetime, mode: Literal["append", "overwrite"]
+    ) -> None:
         logger.debug(f"[{self.full_name}] found {len(items)} new records for {self._table}")
         self._backend.save_table(self.full_name, items, self._klass, mode=mode)
+        if self._history_appender:
+            self._history_appender.append_snapshot(items, run_start_time=crawl_start_time)

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

@@ -0,0 +1,158 @@
+from __future__ import annotations
+import dataclasses
+import datetime as dt
+import json
+import logging
+import os
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from functools import cached_property
+from typing import ClassVar, Protocol, TypeVar
+
+from databricks.labs.lsql.backends import SqlBackend
+from databricks.sdk import WorkspaceClient
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, kw_only=True)
+class HistoricalRecord:
+    workspace_id: int
+    """The identifier of the workspace where this record was generated."""
+
+    run_id: int
+    """The identifier of the workflow run that generated this record."""
+
+    snapshot_id: int
+    """An identifier that is unique to the records produced for a given snapshot."""
+
+    run_start_time: dt.datetime
+    """When this record was generated."""
+
+    object_type: str
+    """The inventory table for which this record was generated."""
+
+    object_type_version: int
+    """Versioning of inventory table, for forward compatibility."""
+
+    object_id: list[str]
+    """The type-specific identifier for this inventory record."""
+
+    object_data: dict[str, str]
+    """Type-specific data of the inventory record. Keys are top-level attributes, values are their JSON-encoded values."""
+
+    failures: list[str]
+    """The list of problems associated with the object that this inventory record covers."""
+
+    owner: str
+    """The identity of the account that created this inventory record."""
+
+
+class DataclassInstance(Protocol):
+    __dataclass_fields__: ClassVar[dict]
+    # TODO: Once all record types provide the property: key_fields: ClassVar[Sequence[str]]
+
+
+Record = TypeVar("Record", bound=DataclassInstance)
+
+
+class HistoryLog:
+    __slots__ = ("_ws", "_backend", "_run_id", "_catalog", "_schema", "_table")
+
+    def __init__(
+        self,
+        ws: WorkspaceClient,
+        backend: SqlBackend,
+        run_id: int,
+        catalog: str,
+        schema: str,
+        table: str,
+    ) -> None:
+        self._ws = ws
+        self._backend = backend
+        self._run_id = run_id
+        self._catalog = catalog
+        self._schema = schema
+        self._table = table
+
+    @property
+    def full_name(self) -> str:
+        return f"{self._catalog}.{self._schema}.{self._table}"
+
+    def _append_history_snapshot(self, object_type: str, snapshot: list[HistoricalRecord]) -> None:
+        logger.debug(f"[{self.full_name}] appending {len(snapshot)} new records for {object_type}")
+        # Concurrent writes do not need to be handled here; appends cannot conflict.
+        # TODO: Although documented as conflict-free, verify that this is truly is the case.
+        self._backend.save_table(self.full_name, snapshot, HistoricalRecord, mode="append")
+
+    class Appender:
+        __slots__ = ("_ws", "_object_type", "_object_type_version", "_key_from", "_run_id", "_persist")
+
+        def __init__(
+            self,
+            ws: WorkspaceClient,
+            run_id: int,
+            klass: type[Record],
+            key_from: Callable[[Record], list[str]],
+            persist: Callable[[str, list[HistoricalRecord]], None],
+        ) -> None:
+            self._ws = ws
+            self._run_id = run_id
+            self._object_type = klass.__name__
+            # Versioning support: if the dataclass has a _ucx_version class attribute that is the current version.
+            self._object_type_version = getattr(klass, "_ucx_version") if hasattr(klass, "_ucx_version") else 0
+            self._key_from = key_from
+            self._persist = persist
+
+        @cached_property
+        def _workspace_id(self) -> int:
+            return self._ws.get_workspace_id()
+
+        @cached_property
+        def _owner(self) -> str:
+            current_user = self._ws.current_user.me()
+            owner = current_user.user_name or current_user.id
+            assert owner
+            return owner
+
+        def append_snapshot(self, records: Sequence[Record], *, run_start_time: dt.datetime) -> None:
+            snapshot_id = int.from_bytes(os.urandom(7), byteorder="big")
+            historical_records = [
+                self._inventory_record_to_historical(record, snapshot_id=snapshot_id, run_start_time=run_start_time)
+                for record in records
+            ]
+            self._persist(self._object_type, historical_records)
+
+        def _inventory_record_to_historical(
+            self,
+            record: Record,
+            *,
+            snapshot_id: int,
+            run_start_time: dt.datetime,
+        ) -> HistoricalRecord:
+            object_id = self._key_from(record)
+            object_as_dict = dataclasses.asdict(record)
+            flattened_object_data = {k: json.dumps(v) for k, v in object_as_dict.items()}
+            # TODO: Get failures.
+            failures: list[str] = []
+            return HistoricalRecord(
+                workspace_id=self._workspace_id,
+                run_id=self._run_id,
+                snapshot_id=snapshot_id,
+                run_start_time=run_start_time,
+                object_type=self._object_type,
+                object_type_version=self._object_type_version,
+                object_id=object_id,
+                object_data=flattened_object_data,
+                failures=failures,
+                owner=self._owner,
+            )
+
+    def appender(self, klass: type[Record]) -> Appender:
+        key_fields = getattr(klass, "key_fields", ())
+
+        def key_from(record: Record) -> list[str]:
+            return [getattr(record, field) for field in key_fields]
+
+        return self.Appender(self._ws, self._run_id, klass, key_from, self._append_history_snapshot)

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

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
+import datetime as dt
 import logging
 import sys
 from collections.abc import Sequence, Iterable
@@ -97,14 +98,14 @@ def __init__(self, backend: SqlBackend, schema: str, table: str):
         """
         super().__init__(backend=backend, catalog="hive_metastore", schema=schema, table=table, klass=DirectFsAccess)
 
-    def dump_all(self, dfsas: Sequence[DirectFsAccess]):
+    def dump_all(self, dfsas: Sequence[DirectFsAccess], *, crawl_start_time: dt.datetime) -> None:
         """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
         """
         try:
             # TODO until we historize data, we append all DFSAs
-            self._update_snapshot(dfsas, mode="append")
+            self._update_snapshot(dfsas, crawl_start_time=crawl_start_time, mode="append")
         except DatabricksError as e:
             logger.error("Failed to store DFSAs", exc_info=e)
 

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

@@ -352,6 +352,7 @@ def __init__(
         self._include_job_ids = include_job_ids
 
     def refresh_report(self, sql_backend: SqlBackend, inventory_database: str):
+        crawl_start_time = datetime.now(tz=timezone.utc)
         tasks = []
         all_jobs = list(self._ws.jobs.list())
         logger.info(f"Preparing {len(all_jobs)} linting tasks...")
@@ -374,7 +375,7 @@ def refresh_report(self, sql_backend: SqlBackend, inventory_database: str):
             JobProblem,
             mode='overwrite',
         )
-        self._directfs_crawler.dump_all(job_dfsas)
+        self._directfs_crawler.dump_all(job_dfsas, crawl_start_time=crawl_start_time)
         if len(errors) > 0:
             raise ManyError(errors)
 

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

@@ -84,7 +84,7 @@ def refresh_report(self, sql_backend: SqlBackend, inventory_database: str):
             )
             for dfsa in all_dfsas
         ]
-        self._directfs_crawler.dump_all(all_dfsas)
+        self._directfs_crawler.dump_all(all_dfsas, crawl_start_time=assessment_start)
 
     def _dashboard_ids_in_scope(self) -> list[str]:
         if self._include_dashboard_ids is not None:  # an empty list is accepted

tests/unit/framework/test_crawlers.py

@@ -1,13 +1,14 @@
 from collections.abc import Iterable
 from dataclasses import dataclass
-from unittest.mock import Mock
+from unittest.mock import ANY, Mock, create_autospec
 
 import pytest
 from databricks.labs.lsql import Row
 from databricks.labs.lsql.backends import MockBackend
 from databricks.sdk.errors import NotFound
 
 from databricks.labs.ucx.framework.crawlers import CrawlerBase, Result, ResultFn
+from databricks.labs.ucx.framework.history import HistoryLog
 
 
 @dataclass
@@ -37,11 +38,12 @@ def __init__(
         schema: str,
         table: str,
         klass: type[Result],
+        history_log: HistoryLog | None = None,
         *,
         fetcher: ResultFn = lambda: [],
         loader: ResultFn = lambda: [],
     ):
-        super().__init__(backend, catalog, schema, table, klass)
+        super().__init__(backend, catalog, schema, table, klass, history_log)
         self._fetcher = fetcher
         self._loader = loader
 
@@ -62,6 +64,14 @@ def test_full_name():
     assert cb.full_name == "a.b.c"
 
 
+def test_history_appender_setup() -> None:
+    """Verify that initializing the crawler also initializes the appender."""
+    mock_history = create_autospec(HistoryLog)
+    _ = _CrawlerFixture(MockBackend(), "a", "b", "c", Baz, history_log=mock_history)
+
+    mock_history.appender.assert_called_once_with(Baz)
+
+
 def test_snapshot_crawls_when_no_prior_crawl() -> None:
     """Check that the crawler is invoked when the fetcher reports that the inventory doesn't exist."""
     mock_backend = MockBackend()
@@ -132,7 +142,7 @@ def test_snapshot_force_refresh_replaces_prior_data() -> None:
     assert [Row(first="second", second=None)] == mock_backend.rows_written_for("a.b.c", mode="overwrite")
 
 
-def test_snapshot_updates_existing_table() -> None:
+def test_snapshot_updates_existing_inventory_table() -> None:
     mock_backend = MockBackend()
     cb = _CrawlerFixture[Baz](mock_backend, "a", "b", "c", Baz, loader=lambda: [Baz(first="first")])
 
@@ -142,7 +152,7 @@ def test_snapshot_updates_existing_table() -> None:
     assert [Row(first="first", second=None)] == mock_backend.rows_written_for("a.b.c", "overwrite")
 
 
-def test_snapshot_updates_new_table() -> None:
+def test_snapshot_updates_new_inventory_table() -> None:
     mock_backend = MockBackend()
 
     def fetcher():
@@ -159,6 +169,19 @@ def fetcher():
     assert [Row(first="first", second=True)] == mock_backend.rows_written_for("a.b.c", "overwrite")
 
 
+def test_snapshot_appends_to_history() -> None:
+    """Verify that when a snapshot is saved, it's also persisted in the history."""
+    mock_history = create_autospec(HistoryLog)
+    cb = _CrawlerFixture[Baz](
+        MockBackend(), "a", "b", "c", Baz, loader=lambda: [Baz(first="first")], history_log=mock_history
+    )
+
+    result = cb.snapshot()
+
+    assert [Baz(first="first")] == result
+    mock_history.appender(Baz).append_snapshot.assert_called_once_with([Baz(first="first")], run_start_time=ANY)
+
+
 def test_snapshot_wrong_error() -> None:
     sql_backend = MockBackend()
 

tests/unit/source_code/test_directfs_access.py

@@ -1,4 +1,4 @@
-from datetime import datetime
+import datetime as dt
 
 from databricks.labs.lsql.backends import MockBackend
 
@@ -12,19 +12,20 @@
 def test_crawler_appends_dfsas():
     backend = MockBackend()
     crawler = DirectFsAccessCrawler.for_paths(backend, "schema")
+    now = dt.datetime.now(tz=dt.timezone.utc)
     dfsas = list(
         DirectFsAccess(
             path=path,
             is_read=False,
             is_write=False,
             source_id="ID",
-            source_timestamp=datetime.now(),
+            source_timestamp=now,
             source_lineage=[LineageAtom(object_type="LINEAGE", object_id="ID")],
-            assessment_start_timestamp=datetime.now(),
-            assessment_end_timestamp=datetime.now(),
+            assessment_start_timestamp=now,
+            assessment_end_timestamp=now,
         )
         for path in ("a", "b", "c")
     )
-    crawler.dump_all(dfsas)
+    crawler.dump_all(dfsas, crawl_start_time=now)
     rows = backend.rows_written_for(crawler.full_name, "append")
     assert len(rows) == 3