Add vulnerability counts and image counts to project and file responses

main.py

@@ -168,6 +168,8 @@ class FileResponse(BaseModel):
     is_active: bool
     created_at: datetime
     updated_at: datetime
+    image_count: Optional[int] = None
+    vulnerability_counts: Optional[dict] = None
 
 
 class ImageResponse(BaseModel):
@@ -181,6 +183,7 @@ class ImageResponse(BaseModel):
     created_at: datetime
     updated_at: datetime
     usage_count: Optional[int] = None
+    vulnerability_counts: Optional[dict] = None
 
 
 class VulnerabilityResponse(BaseModel):
@@ -320,19 +323,37 @@ async def get_projects(
             'total': 0
         }
         
-        for image in project_images:
-            vulnerabilities = db.query(Vulnerability).filter(
-                Vulnerability.image_id == image.id,
-                Vulnerability.is_active == True
-            ).all()
-            
-            for vuln in vulnerabilities:
-                severity = vuln.severity.lower()
-                if severity in vulnerability_counts:
-                    vulnerability_counts[severity] += 1
-                else:
-                    vulnerability_counts['unspecified'] += 1
-                vulnerability_counts['total'] += 1
+        # Count vulnerabilities by severity for all images in this project using SQL COUNT queries
+        if project_images:
+            image_ids = [image.id for image in project_images]
+            vulnerability_counts = {
+                'critical': db.query(Vulnerability).filter(
+                    Vulnerability.image_id.in_(image_ids),
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'critical'
+                ).count(),
+                'high': db.query(Vulnerability).filter(
+                    Vulnerability.image_id.in_(image_ids),
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'high'
+                ).count(),
+                'medium': db.query(Vulnerability).filter(
+                    Vulnerability.image_id.in_(image_ids),
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'medium'
+                ).count(),
+                'low': db.query(Vulnerability).filter(
+                    Vulnerability.image_id.in_(image_ids),
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'low'
+                ).count(),
+                'unspecified': db.query(Vulnerability).filter(
+                    Vulnerability.image_id.in_(image_ids),
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'unspecified'
+                ).count(),
+            }
+            vulnerability_counts['total'] = sum(vulnerability_counts.values())
         
         project_dict = {
             "id": project.id,
@@ -364,13 +385,85 @@ async def get_project_files(
     project_id: int,
     skip: int = 0,
     limit: int = 100,
+    include_image_counts: bool = False,
     db: Session = Depends(get_db)
 ):
     files = db.query(File).filter(
         File.project_id == project_id,
         File.is_active == True
     ).offset(skip).limit(limit).all()
-    return files
+    
+    if not include_image_counts:
+        return files
+    
+    # Add image and vulnerability counts for each file
+    result = []
+    for file in files:
+        # Count distinct images in this file
+        distinct_images = db.query(Image).join(FileImageUsage).filter(
+            FileImageUsage.file_id == file.id,
+            FileImageUsage.is_active == True,
+            Image.is_active == True
+        ).distinct().all()
+        
+        # Count vulnerabilities by severity for all images in this file
+        vulnerability_counts = {
+            'critical': 0,
+            'high': 0,
+            'medium': 0,
+            'low': 0,
+            'unspecified': 0,
+            'total': 0
+        }
+        
+        if distinct_images:
+            # Count vulnerabilities by severity for all images in this file using SQL COUNT queries
+            image_ids = [image.id for image in distinct_images]
+            vulnerability_counts = {
+                'critical': db.query(Vulnerability).filter(
+                    Vulnerability.image_id.in_(image_ids),
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'critical'
+                ).count(),
+                'high': db.query(Vulnerability).filter(
+                    Vulnerability.image_id.in_(image_ids),
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'high'
+                ).count(),
+                'medium': db.query(Vulnerability).filter(
+                    Vulnerability.image_id.in_(image_ids),
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'medium'
+                ).count(),
+                'low': db.query(Vulnerability).filter(
+                    Vulnerability.image_id.in_(image_ids),
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'low'
+                ).count(),
+                'unspecified': db.query(Vulnerability).filter(
+                    Vulnerability.image_id.in_(image_ids),
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'unspecified'
+                ).count(),
+            }
+            vulnerability_counts['total'] = sum(vulnerability_counts.values())
+        
+        file_dict = {
+            "id": file.id,
+            "project_id": file.project_id,
+            "file_path": file.file_path,
+            "branch": file.branch,
+            "file_type": file.file_type,
+            "last_scanned": file.last_scanned,
+            "is_active": file.is_active,
+            "created_at": file.created_at,
+            "updated_at": file.updated_at,
+            "image_count": len(distinct_images),
+            "vulnerability_counts": vulnerability_counts
+        }
+        result.append(file_dict)
+    
+    return result
 
 
 @app.get("/projects/{project_id}/images", response_model=List[ImageResponse])
@@ -378,6 +471,7 @@ async def get_project_images(
     project_id: int,
     skip: int = 0,
     limit: int = 100,
+    include_vulnerability_counts: bool = False,
     db: Session = Depends(get_db)
 ):
     images = db.query(Image).join(FileImageUsage).join(File).filter(
@@ -386,7 +480,7 @@ async def get_project_images(
         FileImageUsage.is_active == True
     ).distinct().offset(skip).limit(limit).all()
     
-    # Add usage count for each image
+    # Add usage count and optionally vulnerability counts for each image
     result = []
     for image in images:
         usage_count = db.query(FileImageUsage).filter(
@@ -406,6 +500,39 @@ async def get_project_images(
             "updated_at": image.updated_at,
             "usage_count": usage_count
         }
+        
+        if include_vulnerability_counts:
+            # Count vulnerabilities by severity for this image using SQL COUNT queries
+            vulnerability_counts = {
+                'critical': db.query(Vulnerability).filter(
+                    Vulnerability.image_id == image.id,
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'critical'
+                ).count(),
+                'high': db.query(Vulnerability).filter(
+                    Vulnerability.image_id == image.id,
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'high'
+                ).count(),
+                'medium': db.query(Vulnerability).filter(
+                    Vulnerability.image_id == image.id,
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'medium'
+                ).count(),
+                'low': db.query(Vulnerability).filter(
+                    Vulnerability.image_id == image.id,
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'low'
+                ).count(),
+                'unspecified': db.query(Vulnerability).filter(
+                    Vulnerability.image_id == image.id,
+                    Vulnerability.is_active == True,
+                    Vulnerability.severity == 'unspecified'
+                ).count(),
+            }
+            vulnerability_counts['total'] = sum(vulnerability_counts.values())
+            image_dict["vulnerability_counts"] = vulnerability_counts
+        
         result.append(image_dict)
     
     return result
@@ -450,7 +577,11 @@ async def get_images(
 
 
 @app.get("/images/{image_id}", response_model=ImageResponse)
-async def get_image(image_id: int, db: Session = Depends(get_db)):
+async def get_image(
+    image_id: int, 
+    include_vulnerability_counts: bool = False,
+    db: Session = Depends(get_db)
+):
     image = db.query(Image).filter(Image.id == image_id).first()
     if not image:
         raise HTTPException(status_code=404, detail="Image not found")
@@ -460,7 +591,7 @@ async def get_image(image_id: int, db: Session = Depends(get_db)):
         FileImageUsage.is_active == True
     ).count()
     
-    return {
+    result = {
         "id": image.id,
         "image_name": image.image_name,
         "tag": image.tag,
@@ -472,6 +603,40 @@ async def get_image(image_id: int, db: Session = Depends(get_db)):
         "updated_at": image.updated_at,
         "usage_count": usage_count
     }
+    
+    if include_vulnerability_counts:
+        # Count vulnerabilities by severity for this image using SQL COUNT queries
+        vulnerability_counts = {
+            'critical': db.query(Vulnerability).filter(
+                Vulnerability.image_id == image.id,
+                Vulnerability.is_active == True,
+                Vulnerability.severity == 'critical'
+            ).count(),
+            'high': db.query(Vulnerability).filter(
+                Vulnerability.image_id == image.id,
+                Vulnerability.is_active == True,
+                Vulnerability.severity == 'high'
+            ).count(),
+            'medium': db.query(Vulnerability).filter(
+                Vulnerability.image_id == image.id,
+                Vulnerability.is_active == True,
+                Vulnerability.severity == 'medium'
+            ).count(),
+            'low': db.query(Vulnerability).filter(
+                Vulnerability.image_id == image.id,
+                Vulnerability.is_active == True,
+                Vulnerability.severity == 'low'
+            ).count(),
+            'unspecified': db.query(Vulnerability).filter(
+                Vulnerability.image_id == image.id,
+                Vulnerability.is_active == True,
+                Vulnerability.severity == 'unspecified'
+            ).count(),
+        }
+        vulnerability_counts['total'] = sum(vulnerability_counts.values())
+        result["vulnerability_counts"] = vulnerability_counts
+    
+    return result
 
 
 @app.get("/images/{image_id}/vulnerabilities", response_model=List[VulnerabilityResponse])
@@ -681,6 +846,7 @@ async def get_scan_job(job_id: int, db: Session = Depends(get_db)):
     }
 
 
+
 @app.get("/scan/status")
 async def get_scan_status(db: Session = Depends(get_db)):
     """Check if there are any running or pending scans"""