diff --git a/cmd/db.go b/cmd/db.go index 5030b86..89d4626 100644 --- a/cmd/db.go +++ b/cmd/db.go @@ -16,7 +16,6 @@ limitations under the License. package cmd import ( - "encoding/json" "fmt" "os" "path" @@ -25,6 +24,7 @@ import ( "github.com/mrsimonemms/gobblr/pkg/drivers" "github.com/mrsimonemms/gobblr/pkg/gobblr" + "github.com/rs/zerolog/log" "github.com/spf13/cobra" "github.com/spf13/viper" ) @@ -47,6 +47,12 @@ var dbCmd = &cobra.Command{ // // There can be only one PersistentPostRun command. + log.Debug(). + Str("data path", dbOpts.DataPath). + Str("driver", dbOpts.Driver.DriverName()). + Int("retries", int(dbOpts.Retries)). + Msg("Ingesting data into database") + inserted, err := gobblr.Execute(dbOpts.DataPath, dbOpts.Driver, dbOpts.Retries) if err != nil { return err @@ -57,11 +63,9 @@ var dbCmd = &cobra.Command{ return gobblr.Serve(dbOpts.DataPath, dbOpts.Driver, dbOpts.Retries, dbOpts.WebPort) } - jsonData, err := json.MarshalIndent(inserted, "", " ") - if err != nil { - return err - } - fmt.Println(string(jsonData)) + log.Info().Fields(map[string]interface{}{ + "inserted": inserted, + }).Msg("Successfully inserted data") return nil }, diff --git a/cmd/root.go b/cmd/root.go index 69adade..cff778f 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -19,16 +19,29 @@ import ( "fmt" "os" + "github.com/rs/zerolog" "github.com/spf13/cobra" "github.com/spf13/viper" ) -var cfgFile string +var ( + cfgFile string + logLevel string +) // rootCmd represents the base command when called without any subcommands var rootCmd = &cobra.Command{ Use: "gobblr", Short: "Easily ingest test data into your development stack", + PersistentPreRunE: func(cmd *cobra.Command, args []string) error { + if level, err := zerolog.ParseLevel(logLevel); err != nil { + return err + } else { + zerolog.SetGlobalLevel(level) + } + + return nil + }, } // Execute adds all child commands to the root command and sets flags appropriately. @@ -43,7 +56,11 @@ func Execute() { func init() { cobra.OnInitialize(initConfig) + bindEnv("log-level") + + viper.SetDefault("log-level", "info") rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.gobblr.yaml)") + rootCmd.PersistentFlags().StringVar(&logLevel, "log-level", viper.GetString("log-level"), `log level - "trace", "debug", "info", "warn", "error", "fatal", "panic" or "disabled"`) } // initConfig reads in config file and ENV variables if set. diff --git a/pkg/drivers/driver.go b/pkg/drivers/driver.go index 1d3e386..d4df04c 100644 --- a/pkg/drivers/driver.go +++ b/pkg/drivers/driver.go @@ -12,6 +12,9 @@ type Driver interface { // Close the database connection and free up resources Close() error + // Name of the driver + DriverName() string + // Insert data in bulk InsertBulk(table string, data []map[string]interface{}) (inserted int, err error) diff --git a/pkg/drivers/mongodb/mongodb.go b/pkg/drivers/mongodb/mongodb.go index 1451823..209ebea 100644 --- a/pkg/drivers/mongodb/mongodb.go +++ b/pkg/drivers/mongodb/mongodb.go @@ -43,6 +43,10 @@ func (db *MongoDB) Close() error { return db.activeConnection.client.Disconnect(context.TODO()) } +func (db *MongoDB) DriverName() string { + return "mongodb" +} + func (db *MongoDB) InsertBulk(collection string, raw []map[string]interface{}) (int, error) { var err error data := make([]interface{}, 0) diff --git a/pkg/drivers/sql/sql.go b/pkg/drivers/sql/sql.go index 7ab2da9..c70763d 100644 --- a/pkg/drivers/sql/sql.go +++ b/pkg/drivers/sql/sql.go @@ -34,6 +34,10 @@ func (db *SQL) Close() error { return nil } +func (db *SQL) DriverName() string { + return db.driver.Name() +} + func (db *SQL) InsertBulk(table string, data []map[string]interface{}) (int, error) { result := db.activeConnection.Table(table).CreateInBatches(data, 100) if err := result.Error; err != nil { diff --git a/pkg/gobblr/execute.go b/pkg/gobblr/execute.go index 1b6a108..9a3b418 100644 --- a/pkg/gobblr/execute.go +++ b/pkg/gobblr/execute.go @@ -5,16 +5,23 @@ import ( "github.com/cenkalti/backoff/v4" "github.com/mrsimonemms/gobblr/pkg/drivers" + "github.com/rs/zerolog" + "github.com/rs/zerolog/log" ) +var logger *zerolog.Logger = &log.Logger + type Inserted struct { Table string `json:"table"` Count int `json:"count"` } func Execute(dataPath string, db drivers.Driver, retries uint64) ([]Inserted, error) { + attempt := 0 return backoff.RetryWithData( func() ([]Inserted, error) { + logger.Debug().Int("attempt", attempt).Msg("Attempting to gobble data") + return retryExecute(dataPath, db) }, backoff.WithMaxRetries(backoff.NewExponentialBackOff(), retries), @@ -25,21 +32,28 @@ func retryExecute(dataPath string, db drivers.Driver) ([]Inserted, error) { inserted := make([]Inserted, 0) var err error + logCtx := (&log.Logger).With().Str("dbType", db.DriverName()).Logger() + // Connect to database + logCtx.Debug().Msg("Authenticating database") if err := db.Auth(); err != nil { + logCtx.Error().Err(err).Msg("Failed to connect to database") return nil, err } // We've finished - let's clear up after ourselves defer func() { + logCtx.Debug().Msg("Closing database connection") err = db.Close() }() // Find the files files, err := FindFiles(dataPath) if err != nil { + logCtx.Error().Err(err).Msg("Failed to connect to database") return nil, err } + logCtx.Debug().Str("path", dataPath).Int("count", len(files)).Msg("Found files to ingest") // Iterate over each file, delete and then ingest data for _, file := range files { diff --git a/pkg/gobblr/server.go b/pkg/gobblr/server.go index 15ab3fd..8c31d2b 100644 --- a/pkg/gobblr/server.go +++ b/pkg/gobblr/server.go @@ -6,21 +6,22 @@ import ( "github.com/gin-contrib/requestid" "github.com/gin-gonic/gin" - logger "github.com/mrsimonemms/gin-structured-logger" + gsl "github.com/mrsimonemms/gin-structured-logger" "github.com/mrsimonemms/gobblr/pkg/drivers" - "github.com/rs/zerolog/log" ) +var runCount = 0 + func Serve(dataPath string, db drivers.Driver, retries uint64, port int) error { gin.SetMode(gin.ReleaseMode) r := gin.New() r.Use( requestid.New(), - logger.New(), + gsl.New(), gin.Recovery(), func(ctx *gin.Context) { - logger.Get(ctx).Debug().Str("path", ctx.Request.URL.Path).Msg("New HTTP call") + gsl.Get(ctx).Debug().Str("path", ctx.Request.URL.Path).Msg("New HTTP call") }, ) @@ -33,7 +34,7 @@ func Serve(dataPath string, db drivers.Driver, retries uint64, port int) error { // Register the routes r.POST("/data/reset", h.ResetData) - (&log.Logger).Info().Int("port", port).Msg("Starting web server") + logger.Info().Int("port", port).Msg("Starting web server") return r.Run(fmt.Sprintf(":%d", port)) } @@ -46,7 +47,10 @@ type handler struct { // ResetData runs the Execute command whenever it receives a call func (h handler) ResetData(c *gin.Context) { - log := logger.Get(c).With().Logger() + // Increment the run count + runCount++ + + log := gsl.Get(c).With().Logger() inserted, err := Execute(h.DataPath, h.Driver, h.Retries) if err != nil { @@ -63,7 +67,7 @@ func (h handler) ResetData(c *gin.Context) { // Log the result log.Info().Fields(map[string]interface{}{ "inserted": inserted, - }).Msg("Successfully inserted data") + }).Int("run count", runCount).Msg("Successfully inserted data") c.JSON(http.StatusOK, inserted) }