From 76d3177e2731599aa98293157f6e334324e7208b Mon Sep 17 00:00:00 2001 From: Steven Danna Date: Mon, 17 Jun 2024 09:17:39 +0100 Subject: [PATCH] roachprod: add database flag to sql and pgurl command Epic: none Release note: None --- pkg/cmd/roachprod/flags.go | 2 ++ pkg/cmd/roachprod/main.go | 8 +++++--- pkg/roachprod/install/cluster_synced.go | 10 +++++----- pkg/roachprod/install/cockroach.go | 24 ++++++++++++++++-------- pkg/roachprod/roachprod.go | 13 ++++++++----- 5 files changed, 36 insertions(+), 21 deletions(-) diff --git a/pkg/cmd/roachprod/flags.go b/pkg/cmd/roachprod/flags.go index 9804ed9afd3b..db2422a6d7cc 100644 --- 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 index 915d1f60483c..5cf35aeb2837 100644 --- 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 index e422ec027896..d934a24847b4 100644 --- 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 index 14f18e102c73..4c0d45d5b228 100644 --- 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 index c2f1ca886b04..9433c7d0a2a3 100644 --- a/pkg/roachprod/roachprod.go +++ b/pkg/roachprod/roachprod.go @@ -474,6 +474,7 @@ func SQL( tenantName string, tenantInstance int, authMode install.PGAuthMode, + database string, cmdArray []string, ) error { c, err := getClusterFromCache(l, clusterName, install.SecureOption(secure)) @@ -481,10 +482,10 @@ func SQL( 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