diff --git a/examples/client.go b/examples/client.go index b2a57e5..dc71f8e 100644 --- a/examples/client.go +++ b/examples/client.go @@ -18,7 +18,6 @@ import ( "flag" "fmt" "io/ioutil" - "log" "os" "time" @@ -95,7 +94,6 @@ func main() { Factor: 2.0, Jitter: 0.1, }), - realis.SetLogger(log.New(os.Stdout, "realis-debug: ", log.Ldate)), realis.Debug(), } diff --git a/logger.go b/logger.go index 4597dde..05a3e85 100644 --- a/logger.go +++ b/logger.go @@ -33,6 +33,10 @@ type LevelLogger struct { debug bool } +func (l *LevelLogger) EnableDebug(enable bool) { + l.debug = enable +} + func (l LevelLogger) DebugPrintf(format string, a ...interface{}) { if l.debug { l.Print("[DEBUG] ") diff --git a/realis.go b/realis.go index c9291d1..82a5eac 100644 --- a/realis.go +++ b/realis.go @@ -26,9 +26,9 @@ import ( "net/http/cookiejar" "os" "path/filepath" - "time" - + "strings" "sync" + "time" "git.apache.org/thrift.git/lib/go/thrift" "github.com/paypal/gorealis/gen-go/apache/aurora" @@ -96,12 +96,13 @@ type RealisConfig struct { backoff Backoff transport thrift.TTransport protoFactory thrift.TProtocolFactory - logger Logger + logger *LevelLogger InsecureSkipVerify bool certspath string clientkey, clientcert string options []ClientOption debug bool + zkOptions []ZKOpt } var defaultBackoff = Backoff{ @@ -140,8 +141,15 @@ func ZKCluster(cluster *Cluster) ClientOption { } func ZKUrl(url string) ClientOption { + + opts := []ZKOpt{ZKEndpoints(strings.Split(url, ",")...), ZKPath("/aurora/scheduler")} + return func(config *RealisConfig) { - config.cluster = GetDefaultClusterFromZKUrl(url) + if config.zkOptions == nil { + config.zkOptions = opts + } else { + config.zkOptions = append(config.zkOptions, opts...) + } } } @@ -187,10 +195,18 @@ func ClientCerts(clientKey, clientCert string) ClientOption { } } +// Use this option if you'd like to override default settings for connecting to Zookeeper. +// See zk.go for what is possible to set as an option. +func ZookeeperOptions(opts ...ZKOpt) ClientOption { + return func(config *RealisConfig) { + config.zkOptions = opts + } +} + // Using the word set to avoid name collision with Interface. func SetLogger(l Logger) ClientOption { return func(config *RealisConfig) { - config.logger = l + config.logger = &LevelLogger{l, false} } } @@ -232,7 +248,7 @@ func NewRealisClient(options ...ClientOption) (Realis, error) { // Default configs config.timeoutms = 10000 config.backoff = defaultBackoff - config.logger = LevelLogger{NoopLogger{}, false} + config.logger = &LevelLogger{log.New(os.Stdout, "realis: ", log.Ltime|log.Ldate|log.LUTC), false} // Save options to recreate client if a connection error happens config.options = options @@ -242,13 +258,23 @@ func NewRealisClient(options ...ClientOption) (Realis, error) { opt(config) } - config.logger.Println("Number of options applied to config: ", len(options)) + // TODO(rdelvalle): Move this logic to it's own function to make initialization code easier to read. + + // Turn off all logging (including debug) + if config.logger == nil { + config.logger = &LevelLogger{NoopLogger{}, 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 = log.New(os.Stdout, "realis: ", log.Ltime|log.Ldate|log.LUTC) + config.logger = &LevelLogger{log.New(os.Stdout, "realis: ", log.Ltime|log.Ldate|log.LUTC), true} } + // Note, by this point, a LevelLogger should have been created. + config.logger.EnableDebug(config.debug) + + config.logger.DebugPrintln("Number of options applied to config: ", len(options)) + //Set default Transport to JSON if needed. if !config.jsonTransport && !config.binTransport { config.jsonTransport = true @@ -257,9 +283,16 @@ func NewRealisClient(options ...ClientOption) (Realis, error) { var url string var err error - // Determine how to get information to connect to the scheduler. - // Prioritize getting leader from ZK over using a direct URL. - if config.cluster != nil { + // Find the leader using custom Zookeeper options if options are provided + if config.zkOptions != nil { + url, err = LeaderFromZKOpts(config.zkOptions...) + if err != nil { + return nil, NewTemporaryError(errors.Wrap(err, "LeaderFromZK error")) + } + config.logger.Println("Scheduler URL from ZK: ", url) + } else if config.cluster != nil { + // Determine how to get information to connect to the scheduler. + // Prioritize getting leader from ZK over using a direct URL. url, err = LeaderFromZK(*config.cluster) // If ZK is configured, throw an error if the leader is unable to be determined if err != nil { @@ -270,7 +303,7 @@ func NewRealisClient(options ...ClientOption) (Realis, error) { url = config.url config.logger.Println("Scheduler URL: ", url) } else { - return nil, errors.New("Incomplete Options -- url or cluster required") + return nil, errors.New("Incomplete Options -- url, cluster.json, or Zookeeper address required") } if config.jsonTransport { diff --git a/realis_e2e_test.go b/realis_e2e_test.go index d27bc11..342dbb1 100644 --- a/realis_e2e_test.go +++ b/realis_e2e_test.go @@ -40,8 +40,7 @@ func TestMain(m *testing.M) { // New configuration to connect to Vagrant image r, err = realis.NewRealisClient(realis.SchedulerUrl("http://192.168.33.7:8081"), realis.BasicAuth("aurora", "secret"), - realis.TimeoutMS(20000), - realis.Debug()) + realis.TimeoutMS(20000)) if err != nil { fmt.Println("Please run vagrant box before running test suite")