Refactor security worker

data/model/image.py

@@ -1,5 +1,6 @@
 import logging
 import dateutil.parser
+import random
 
 from peewee import JOIN_LEFT_OUTER, fn, SQL
 from datetime import datetime
@@ -7,7 +8,8 @@ from datetime import datetime
 from data.model import (DataModelException, db_transaction, _basequery, storage,
                         InvalidImageException, config)
 from data.database import (Image, Repository, ImageStoragePlacement, Namespace, ImageStorage,
-                           ImageStorageLocation, RepositoryPermission, db_for_update)
+                           ImageStorageLocation, RepositoryPermission, db_for_update,
+                           db_random_func)
 
 
 logger = logging.getLogger(__name__)
@@ -459,3 +461,74 @@ def ensure_image_locations(*names):
 
     data = [{'name': name} for name in insert_names]
     ImageStorageLocation.insert_many(data).execute()
+
+def get_secscan_candidates(engine_version, batch_size):
+  Parent = Image.alias()
+  ParentImageStorage = ImageStorage.alias()
+  rimages = []
+
+  # Collect the images without parents
+  candidates = (Image
+                .select(Image.id)
+                .join(ImageStorage)
+                .where(Image.security_indexed_engine < engine_version,
+                       Image.parent >> None,
+                       ImageStorage.uploading == False)
+                .limit(batch_size*10))
+
+  images = (Image
+            .select(Image, ImageStorage)
+            .join(ImageStorage)
+            .where(Image.id << candidates)
+            .order_by(db_random_func())
+            .limit(batch_size))
+
+  for image in images:
+    rimages.append(image)
+
+  # Collect the images with analyzed parents.
+  candidates = (Image
+                .select(Image.id)
+                .join(Parent, on=(Image.parent == Parent.id))
+                .switch(Image)
+                .join(ImageStorage)
+                .where(Image.security_indexed_engine < engine_version,
+                       Parent.security_indexed == True,
+                       Parent.security_indexed_engine >= engine_version,
+                       ImageStorage.uploading == False)
+                .limit(batch_size*10))
+
+  images = (Image
+            .select(Image, ImageStorage, Parent, ParentImageStorage)
+            .join(Parent, on=(Image.parent == Parent.id))
+            .join(ParentImageStorage, on=(ParentImageStorage.id == Parent.storage))
+            .switch(Image)
+            .join(ImageStorage)
+            .where(Image.id << candidates)
+            .order_by(db_random_func())
+            .limit(batch_size))
+
+  for image in images:
+    rimages.append(image)
+
+  # Shuffle the images, otherwise the images without parents will always be on the top
+  random.shuffle(rimages)
+
+  return rimages
+
+
+def set_secscan_status(image, indexed, version):
+  query = (Image
+           .select()
+           .join(ImageStorage)
+           .where(Image.docker_image_id == image.docker_image_id,
+                  ImageStorage.uuid == image.storage.uuid))
+
+  ids_to_update = [row.id for row in query]
+  if not ids_to_update:
+    return
+
+  (Image
+   .update(security_indexed=indexed, security_indexed_engine=version)
+   .where(Image.id << ids_to_update)
+   .execute())