From ecc15a2f71752c3b980b5855d57c75fb5f52b319 Mon Sep 17 00:00:00 2001
From: Selene <selenepinillos@gmail.com>
Date: Thu, 23 Jun 2022 11:12:07 +0200
Subject: [PATCH] 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             |  6 ++++
 pkg/infra/kvstore/kvstore_test.go        | 45 ++++++++++++++++++++++++
 pkg/infra/kvstore/sql.go                 | 25 +++++++++++++
 pkg/services/ngalert/notifier/testing.go |  4 +++
 4 files changed, 80 insertions(+)

diff --git a/pkg/infra/kvstore/kvstore.go b/pkg/infra/kvstore/kvstore.go
index 0b7e40b0a39..f9d419f76a3 100644
--- a/pkg/infra/kvstore/kvstore.go
+++ b/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)
+}
diff --git a/pkg/infra/kvstore/kvstore_test.go b/pkg/infra/kvstore/kvstore_test.go
index 78350342ccf..c35ded9dc21 100644
--- a/pkg/infra/kvstore/kvstore_test.go
+++ b/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())
+		}
+	})
+}
diff --git a/pkg/infra/kvstore/sql.go b/pkg/infra/kvstore/sql.go
index cf4894338bf..eae7e808c7d 100644
--- a/pkg/infra/kvstore/sql.go
+++ b/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
+}
diff --git a/pkg/services/ngalert/notifier/testing.go b/pkg/services/ngalert/notifier/testing.go
index a7cd451b11c..3415569bd27 100644
--- a/pkg/services/ngalert/notifier/testing.go
+++ b/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
 }