From 9f19a4931701f3982bc00e975c094366e222ba5c Mon Sep 17 00:00:00 2001
From: Denis Arh <denis.arh@gmail.com>
Date: Wed, 17 Jul 2019 10:20:41 +0200
Subject: [PATCH] Improve DB conn. code, make retry/timeout configurable

---
 internal/db/connector.go | 40 ++++++++++++++++++++++------------------
 pkg/cli/options/db.go    | 14 ++++++++++++--
 pkg/cli/runner.go        |  2 +-
 3 files changed, 35 insertions(+), 21 deletions(-)

diff --git a/internal/db/connector.go b/internal/db/connector.go
index 124aa122a..a66f47c68 100644
--- a/internal/db/connector.go
+++ b/internal/db/connector.go
@@ -2,37 +2,41 @@ package db
 
 import (
 	"context"
+	"regexp"
 	"time"
 
-	sentry "github.com/getsentry/sentry-go"
+	"github.com/getsentry/sentry-go"
 	"github.com/pkg/errors"
 	"github.com/titpetric/factory"
 	"go.uber.org/zap"
+
+	"github.com/cortezaproject/corteza-server/pkg/cli/options"
 )
 
-const (
-	maxTries = 100
-	delay    = 5 * time.Second
-	timeout  = 1 * time.Minute
+var (
+	dsnMasker = regexp.MustCompile("(.)(?:.*)(.):(.)(?:.*)(.)@")
 )
 
-func TryToConnect(ctx context.Context, log *zap.Logger, name, dsn, profiler string) (db *factory.DB, err error) {
-	factory.Database.Add(name, dsn)
+func TryToConnect(ctx context.Context, log *zap.Logger, name string, opt options.DBOpt) (db *factory.DB, err error) {
+	factory.Database.Add(name, opt.DSN)
 
 	var (
 		connErrCh = make(chan error, 1)
 	)
 
+	// We'll not add this to the general log because we do not want to carry it with us for every query.
+	dsnField := zap.String("dsn", dsnMasker.ReplaceAllString(opt.DSN, "$1****$2:$3****$4@"))
+
 	// This logger is also used inside profiler
 	log = log.Named("database").With(zap.String("name", name))
 
 	defer close(connErrCh)
 
 	log.Debug("connecting to the database",
-		zap.String("dsn", dsn),
-		zap.Int("tries", maxTries),
-		zap.Duration("delay", delay),
-		zap.Duration("timeout", timeout))
+		dsnField,
+		zap.Int("tries", opt.MaxTries),
+		zap.Duration("delay", opt.Delay),
+		zap.Duration("timeout", opt.Timeout))
 
 	go func() {
 		defer sentry.Recover()
@@ -44,7 +48,7 @@ func TryToConnect(ctx context.Context, log *zap.Logger, name, dsn, profiler stri
 		for {
 			try++
 
-			if maxTries <= try {
+			if opt.MaxTries <= try {
 				err = errors.Errorf("could not connect to %q, in %d tries", name, try)
 				return
 			}
@@ -55,21 +59,21 @@ func TryToConnect(ctx context.Context, log *zap.Logger, name, dsn, profiler stri
 					"could not connect to the database",
 					zap.Error(err),
 					zap.Int("try", try),
-					zap.String("dsn", dsn),
-					zap.Float64("delay", delay.Seconds()),
+					dsnField,
+					zap.Float64("delay", opt.Delay.Seconds()),
 				)
 
 				select {
 				case <-ctx.Done():
 					// Forced break
 					break
-				case <-time.After(delay):
+				case <-time.After(opt.Delay):
 					// Wait before next try
 					continue
 				}
 			}
 
-			log.Info("connected to the database", zap.String("dsn", dsn))
+			log.Info("connected to the database", dsnField)
 
 			// Connected
 			break
@@ -82,14 +86,14 @@ func TryToConnect(ctx context.Context, log *zap.Logger, name, dsn, profiler stri
 	select {
 	case err = <-connErrCh:
 		break
-	case <-time.After(timeout):
+	case <-time.After(opt.Timeout):
 		// Wait before next try
 		return nil, errors.Errorf("db init for %q timedout", name)
 	case <-ctx.Done():
 		return nil, errors.Errorf("db connection for %q cancelled", name)
 	}
 
-	switch profiler {
+	switch opt.Profiler {
 	case "stdout":
 		db.Profiler = &factory.Database.ProfilerStdout
 	case "logger":
diff --git a/pkg/cli/options/db.go b/pkg/cli/options/db.go
index f0508dc4b..7fb3a805a 100644
--- a/pkg/cli/options/db.go
+++ b/pkg/cli/options/db.go
@@ -1,9 +1,16 @@
 package options
 
+import (
+	"time"
+)
+
 type (
 	DBOpt struct {
-		DSN      string `env:"DB_DSN"`
-		Profiler string `env:"DB_PROFILER"`
+		DSN      string        `env:"DB_DSN"`
+		Profiler string        `env:"DB_PROFILER"`
+		MaxTries int           `env:"DB_MAX_TRIES"`
+		Delay    time.Duration `env:"DB_CONN_ERR_DELAY"`
+		Timeout  time.Duration `env:"DB_CONN_TIMEOUT"`
 	}
 )
 
@@ -11,6 +18,9 @@ func DB(pfix string) (o *DBOpt) {
 	o = &DBOpt{
 		DSN:      "corteza:corteza@tcp(db:3306)/corteza?collation=utf8mb4_general_ci",
 		Profiler: "none",
+		MaxTries: 100,
+		Delay:    5 * time.Second,
+		Timeout:  1 * time.Minute,
 	}
 
 	fill(o, pfix)
diff --git a/pkg/cli/runner.go b/pkg/cli/runner.go
index 7811f3370..3021e27da 100644
--- a/pkg/cli/runner.go
+++ b/pkg/cli/runner.go
@@ -190,7 +190,7 @@ func (c *Config) Init() {
 	if c.RootCommandDBSetup == nil {
 		c.RootCommandDBSetup = Runners{func(ctx context.Context, cmd *cobra.Command, c *Config) (err error) {
 			if c.DbOpt != nil {
-				_, err = db.TryToConnect(ctx, c.Log, c.DatabaseName, c.DbOpt.DSN, c.DbOpt.Profiler)
+				_, err = db.TryToConnect(ctx, c.Log, c.DatabaseName, *c.DbOpt)
 				if err != nil {
 					return errors.Wrap(err, "could not connect to database")
 				}