Read and write checks to/from attestable data

Author: alitheg

atr/attestable.py

@@ -28,6 +28,7 @@
 import atr.log as log
 import atr.models.attestable as models
 import atr.util as util
+from atr.models.attestable import AttestableChecksV1
 
 if TYPE_CHECKING:
     import pathlib
@@ -45,6 +46,10 @@ def github_tp_payload_path(project_name: str, version_name: str, revision_number
     return util.get_attestable_dir() / project_name / version_name / f"{revision_number}.github-tp.json"
 
 
+def attestable_checks_path(project_name: str, version_name: str, revision_number: str) -> pathlib.Path:
+    return util.get_attestable_dir() / project_name / version_name / f"{revision_number}.checks.json"
+
+
 async def github_tp_payload_write(
     project_name: str, version_name: str, revision_number: str, github_payload: dict[str, Any]
 ) -> None:
@@ -89,6 +94,22 @@ async def load_paths(
     return None
 
 
+async def load_checks(
+    project_name: str,
+    version_name: str,
+    revision_number: str,
+) -> list[int] | None:
+    file_path = attestable_checks_path(project_name, version_name, revision_number)
+    if await aiofiles.os.path.isfile(file_path):
+        try:
+            async with aiofiles.open(file_path, encoding="utf-8") as f:
+                data = json.loads(await f.read())
+            return models.AttestableChecksV1.model_validate(data).checks
+        except (json.JSONDecodeError, pydantic.ValidationError) as e:
+            log.warning(f"Could not parse {file_path}: {e}")
+    return []
+
+
 def migrate_to_paths_files() -> int:
     attestable_dir = util.get_attestable_dir()
     if not attestable_dir.is_dir():
@@ -135,7 +156,7 @@ async def paths_to_hashes_and_sizes(directory: pathlib.Path) -> tuple[dict[str,
     return path_to_hash, path_to_size
 
 
-async def write(
+async def write_files_data(
     project_name: str,
     version_name: str,
     revision_number: str,
@@ -145,12 +166,37 @@ async def write(
     path_to_hash: dict[str, str],
     path_to_size: dict[str, int],
 ) -> None:
-    result = _generate(path_to_hash, path_to_size, revision_number, release_policy, uploader_uid, previous)
+    result = _generate_files_data(path_to_hash, path_to_size, revision_number, release_policy, uploader_uid, previous)
     file_path = attestable_path(project_name, version_name, revision_number)
     await util.atomic_write_file(file_path, result.model_dump_json(indent=2))
     paths_result = models.AttestablePathsV1(paths=result.paths)
     paths_file_path = attestable_paths_path(project_name, version_name, revision_number)
     await util.atomic_write_file(paths_file_path, paths_result.model_dump_json(indent=2))
+    checks_file_path = attestable_checks_path(project_name, version_name, revision_number)
+    if not checks_file_path.exists():
+        async with aiofiles.open(checks_file_path, "w", encoding="utf-8") as f:
+            await f.write(models.AttestableChecksV1().model_dump_json(indent=2))
+
+
+async def write_checks_data(
+    project_name: str,
+    version_name: str,
+    revision_number: str,
+    checks: list[int],
+) -> None:
+    log.info(f"Writing checks for {project_name}/{version_name}/{revision_number}: {checks}")
+
+    def modify(content: str) -> str:
+        try:
+            current = AttestableChecksV1.model_validate_json(content).checks
+        except pydantic.ValidationError:
+            current = []
+        new_checks = set(current or [])
+        new_checks.update(checks)
+        result = models.AttestableChecksV1(checks=sorted(new_checks))
+        return result.model_dump_json(indent=2)
+
+    await util.atomic_modify_file(attestable_checks_path(project_name, version_name, revision_number), modify)
 
 
 def _compute_hashes_with_attribution(
@@ -188,7 +234,7 @@ def _compute_hashes_with_attribution(
     return new_hashes
 
 
-def _generate(
+def _generate_files_data(
     path_to_hash: dict[str, str],
     path_to_size: dict[str, int],
     revision_number: str,

atr/db/__init__.py

@@ -155,6 +155,7 @@ async def begin_immediate(self) -> None:
     def check_result(
         self,
         id: Opt[int] = NOT_SET,
+        id_in: Opt[list[int]] = NOT_SET,
         release_name: Opt[str] = NOT_SET,
         revision_number: Opt[str] = NOT_SET,
         checker: Opt[str] = NOT_SET,
@@ -169,8 +170,12 @@ def check_result(
     ) -> Query[sql.CheckResult]:
         query = sqlmodel.select(sql.CheckResult)
 
+        via = sql.validate_instrumented_attribute
+
         if is_defined(id):
             query = query.where(sql.CheckResult.id == id)
+        if is_defined(id_in):
+            query = query.where(via(sql.CheckResult.id).in_(id_in))
         if is_defined(release_name):
             query = query.where(sql.CheckResult.release_name == release_name)
         if is_defined(revision_number):

atr/models/attestable.py

@@ -27,6 +27,11 @@ class HashEntry(schema.Strict):
     uploaders: list[Annotated[tuple[str, str], pydantic.BeforeValidator(tuple)]]
 
 
+class AttestableChecksV1(schema.Strict):
+    version: Literal[1] = 1
+    checks: list[int] = schema.factory(list)
+
+
 class AttestablePathsV1(schema.Strict):
     version: Literal[1] = 1
     paths: dict[str, str] = schema.factory(dict)

atr/storage/readers/checks.py

@@ -21,6 +21,7 @@
 import importlib
 from typing import TYPE_CHECKING
 
+import atr.attestable as attestable
 import atr.db as db
 import atr.hashes as hashes
 import atr.models.sql as sql
@@ -44,22 +45,22 @@ async def _filter_check_results_by_hash(
     if release.latest_revision_number is None:
         raise ValueError("Release has no revision - Invalid state")
     for cr in all_check_results:
-        module_path = cr.checker.rsplit(".", 1)[0]
-        if module_path not in input_hash_by_module:
+        if cr.checker not in input_hash_by_module:
+            module_path = cr.checker.rsplit(".", 1)[0]
             try:
                 module = importlib.import_module(module_path)
                 policy_keys: list[str] = module.INPUT_POLICY_KEYS
                 extra_arg_names: list[str] = getattr(module, "INPUT_EXTRA_ARGS", [])
             except (ImportError, AttributeError):
                 policy_keys = []
                 extra_arg_names = []
-            extra_args = checks.resolve_extra_args(extra_arg_names, release)
+            extra_args = await checks.resolve_extra_args(extra_arg_names, release)
             cache_key = await checks.resolve_cache_key(
                 cr.checker, policy_keys, release, release.latest_revision_number, extra_args, file=rel_path.name
             )
-            input_hash_by_module[module_path] = hashes.compute_dict_hash(cache_key) if cache_key else None
+            input_hash_by_module[cr.checker] = hashes.compute_dict_hash(cache_key) if cache_key else None
 
-        if cr.inputs_hash == input_hash_by_module[module_path]:
+        if cr.inputs_hash == input_hash_by_module[cr.checker]:
             filtered_check_results.append(cr)
     return filtered_check_results
 
@@ -81,31 +82,26 @@ async def by_release_path(self, release: sql.Release, rel_path: pathlib.Path) ->
         if release.latest_revision_number is None:
             raise ValueError("Release has no revision - Invalid state")
 
-        # TODO: Is this potentially too much data? Within a revision I hope it's not too bad?
-
-        query = self.__data.check_result(
-            release_name=release.name,
-            primary_rel_path=str(rel_path),
-        ).order_by(
-            sql.validate_instrumented_attribute(sql.CheckResult.checker).asc(),
-            sql.validate_instrumented_attribute(sql.CheckResult.created).desc(),
-        )
-        all_check_results = await query.all()
-
-        # Filter to checks for the current file version / policy
-        # Cache the computed input hash per checker module, since all results here share the same file and release
-        input_hash_by_module: dict[str, str | None] = {}
-        # TODO: This has a bug - create an archive, it'll scan with a hash and show missing checksum.
-        # Then generate a checksum. It'll re-scan the file with the same hash, but now has one. Two checks shown.
-        filtered_check_results = await _filter_check_results_by_hash(
-            all_check_results, rel_path, input_hash_by_module, release
+        check_ids = await attestable.load_checks(release.project_name, release.version, release.latest_revision_number)
+        all_check_results = (
+            [
+                a
+                for a in await self.__data.check_result(id_in=check_ids)
+                .order_by(
+                    sql.validate_instrumented_attribute(sql.CheckResult.checker).asc(),
+                    sql.validate_instrumented_attribute(sql.CheckResult.created).desc(),
+                )
+                .all()
+            ]
+            if check_ids
+            else []
         )
 
         # Filter out any results that are ignored
         unignored_checks = []
         ignored_checks = []
         match_ignore = await self.ignores_matcher(release.project_name)
-        for cr in filtered_check_results:
+        for cr in all_check_results:
             if not match_ignore(cr):
                 unignored_checks.append(cr)
             else:

atr/storage/readers/releases.py

@@ -21,6 +21,7 @@
 import dataclasses
 import pathlib
 
+import atr.attestable as attestable
 import atr.classify as classify
 import atr.db as db
 import atr.models.sql as sql
@@ -122,10 +123,11 @@ async def __successes_errors_warnings(
         self, release: sql.Release, latest_revision_number: str, info: types.PathInfo
     ) -> None:
         match_ignore = await self.__read_as.checks.ignores_matcher(release.project_name)
+        check_ids = await attestable.load_checks(release.project_name, release.version, latest_revision_number)
+        attestable_checks = [a for a in await self.__data.check_result(id_in=check_ids).all()] if check_ids else []
 
         cs = types.ChecksSubset(
-            release=release,
-            latest_revision_number=latest_revision_number,
+            checks=attestable_checks,
             info=info,
             match_ignore=match_ignore,
         )
@@ -137,23 +139,13 @@ async def __successes_errors_warnings(
         await self.__blocker(cs)
 
     async def __blocker(self, cs: types.ChecksSubset) -> None:
-        blocker = await self.__data.check_result(
-            release_name=cs.release.name,
-            revision_number=cs.latest_revision_number,
-            member_rel_path=None,
-            status=sql.CheckResultStatus.BLOCKER,
-        ).all()
+        blocker = [cr for cr in cs.checks if cr.status == sql.CheckResultStatus.BLOCKER]
         for result in blocker:
             if primary_rel_path := result.primary_rel_path:
                 cs.info.errors.setdefault(pathlib.Path(primary_rel_path), []).append(result)
 
     async def __errors(self, cs: types.ChecksSubset) -> None:
-        errors = await self.__data.check_result(
-            release_name=cs.release.name,
-            revision_number=cs.latest_revision_number,
-            member_rel_path=None,
-            status=sql.CheckResultStatus.FAILURE,
-        ).all()
+        errors = [cr for cr in cs.checks if cr.status == sql.CheckResultStatus.FAILURE]
         for error in errors:
             if cs.match_ignore(error):
                 cs.info.ignored_errors.append(error)
@@ -162,24 +154,14 @@ async def __errors(self, cs: types.ChecksSubset) -> None:
                 cs.info.errors.setdefault(pathlib.Path(primary_rel_path), []).append(error)
 
     async def __successes(self, cs: types.ChecksSubset) -> None:
-        successes = await self.__data.check_result(
-            release_name=cs.release.name,
-            revision_number=cs.latest_revision_number,
-            member_rel_path=None,
-            status=sql.CheckResultStatus.SUCCESS,
-        ).all()
+        successes = [cr for cr in cs.checks if cr.status == sql.CheckResultStatus.SUCCESS]
         for success in successes:
             # Successes cannot be ignored
             if primary_rel_path := success.primary_rel_path:
                 cs.info.successes.setdefault(pathlib.Path(primary_rel_path), []).append(success)
 
     async def __warnings(self, cs: types.ChecksSubset) -> None:
-        warnings = await self.__data.check_result(
-            release_name=cs.release.name,
-            revision_number=cs.latest_revision_number,
-            member_rel_path=None,
-            status=sql.CheckResultStatus.WARNING,
-        ).all()
+        warnings = [cr for cr in cs.checks if cr.status == sql.CheckResultStatus.WARNING]
         for warning in warnings:
             if cs.match_ignore(warning):
                 cs.info.ignored_warnings.append(warning)

atr/storage/types.py

@@ -75,8 +75,7 @@ class PathInfo(schema.Strict):
 
 @dataclasses.dataclass
 class ChecksSubset:
-    release: sql.Release
-    latest_revision_number: str
+    checks: list[sql.CheckResult]
     info: PathInfo
     match_ignore: Callable[[sql.CheckResult], bool]
 

atr/storage/writers/revision.py

@@ -245,7 +245,7 @@ async def create_and_manage(  # noqa: C901
 
             policy = release.release_policy or release.project.release_policy
 
-            await attestable.write(
+            await attestable.write_files_data(
                 project_name,
                 version_name,
                 new_revision.number,