Hash v1 uploads for torrent chunks

data/model/_basequery.py

@@ -1,8 +1,9 @@
-from peewee import JOIN_LEFT_OUTER, Clause, SQL
+from peewee import Clause, SQL, fn
 from cachetools import lru_cache
 
+from data.model import DataModelException
 from data.database import (Repository, User, Team, TeamMember, RepositoryPermission, TeamRole,
-                           Namespace, Visibility, db_for_update)
+                           Namespace, Visibility, ImageStorage, Image, db_for_update)
 
 
 def prefix_search(field, prefix_query):
@@ -97,3 +98,30 @@ def get_user_organizations(username):
           .join(TeamMember)
           .join(UserAlias, on=(UserAlias.id == TeamMember.user))
           .where(User.organization == True, UserAlias.username == username))
+
+
+def calculate_image_aggregate_size(ancestors_str, image_size, parent_image):
+  ancestors = ancestors_str.split('/')[1:-1]
+  if not ancestors:
+    return image_size
+
+  if parent_image is None:
+    raise DataModelException('Could not load parent image')
+
+  ancestor_size = parent_image.aggregate_size
+  if ancestor_size is not None:
+    return ancestor_size + image_size
+
+  # Fallback to a slower path if the parent doesn't have an aggregate size saved.
+  # TODO: remove this code if/when we do a full backfill.
+  ancestor_size = (ImageStorage
+                   .select(fn.Sum(ImageStorage.image_size))
+                   .join(Image)
+                   .where(Image.id << ancestors)
+                   .scalar())
+  if ancestor_size is None:
+    return None
+
+  return ancestor_size + image_size
+
+