Update check caching to use hash keys of inputs

Author: alitheg

atr/attestable.py

@@ -18,22 +18,21 @@
 from __future__ import annotations
 
 import json
-from typing import TYPE_CHECKING, Any, Final
+from typing import TYPE_CHECKING, Any
 
 import aiofiles
 import aiofiles.os
-import blake3
 import pydantic
 
+import atr.hashes as hashes
 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
 
-_HASH_CHUNK_SIZE: Final[int] = 4 * 1024 * 1024
-
 
 def attestable_path(project_name: str, version_name: str, revision_number: str) -> pathlib.Path:
     return util.get_attestable_dir() / project_name / version_name / f"{revision_number}.json"
@@ -43,18 +42,14 @@ def attestable_paths_path(project_name: str, version_name: str, revision_number:
     return util.get_attestable_dir() / project_name / version_name / f"{revision_number}.paths.json"
 
 
-async def compute_file_hash(path: pathlib.Path) -> str:
-    hasher = blake3.blake3()
-    async with aiofiles.open(path, "rb") as f:
-        while chunk := await f.read(_HASH_CHUNK_SIZE):
-            hasher.update(chunk)
-    return f"blake3:{hasher.hexdigest()}"
-
-
 def github_tp_payload_path(project_name: str, version_name: str, revision_number: str) -> pathlib.Path:
     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:
@@ -99,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():
@@ -140,26 +151,52 @@ async def paths_to_hashes_and_sizes(directory: pathlib.Path) -> tuple[dict[str,
         if "\\" in path_key:
             # TODO: We should centralise this, and forbid some other characters too
             raise ValueError(f"Backslash in path is forbidden: {path_key}")
-        path_to_hash[path_key] = await compute_file_hash(full_path)
+        path_to_hash[path_key] = await hashes.compute_file_hash(full_path)
         path_to_size[path_key] = (await aiofiles.os.stat(full_path)).st_size
     return path_to_hash, path_to_size
 
 
-async def write(
+async def write_files_data(
     project_name: str,
     version_name: str,
     revision_number: str,
+    release_policy: dict[str, Any] | None,
     uploader_uid: str,
     previous: models.AttestableV1 | None,
     path_to_hash: dict[str, str],
     path_to_size: dict[str, int],
 ) -> None:
-    result = _generate(path_to_hash, path_to_size, revision_number, 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(
@@ -197,10 +234,11 @@ 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,
+    release_policy: dict[str, Any] | None,
     uploader_uid: str,
     previous: models.AttestableV1 | None,
 ) -> models.AttestableV1:
@@ -215,4 +253,5 @@ def _generate(
     return models.AttestableV1(
         paths=dict(path_to_hash),
         hashes=dict(new_hashes),
+        policy=release_policy or {},
     )

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,
@@ -164,12 +165,17 @@ def check_result(
         status: Opt[sql.CheckResultStatus] = NOT_SET,
         message: Opt[str] = NOT_SET,
         data: Opt[Any] = NOT_SET,
+        inputs_hash: Opt[str] = NOT_SET,
         _release: bool = False,
     ) -> 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):
@@ -188,6 +194,8 @@ def check_result(
             query = query.where(sql.CheckResult.message == message)
         if is_defined(data):
             query = query.where(sql.CheckResult.data == data)
+        if is_defined(inputs_hash):
+            query = query.where(sql.CheckResult.inputs_hash == inputs_hash)
 
         if _release:
             query = query.options(joined_load(sql.CheckResult.release))

atr/docs/checks.md

@@ -52,7 +52,7 @@ This check records separate checker keys for errors, warnings, and success. Use
 
 For each `.sha256` or `.sha512` file, ATR computes the hash of the referenced artifact and compares it with the expected value. It supports files that contain just the hash as well as files that include a filename and hash on the same line. If the suffix does not indicate `sha256` or `sha512`, the check fails.
 
-The checker key is `atr.tasks.checks.hashing.check`.
+The checker key is `atr.tasks.checks.file_hash.check`.
 
 ### Signature verification
 

atr/docs/tasks.md

@@ -41,7 +41,7 @@ In `atr/tasks/checks` you will find several modules that perform these check tas
 In `atr/tasks/__init__.py` you will see imports for existing modules where you can add an import for new check task, for example:
 
 ```python
-import atr.tasks.checks.hashing as hashing
+import atr.tasks.checks.hashing as file_hash
 import atr.tasks.checks.license as license
 ```
 

atr/get/report.py

@@ -40,10 +40,17 @@ async def selected_path(session: web.Committer, project_name: str, version_name:
 
     # If the draft is not found, we try to get the release candidate
     try:
-        release = await session.release(project_name, version_name, with_committee=True)
+        release = await session.release(
+            project_name, version_name, with_committee=True, with_release_policy=True, with_project_release_policy=True
+        )
     except base.ASFQuartException:
         release = await session.release(
-            project_name, version_name, phase=sql.ReleasePhase.RELEASE_CANDIDATE, with_committee=True
+            project_name,
+            version_name,
+            phase=sql.ReleasePhase.RELEASE_CANDIDATE,
+            with_committee=True,
+            with_release_policy=True,
+            with_project_release_policy=True,
         )
 
     if release.committee is None:

atr/merge.py

@@ -24,6 +24,7 @@
 import aiofiles.os
 
 import atr.attestable as attestable
+import atr.hashes as hashes
 import atr.util as util
 
 if TYPE_CHECKING:
@@ -131,7 +132,7 @@ async def _add_from_prior(
     if (prior_hashes is not None) and (path in prior_hashes):
         n_hashes[path] = prior_hashes[path]
     else:
-        n_hashes[path] = await attestable.compute_file_hash(target)
+        n_hashes[path] = await hashes.compute_file_hash(target)
     stat_result = await aiofiles.os.stat(target)
     n_sizes[path] = stat_result.st_size
     return prior_hashes
@@ -211,7 +212,7 @@ async def _merge_all_present(
         if (prior_hashes is not None) and (path in prior_hashes):
             p_hash = prior_hashes[path]
         else:
-            p_hash = await attestable.compute_file_hash(prior_dir / path)
+            p_hash = await hashes.compute_file_hash(prior_dir / path)
         if p_hash != b_hash:
             # Case 11 via hash: base and new have the same content but prior differs
             return await _replace_with_prior(
@@ -250,7 +251,7 @@ async def _replace_with_prior(
     if (prior_hashes is not None) and (path in prior_hashes):
         n_hashes[path] = prior_hashes[path]
     else:
-        n_hashes[path] = await attestable.compute_file_hash(file_path)
+        n_hashes[path] = await hashes.compute_file_hash(file_path)
     stat_result = await aiofiles.os.stat(file_path)
     n_sizes[path] = stat_result.st_size
     return prior_hashes

atr/models/attestable.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from typing import Annotated, Literal
+from typing import Annotated, Any, Literal
 
 import pydantic
 
@@ -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)
@@ -36,3 +41,4 @@ class AttestableV1(schema.Strict):
     version: Literal[1] = 1
     paths: dict[str, str] = schema.factory(dict)
     hashes: dict[str, HashEntry] = schema.factory(dict)
+    policy: dict[str, Any] = schema.factory(dict)

atr/models/sql.py

@@ -946,7 +946,7 @@ class CheckResult(sqlmodel.SQLModel, table=True):
     data: Any = sqlmodel.Field(
         sa_column=sqlalchemy.Column(sqlalchemy.JSON), **example({"expected": "...", "found": "..."})
     )
-    input_hash: str | None = sqlmodel.Field(default=None, index=True, **example("blake3:7f83b1657ff1fc..."))
+    inputs_hash: str | None = sqlmodel.Field(default=None, index=True, **example("blake3:7f83b1657ff1fc..."))
     cached: bool = sqlmodel.Field(default=False, **example(False))
 
 

atr/storage/readers/checks.py

@@ -18,17 +18,51 @@
 # Removing this will cause circular imports
 from __future__ import annotations
 
+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
 import atr.storage as storage
 import atr.storage.types as types
+import atr.tasks.checks as checks
 import atr.util as util
 
 if TYPE_CHECKING:
     import pathlib
-    from collections.abc import Callable
+    from collections.abc import Callable, Sequence
+
+
+async def _filter_check_results_by_hash(
+    all_check_results: Sequence[sql.CheckResult],
+    rel_path: pathlib.Path,
+    input_hash_by_module: dict[str, str | None],
+    release: sql.Release,
+) -> Sequence[sql.CheckResult]:
+    filtered_check_results = []
+    if release.latest_revision_number is None:
+        raise ValueError("Release has no revision - Invalid state")
+    for cr in all_check_results:
+        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 = 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[cr.checker] = hashes.compute_dict_hash(cache_key) if cache_key else None
+
+        if cr.inputs_hash == input_hash_by_module[cr.checker]:
+            filtered_check_results.append(cr)
+    return filtered_check_results
 
 
 class GeneralPublic:
@@ -48,15 +82,20 @@ 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")
 
-        query = self.__data.check_result(
-            release_name=release.name,
-            revision_number=release.latest_revision_number,
-            primary_rel_path=str(rel_path),
-        ).order_by(
-            sql.validate_instrumented_attribute(sql.CheckResult.checker).asc(),
-            sql.validate_instrumented_attribute(sql.CheckResult.created).desc(),
+        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 []
         )
-        all_check_results = await query.all()
 
         # Filter out any results that are ignored
         unignored_checks = []