Use logrus as the new logger instead of fmt.

Author: clems4ever

cmd/datasource-csv/main.go

@@ -8,6 +8,7 @@ import (
 	"os"
 
 	"github.com/clems4ever/go-graphkb/graphkb"
+	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 )
@@ -24,7 +25,7 @@ func NewCSVSource(graphAPI *graphkb.GraphAPI) *CSVSource {
 	csvSource.dataPath = viper.GetString("path")
 
 	if csvSource.dataPath == "" {
-		log.Fatal(fmt.Errorf("Unable to detect CSV file path in configuration. Check patch configuration is provided"))
+		logrus.Fatal("Unable to detect CSV file path in configuration. Check patch configuration is provided")
 	}
 
 	csvSource.graphAPI = graphAPI
@@ -74,7 +75,7 @@ func (cs *CSVSource) Publish() error {
 
 	err = tx.Commit()
 	if err == nil {
-		fmt.Println("CSV data has been sent successfully")
+		logrus.Info("CSV data has been sent successfully")
 	}
 	return err
 }
@@ -118,6 +119,6 @@ func main() {
 		"Provide the path to the configuration file (required)")
 
 	if err := rootCmd.Execute(); err != nil {
-		log.Fatal(err)
+		logrus.Fatal(err)
 	}
 }

cmd/go-graphkb/main.go

@@ -3,13 +3,13 @@ package main
 import (
 	"context"
 	"fmt"
-	"log"
 	"strings"
 	"time"
 
 	"github.com/clems4ever/go-graphkb/internal/database"
 	"github.com/clems4ever/go-graphkb/internal/knowledge"
 	"github.com/clems4ever/go-graphkb/internal/server"
+	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 )
@@ -20,10 +20,10 @@ var Database *database.MariaDB
 // ConfigPath string
 var ConfigPath string
 
