Merge 1bf341b815 into 8ca953f925

realis.go

@@ -26,6 +26,7 @@ import (
 	"net/http/cookiejar"
 	"os"
 	"path/filepath"
+	"strings"
 	"time"
 
 	"sync"
@@ -102,6 +103,7 @@ type RealisConfig struct {
 	clientkey, clientcert       string
 	options                     []ClientOption
 	debug                       bool
+	zkOptions                   []ZKOpt
 }
 
 var defaultBackoff = Backoff{
@@ -140,8 +142,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,6 +196,14 @@ 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) {
@@ -257,9 +274,16 @@ func NewRealisClient(options ...ClientOption) (Realis, error) {
 	var url string
 	var err error
 
-	// Determine how to get information to connect to the scheduler.
+	// Find the leader using custom Zookeeper options if options are provided
-	// Prioritize getting leader from ZK over using a direct URL.
+	if config.zkOptions != nil {
-	if config.cluster != 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 +294,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 {