refactor: prelculate advisory and vulnerability average scores and severities

common/src/db/multi_model.rs

@@ -1,3 +1,4 @@
+use migration::IntoIden;
 use sea_orm::{
     ColumnTrait, DbErr, EntityTrait, FromQueryResult, IntoIdentity, IntoSimpleExpr, Iterable,
     QueryResult, QuerySelect, Select, SelectModel, Selector,
@@ -44,6 +45,12 @@ impl<T: QuerySelect> ColumnsPrefixed for T {
 pub trait SelectIntoMultiModel: Sized {
     fn try_model_columns<E: EntityTrait>(self, entity: E) -> Result<Self, DbErr>;
 
+    fn try_model_columns_excluding<O: EntityTrait>(
+        self,
+        entity: O,
+        excluded: &[O::Column],
+    ) -> Result<Self, DbErr>;
+
     fn try_model_columns_from_alias<E: EntityTrait>(
         self,
         entity: E,
@@ -62,6 +69,27 @@ impl<E: EntityTrait> SelectIntoMultiModel for Select<E> {
         self.try_columns_prefixed(&prefix, O::Column::iter())
     }
 
+    fn try_model_columns_excluding<O: EntityTrait>(
+        self,
+        entity: O,
+        excluded: &[O::Column],
+    ) -> Result<Self, DbErr> {
+        let excluded_names: Vec<_> = excluded
+            .iter()
+            .map(|col| (*col).into_iden().to_string())
+            .collect();
+
+        let columns: Vec<_> = O::Column::iter()
+            .filter(|col| {
+                let col_name = (*col).into_iden().to_string();
+                !excluded_names.contains(&col_name)
+            })
+            .collect();
+
+        let prefix = format!("{}$", entity.module_name());
+        self.try_columns_prefixed(&prefix, columns)
+    }
+
     fn try_model_columns_from_alias<O: EntityTrait>(
         mut self,
         _entity: O,

entity/src/advisory.rs

@@ -8,7 +8,7 @@ use trustify_common::{
     id::{Id, IdError, TryFilterForId},
 };
 
-#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel, SimpleObject)]
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, SimpleObject)]
 #[graphql(complex)]
 #[graphql(concrete(name = "Advisory", params()))]
 #[sea_orm(table_name = "advisory")]
@@ -30,6 +30,8 @@ pub struct Model {
     pub title: Option<String>,
     pub labels: Labels,
     pub source_document_id: Option<Uuid>,
+    pub average_score: Option<f64>,
+    pub average_severity: Option<super::cvss3::Severity>,
 }
 
 #[ComplexObject]

entity/src/cvss3.rs

@@ -1,4 +1,5 @@
 use crate::{advisory, advisory_vulnerability, vulnerability};
+use async_graphql::Enum;
 use sea_orm::entity::prelude::*;
 use std::fmt::{Display, Formatter};
 use trustify_cvss::cvss3;
@@ -331,7 +332,7 @@ impl From<cvss3::Availability> for Availability {
     }
 }
 
-#[derive(Debug, Copy, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum)]
+#[derive(Debug, Copy, Clone, PartialEq, Eq, EnumIter, DeriveActiveEnum, Enum)]
 #[sea_orm(rs_type = "String", db_type = "Enum", enum_name = "cvss3_severity")]
 pub enum Severity {
     #[sea_orm(string_value = "none")]
@@ -391,3 +392,18 @@ impl From<Severity> for cvss3::severity::Severity {
         }
     }
 }
+
+impl std::str::FromStr for Severity {
+    type Err = ();
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s {
+            "none" => Ok(Severity::None),
+            "low" => Ok(Severity::Low),
+            "medium" => Ok(Severity::Medium),
+            "high" => Ok(Severity::High),
+            "critical" => Ok(Severity::Critical),
+            _ => Err(()),
+        }
+    }
+}

entity/src/vulnerability.rs

