From 6babb9c2af84cf66f46254e88a173c77718b0cc7 Mon Sep 17 00:00:00 2001
From: Julien Riou <julien@riou.xyz>
Date: Sun, 9 Oct 2022 10:58:32 +0200
Subject: [PATCH] feat: Add logging library

Signed-off-by: Julien Riou <julien@riou.xyz>
---
 go.mod  |  5 +++++
 go.sum  | 10 ++++++++++
 main.go | 26 +++++++++++++++++++++++---
 3 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/go.mod b/go.mod
index c59b797..75e1b51 100644
--- a/go.mod
+++ b/go.mod
@@ -8,3 +8,8 @@ require (
 )
 
 require gopkg.in/yaml.v2 v2.4.0
+
+require (
+	github.com/sirupsen/logrus v1.9.0 // indirect
+	golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 // indirect
+)
diff --git a/go.sum b/go.sum
index 187a02a..c96922f 100644
--- a/go.sum
+++ b/go.sum
@@ -1,8 +1,18 @@
+github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/go-sql-driver/mysql v1.6.0 h1:BCTh4TKNUYmOmMUcQ3IipzF5prigylS7XXjEkfCHuOE=
 github.com/go-sql-driver/mysql v1.6.0/go.mod h1:DCzpHaOWr8IXmIStZouvnhqoel9Qv2LBy8hT2VhHyBg=
 github.com/lib/pq v1.10.7 h1:p7ZhMD+KsSRozJr34udlUrhboJwWAgCg34+/ZZNvZZw=
 github.com/lib/pq v1.10.7/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/sirupsen/logrus v1.9.0 h1:trlNQbNUG3OdDrDil03MCb1H2o9nJ1x4/5LYw7byDE0=
+github.com/sirupsen/logrus v1.9.0/go.mod h1:naHLuLoDiP4jHNo9R0sCBMtWGeIprob74mVsIT4qYEQ=
+github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
+github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
+golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8 h1:0A+M6Uqn+Eje4kHMK80dtF3JCXC4ykBgQG4Fe06QRhQ=
+golang.org/x/sys v0.0.0-20220715151400-c0bba94af5f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
 gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
+gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
diff --git a/main.go b/main.go
index f512b8b..9558971 100644
--- a/main.go
+++ b/main.go
@@ -3,9 +3,11 @@ package main
 import (
 	"flag"
 	"fmt"
-	"log"
+	"os"
 	"time"
 
+	log "github.com/sirupsen/logrus"
+
 	_ "github.com/go-sql-driver/mysql"
 	_ "github.com/lib/pq"
 )
@@ -22,11 +24,18 @@ var GitCommit string
 // GoVersion to set Go version at compilation time
 var GoVersion string
 
+func init() {
+	log.SetOutput(os.Stdout)
+}
+
 func main() {
 
 	config := NewConfig()
 
 	version := flag.Bool("version", false, "Print version and exit")
+	quiet := flag.Bool("quiet", false, "Log errors only")
+	verbose := flag.Bool("verbose", false, "Print more logs")
+	debug := flag.Bool("debug", false, "Print even more logs")
 	configFile := flag.String("config", "", "Configuration file")
 	flag.StringVar(&config.Driver, "driver", "postgres", "Database driver (postgres or mysql)")
 	flag.IntVar(&config.Connections, "connections", 1, "Number of concurrent connections to the database")
@@ -47,6 +56,17 @@ func main() {
 		return
 	}
 
+	log.SetLevel(log.WarnLevel)
+	if *debug {
+		log.SetLevel(log.DebugLevel)
+	}
+	if *verbose {
+		log.SetLevel(log.InfoLevel)
+	}
+	if *quiet {
+		log.SetLevel(log.ErrorLevel)
+	}
+
 	if *configFile != "" {
 		err := config.Read(*configFile)
 		if err != nil {
@@ -61,8 +81,8 @@ func main() {
 		log.Fatalf("Cannot perform benchmark: %v", err)
 	}
 	benchmark.Run()
-	log.Printf("Queries: %.0f", benchmark.Queries())
-	log.Printf("Queries per second: %.0f", benchmark.QueriesPerSecond())
+	fmt.Printf("Queries: %.0f\n", benchmark.Queries())
+	fmt.Printf("Queries per second: %.0f\n", benchmark.QueriesPerSecond())
 }
 
 func showVersion() {