From 6644e5e6768eadc81fceb8ca0d41e3531a664934 Mon Sep 17 00:00:00 2001
From: Matthew Jacobson <matthew.jacobson@grafana.com>
Date: Fri, 1 Dec 2023 15:18:41 -0500
Subject: [PATCH] Alerting: Fix migration that is brittle to version downgrades
 (#78976)

---
 .../ualert/org_upgrade_state_mig.go           | 59 +++++++++++++++++++
 1 file changed, 59 insertions(+)

diff --git a/pkg/services/sqlstore/migrations/ualert/org_upgrade_state_mig.go b/pkg/services/sqlstore/migrations/ualert/org_upgrade_state_mig.go
index 5573888be03..7c4c26e4487 100644
--- a/pkg/services/sqlstore/migrations/ualert/org_upgrade_state_mig.go
+++ b/pkg/services/sqlstore/migrations/ualert/org_upgrade_state_mig.go
@@ -2,6 +2,7 @@ package ualert
 
 import (
 	"fmt"
+	"time"
 
 	"xorm.io/xorm"
 
@@ -30,6 +31,11 @@ const (
 // detect transitions between Legacy and UnifiedAlerting by comparing to the desired type in the configuration.
 func CreateOrgMigratedKVStoreEntries(mg *migrator.Migrator) {
 	mg.AddMigration("copy kvstore migration status to each org", &createOrgMigratedKVStoreEntries{})
+
+	// Adds back the entry for orgId=0 if it doesn't exist. This is not strictly necessary, but
+	// helps to cover some issues with createOrgMigratedKVStoreEntries if a grafana instance on UA is downgraded to before
+	// the split for org-specific migrated statuses.
+	mg.AddMigration("add back entry for orgid=0 migrated status", &addBackAnyOrgMigratedEntry{})
 }
 
 type createOrgMigratedKVStoreEntries struct {
@@ -100,3 +106,56 @@ func (c createOrgMigratedKVStoreEntries) Exec(sess *xorm.Session, mg *migrator.M
 
 	return nil
 }
+
+type addBackAnyOrgMigratedEntry struct {
+	migrator.MigrationBase
+}
+
+func (c addBackAnyOrgMigratedEntry) SQL(migrator.Dialect) string {
+	return codeMigration
+}
+
+func (c addBackAnyOrgMigratedEntry) Exec(sess *xorm.Session, mg *migrator.Migrator) error {
+	var anyOrg int64 = 0
+	migratedEntries := make([]kvStoreV1Entry, 0)
+	cond := kvStoreV1Entry{
+		Namespace: &KVNamespace,
+		Key:       &migratedKey,
+	}
+	err := sess.Table("kv_store").Find(&migratedEntries, &cond)
+	if err != nil {
+		return err
+	}
+
+	if len(migratedEntries) == 0 {
+		mg.Logger.Debug("No migrated orgs in kvstore, nothing to set")
+		return nil
+	}
+
+	migratedStatus := "false"
+	for _, migrated := range migratedEntries {
+		if migrated.OrgID != nil && *migrated.OrgID == anyOrg {
+			mg.Logger.Debug("Already have migrated status for orgid=0, nothing to set")
+			return nil
+		}
+		// All org entries should have the same value. If, somehow, they don't, we set it to true if any orgs are migrated.
+		// This is because assuming an org is migrated when it isn't is less bad as force_migration check exists.
+		if migrated.Value == "true" {
+			migratedStatus = "true"
+		}
+	}
+
+	now := time.Now()
+	if _, err := sess.Table("kv_store").Insert(&kvStoreV1Entry{
+		OrgID:     &anyOrg,
+		Namespace: &KVNamespace,
+		Key:       &migratedKey,
+		Value:     migratedStatus,
+		Created:   now,
+		Updated:   now,
+	}); err != nil {
+		mg.Logger.Error("failed to insert orgid=0 migrated status in kvstore", "err", err)
+		return err
+	}
+	return nil
+}