roachprod: load balancer from pgurl command

pkg/cmd/roachprod/main.go

@@ -1021,6 +1021,13 @@ var pgurlCmd = &cobra.Command{
 	Short: "generate pgurls for the nodes in a cluster",
 	Long: `Generate pgurls for the nodes in a cluster.
 
+The command generates postgres urls for the specified nodes in a cluster.
+Both the nodes or a load balancer can be specified as the target of the pgurl.
+
+Examples of <cluster>:
+	cluster-name:1-3
+	cluster-name:lb
+
 --auth-mode specifies the method of authentication unless --insecure is passed.
 Defaults to root if not passed. Available auth-modes are:
 

pkg/roachprod/install/expander.go

@@ -23,8 +23,8 @@ import (
 )
 
 var parameterRe = regexp.MustCompile(`{[^{}]*}`)
-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 pgURLRe = regexp.MustCompile(`{pgurl(:[-,0-9]+|:(?i)lb)?(:[a-z0-9\-]+)?(:[0-9]+)?}`)
+var pgHostRe = regexp.MustCompile(`{pghost(:[-,0-9]+|:(?i)lb)?(:[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]+)?}`)
@@ -158,8 +158,8 @@ func (e *expander) maybeExpandPgURL(
 	if err != nil {
 		return "", false, err
 	}
-	switch m[1] {
-	case ":L":
+	switch strings.ToLower(m[1]) {
+	case ":lb":
 		url, err := c.loadBalancerURL(ctx, l, virtualClusterName, sqlInstance, DefaultAuthMode)
 		return url, url != "", err
 	default:
@@ -187,8 +187,8 @@ func (e *expander) maybeExpandPgHost(
 		return "", false, err
 	}
 
-	switch m[1] {
-	case ":L":
+	switch strings.ToLower(m[1]) {
+	case ":lb":
 		services, err := c.DiscoverServices(ctx, virtualClusterName, ServiceTypeSQL, ServiceInstancePredicate(sqlInstance))
 		if err != nil {
 			return "", false, err

pkg/roachprod/install/nodes.go

@@ -44,7 +44,9 @@ func ListNodes(s string, numNodesInCluster int) (Nodes, error) {
 		return nil, errors.AssertionFailedf("invalid number of nodes %d", numNodesInCluster)
 	}
 
-	if s == "all" {
+	// "lb" is a special value that also returns all nodes, but is used to
+	// indicate that a load balancer should be used.
+	if s == "all" || strings.ToLower(s) == "lb" {
 		return allNodes(numNodesInCluster), nil
 	}
 

pkg/roachprod/roachprod.go

@@ -185,6 +185,13 @@ func newCluster(
 	return c, nil
 }
 
+// shouldUseLoadBalancer determines if the node selector section in the cluster
+// name is set to select a load balancer.
+func shouldUseLoadBalancer(name string) bool {
+	parts := strings.Split(name, ":")
+	return len(parts) == 2 && strings.ToLower(parts[1]) == "lb"
+}
+
 // userClusterNameRegexp returns a regexp that matches all clusters owned by the
 // current user.
 func userClusterNameRegexp(l *logger.Logger) (*regexp.Regexp, error) {
@@ -929,6 +936,28 @@ func PgURL(
 	if err != nil {
 		return nil, err
 	}
+
+	if shouldUseLoadBalancer(clusterName) {
+		services, err := c.DiscoverServices(ctx, opts.VirtualClusterName, install.ServiceTypeSQL,
+			install.ServiceInstancePredicate(opts.SQLInstance))
+		if err != nil {
+			return nil, err
+		}
+		port := config.DefaultSQLPort
+		serviceMode := install.ServiceModeExternal
+		if len(services) > 0 {
+			port = services[0].Port
+			serviceMode = services[0].ServiceMode
+		} else {
+			l.Printf("no services found, searching for load balancer on default port %d", port)
+		}
+		addr, err := c.FindLoadBalancer(l, port)
+		if err != nil {
+			return nil, err
+		}
+		return []string{c.NodeURL(addr.IP, port, opts.VirtualClusterName, serviceMode, opts.Auth)}, nil
+	}
+
 	nodes := c.TargetNodes()
 	ips := make([]string, len(nodes))