Compute attestable metadata before starting a database transaction

Author: sbp

atr/attestable.py

@@ -131,18 +131,30 @@ def migrate_to_paths_files() -> int:
     return count
 
 
+async def paths_to_hashes_and_sizes(directory: pathlib.Path) -> tuple[dict[str, str], dict[str, int]]:
+    path_to_hash: dict[str, str] = {}
+    path_to_size: dict[str, int] = {}
+    async for rel_path in util.paths_recursive(directory):
+        full_path = directory / rel_path
+        path_key = str(rel_path)
+        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_size[path_key] = (await aiofiles.os.stat(full_path)).st_size
+    return path_to_hash, path_to_size
+
+
 async def write(
-    release_directory: pathlib.Path,
     project_name: str,
     version_name: str,
     revision_number: str,
     uploader_uid: str,
-    parent_revision_number: str | None,
+    previous: models.AttestableV1 | None,
+    path_to_hash: dict[str, str],
+    path_to_size: dict[str, int],
 ) -> None:
-    previous: models.AttestableV1 | None = None
-    if parent_revision_number is not None:
-        previous = await load(project_name, version_name, parent_revision_number)
-    result = await _generate(release_directory, revision_number, uploader_uid, previous)
+    result = _generate(path_to_hash, path_to_size, revision_number, 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)
@@ -185,34 +197,22 @@ def _compute_hashes_with_attribution(
     return new_hashes
 
 
-async def _generate(
-    directory: pathlib.Path,
+def _generate(
+    path_to_hash: dict[str, str],
+    path_to_size: dict[str, int],
     revision_number: str,
     uploader_uid: str,
     previous: models.AttestableV1 | None,
 ) -> models.AttestableV1:
-    current_path_to_hash: dict[str, str] = {}
     current_hash_to_paths: dict[str, set[str]] = {}
-    path_to_size: dict[str, int] = {}
-
-    async for rel_path in util.paths_recursive(directory):
-        full_path = directory / rel_path
-        path_key = str(rel_path)
-        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}")
-        hash_ref = await compute_file_hash(full_path)
-        file_size = (await aiofiles.os.stat(full_path)).st_size
-
-        current_path_to_hash[path_key] = hash_ref
-        path_to_size[path_key] = file_size
+    for path_key, hash_ref in path_to_hash.items():
         current_hash_to_paths.setdefault(hash_ref, set()).add(path_key)
 
     new_hashes = _compute_hashes_with_attribution(
         current_hash_to_paths, path_to_size, previous, uploader_uid, revision_number
     )
 
     return models.AttestableV1(
-        paths=dict(current_path_to_hash),
+        paths=dict(path_to_hash),
         hashes=dict(new_hashes),
     )

atr/storage/writers/revision.py

@@ -105,7 +105,7 @@ def __init__(
         self.__committee_name = committee_name
 
     @contextlib.asynccontextmanager
-    async def create_and_manage(
+    async def create_and_manage(  # noqa: C901
         self,
         project_name: str,
         version_name: str,
@@ -165,6 +165,16 @@ async def create_and_manage(
             await aioshutil.rmtree(temp_dir)
             raise
 
+        try:
+            path_to_hash, path_to_size = await attestable.paths_to_hashes_and_sizes(temp_dir_path)
+            parent_revision_number = old_revision.number if old_revision else None
+            previous_attestable = None
+            if parent_revision_number is not None:
+                previous_attestable = await attestable.load(project_name, version_name, parent_revision_number)
+        except Exception:
+            await aioshutil.rmtree(temp_dir)
+            raise
+
         async with SafeSession(temp_dir) as data:
             try:
                 # This is the only place where models.Revision is constructed
@@ -206,9 +216,14 @@ async def create_and_manage(
                 await aioshutil.rmtree(temp_dir)
                 raise
 
-            parent_revision_number = old_revision.number if old_revision else None
             await attestable.write(
-                new_revision_dir, project_name, version_name, new_revision.number, asf_uid, parent_revision_number
+                project_name,
+                version_name,
+                new_revision.number,
+                asf_uid,
+                previous_attestable,
+                path_to_hash,
+                path_to_size,
             )
 
             # Commit to end the transaction started by data.begin_immediate