From 583c1cce79c0a6a29887a4fdfd3111ef2b116b8c Mon Sep 17 00:00:00 2001
From: Igor German <igor.german@synesis.ru>
Date: Wed, 5 Jun 2019 17:01:08 +0300
Subject: [PATCH] Add ability to configure etcd store default limit

---
 cmd/regula/cli/cli.go       |  6 ++++--
 cmd/regula/main.go          |  7 ++++---
 store/etcd/rulesets.go      | 13 ++++++++-----
 store/etcd/rulesets_test.go | 18 +++++++++++++-----
 4 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/cmd/regula/cli/cli.go b/cmd/regula/cli/cli.go
index 933f460..279239d 100644
--- a/cmd/regula/cli/cli.go
+++ b/cmd/regula/cli/cli.go
@@ -22,8 +22,9 @@ import (
 // Config holds the server configuration.
 type Config struct {
 	Etcd struct {
-		Endpoints []string `config:"etcd-endpoints"`
-		Namespace string   `config:"etcd-namespace"`
+		Endpoints    []string `config:"etcd-endpoints"`
+		DefaultLimit int      `config:"etcd-default-limit"`
+		Namespace    string   `config:"etcd-namespace"`
 	}
 	Server struct {
 		Address      string        `config:"addr"`
@@ -40,6 +41,7 @@ func LoadConfig(args []string) (*Config, error) {
 	var cfg Config
 	flag := stdflag.NewFlagSet("", stdflag.ContinueOnError)
 	flag.StringVar(&cfg.Etcd.Namespace, "etcd-namespace", "", "etcd namespace to use")
+	flag.IntVar(&cfg.Etcd.DefaultLimit, "etcd-default-limit", 50, "etcd default limit for Get requests")
 	flag.StringVar(&cfg.LogLevel, "log-level", zerolog.DebugLevel.String(), "debug level")
 	cfg.Etcd.Endpoints = []string{"127.0.0.1:2379"}
 	flag.Var(commaSeparatedFlag{&cfg.Etcd.Endpoints}, "etcd-endpoints", "comma separated etcd endpoints")
diff --git a/cmd/regula/main.go b/cmd/regula/main.go
index 1df5d5b..67df4de 100644
--- a/cmd/regula/main.go
+++ b/cmd/regula/main.go
@@ -34,9 +34,10 @@ func main() {
 	defer etcdCli.Close()
 
 	service := etcd.RulesetService{
-		Client:    etcdCli,
-		Namespace: cfg.Etcd.Namespace,
-		Logger:    logger.With().Str("service", "etcd").Logger(),
+		Client:       etcdCli,
+		DefaultLimit: cfg.Etcd.DefaultLimit,
+		Namespace:    cfg.Etcd.Namespace,
+		Logger:       logger.With().Str("service", "etcd").Logger(),
 	}
 
 	srv := server.New(&service, server.Config{
diff --git a/store/etcd/rulesets.go b/store/etcd/rulesets.go
index 4a0a947..4631f20 100644
--- a/store/etcd/rulesets.go
+++ b/store/etcd/rulesets.go
@@ -21,11 +21,14 @@ import (
 	"github.com/segmentio/ksuid"
 )
 
+const maxLimit = 100
+
 // RulesetService manages the rulesets using etcd.
 type RulesetService struct {
-	Client    *clientv3.Client
-	Logger    zerolog.Logger
-	Namespace string
+	Client       *clientv3.Client
+	DefaultLimit int
+	Logger       zerolog.Logger
+	Namespace    string
 }
 
 // List returns all the rulesets entries under the given prefix.
@@ -34,8 +37,8 @@ func (s *RulesetService) List(ctx context.Context, prefix string, limit int, con
 
 	var key string
 
-	if limit < 0 || limit > 100 {
-		limit = 50 // TODO(asdine): make this configurable in future releases.
+	if limit < 0 || limit > maxLimit {
+		limit = s.DefaultLimit
 	}
 
 	if continueToken != "" {
diff --git a/store/etcd/rulesets_test.go b/store/etcd/rulesets_test.go
index 17b1c3d..1c0ddac 100644
--- a/store/etcd/rulesets_test.go
+++ b/store/etcd/rulesets_test.go
@@ -24,11 +24,12 @@ var (
 )
 
 var (
-	dialTimeout = 5 * time.Second
-	endpoints   = []string{"localhost:2379", "etcd:2379"}
+	dialTimeout  = 5 * time.Second
+	endpoints    = []string{"localhost:2379", "etcd:2379"}
+	defaultLimit = 50
 )
 
-func Init() {
+func init() {
 	rand.Seed(time.Now().Unix())
 }
 
@@ -42,8 +43,9 @@ func newEtcdRulesetService(t *testing.T) (*etcd.RulesetService, func()) {
 	require.NoError(t, err)
 
 	s := etcd.RulesetService{
-		Client:    cli,
-		Namespace: fmt.Sprintf("regula-store-tests-%d/", rand.Int()),
+		Client:       cli,
+		DefaultLimit: defaultLimit,
+		Namespace:    fmt.Sprintf("regula-store-tests-%d/", rand.Int()),
 	}
 
 	return &s, func() {
@@ -150,6 +152,12 @@ func TestList(t *testing.T) {
 		entries, err = s.List(context.Background(), "y", -10, "")
 		require.NoError(t, err)
 		require.Len(t, entries.Entries, 5)
+
+		// change default limit and check whether it's applied
+		s.DefaultLimit = 3
+		entries, err = s.List(context.Background(), "y", -10, "")
+		require.NoError(t, err)
+		require.Len(t, entries.Entries, 3)
 	})
 }