Alerting: delete orphaned records from kvstore (#40337)

pkg/infra/kvstore/kvstore.go

@@ -7,6 +7,11 @@ import (
 	"github.com/grafana/grafana/pkg/services/sqlstore"
 )
 
+const (
+	// Wildcard to query all organizations
+	AllOrganizations = -1
+)
+
 func ProvideService(sqlStore *sqlstore.SQLStore) KVStore {
 	return &kvStoreSQL{
 		sqlStore: sqlStore,
@@ -19,6 +24,7 @@ type KVStore interface {
 	Get(ctx context.Context, orgId int64, namespace string, key string) (string, bool, error)
 	Set(ctx context.Context, orgId int64, namespace string, key string, value string) error
 	Del(ctx context.Context, orgId int64, namespace string, key string) error
+	Keys(ctx context.Context, orgId int64, namespace string, keyPrefix string) ([]Key, error)
 }
 
 // WithNamespace returns a kvstore wrapper with fixed orgId and namespace.
@@ -48,3 +54,7 @@ func (kv *NamespacedKVStore) Set(ctx context.Context, key string, value string)
 func (kv *NamespacedKVStore) Del(ctx context.Context, key string) error {
 	return kv.kvStore.Del(ctx, kv.orgId, kv.namespace, key)
 }
+
+func (kv *NamespacedKVStore) Keys(ctx context.Context, keyPrefix string) ([]Key, error) {
+	return kv.kvStore.Keys(ctx, kv.orgId, kv.namespace, keyPrefix)
+}

pkg/infra/kvstore/kvstore_test.go

@@ -165,4 +165,77 @@ func TestKVStore(t *testing.T) {
 			require.False(t, ok, "all keys should be deleted at this point")
 		}
 	})
+
+	t.Run("listing existing keys", func(t *testing.T) {
+		kv := createTestableKVStore(t)
+
+		ctx := context.Background()
+
+		namespace, key := "listtest", "listtest"
+
+		testCases := []*TestCase{
+			{
+				OrgId:     1,
+				Namespace: namespace,
+				Key:       key + "_1",
+			},
+			{
+				OrgId:     2,
+				Namespace: namespace,
+				Key:       key + "_2",
+			},
+			{
+				OrgId:     3,
+				Namespace: namespace,
+				Key:       key + "_3",
+			},
+			{
+				OrgId:     4,
+				Namespace: namespace,
+				Key:       key + "_4",
+			},
+			{
+				OrgId:     1,
+				Namespace: namespace,
+				Key:       "other_key",
+			},
+			{
+				OrgId:     4,
+				Namespace: namespace,
+				Key:       "another_one",
+			},
+		}
+
+		for _, tc := range testCases {
+			err := kv.Set(ctx, tc.OrgId, tc.Namespace, tc.Key, tc.Value())
+			require.NoError(t, err)
+		}
+
+		keys, err := kv.Keys(ctx, AllOrganizations, namespace, key[0:6])
+
+		require.NoError(t, err)
+		require.Len(t, keys, 4)
+
+		found := 0
+
+		for _, key := range keys {
+			for _, tc := range testCases {
+				if key.Key == tc.Key {
+					found++
+					break
+				}
+			}
+		}
+
+		require.Equal(t, 4, found, "querying with the wildcard should return 4 records")
+
+		keys, err = kv.Keys(ctx, 1, namespace, key[0:6])
+
+		require.NoError(t, err)
+		require.Len(t, keys, 1, "querying for a specific org should return 1 record")
+
+		keys, err = kv.Keys(ctx, AllOrganizations, "not_existing_namespace", "not_existing_key")
+		require.NoError(t, err, "querying a not existing namespace and key should not throw an error")
+		require.Len(t, keys, 0, "querying a not existing namespace and key should return an empty slice")
+	})
 }

pkg/infra/kvstore/model.go

@@ -19,3 +19,13 @@ type Item struct {
 func (i *Item) TableName() string {
 	return "kv_store"
 }
+
+type Key struct {
+	OrgId     int64
+	Namespace string
+	Key       string
+}
+
+func (i *Key) TableName() string {
+	return "kv_store"
+}

pkg/infra/kvstore/sql.go

@@ -93,3 +93,17 @@ func (kv *kvStoreSQL) Del(ctx context.Context, orgId int64, namespace string, ke
 	})
 	return err
 }
+
+// Keys get all keys for a given namespace and keyPrefix. To query for all
+// organizations the constant 'kvstore.AllOrganizations' can be passed as orgId.
+func (kv *kvStoreSQL) Keys(ctx context.Context, orgId int64, namespace string, keyPrefix string) ([]Key, error) {
+	var keys []Key
+	err := kv.sqlStore.WithDbSession(ctx, func(dbSession *sqlstore.DBSession) error {
+		query := dbSession.Where("namespace = ?", namespace).And("key LIKE ?", keyPrefix+"%")
+		if orgId != AllOrganizations {
+			query.And("org_id = ?", orgId)
+		}
+		return query.Find(&keys)
+	})
+	return keys, err
+}