From 787725e37d786cbc2acbb369ecbb9bef961381f0 Mon Sep 17 00:00:00 2001 From: Herko Lategan Date: Fri, 5 Apr 2024 14:09:55 +0100 Subject: [PATCH 1/4] roachprod: resolve load balancer Add the ability to find a load balancer given the service port it serves. If no port (0) is specified the first load balancer found will be returned. Epic: None Release Note: None --- pkg/roachprod/install/services.go | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/pkg/roachprod/install/services.go b/pkg/roachprod/install/services.go index 014622c67033..daaf56724e60 100644 --- a/pkg/roachprod/install/services.go +++ b/pkg/roachprod/install/services.go @@ -507,3 +507,19 @@ func (c *SyncedCluster) TargetDNSName(node Node) string { // Targets always end with a period as per SRV record convention. return fmt.Sprintf("%s.%s", cVM.PublicDNS, postfix) } + +// FindLoadBalancer returns the first load balancer address that matches the +// given port. If no load balancer is found, an error is returned. +func (c *SyncedCluster) FindLoadBalancer(l *logger.Logger, port int) (*vm.ServiceAddress, error) { + addresses, err := c.ListLoadBalancers(l) + if err != nil { + return nil, err + } + // Find the load balancer with the matching port. + for _, a := range addresses { + if a.Port == port { + return &a, nil + } + } + return nil, errors.Newf("no load balancer found for port %d", port) +} From 0f56ecdc3bf4caf061314bdf1a6f157b1d5edad3 Mon Sep 17 00:00:00 2001 From: Herko Lategan Date: Fri, 5 Apr 2024 14:12:10 +0100 Subject: [PATCH 2/4] roachprod: load balancer postgres URL Add ability to get the postgres URL for a load balancer given a service. Epic: None Release Note: None --- pkg/roachprod/install/cluster_synced.go | 27 +++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/pkg/roachprod/install/cluster_synced.go b/pkg/roachprod/install/cluster_synced.go index 08b2840ebb3a..053c414e5f88 100644 --- a/pkg/roachprod/install/cluster_synced.go +++ b/pkg/roachprod/install/cluster_synced.go @@ -2654,6 +2654,33 @@ func (c *SyncedCluster) pghosts( return m, nil } +// resolveLoadBalancerURL resolves the load balancer postgres URL for the given +// virtual cluster and SQL instance. Returns an empty string if a load balancer +// is not found. +func (c *SyncedCluster) loadBalancerURL( + ctx context.Context, l *logger.Logger, virtualClusterName string, sqlInstance int, +) (string, error) { + services, err := c.DiscoverServices(ctx, virtualClusterName, ServiceTypeSQL) + if err != nil { + return "", err + } + var port int + serviceMode := ServiceModeShared + for _, service := range services { + if service.VirtualClusterName == virtualClusterName && service.Instance == sqlInstance { + serviceMode = service.ServiceMode + port = service.Port + break + } + } + address, err := c.FindLoadBalancer(l, port) + if err != nil { + return "", err + } + loadBalancerURL := c.NodeURL(address.IP, address.Port, virtualClusterName, serviceMode, AuthRootCert) + return loadBalancerURL, nil +} + // SSH creates an interactive shell connecting the caller to the first // node on the cluster (or to all nodes in an iterm2 split screen if // supported). From 4edacf953056c88b27c3df2c091721496a5b523e Mon Sep 17 00:00:00 2001 From: Herko Lategan Date: Fri, 5 Apr 2024 14:16:27 +0100 Subject: [PATCH 3/4] roachprod: load balancer expanders Add "L" to the node expansion in `roachprod` (e.g., `test-cluster:L` vs. `test-cluster:1`). This allows users to pass "L" instead of a node range / number. It will return the load balancer IP or postgres URL (depending on the expander used) and will enable `workload` and other utilities to target the load balancer IP rather than a specific node. Epic: None Release Note: None --- pkg/roachprod/install/expander.go | 55 +++++++++++++++++++++++-------- 1 file changed, 41 insertions(+), 14 deletions(-) diff --git a/pkg/roachprod/install/expander.go b/pkg/roachprod/install/expander.go index b6e123808daf..6679dd1195be 100644 --- a/pkg/roachprod/install/expander.go +++ b/pkg/roachprod/install/expander.go @@ -22,8 +22,8 @@ import ( ) var parameterRe = regexp.MustCompile(`{[^{}]*}`) -var pgURLRe = regexp.MustCompile(`{pgurl(:[-,0-9]+)?(:[a-z0-9\-]+)?(:[0-9]+)?}`) -var pgHostRe = regexp.MustCompile(`{pghost(:[-,0-9]+)?}`) +var pgURLRe = regexp.MustCompile(`{pgurl(:[-,0-9]+|:L)?(:[a-z0-9\-]+)?(:[0-9]+)?}`) +var pgHostRe = regexp.MustCompile(`{pghost(:[-,0-9]+|:L)?(:[a-z0-9\-]+)?(:[0-9]+)?}`) var pgPortRe = regexp.MustCompile(`{pgport(:[-,0-9]+)?(:[a-z0-9\-]+)?(:[0-9]+)?}`) var uiPortRe = regexp.MustCompile(`{uiport(:[-,0-9]+)}`) var storeDirRe = regexp.MustCompile(`{store-dir(:[0-9]+)?}`) @@ -157,14 +157,20 @@ func (e *expander) maybeExpandPgURL( if err != nil { return "", false, err } - if e.pgURLs[virtualClusterName] == nil { - e.pgURLs[virtualClusterName], err = c.pgurls(ctx, l, allNodes(len(c.VMs)), virtualClusterName, sqlInstance) - if err != nil { - return "", false, err + switch m[1] { + case ":L": + url, err := c.loadBalancerURL(ctx, l, virtualClusterName, sqlInstance) + return url, url != "", err + default: + if e.pgURLs[virtualClusterName] == nil { + e.pgURLs[virtualClusterName], err = c.pgurls(ctx, l, allNodes(len(c.VMs)), virtualClusterName, sqlInstance) + if err != nil { + return "", false, err + } } + s, err = e.maybeExpandMap(c, e.pgURLs[virtualClusterName], m[1]) + return s, err == nil, err } - s, err = e.maybeExpandMap(c, e.pgURLs[virtualClusterName], m[1]) - return s, err == nil, err } // maybeExpandPgHost is an expanderFunc for {pghost:} @@ -175,17 +181,38 @@ func (e *expander) maybeExpandPgHost( if m == nil { return s, false, nil } + virtualClusterName, sqlInstance, err := extractVirtualClusterInfo(m[2:]) + if err != nil { + return "", false, err + } - if e.pgHosts == nil { - var err error - e.pgHosts, err = c.pghosts(ctx, l, allNodes(len(c.VMs))) + switch m[1] { + case ":L": + services, err := c.DiscoverServices(ctx, virtualClusterName, ServiceTypeSQL, ServiceInstancePredicate(sqlInstance)) if err != nil { return "", false, err } + for _, svc := range services { + if svc.VirtualClusterName == virtualClusterName && svc.Instance == sqlInstance { + addr, err := c.FindLoadBalancer(l, svc.Port) + if err != nil { + return "", false, err + } + return addr.IP, true, nil + } + } + return "", false, err + default: + if e.pgHosts == nil { + var err error + e.pgHosts, err = c.pghosts(ctx, l, allNodes(len(c.VMs))) + if err != nil { + return "", false, err + } + } + s, err := e.maybeExpandMap(c, e.pgHosts, m[1]) + return s, err == nil, err } - - s, err := e.maybeExpandMap(c, e.pgHosts, m[1]) - return s, err == nil, err } // maybeExpandPgURL is an expanderFunc for {pgport:} From 8cb12bddda9acd769aafba6cbc2803de40bcaa92 Mon Sep 17 00:00:00 2001 From: Herko Lategan Date: Wed, 24 Apr 2024 09:15:49 +0100 Subject: [PATCH 4/4] roachprod: pass `PGAuthMode` to `loadBalancerURL` Epic: None Release Note: None --- pkg/roachprod/install/cluster_synced.go | 8 ++++++-- pkg/roachprod/install/expander.go | 2 +- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/pkg/roachprod/install/cluster_synced.go b/pkg/roachprod/install/cluster_synced.go index 053c414e5f88..c58252fd99b2 100644 --- a/pkg/roachprod/install/cluster_synced.go +++ b/pkg/roachprod/install/cluster_synced.go @@ -2658,7 +2658,11 @@ func (c *SyncedCluster) pghosts( // virtual cluster and SQL instance. Returns an empty string if a load balancer // is not found. func (c *SyncedCluster) loadBalancerURL( - ctx context.Context, l *logger.Logger, virtualClusterName string, sqlInstance int, + ctx context.Context, + l *logger.Logger, + virtualClusterName string, + sqlInstance int, + auth PGAuthMode, ) (string, error) { services, err := c.DiscoverServices(ctx, virtualClusterName, ServiceTypeSQL) if err != nil { @@ -2677,7 +2681,7 @@ func (c *SyncedCluster) loadBalancerURL( if err != nil { return "", err } - loadBalancerURL := c.NodeURL(address.IP, address.Port, virtualClusterName, serviceMode, AuthRootCert) + loadBalancerURL := c.NodeURL(address.IP, address.Port, virtualClusterName, serviceMode, auth) return loadBalancerURL, nil } diff --git a/pkg/roachprod/install/expander.go b/pkg/roachprod/install/expander.go index 6679dd1195be..b0ea40b3a5a0 100644 --- a/pkg/roachprod/install/expander.go +++ b/pkg/roachprod/install/expander.go @@ -159,7 +159,7 @@ func (e *expander) maybeExpandPgURL( } switch m[1] { case ":L": - url, err := c.loadBalancerURL(ctx, l, virtualClusterName, sqlInstance) + url, err := c.loadBalancerURL(ctx, l, virtualClusterName, sqlInstance, AuthUserCert) return url, url != "", err default: if e.pgURLs[virtualClusterName] == nil {