diff --git a/zk.go b/zk.go index 838714c..f930bd9 100644 --- a/zk.go +++ b/zk.go @@ -42,6 +42,7 @@ type zkConfig struct { timeout time.Duration logger Logger auroraSchemeOverride *string + auroraPortOverride *int } type ZKOpt func(z *zkConfig) @@ -82,6 +83,12 @@ func ZKAuroraSchemeOverride(scheme string) ZKOpt { } } +func ZKAuroraPortOverride(port int) ZKOpt { + return func(z *zkConfig) { + z.auroraPortOverride = &port + } +} + // Retrieves current Aurora leader from ZK. func LeaderFromZK(cluster Cluster) (string, error) { return LeaderFromZKOpts(ZKEndpoints(strings.Split(cluster.ZK, ",")...), ZKPath(cluster.SchedZKPath)) @@ -167,7 +174,11 @@ func LeaderFromZKOpts(options ...ZKOpt) (string, error) { host = v.Host - port = strconv.Itoa(v.Port) + if config.auroraPortOverride == nil { + port = strconv.Itoa(v.Port) + } else { + port = strconv.Itoa(*config.auroraPortOverride) + } } leaderURL = scheme + "://" + host + ":" + port