1 files changed, 86 insertions, 0 deletions
diff --git a/db/migrate/20220712025712_update_insert_or_update_vulnerability_reads_function.rb b/db/migrate/20220712025712_update_insert_or_update_vulnerability_reads_function.rb
new file mode 100644
index 00000000000..55aeacfbecd
--- /dev/null
+++ b/db/migrate/20220712025712_update_insert_or_update_vulnerability_reads_function.rb
@@ -0,0 +1,86 @@
+# frozen_string_literal: true
+
+class UpdateInsertOrUpdateVulnerabilityReadsFunction < Gitlab::Database::Migration[2.0]
+  FUNCTION_NAME = 'insert_or_update_vulnerability_reads'
+
+  enable_lock_retries!
+
+  def up
+    execute(<<~SQL)
+      CREATE OR REPLACE FUNCTION #{FUNCTION_NAME}() RETURNS trigger
+          LANGUAGE plpgsql
+          AS $$
+      DECLARE
+        severity smallint;
+        state smallint;
+        report_type smallint;
+        resolved_on_default_branch boolean;
+        present_on_default_branch boolean;
+      BEGIN
+        IF (NEW.vulnerability_id IS NULL AND (TG_OP = 'INSERT' OR TG_OP = 'UPDATE')) THEN
+          RETURN NULL;
+        END IF;
+
+        IF (TG_OP = 'UPDATE' AND OLD.vulnerability_id IS NOT NULL AND NEW.vulnerability_id IS NOT NULL) THEN
+          RETURN NULL;
+        END IF;
+
+        SELECT
+          vulnerabilities.severity, vulnerabilities.state, vulnerabilities.report_type, vulnerabilities.resolved_on_default_branch, vulnerabilities.present_on_default_branch
+        INTO
+          severity, state, report_type, resolved_on_default_branch, present_on_default_branch
+        FROM
+           vulnerabilities
+        WHERE
+          vulnerabilities.id = NEW.vulnerability_id;
+
+        IF present_on_default_branch IS NOT true THEN
+          RETURN NULL;
+        END IF;
+
+        INSERT INTO vulnerability_reads (vulnerability_id, project_id, scanner_id, report_type, severity, state, resolved_on_default_branch, uuid, location_image, cluster_agent_id, casted_cluster_agent_id)
+          VALUES (NEW.vulnerability_id, NEW.project_id, NEW.scanner_id, report_type, severity, state, resolved_on_default_branch, NEW.uuid::uuid, NEW.location->>'image', NEW.location->'kubernetes_resource'->>'agent_id', CAST(NEW.location->'kubernetes_resource'->>'agent_id' AS bigint))
+          ON CONFLICT(vulnerability_id) DO NOTHING;
+        RETURN NULL;
+      END
+      $$;
+    SQL
+  end
+
+  def down
+    execute(<<~SQL)
+      CREATE OR REPLACE FUNCTION #{FUNCTION_NAME}() RETURNS trigger
+          LANGUAGE plpgsql
+          AS $$
+      DECLARE
+        severity smallint;
+        state smallint;
+        report_type smallint;
+        resolved_on_default_branch boolean;
+      BEGIN
+        IF (NEW.vulnerability_id IS NULL AND (TG_OP = 'INSERT' OR TG_OP = 'UPDATE')) THEN
+          RETURN NULL;
+        END IF;
+
+        IF (TG_OP = 'UPDATE' AND OLD.vulnerability_id IS NOT NULL AND NEW.vulnerability_id IS NOT NULL) THEN
+          RETURN NULL;
+        END IF;
+
+        SELECT
+          vulnerabilities.severity, vulnerabilities.state, vulnerabilities.report_type, vulnerabilities.resolved_on_default_branch
+        INTO
+          severity, state, report_type, resolved_on_default_branch
+        FROM
+           vulnerabilities
+        WHERE
+          vulnerabilities.id = NEW.vulnerability_id;
+
+        INSERT INTO vulnerability_reads (vulnerability_id, project_id, scanner_id, report_type, severity, state, resolved_on_default_branch, uuid, location_image, cluster_agent_id, casted_cluster_agent_id)
+          VALUES (NEW.vulnerability_id, NEW.project_id, NEW.scanner_id, report_type, severity, state, resolved_on_default_branch, NEW.uuid::uuid, NEW.location->>'image', NEW.location->'kubernetes_resource'->>'agent_id', CAST(NEW.location->'kubernetes_resource'->>'agent_id' AS bigint))
+          ON CONFLICT(vulnerability_id) DO NOTHING;
+        RETURN NULL;
+      END
+      $$;
+    SQL
+  end
+end