@@ -1,9 +1,10 @@
-use crate::{advisory, advisory_vulnerability, cvss3, vulnerability_description};
+use crate::cvss3;
+use crate::{advisory, advisory_vulnerability, vulnerability_description};
 use async_graphql::SimpleObject;
 use sea_orm::entity::prelude::*;
 use time::OffsetDateTime;
 
-#[derive(Clone, Debug, PartialEq, Eq, DeriveEntityModel, SimpleObject)]
+#[derive(Clone, Debug, PartialEq, DeriveEntityModel, SimpleObject)]
 #[sea_orm(table_name = "vulnerability")]
 #[graphql(concrete(name = "Vulnerability", params()))]
 pub struct Model {
@@ -15,6 +16,8 @@ pub struct Model {
     pub modified: Option<OffsetDateTime>,
     pub withdrawn: Option<OffsetDateTime>,
     pub cwes: Option<Vec<String>>,
+    pub average_score: Option<f64>,
+    pub average_severity: Option<super::cvss3::Severity>,
 }
 
 #[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]

migration/src/lib.rs

@@ -14,6 +14,7 @@ mod m0000070_perf_adv_vuln4;
 mod m0000080_get_purl_refactor;
 mod m0000090_release_perf;
 mod m0000100_perf_adv_vuln5;
+mod m0000110_alter_aggregate_scores;
 mod m0000970_alter_importer_add_heartbeat;
 
 pub struct Migrator;
@@ -33,6 +34,7 @@ impl MigratorTrait for Migrator {
             Box::new(m0000080_get_purl_refactor::Migration),
             Box::new(m0000090_release_perf::Migration),
             Box::new(m0000100_perf_adv_vuln5::Migration),
+            Box::new(m0000110_alter_aggregate_scores::Migration),
         ]
     }
 }

migration/src/m0000110_alter_aggregate_scores.rs

@@ -0,0 +1,111 @@
+use sea_orm_migration::prelude::*;
+
+#[derive(DeriveMigrationName)]
+pub struct Migration;
+
+#[async_trait::async_trait]
+impl MigrationTrait for Migration {
+    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        // Add columns to advisory
+        manager
+            .alter_table(
+                Table::alter()
+                    .table(Advisory::Table)
+                    .add_column(ColumnDef::new(Advisory::AverageScore).double())
+                    .add_column(
+                        ColumnDef::new(Advisory::AverageSeverity)
+                            .custom(Alias::new("cvss3_severity")),
+                    )
+                    .to_owned(),
+            )
+            .await?;
+
+        // Add columns to vulnerability
+        manager
+            .alter_table(
+                Table::alter()
+                    .table(Vulnerability::Table)
+                    .add_column(ColumnDef::new(Vulnerability::AverageScore).double())
+                    .add_column(
+                        ColumnDef::new(Vulnerability::AverageSeverity)
+                            .custom(Alias::new("cvss3_severity")),
+                    )
+                    .to_owned(),
+            )
+            .await?;
+
+        manager
+            .get_connection()
+            .execute_unprepared(include_str!(
+                "m000080_alter_aggregate_scores_fns/recalculate_cvss_aggregates.sql"
+            ))
+            .await
+            .map(|_| ())?;
+
+        manager
+            .get_connection()
+            .execute_unprepared(include_str!(
+                "m000080_alter_aggregate_scores_fns/update_cvss_aggregates_on_change.sql"
+            ))
+            .await
+            .map(|_| ())?;
+        Ok(())
+    }
+
+    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        // Drop trigger
+        manager
+            .get_connection()
+            .execute_unprepared("DROP TRIGGER IF EXISTS cvss3_insert_update_trigger ON cvss3")
+            .await?;
+
+        // Drop functions
+        manager
+            .get_connection()
+            .execute_unprepared("DROP FUNCTION IF EXISTS update_cvss_aggregates_on_change")
+            .await?;
+
+        manager
+            .get_connection()
+            .execute_unprepared("DROP FUNCTION IF EXISTS recalculate_cvss_aggregates")
+            .await?;
+
+        // Drop columns from vulnerability
+        manager
+            .alter_table(
+                Table::alter()
+                    .table(Vulnerability::Table)
+                    .drop_column(Vulnerability::AverageScore)
+                    .drop_column(Vulnerability::AverageSeverity)
+                    .to_owned(),
+            )
+            .await?;
+
+        // Drop columns from advisory
+        manager
+            .alter_table(
+                Table::alter()
+                    .table(Advisory::Table)
+                    .drop_column(Advisory::AverageScore)
+                    .drop_column(Advisory::AverageSeverity)
+                    .to_owned(),
+            )
+            .await?;
+
+        Ok(())
+    }
+}
+
+#[derive(Iden)]
+enum Advisory {
+    Table,
+    AverageScore,
+    AverageSeverity,
+}
+
+#[derive(Iden)]
+enum Vulnerability {
+    Table,
+    AverageScore,
+    AverageSeverity,
+}

