[coordinator] Update ApplyCustomRuleStore interface (#3000)

m3db · Dec 10, 2020 · d497d95 · d497d95
1 parent 76e5e21
commit d497d95
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 8 deletions.
diff --git a/src/cmd/services/m3coordinator/downsample/options.go b/src/cmd/services/m3coordinator/downsample/options.go
@@ -103,7 +103,7 @@ var (
 )
 
 // CustomRuleStoreFn is a function to swap the backend used for the rule stores.
-type CustomRuleStoreFn func(kv.Store) (kv.Store, error)
+type CustomRuleStoreFn func(clusterclient.Client) (kv.TxnStore, error)
 
 // DownsamplerOptions is a set of required downsampler options.
 type DownsamplerOptions struct {

diff --git a/src/query/server/query.go b/src/query/server/query.go
@@ -33,6 +33,7 @@ import (
 	"github.com/m3db/m3/src/aggregator/server"
 	clusterclient "github.com/m3db/m3/src/cluster/client"
 	etcdclient "github.com/m3db/m3/src/cluster/client/etcd"
+	"github.com/m3db/m3/src/cluster/kv"
 	"github.com/m3db/m3/src/cmd/services/m3aggregator/serve"
 	"github.com/m3db/m3/src/cmd/services/m3coordinator/downsample"
 	"github.com/m3db/m3/src/cmd/services/m3coordinator/ingest"
@@ -742,13 +743,17 @@ func newDownsampler(
 			"must set this config for downsampler")
 	}
 
-	kvStore, err := clusterManagementClient.KV()
-	if err != nil {
-		return nil, errors.Wrap(err, "unable to create KV store from the "+
-			"cluster management config client")
-	}
-	if applyCustomRuleStore != nil {
-		kvStore, err = applyCustomRuleStore(kvStore)
+	var kvStore kv.Store
+	var err error
+
+	if applyCustomRuleStore == nil {
+		kvStore, err = clusterManagementClient.KV()
+		if err != nil {
+			return nil, errors.Wrap(err, "unable to create KV store from the "+
+				"cluster management config client")
+		}
+	} else {
+		kvStore, err = applyCustomRuleStore(clusterManagementClient)
 		if err != nil {
 			return nil, errors.Wrap(err, "unable to apply custom rule store")
 		}