Merge #123206 #123458

123206: roachprod: load balancer from pgurl command r=srosenberg,renatolabs a=herkolategan

Previously, the load balancer IP or URL can only be retrieved using the expanders. This change adds the ability to also get the URL from the `roachprod` pgurl command by using the `lb` node selector similar to how it is done with expansion.

Epic: None
Release Note: None

123458: roachprod: fix `start.sh` profile check r=srosenberg,renatolabs a=herkolategan

Previously, multiple entries would be created in "${HOME}/.profile" since the check for "${HOME}/.profile-cockroach" would always fail. This updates the check to look for the correct file based on the virtual cluster label.

Fixes: #121267

Epic: None
Release Note: None

Co-authored-by: Herko Lategan <herko@cockroachlabs.com>

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/install/scripts/start.sh

@@ -77,7 +77,7 @@ sudo systemctl reset-failed "${VIRTUAL_CLUSTER_LABEL}" 2>/dev/null || true
 
 # The first time we run, install a small script that shows some helpful
 # information when we ssh in.
-if [ ! -e "${HOME}/.profile-cockroach" ]; then
+if [ ! -e "${HOME}/.profile-${VIRTUAL_CLUSTER_LABEL}" ]; then
   cat > "${HOME}/.profile-${VIRTUAL_CLUSTER_LABEL}" <<EOQ
 echo ""
 if systemctl is-active -q ${VIRTUAL_CLUSTER_LABEL}; then

pkg/roachprod/install/testdata/start/start.txt

@@ -81,7 +81,7 @@ sudo systemctl reset-failed "${VIRTUAL_CLUSTER_LABEL}" 2>/dev/null || true
 
 # The first time we run, install a small script that shows some helpful
 # information when we ssh in.
-if [ ! -e "${HOME}/.profile-cockroach" ]; then
+if [ ! -e "${HOME}/.profile-${VIRTUAL_CLUSTER_LABEL}" ]; then
   cat > "${HOME}/.profile-${VIRTUAL_CLUSTER_LABEL}" <<EOQ
 echo ""
 if systemctl is-active -q ${VIRTUAL_CLUSTER_LABEL}; then

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))