From 0dec820951854b5a12eb0917da3b200643857bea Mon Sep 17 00:00:00 2001
From: Renan DelValle <renanidelvalle+noreply@gmail.com>
Date: Thu, 12 Apr 2018 15:12:36 -0700
Subject: [PATCH] Removing port override as it is not needed

---
 examples/client.go |  3 ++-
 realis.go          | 36 +++++++++++++++++++++++++++++++-----
 zk.go              | 26 ++++++++++++++++++++------
 3 files changed, 53 insertions(+), 12 deletions(-)

diff --git a/examples/client.go b/examples/client.go
index 7162408..9e83d9a 100644
--- a/examples/client.go
+++ b/examples/client.go
@@ -99,7 +99,8 @@ func main() {
 		fmt.Println("zkUrl: ", zkUrl)
 		clientOptions = append(clientOptions, realis.ZKUrl(zkUrl))
 	} else {
-		clientOptions = append(clientOptions, realis.SchedulerUrl(url))
+		clientOptions = append(clientOptions, realis.SchedulerUrl(url),
+			realis.ZookeeperOptions(realis.ZKAuroraPortOverride(2343), realis.ZKAuroraSchemeOverride("https")))
 	}
 
 	r, err = realis.NewRealisClient(clientOptions...)
diff --git a/realis.go b/realis.go
index c9291d1..e94e1e7 100644
--- a/realis.go
+++ b/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,16 @@ func ClientCerts(clientKey, clientCert string) ClientOption {
 	}
 }
 
+// Use this option if you'd like to override default settings for connecting to Zookeeper.
+// For example, this can be used to override the port on which to communicate with Aurora.
+// This may be helpful if Aurora is behind another service and running on a port that is different
+// what is advertised in Zookeeper.
+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 +276,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 +296,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/zk.go b/zk.go
index dd711e0..838714c 100644
--- a/zk.go
+++ b/zk.go
@@ -36,11 +36,12 @@ type ServiceInstance struct {
 }
 
 type zkConfig struct {
-	endpoints []string
-	path      string
-	backoff   Backoff
-	timeout   time.Duration
-	logger    Logger
+	endpoints            []string
+	path                 string
+	backoff              Backoff
+	timeout              time.Duration
+	logger               Logger
+	auroraSchemeOverride *string
 }
 
 type ZKOpt func(z *zkConfig)
@@ -75,6 +76,12 @@ func ZKLogger(l Logger) ZKOpt {
 	}
 }
 
+func ZKAuroraSchemeOverride(scheme string) ZKOpt {
+	return func(z *zkConfig) {
+		z.auroraSchemeOverride = &scheme
+	}
+}
+
 // Retrieves current Aurora leader from ZK.
 func LeaderFromZK(cluster Cluster) (string, error) {
 	return LeaderFromZKOpts(ZKEndpoints(strings.Split(cluster.ZK, ",")...), ZKPath(cluster.SchedZKPath))
@@ -151,8 +158,15 @@ func LeaderFromZKOpts(options ...ZKOpt) (string, error) {
 
 				var scheme, host, port string
 				for k, v := range serviceInst.AdditionalEndpoints {
-					scheme = k
+
+					if config.auroraSchemeOverride == nil {
+						scheme = k
+					} else {
+						scheme = *config.auroraSchemeOverride
+					}
+
 					host = v.Host
+
 					port = strconv.Itoa(v.Port)
 				}