From f67cda7c9dedf6e89e36aa65c80a772d81f9ab8d Mon Sep 17 00:00:00 2001
From: Lloyd Dakin <taidakin99lol@gmail.com>
Date: Tue, 28 Nov 2023 13:10:52 -0800
Subject: [PATCH] fix logic filtering smart stacks prior implementation was
 incorrect and after testing learned SMARTSTACK represents an unique id for
 smart stack images. Changed logic to mimic frontend implementation of
 removing intermediates. Creating a dict of max stknum smart stacks and then
 using list comprehension to build a new list of only the max stknums and non
 smart stack images. Tested and working properly.

---
 api/db.py | 33 +++++++++++++++++++++++++--------
 1 file changed, 25 insertions(+), 8 deletions(-)

diff --git a/api/db.py b/api/db.py
index e158453..3eb91c8 100644
--- a/api/db.py
+++ b/api/db.py
@@ -288,17 +288,34 @@ def filtered_images_query(db_address: str, query_filters: list):
             .all()
         session.expunge_all()
     image_pkgs = [image.get_image_pkg() for image in images]
-    image_pkgs = list(filter(filter_img_pkgs_final_sstack, image_pkgs))
+    image_pkgs = filter_img_pkgs_final_sstack(image_pkgs)
     return image_pkgs
 
 # Filter smart stacks to reduce the size of ui payload
-def filter_img_pkgs_final_sstack(img_pkg):
-    try:
-        if img_pkg["SMARTSTK"] == 'no' or img_pkg["SMARTSTK"] == img_pkg["SSTKNUM"]:
-            return True
-        return False
-    except KeyError :
-        return True
+def filter_img_pkgs_final_sstack(image_pkgs):
+    # dict containing max stknum for each unique smartstack
+    max_sstk_nums = {}
+
+    for img_pkg in image_pkgs:
+        smart_stk = img_pkg.get('SMARTSTK')
+        sstk_num = img_pkg.get('SSTKNUM')
+
+        if smart_stk is None or smart_stk == 'no':
+            continue
+        
+        if smart_stk not in max_sstk_nums or sstk_num > max_sstk_nums[smart_stk]:
+            max_sstk_nums[smart_stk] = sstk_num
+    
+    # keep non intermediate img_pkgs, non smart stack img_pkgs, and img_pkgs missing smart stack keys
+    filtered_arr = [
+        img_pkg for img_pkg in image_pkgs
+        if img_pkg.get('SMARTSTK') is None
+        or img_pkg.get('SSTKNUM') is None
+        or img_pkg.get('SMARTSTK') == 'no'
+        or img_pkg['SSTKNUM'] >= max_sstk_nums.get(img_pkg.get('SMARTSTK'), 0)
+    ]
+
+    return filtered_arr
 
 def get_image_by_filename(db_address, base_filename):
     """Gets the image package for the image specified by the filename.