migration/src/m000080_alter_aggregate_scores_fns/recalculate_cvss_aggregates.sql

@@ -0,0 +1,28 @@
+CREATE OR REPLACE FUNCTION recalculate_cvss_aggregates()
+RETURNS void AS $$
+BEGIN
+    -- Update advisories
+    UPDATE advisory SET
+        average_score = sub.avg_score,
+        average_severity = cvss3_severity(sub.avg_score)
+    FROM (
+        SELECT advisory_id, AVG(score) AS avg_score
+        FROM cvss3
+        GROUP BY advisory_id
+    ) AS sub
+    WHERE advisory.id = sub.advisory_id;
+
+    -- Update vulnerabilities
+    UPDATE vulnerability SET
+        average_score = sub.avg_score,
+        average_severity = cvss3_severity(sub.avg_score)
+    FROM (
+        SELECT vulnerability_id, AVG(score) AS avg_score
+        FROM cvss3
+        GROUP BY vulnerability_id
+    ) AS sub
+    WHERE vulnerability.id = sub.vulnerability_id;
+END;
+$$ LANGUAGE plpgsql;
+
+SELECT recalculate_cvss_aggregates();

migration/src/m000080_alter_aggregate_scores_fns/update_cvss_aggregates_on_change.sql

@@ -0,0 +1,38 @@
+CREATE OR REPLACE FUNCTION update_cvss_aggregates_on_change()
+RETURNS trigger AS $$
+BEGIN
+    -- Update advisory aggregate
+    IF NEW.advisory_id IS NOT NULL THEN
+        UPDATE advisory SET
+            average_score = sub.avg_score,
+            average_severity = cvss3_severity(sub.avg_score)
+        FROM (
+            SELECT AVG(score) AS avg_score
+            FROM cvss3
+            WHERE advisory_id = NEW.advisory_id
+        ) AS sub
+        WHERE advisory.id = NEW.advisory_id;
+    END IF;
+
+    -- Update vulnerability aggregate
+    IF NEW.vulnerability_id IS NOT NULL THEN
+        UPDATE vulnerability SET
+            average_score = sub.avg_score,
+            average_severity = cvss3_severity(sub.avg_score)
+        FROM (
+            SELECT AVG(score) AS avg_score
+            FROM cvss3
+            WHERE vulnerability_id = NEW.vulnerability_id
+        ) AS sub
+        WHERE vulnerability.id = NEW.vulnerability_id;
+    END IF;
+
+    RETURN NULL;
+END;
+$$ LANGUAGE plpgsql
+PARALLEL SAFE;
+
+CREATE TRIGGER cvss3_insert_update_trigger
+AFTER INSERT OR UPDATE OR DELETE ON cvss3
+FOR EACH ROW
+EXECUTE FUNCTION update_cvss_aggregates_on_change();