-func main() {
-	// Display the code line where log.Fatal appeared for troubleshooting
-	log.SetFlags(log.LstdFlags | log.Lshortfile)
+// LogLevel the log level
+var LogLevel string
 
+func main() {
 	// Also read env variables prefixed with GRAPHKB_.
 	viper.SetEnvPrefix("GRAPHKB")
 	viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
@@ -61,15 +61,32 @@ func main() {
 	}
 
 	rootCmd.PersistentFlags().StringVar(&ConfigPath, "config", "config.yml", "Provide the path to the configuration file (required)")
+	rootCmd.PersistentFlags().StringVar(&LogLevel, "log-level", "info", "The log level among 'debug', 'info', 'warn', 'error'")
 
 	cobra.OnInitialize(onInit)
 
 	rootCmd.AddCommand(cleanCmd, listenCmd, countCmd, readCmd, queryCmd)
 	if err := rootCmd.Execute(); err != nil {
-		log.Fatal(err)
+		logrus.Fatal(err)
 	}
 }
 
+func logLevelParamToSeverity(level string) logrus.Level {
+	switch level {
+	case "debug":
+		return logrus.DebugLevel
+	case "info":
+		return logrus.InfoLevel
+	case "warn":
+		return logrus.WarnLevel
+	case "error":
+		return logrus.ErrorLevel
+	}
+	logrus.Fatal("Provided level %s is not a valid option")
+	// This should never be reached but needed by the compiler
+	return logrus.InfoLevel
+}
+
 func onInit() {
 	viper.SetConfigFile(ConfigPath)
 	viper.SetConfigType("yaml")
@@ -78,11 +95,13 @@ func onInit() {
 		panic(fmt.Errorf("Cannot read configuration file from %s", ConfigPath))
 	}
 
-	fmt.Println("Using config file:", viper.ConfigFileUsed())
+	logrus.Info("Using config file: ", viper.ConfigFileUsed())
+	logrus.SetLevel(logLevelParamToSeverity(LogLevel))
+	logrus.Info("Using log severity: ", LogLevel)
 
 	dbName := viper.GetString("mariadb_database")
 	if dbName == "" {
-		log.Fatal("Please provide database_name option in your configuration file")
+		logrus.Fatal("Please provide database_name option in your configuration file")
 	}
 	Database = database.NewMariaDB(
 		viper.GetString("mariadb_username"),
@@ -95,26 +114,26 @@ func onInit() {
 func count(cmd *cobra.Command, args []string) {
 	countAssets, err := Database.CountAssets()
 	if err != nil {
-		log.Fatal(err)
+		logrus.Fatal(err)
 	}
 
 	countRelations, err := Database.CountRelations()
 	if err != nil {
-		log.Fatal(err)
+		logrus.Fatal(err)
 	}
 	fmt.Printf("%d assets\n%d relations\n", countAssets, countRelations)
 }
 
 func flush(cmd *cobra.Command, args []string) {
 	if err := Database.FlushAll(); err != nil {
-		log.Fatal(err)
+		logrus.Fatal(err)
 	}
-	fmt.Println("Successul flush")
+	logrus.Info("Successul flush")
 }
 
 func listen(cmd *cobra.Command, args []string) {
 	if err := Database.InitializeSchema(); err != nil {
-		log.Fatal(err)
+		logrus.Fatal(err)
 	}
 
 	listenInterface := viper.GetString("server_listen")
@@ -130,7 +149,7 @@ func read(cmd *cobra.Command, args []string) {
 	g := knowledge.NewGraph()
 	err := Database.ReadGraph(args[0], g)
 	if err != nil {
-		log.Fatal(err)
+		logrus.Fatal(err)
 	}
 
 	fmt.Printf("assets = %d\nrelations = %d\n", len(g.Assets()), len(g.Relations()))
@@ -144,15 +163,15 @@ func queryFunc(cmd *cobra.Command, args []string) {
 
 	r, err := q.Query(ctx, args[0])
 	if err != nil {
-		log.Fatal(err)
+		logrus.Fatal(err)
 	}
 
 	resultsCount := 0
 	for r.Cursor.HasMore() {
 		var m interface{}
 		err := r.Cursor.Read(context.Background(), &m)
 		if err != nil {
-			log.Fatal(err)
+			logrus.Fatal(err)
 		}
 
 		doc := m.([]interface{})

internal/client/graph_api.go

@@ -5,6 +5,7 @@ import (
 	"time"
 
 	"github.com/clems4ever/go-graphkb/internal/knowledge"
+	"github.com/sirupsen/logrus"
 )
 
 // GraphAPI represent the graph API from a data source point of view
@@ -52,7 +53,7 @@ func NewGraphAPI(options GraphAPIOptions) *GraphAPI {
 // with previous version of it.
 func (gapi *GraphAPI) CreateTransaction() (*Transaction, error) {
 	if gapi.currentGraph == nil {
-		fmt.Println("transaction: fetching remote graph")
+		logrus.Debug("transaction: fetching remote graph")
 		g, err := gapi.ReadCurrentGraph()
 		if err != nil {
 			return nil, fmt.Errorf("create transaction: %w", err)
@@ -100,7 +101,7 @@ func (gapi *GraphAPI) CreateTransaction() (*Transaction, error) {
 	transaction.onError = func(err error) {
 		// there was an error, we don't know the remote graph state.
 		// we clear the cached copy to refresh in on the next run.
-		fmt.Println("transaction: clearing graph cache because of error:", err)
+		logrus.Debug("transaction: clearing graph cache because of error:", err)
 		gapi.currentGraph = nil
 	}
 

internal/client/transaction.go

@@ -11,6 +11,7 @@ import (
 	"github.com/clems4ever/go-graphkb/internal/knowledge"
 	"github.com/clems4ever/go-graphkb/internal/schema"
 	"github.com/clems4ever/go-graphkb/internal/utils"
+	"github.com/sirupsen/logrus"
 )
 
 // Transaction represent a transaction generating updates by diffing the provided graph against
@@ -62,7 +63,7 @@ func withRetryOnTooManyRequests(fn func() error, backoffFactor float64, maxRetri
 		err := fn()
 		if err != nil {
 			backoffTime := time.Duration(int(math.Pow(backoffFactor, float64(trials)))) * delay
-			fmt.Printf("Sleeping for %d seconds\n", backoffTime/time.Second)
+			logrus.Info("Sleeping for %f seconds\n", backoffTime/time.Second)
 			time.Sleep(backoffTime)
 		} else {
 			return nil
@@ -78,18 +79,18 @@ func withRetryOnTooManyRequests(fn func() error, backoffFactor float64, maxRetri
 func (cgt *Transaction) Commit() error {
 	sg := cgt.newGraph.ExtractSchema()
 
-	fmt.Println("Start uploading the schema of the graph...")
+	logrus.Debug("Start uploading the schema of the graph...")
 	if err := cgt.client.UpdateSchema(sg); err != nil {
 		err := fmt.Errorf("Unable to update the schema of the graph: %v", err)
 		cgt.onError(err)
 		return err
 	}
 
-	fmt.Println("Finished uploading the schema of the graph...")
+	logrus.Debug("Finished uploading the schema of the graph...")
 
 	bulk := knowledge.GenerateGraphUpdatesBulk(cgt.currentGraph, cgt.newGraph)
 
-	fmt.Println("Start uploading the graph...")
+	logrus.Debug("Start uploading the graph...")
 
 	progress := pb.New(len(bulk.GetAssetRemovals()) + len(bulk.GetAssetUpserts()) + len(bulk.GetRelationRemovals()) + len(bulk.GetRelationUpserts()))
 
@@ -103,7 +104,7 @@ func (cgt *Transaction) Commit() error {
 
 	chunkSize := cgt.chunkSize
 
-	fmt.Printf("Assets to be inserted=%d removed=%d, Relations to be inserted=%d removed=%d\n",
+	logrus.Debug("Assets to be inserted=%d removed=%d, Relations to be inserted=%d removed=%d\n",
 		len(bulk.GetAssetUpserts()), len(bulk.GetAssetRemovals()),
 		len(bulk.GetRelationUpserts()), len(bulk.GetRelationRemovals()))
 
@@ -192,7 +193,7 @@ func (cgt *Transaction) Commit() error {
 		}
 	}
 
-	fmt.Println("Finished uploading the graph...")
+	logrus.Debug("Finished uploading the graph...")
 
 	cgt.onSuccess(cgt.newGraph)
 	cgt.newGraph = knowledge.NewGraph()

internal/database/mariadb.go

@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"hash/fnv"
-	"log"
 	"reflect"
 	"strconv"
 	"time"
@@ -31,7 +30,7 @@ func NewMariaDB(username string, password string, host string, databaseName stri
 	db, err := sql.Open("mysql", fmt.Sprintf("%s:%s@(%s)/%s?allowCleartextPasswords=%s", username, password,
 		host, databaseName, strconv.FormatBool(allowCleartextPassword)))
 	if err != nil {
-		log.Fatal(err)
+		logrus.Fatal(err)
 	}
 	db.SetMaxIdleConns(10)
 	db.SetMaxOpenConns(10)
@@ -426,7 +425,7 @@ func (m *MariaDB) RemoveRelations(source string, relations []knowledge.Relation)
 
 // ReadGraph read source subgraph
 func (m *MariaDB) ReadGraph(sourceName string, graph *knowledge.Graph) error {
-	fmt.Printf("Start reading graph of data source with name %s\n", sourceName)
+	logrus.Debugf("Start reading graph of data source with name %s", sourceName)
 	sourceID, err := m.resolveSourceID(sourceName)
 	if err != nil {
 		return fmt.Errorf("Unable to resolve source ID from name %s: %v", sourceName, err)
@@ -507,7 +506,7 @@ WHERE abs.source_id = ?
 	}
 
 	elapsed := time.Since(now)
-	fmt.Printf("Read graph of data source with name %s in %fs\n", sourceName, elapsed.Seconds())
+	logrus.Debugf("Read graph of data source with name %s in %fs\n", sourceName, elapsed.Seconds())
 	return nil
 }
 
@@ -607,7 +606,7 @@ func (m *MariaDB) Query(ctx context.Context, sql knowledge.SQLTranslation) (*kno
 		// Query can take 35 seconds max before being aborted...
 		sql.Query = fmt.Sprintf("SET STATEMENT max_statement_time=%f FOR %s", time.Until(deadline).Seconds()+5, sql.Query)
 	}
-	fmt.Println(sql.Query)
+	logrus.Debug("Query to be executed: ", sql.Query)
 
 	rows, err := m.db.QueryContext(ctx, sql.Query)
 	if err != nil {

internal/handlers/handler_query.go

@@ -3,7 +3,6 @@ package handlers
 import (
 	"context"
 	"encoding/json"
-	"fmt"
 	"net/http"
 	"time"
 
@@ -38,8 +37,7 @@ func PostQuery(database knowledge.GraphDB, queryHistorizer history.Historizer) h
 
 		if requestBody.Query == "" {
 			w.WriteHeader(http.StatusBadRequest)
-			fmt.Println("Empty query parameter")
-			_, err = w.Write([]byte("Empty query parameter"))
+			_, err = w.Write([]byte("Empty 'query' parameter"))
 			if err != nil {
 				ReplyWithInternalError(w, err)
 			}

internal/handlers/replies.go

@@ -2,10 +2,10 @@ package handlers
 
 import (
 	"encoding/json"
-	"fmt"
 	"net/http"
 
 	"github.com/clems4ever/go-graphkb/internal/schema"
+	"github.com/sirupsen/logrus"
 )
 
 // ReplyWithSourceGraph sends a reply containing the schema graph
@@ -23,11 +23,21 @@ func ReplyWithSourceGraph(w http.ResponseWriter, sg *schema.SchemaGraph) {
 
 // ReplyWithInternalError send response with internal error.
 func ReplyWithInternalError(w http.ResponseWriter, err error) {
-	fmt.Println(err)
+	logrus.Error(err)
 	w.WriteHeader(http.StatusInternalServerError)
 	_, werr := w.Write([]byte(err.Error()))
 	if werr != nil {
-		fmt.Println(werr)
+		logrus.Error(werr)
+	}
+}
+
+// ReplyWithBadRequest send response with bad request.
+func ReplyWithBadRequest(w http.ResponseWriter, err error) {
+	logrus.Error(err)
+	w.WriteHeader(http.StatusBadRequest)
+	_, werr := w.Write([]byte(err.Error()))
+	if werr != nil {
+		logrus.Error(werr)
 	}
 }
 
@@ -36,7 +46,7 @@ func ReplyWithUnauthorized(w http.ResponseWriter) {
 	w.WriteHeader(http.StatusUnauthorized)
 	_, werr := w.Write([]byte("Unauthorized"))
 	if werr != nil {
-		fmt.Println(werr)
+		logrus.Error(werr)
 	}
 }
 
@@ -45,6 +55,6 @@ func ReplyWithTooManyRequests(w http.ResponseWriter) {
 	w.WriteHeader(http.StatusTooManyRequests)
 	_, werr := w.Write([]byte("Too Many Requests. Retry later."))
 	if werr != nil {
-		fmt.Println(werr)
+		logrus.Error(werr)
 	}
 }