1 files changed, 81 insertions, 0 deletions
diff --git a/db/migrate/20230803125434_add_has_merge_request_on_vulnerability_reads_trigger.rb b/db/migrate/20230803125434_add_has_merge_request_on_vulnerability_reads_trigger.rb
new file mode 100644
index 00000000000..dc6e4b94d9c
--- /dev/null
+++ b/db/migrate/20230803125434_add_has_merge_request_on_vulnerability_reads_trigger.rb
@@ -0,0 +1,81 @@
+# frozen_string_literal: true
+
+class AddHasMergeRequestOnVulnerabilityReadsTrigger < Gitlab::Database::Migration[2.1]
+  enable_lock_retries!
+
+  include Gitlab::Database::SchemaHelpers
+
+  TRIGGER_ON_INSERT = 'trigger_update_has_merge_request_on_vulnerability_mr_links_update'
+  INSERT_FUNCTION_NAME = 'set_has_merge_request_on_vulnerability_reads'
+
+  TRIGGER_ON_DELETE = 'trigger_update_has_merge_request_on_vulnerability_mr_links_delete'
+  DELETE_FUNCTION_NAME = 'unset_has_merge_request_on_vulnerability_reads'
+
+  def up
+    create_trigger_function(INSERT_FUNCTION_NAME, replace: true) do
+      <<~SQL
+        UPDATE
+          vulnerability_reads
+        SET
+          has_merge_request = true
+        WHERE
+          vulnerability_id = NEW.vulnerability_id AND has_merge_request IS FALSE;
+        RETURN NULL;
+      SQL
+    end
+
+    execute(<<~SQL)
+      CREATE OR REPLACE FUNCTION #{DELETE_FUNCTION_NAME}()
+      RETURNS TRIGGER
+      LANGUAGE plpgsql
+      AS $$
+      DECLARE
+        has_merge_request_links integer;
+      BEGIN
+        PERFORM 1
+        FROM
+          vulnerability_reads
+        WHERE
+          vulnerability_id = OLD.vulnerability_id
+        FOR UPDATE;
+
+        SELECT 1 INTO has_merge_request_links FROM vulnerability_merge_request_links WHERE vulnerability_id = OLD.vulnerability_id LIMIT 1;
+
+        IF (has_merge_request_links = 1) THEN
+          RETURN NULL;
+        END IF;
+
+        UPDATE
+          vulnerability_reads
+        SET
+          has_merge_request = false
+        WHERE
+          vulnerability_id = OLD.vulnerability_id;
+
+        RETURN NULL;
+      END
+      $$;
+    SQL
+
+    execute(<<~SQL)
+      CREATE TRIGGER #{TRIGGER_ON_INSERT}
+      AFTER INSERT ON vulnerability_merge_request_links
+      FOR EACH ROW
+      EXECUTE FUNCTION #{INSERT_FUNCTION_NAME}();
+    SQL
+
+    execute(<<~SQL)
+      CREATE TRIGGER #{TRIGGER_ON_DELETE}
+      AFTER DELETE ON vulnerability_merge_request_links
+      FOR EACH ROW
+      EXECUTE FUNCTION #{DELETE_FUNCTION_NAME}();
+    SQL
+  end
+
+  def down
+    drop_trigger(:vulnerability_merge_request_links, TRIGGER_ON_INSERT)
+    drop_function(INSERT_FUNCTION_NAME)
+    drop_trigger(:vulnerability_merge_request_links, TRIGGER_ON_DELETE)
+    drop_function(DELETE_FUNCTION_NAME)
+  end
+end