modules/fundamental/src/advisory/model/details/advisory_vulnerability.rs

@@ -2,8 +2,8 @@ use crate::{Error, vulnerability::model::VulnerabilityHead};
 use sea_orm::{ColumnTrait, ConnectionTrait, EntityTrait, LoaderTrait, QueryFilter};
 use serde::{Deserialize, Serialize};
 use trustify_common::memo::Memo;
+use trustify_cvss::cvss3::Cvss3Base;
 use trustify_cvss::cvss3::severity::Severity;
-use trustify_cvss::{cvss3::Cvss3Base, cvss3::score::Score};
 use trustify_entity::{advisory, advisory_vulnerability, cvss3, vulnerability};
 use utoipa::ToSchema;
 
@@ -34,14 +34,6 @@ impl AdvisoryVulnerabilityHead {
         vulnerability: &vulnerability::Model,
         tx: &C,
     ) -> Result<Self, Error> {
-        let cvss3 = cvss3::Entity::find()
-            .filter(cvss3::Column::AdvisoryId.eq(advisory.id))
-            .filter(cvss3::Column::VulnerabilityId.eq(&vulnerability.id))
-            .all(tx)
-            .await?;
-
-        let score = Score::from_iter(cvss3.iter().map(Cvss3Base::from));
-
         let advisory_vuln = advisory_vulnerability::Entity::find()
             .filter(advisory_vulnerability::Column::AdvisoryId.eq(advisory.id))
             .filter(advisory_vulnerability::Column::VulnerabilityId.eq(&vulnerability.id))
@@ -57,8 +49,11 @@ impl AdvisoryVulnerabilityHead {
             };
             Ok(AdvisoryVulnerabilityHead {
                 head,
-                severity: score.severity(),
-                score: score.value(),
+                severity: advisory
+                    .average_severity
+                    .map(|sev| sev.into())
+                    .unwrap_or(Severity::None),
+                score: advisory.average_score.unwrap_or(0.0),
             })
         } else {
             Err(Error::Data(
@@ -72,34 +67,24 @@ impl AdvisoryVulnerabilityHead {
         vulnerabilities: &[vulnerability::Model],
         tx: &C,
     ) -> Result<Vec<Self>, Error> {
-        let cvss3s = vulnerabilities
-            .load_many(
-                cvss3::Entity::find().filter(cvss3::Column::AdvisoryId.eq(advisory.id)),
-                tx,
-            )
-            .await?;
-
         let mut heads = Vec::new();
 
-        for (vuln, cvss3) in vulnerabilities.iter().zip(cvss3s.iter()) {
-            let score = Score::from_iter(cvss3.iter().map(Cvss3Base::from));
-
+        for vuln in vulnerabilities.iter() {
             let advisory_vuln = advisory_vulnerability::Entity::find()
                 .filter(advisory_vulnerability::Column::AdvisoryId.eq(advisory.id))
                 .filter(advisory_vulnerability::Column::VulnerabilityId.eq(&vuln.id))
                 .one(tx)
                 .await?;
             if let Some(advisory_vuln) = advisory_vuln {
-                let head = if vuln.title.is_some() {
-                    VulnerabilityHead::from_vulnerability_entity(vuln, Memo::NotProvided, tx)
-                        .await?
-                } else {
-                    VulnerabilityHead::from_advisory_vulnerability_entity(&advisory_vuln, vuln)
-                };
+                let head =
+                    VulnerabilityHead::from_advisory_vulnerability_entity(&advisory_vuln, vuln);
                 heads.push(AdvisoryVulnerabilityHead {
                     head,
-                    severity: score.severity(),
-                    score: score.value(),
+                    severity: advisory
+                        .average_severity
+                        .map(|sev| sev.into())
+                        .unwrap_or(Severity::None),
+                    score: advisory.average_score.unwrap_or(0.0),
                 });
             }
         }