From 24e63d69546e1ed0861efaf40294f1e9185b31f9 Mon Sep 17 00:00:00 2001
From: Renan DelValle <renanidelvalle+noreply@gmail.com>
Date: Tue, 8 Jan 2019 16:56:36 -0800
Subject: [PATCH] Add trace level to print out reponse thrift objects. Allows
 user to control wether these are printed or not to avoid pollution.

---
 logger.go |  8 ++++++++
 realis.go | 22 +++++++++++++++++-----
 retry.go  |  2 +-
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/logger.go b/logger.go
index 54a901f..29829b3 100644
--- a/logger.go
+++ b/logger.go
@@ -31,6 +31,7 @@ func (NoopLogger) Println(a ...interface{}) {}
 type LevelLogger struct {
 	Logger
 	debug bool
+	trace bool
 }
 
 func (l *LevelLogger) EnableDebug(enable bool) {
@@ -44,6 +45,13 @@ func (l LevelLogger) DebugPrintf(format string, a ...interface{}) {
 	}
 }
 
+func (l LevelLogger) TracePrintf(format string, a ...interface{}) {
+	if l.debug {
+		l.Print("[TRACE] ")
+		l.Printf(format, a...)
+	}
+}
+
 func (l LevelLogger) DebugPrint(a ...interface{}) {
 	if l.debug {
 		l.Print("[DEBUG] ")
diff --git a/realis.go b/realis.go
index 2597738..69a95bc 100644
--- a/realis.go
+++ b/realis.go
@@ -112,6 +112,7 @@ type RealisConfig struct {
 	clientkey, clientcert       string
 	options                     []ClientOption
 	debug                       bool
+	trace                       bool
 	zkOptions                   []ZKOpt
 }
 
@@ -216,7 +217,7 @@ func ZookeeperOptions(opts ...ZKOpt) ClientOption {
 // Using the word set to avoid name collision with Interface.
 func SetLogger(l Logger) ClientOption {
 	return func(config *RealisConfig) {
-		config.logger = &LevelLogger{l, false}
+		config.logger = &LevelLogger{l, false, false}
 	}
 }
 
@@ -227,6 +228,13 @@ func Debug() ClientOption {
 	}
 }
 
+// Enable debug statements.
+func Trace() ClientOption {
+	return func(config *RealisConfig) {
+		config.trace = true
+	}
+}
+
 func newTJSONTransport(url string, timeout int, config *RealisConfig) (thrift.TTransport, error) {
 	trans, err := defaultTTransport(url, timeout, config)
 	if err != nil {
@@ -262,7 +270,7 @@ func NewRealisClient(options ...ClientOption) (Realis, error) {
 	// Default configs
 	config.timeoutms = 10000
 	config.backoff = defaultBackoff
-	config.logger = &LevelLogger{log.New(os.Stdout, "realis: ", log.Ltime|log.Ldate|log.LUTC), false}
+	config.logger = &LevelLogger{log.New(os.Stdout, "realis: ", log.Ltime|log.Ldate|log.LUTC), false, false}
 
 	// Save options to recreate client if a connection error happens
 	config.options = options
@@ -276,14 +284,17 @@ func NewRealisClient(options ...ClientOption) (Realis, error) {
 
 	// Turn off all logging (including debug)
 	if config.logger == nil {
-		config.logger = &LevelLogger{NoopLogger{}, false}
+		config.logger = &LevelLogger{NoopLogger{}, false, false}
 	}
 
 	// Set a logger if debug has been set to true but no logger has been set
 	if config.logger == nil && config.debug {
-		config.logger = &LevelLogger{log.New(os.Stdout, "realis: ", log.Ltime|log.Ldate|log.LUTC), true}
+		config.logger = &LevelLogger{log.New(os.Stdout, "realis: ", log.Ltime|log.Ldate|log.LUTC), true, false}
 	}
 
+	config.logger.debug = config.debug
+	config.logger.trace = config.trace
+
 	// Note, by this point, a LevelLogger should have been created.
 	config.logger.EnableDebug(config.debug)
 
@@ -350,7 +361,7 @@ func NewRealisClient(options ...ClientOption) (Realis, error) {
 		client:         aurora.NewAuroraSchedulerManagerClientFactory(config.transport, config.protoFactory),
 		readonlyClient: aurora.NewReadOnlySchedulerClientFactory(config.transport, config.protoFactory),
 		adminClient:    aurora.NewAuroraAdminClientFactory(config.transport, config.protoFactory),
-		logger:         LevelLogger{config.logger, config.debug},
+		logger:         LevelLogger{config.logger, config.debug, config.trace},
 		lock:           &sync.Mutex{}}, nil
 }
 
@@ -546,6 +557,7 @@ func (r *realisClient) GetInstanceIds(key *aurora.JobKey, states map[aurora.Sche
 }
 
 func (r *realisClient) GetJobUpdateSummaries(jobUpdateQuery *aurora.JobUpdateQuery) (*aurora.Response, error) {
+
 	r.logger.DebugPrintf("GetJobUpdateSummaries Thrift Payload: %+v\n", jobUpdateQuery)
 
 	resp, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) {
diff --git a/retry.go b/retry.go
index 17b5c1a..e442ee3 100644
--- a/retry.go
+++ b/retry.go
@@ -148,7 +148,7 @@ func (r *realisClient) thriftCallWithRetries(thriftCall auroraThriftCall) (*auro
 
 			resp, clientErr = thriftCall()
 
-			r.logger.DebugPrintf("Aurora Thrift Call ended resp: %v clientErr: %v\n", resp, clientErr)
+			r.logger.TracePrintf("Aurora Thrift Call ended resp: %v clientErr: %v\n", resp, clientErr)
 		}()
 
 		// Check if our thrift call is returning an error. This is a retriable event as we don't know