roachprod: add database flag to sql and pgurl command

Epic: none Release note: None
prdk0 · Jun 17, 2024 · 76d3177 · 76d3177
1 parent efc894b
commit 76d3177
Show file tree

Hide file tree

Showing 5 changed files with 36 additions and 21 deletions.
diff --git a/pkg/cmd/roachprod/flags.go b/pkg/cmd/roachprod/flags.go
@@ -37,6 +37,7 @@ var (
 	numNodes              int
 	numRacks              int
 	username              string
+	database              string
 	dryrun                bool
 	destroyAllMine        bool
 	destroyAllLocal       bool
@@ -217,6 +218,7 @@ func initFlags() {
 	for _, cmd := range []*cobra.Command{pgurlCmd, sqlCmd, loadBalancerPGUrl} {
 		cmd.Flags().StringVar(&authMode,
 			"auth-mode", defaultAuthMode, fmt.Sprintf("form of authentication to use, valid auth-modes: %v", maps.Keys(pgAuthModes)))
+		cmd.Flags().StringVar(&database, "database", "", "database to use")
 	}
 
 	pprofCmd.Flags().DurationVar(&pprofOpts.Duration,

diff --git a/pkg/cmd/roachprod/main.go b/pkg/cmd/roachprod/main.go
@@ -542,6 +542,7 @@ var loadBalancerPGUrl = &cobra.Command{
 			return err
 		}
 		url, err := roachprod.LoadBalancerPgURL(context.Background(), config.Logger, args[0], pgurlCertsDir, roachprod.PGURLOptions{
+			Database:           database,
 			External:           external,
 			Secure:             isSecure,
 			VirtualClusterName: virtualClusterName,
@@ -629,8 +630,8 @@ var updateTargetsCmd = &cobra.Command{
 	Short: "update prometheus target configurations for a cluster",
 	Long: `Update prometheus target configurations of each node of a cluster.
 
-The "start" command updates the prometheus target configuration every time. But, in case of any  
-failure, this command can be used to update the configurations. 
+The "start" command updates the prometheus target configuration every time. But, in case of any
+failure, this command can be used to update the configurations.
 
 The default prometheus url is https://grafana.testeng.crdb.io/. This can be overwritten by using the
 environment variable COCKROACH_PROM_HOST_URL
@@ -1144,7 +1145,7 @@ var sqlCmd = &cobra.Command{
 			return errors.Newf("unsupported auth-mode %s, valid auth-modes: %v", authMode, maps.Keys(pgAuthModes))
 		}
 
-		return roachprod.SQL(context.Background(), config.Logger, args[0], isSecure, virtualClusterName, sqlInstance, auth, args[1:])
+		return roachprod.SQL(context.Background(), config.Logger, args[0], isSecure, virtualClusterName, sqlInstance, auth, database, args[1:])
 	}),
 }
 
@@ -1162,6 +1163,7 @@ var pgurlCmd = &cobra.Command{
 			return err
 		}
 		urls, err := roachprod.PgURL(context.Background(), config.Logger, args[0], pgurlCertsDir, roachprod.PGURLOptions{
+			Database:           database,
 			External:           external,
 			Secure:             isSecure,
 			VirtualClusterName: virtualClusterName,

diff --git a/pkg/roachprod/install/cluster_synced.go b/pkg/roachprod/install/cluster_synced.go
@@ -444,9 +444,9 @@ func (c *SyncedCluster) Stop(
 		}
 		return c.kill(ctx, l, "stop", display, sig, wait, maxWait, virtualClusterLabel)
 	} else {
-		res, err := c.ExecSQL(ctx, l, c.Nodes[:1], "", 0, DefaultAuthMode, []string{
-			"-e", fmt.Sprintf("ALTER TENANT '%s' STOP SERVICE", virtualClusterName),
-		})
+		cmd := fmt.Sprintf("ALTER TENANT '%s' STOP SERVICE", virtualClusterName)
+		res, err := c.ExecSQL(ctx, l, c.Nodes[:1], "", 0, DefaultAuthMode, "", /* database */
+			[]string{"-e", cmd})
 		if err != nil || res[0].Err != nil {
 			if len(res) > 0 {
 				return errors.CombineErrors(err, res[0].Err)
@@ -2652,7 +2652,7 @@ func (c *SyncedCluster) pgurls(
 		if err != nil {
 			return nil, err
 		}
-		m[node] = c.NodeURL(host, desc.Port, virtualClusterName, desc.ServiceMode, AuthUserCert)
+		m[node] = c.NodeURL(host, desc.Port, virtualClusterName, desc.ServiceMode, AuthUserCert, "" /* database */)
 	}
 	return m, nil
 }
@@ -2700,7 +2700,7 @@ func (c *SyncedCluster) loadBalancerURL(
 	if err != nil {
 		return "", err
 	}
-	loadBalancerURL := c.NodeURL(address.IP, address.Port, virtualClusterName, serviceMode, auth)
+	loadBalancerURL := c.NodeURL(address.IP, address.Port, virtualClusterName, serviceMode, auth, "" /* database */)
 	return loadBalancerURL, nil
 }
 

diff --git a/pkg/roachprod/install/cockroach.go b/pkg/roachprod/install/cockroach.go
@@ -592,12 +592,18 @@ const (
 // be used as the virtual cluster name in the URL. This is used to connect to a
 // shared process running services for multiple virtual clusters.
 func (c *SyncedCluster) NodeURL(
-	host string, port int, virtualClusterName string, serviceMode ServiceMode, auth PGAuthMode,
+	host string,
+	port int,
+	virtualClusterName string,
+	serviceMode ServiceMode,
+	auth PGAuthMode,
+	database string,
 ) string {
 	var u url.URL
 	u.Scheme = "postgres"
 	u.User = url.User("root")
 	u.Host = fmt.Sprintf("%s:%d", host, port)
+	u.Path = database
 	v := url.Values{}
 	if c.Secure {
 		user := DefaultUser
@@ -682,6 +688,7 @@ func (c *SyncedCluster) ExecOrInteractiveSQL(
 	virtualClusterName string,
 	sqlInstance int,
 	authMode PGAuthMode,
+	database string,
 	args []string,
 ) error {
 	if len(c.Nodes) != 1 {
@@ -691,7 +698,7 @@ func (c *SyncedCluster) ExecOrInteractiveSQL(
 	if err != nil {
 		return err
 	}
-	url := c.NodeURL("localhost", desc.Port, virtualClusterName, desc.ServiceMode, authMode)
+	url := c.NodeURL("localhost", desc.Port, virtualClusterName, desc.ServiceMode, authMode, database)
 	binary := cockroachNodeBinary(c, c.Nodes[0])
 	allArgs := []string{binary, "sql", "--url", url}
 	allArgs = append(allArgs, ssh.Escape(args))
@@ -708,6 +715,7 @@ func (c *SyncedCluster) ExecSQL(
 	virtualClusterName string,
 	sqlInstance int,
 	authMode PGAuthMode,
+	database string,
 	args []string,
 ) ([]*RunResultDetails, error) {
 	display := fmt.Sprintf("%s: executing sql", c.Name)
@@ -722,7 +730,7 @@ func (c *SyncedCluster) ExecSQL(
 				cmd = fmt.Sprintf(`cd %s ; `, c.localVMDir(node))
 			}
 			cmd += cockroachNodeBinary(c, node) + " sql --url " +
-				c.NodeURL("localhost", desc.Port, virtualClusterName, desc.ServiceMode, authMode) + " " +
+				c.NodeURL("localhost", desc.Port, virtualClusterName, desc.ServiceMode, authMode, database) + " " +
 				ssh.Escape(args)
 			return c.runCmdOnSingleNode(ctx, l, node, cmd, defaultCmdOpts("run-sql"))
 		})
@@ -1195,7 +1203,7 @@ func (c *SyncedCluster) createAdminUserForSecureCluster(
 		// We use the first node in the virtual cluster to create the user.
 		firstNode := c.TargetNodes()[0]
 		results, err := c.ExecSQL(
-			ctx, l, Nodes{firstNode}, virtualClusterName, sqlInstance, AuthRootCert,
+			ctx, l, Nodes{firstNode}, virtualClusterName, sqlInstance, AuthRootCert, "", /* database */
 			[]string{"-e", stmts})
 
 		if err != nil || results[0].Err != nil {
@@ -1281,7 +1289,7 @@ func (c *SyncedCluster) generateClusterSettingCmd(
 	if err != nil {
 		return "", err
 	}
-	url := c.NodeURL("localhost", port, SystemInterfaceName /* virtualClusterName */, ServiceModeShared, AuthRootCert)
+	url := c.NodeURL("localhost", port, SystemInterfaceName /* virtualClusterName */, ServiceModeShared, AuthRootCert, "" /* database */)
 
 	// We use `mkdir -p` here since the directory may not exist if an in-memory
 	// store is used.
@@ -1303,7 +1311,7 @@ func (c *SyncedCluster) generateInitCmd(ctx context.Context, node Node) (string,
 	if err != nil {
 		return "", err
 	}
-	url := c.NodeURL("localhost", port, SystemInterfaceName /* virtualClusterName */, ServiceModeShared, AuthRootCert)
+	url := c.NodeURL("localhost", port, SystemInterfaceName /* virtualClusterName */, ServiceModeShared, AuthRootCert, "" /* database */)
 	binary := cockroachNodeBinary(c, node)
 	initCmd += fmt.Sprintf(`
 		if ! test -e %[1]s ; then
@@ -1380,7 +1388,7 @@ func (c *SyncedCluster) upsertVirtualClusterMetadata(
 ) (int, error) {
 	runSQL := func(stmt string) (string, error) {
 		results, err := startOpts.StorageCluster.ExecSQL(
-			ctx, l, startOpts.StorageCluster.Nodes[:1], "", 0, DefaultAuthMode,
+			ctx, l, startOpts.StorageCluster.Nodes[:1], "", 0, DefaultAuthMode, "", /* database */
 			[]string{"--format", "csv", "-e", stmt})
 		if err != nil {
 			return "", err
@@ -1504,7 +1512,7 @@ func (c *SyncedCluster) createFixedBackupSchedule(
 		serviceMode = ServiceModeExternal
 	}
 
-	url := c.NodeURL("localhost", port, startOpts.VirtualClusterName, serviceMode, AuthRootCert)
+	url := c.NodeURL("localhost", port, startOpts.VirtualClusterName, serviceMode, AuthRootCert, "" /* database */)
 	fullCmd := fmt.Sprintf(`COCKROACH_CONNECT_TIMEOUT=%d %s sql --url %s -e %q`,
 		startSQLTimeout, binary, url, createScheduleCmd)
 	// Instead of using `c.ExecSQL()`, use `c.runCmdOnSingleNode()`, which allows us to

diff --git a/pkg/roachprod/roachprod.go b/pkg/roachprod/roachprod.go
@@ -474,17 +474,18 @@ func SQL(
 	tenantName string,
 	tenantInstance int,
 	authMode install.PGAuthMode,
+	database string,
 	cmdArray []string,
 ) error {
 	c, err := getClusterFromCache(l, clusterName, install.SecureOption(secure))
 	if err != nil {
 		return err
 	}
 	if len(c.Nodes) == 1 {
-		return c.ExecOrInteractiveSQL(ctx, l, tenantName, tenantInstance, authMode, cmdArray)
+		return c.ExecOrInteractiveSQL(ctx, l, tenantName, tenantInstance, authMode, database, cmdArray)
 	}
 
-	results, err := c.ExecSQL(ctx, l, c.Nodes, tenantName, tenantInstance, authMode, cmdArray)
+	results, err := c.ExecSQL(ctx, l, c.Nodes, tenantName, tenantInstance, authMode, database, cmdArray)
 	if err != nil {
 		return err
 	}
@@ -1032,6 +1033,7 @@ func Get(ctx context.Context, l *logger.Logger, clusterName, src, dest string) e
 }
 
 type PGURLOptions struct {
+	Database           string
 	Secure             bool
 	External           bool
 	VirtualClusterName string
@@ -1071,7 +1073,7 @@ func PgURL(
 		if ip == "" {
 			return nil, errors.Errorf("empty ip: %v", ips)
 		}
-		urls = append(urls, c.NodeURL(ip, desc.Port, opts.VirtualClusterName, desc.ServiceMode, opts.Auth))
+		urls = append(urls, c.NodeURL(ip, desc.Port, opts.VirtualClusterName, desc.ServiceMode, opts.Auth, opts.Database))
 	}
 	if len(urls) != len(nodes) {
 		return nil, errors.Errorf("have nodes %v, but urls %v from ips %v", nodes, urls, ips)
@@ -2224,7 +2226,8 @@ func StartJaeger(
 			return err
 		}
 		_, err = c.ExecSQL(
-			ctx, l, nodes, virtualClusterName, 0, install.DefaultAuthMode, []string{"-e", setupStmt},
+			ctx, l, nodes, virtualClusterName, 0, install.DefaultAuthMode, "", /* database */
+			[]string{"-e", setupStmt},
 		)
 		if err != nil {
 			return err
@@ -2650,7 +2653,7 @@ func LoadBalancerPgURL(
 	if err != nil {
 		return "", err
 	}
-	return c.NodeURL(addr.IP, port, opts.VirtualClusterName, serviceMode, opts.Auth), nil
+	return c.NodeURL(addr.IP, port, opts.VirtualClusterName, serviceMode, opts.Auth, opts.Database), nil
 }
 
 // LoadBalancerIP resolves the IP of a load balancer serving the