KVStore: Extend kvstore to retrieve all items (#50848)

* Extend kvstore to retrieve all items

* Fix comment

* Fix tests

* Change test order

* Move test outside to avoid order conditions

* Update Items to GetAll function and return a map

* Add explanation of map result

* Add description comment

Co-authored-by: Tania B <yalyna.ts@gmail.com>

pkg/infra/kvstore/kvstore.go

@@ -25,6 +25,7 @@ type KVStore interface {
 	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)
+	GetAll(ctx context.Context, orgId int64, namespace string) (map[int64]map[string]string, error)
 }
 
 // WithNamespace returns a kvstore wrapper with fixed orgId and namespace.
@@ -58,3 +59,8 @@ func (kv *NamespacedKVStore) Del(ctx context.Context, key string) error {
 func (kv *NamespacedKVStore) Keys(ctx context.Context, keyPrefix string) ([]Key, error) {
 	return kv.kvStore.Keys(ctx, kv.orgId, kv.namespace, keyPrefix)
 }
+
+// GetAll returns all the keys and values stored per organization. It returns a map of org -> key -> value.
+func (kv *NamespacedKVStore) GetAll(ctx context.Context) (map[int64]map[string]string, error) {
+	return kv.kvStore.GetAll(ctx, kv.orgId, kv.namespace)
+}

pkg/infra/kvstore/kvstore_test.go

@@ -242,3 +242,48 @@ func TestIntegrationKVStore(t *testing.T) {
 		require.Len(t, keys, 0, "querying a not existing namespace and key should return an empty slice")
 	})
 }
+
+func TestGetItems(t *testing.T) {
+	kv := createTestableKVStore(t)
+
+	ctx := context.Background()
+
+	testCases := []*TestCase{
+		{
+			OrgId:     1,
+			Namespace: "testing1",
+			Key:       "key1",
+		},
+		{
+			OrgId:     2,
+			Namespace: "testing1",
+			Key:       "key1",
+		},
+		{
+			OrgId:     2,
+			Namespace: "testing1",
+			Key:       "key2",
+		},
+	}
+
+	for _, tc := range testCases {
+		err := kv.Set(ctx, tc.OrgId, tc.Namespace, tc.Key, tc.Value())
+		require.NoError(t, err)
+	}
+
+	t.Run("Get all values per org", func(t *testing.T) {
+		for _, tc := range testCases {
+			items, err := kv.GetAll(ctx, tc.OrgId, tc.Namespace)
+			require.NoError(t, err)
+			require.Equal(t, items[tc.OrgId][tc.Key], tc.Value())
+		}
+	})
+
+	t.Run("Get all values for all orgs", func(t *testing.T) {
+		items, err := kv.GetAll(ctx, AllOrganizations, "testing1")
+		require.NoError(t, err)
+		for _, tc := range testCases {
+			require.Equal(t, items[tc.OrgId][tc.Key], tc.Value())
+		}
+	})
+}

pkg/infra/kvstore/sql.go

@@ -109,3 +109,28 @@ func (kv *kvStoreSQL) Keys(ctx context.Context, orgId int64, namespace string, k
 	})
 	return keys, err
 }
+
+// GetAll get all items a given namespace and org. To query for all
+// organizations the constant 'kvstore.AllOrganizations' can be passed as orgId.
+// The map result is like map[orgId]map[key]value
+func (kv *kvStoreSQL) GetAll(ctx context.Context, orgId int64, namespace string) (map[int64]map[string]string, error) {
+	var results []Item
+	err := kv.sqlStore.WithDbSession(ctx, func(dbSession *sqlstore.DBSession) error {
+		query := dbSession.Where("namespace = ?", namespace)
+		if orgId != AllOrganizations {
+			query.And("org_id = ?", orgId)
+		}
+
+		return query.Find(&results)
+	})
+
+	items := map[int64]map[string]string{}
+	for _, r := range results {
+		if _, ok := items[*r.OrgId]; !ok {
+			items[*r.OrgId] = map[string]string{}
+		}
+		items[*r.OrgId][*r.Key] = r.Value
+	}
+
+	return items, err
+}

pkg/services/ngalert/notifier/testing.go

@@ -201,6 +201,10 @@ func (fkv *FakeKVStore) Keys(ctx context.Context, orgID int64, namespace string,
 	return keys, nil
 }
 
+func (fkv *FakeKVStore) GetAll(ctx context.Context, orgId int64, namespace string) (map[int64]map[string]string, error) {
+	return nil, nil
+}
+
 type fakeState struct {
 	data string
 }