cli: allow internal usage for cli tools

Various cockroach cli tools query internal tables as regular operations.
This includes things like running tsdumps, debug zips, getting statement
diagnostics bundles, etc. These tools are expected to be able to query
and access internal tables without errors or audit logging.

To do this, the app name used for these tools are being changed to have
the prefix: `$ internal `. This allows for the CheckInternalAccess API
to skip the check.

Note: the sql cli does not set this internal app name. Certain actions
and internal checks / queries have been updated to temporarily set
the application name to an internal one in order to run internal queries.

Resolves: CRDB-55062
Epic: CRDB-55276
Release note: None

Author: kyle-a-wong

pkg/cli/auth.go

@@ -17,13 +17,16 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cli/clisqlexec"
 	"github.com/cockroachdb/cockroach/pkg/server/authserver"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	isatty "github.com/mattn/go-isatty"
 	"github.com/spf13/cobra"
 )
 
+const authSessionAppName = catconstants.InternalAppNamePrefix + " cockroach auth-session"
+
 var loginCmd = &cobra.Command{
 	Use:   "login [options] <session-username>",
 	Short: "create a HTTP session and token for the given user",
@@ -89,7 +92,7 @@ func createAuthSessionToken(
 	username string,
 ) (sessionID int64, httpCookie *http.Cookie, resErr error) {
 	ctx := context.Background()
-	sqlConn, err := makeSQLClient(ctx, "cockroach auth-session login", useSystemDb)
+	sqlConn, err := makeSQLClient(ctx, authSessionAppName+" login", useSystemDb)
 	if err != nil {
 		return -1, nil, err
 	}
@@ -177,7 +180,7 @@ The user for which the HTTP sessions are revoked can be arbitrary.
 func runLogout(cmd *cobra.Command, args []string) (resErr error) {
 	username := tree.Name(args[0]).Normalize()
 	ctx := context.Background()
-	sqlConn, err := makeSQLClient(ctx, "cockroach auth-session logout", useSystemDb)
+	sqlConn, err := makeSQLClient(ctx, authSessionAppName+" logout", useSystemDb)
 	if err != nil {
 		return err
 	}
@@ -209,7 +212,7 @@ The user invoking the 'list' CLI command must be an admin on the cluster.
 
 func runAuthList(cmd *cobra.Command, args []string) (resErr error) {
 	ctx := context.Background()
-	sqlConn, err := makeSQLClient(ctx, "cockroach auth-session list", useSystemDb)
+	sqlConn, err := makeSQLClient(ctx, authSessionAppName+" list", useSystemDb)
 	if err != nil {
 		return err
 	}

pkg/cli/clisqlclient/api.go

@@ -101,6 +101,11 @@ type Conn interface {
 	// the fields may be empty if there were errors while retrieving
 	// them when the connection was established.
 	GetServerInfo() ServerInfo
+
+	// AllowExecuteInternal allows internal statements to be executed, regardless
+	// of the allow_unsafe_internals session variable. The returned function should
+	// be called after the internal statements have been executed.
+	AllowExecuteInternal(context.Context) func()
 }
 
 // ServerInfo describes the remote server.

pkg/cli/clisqlclient/conn.go

@@ -28,6 +28,8 @@ import (
 	"github.com/otan/gopgkrb5"
 )
 
+const InternalSqlAppName = "$ internal cockroach sql"
+
 func init() {
 	// Ensure that the CLI client commands can use GSSAPI authentication.
 	pgconn.RegisterGSSProvider(func() (pgconn.GSS, error) { return gopgkrb5.NewGSS() })
@@ -436,22 +438,7 @@ func (c *sqlConn) checkServerMetadata(ctx context.Context) error {
 	defer func(prev bool) { c.alwaysInferResultTypes = prev }(c.alwaysInferResultTypes)
 	c.alwaysInferResultTypes = false
 
-	// Metadata checks require allow_unsafe_internals to be on.
-	allowUnsafeInternals, err := c.getSessionVariable(ctx, "allow_unsafe_internals")
-	if err != nil {
-		fmt.Fprintf(c.errw, "warning: unable to retrieve allow_unsafe_internals setting: %v\n", err)
-	} else if allowUnsafeInternals == "off" {
-		// Temporarily turn on allow_unsafe_internals.
-		if err := c.Exec(ctx, "SET allow_unsafe_internals = on"); err != nil {
-			fmt.Fprintf(c.errw, "warning: unable to set allow_unsafe_internals to true: %v\n", err)
-		} else {
-			defer func() {
-				if resetErr := c.Exec(ctx, "SET allow_unsafe_internals = off"); resetErr != nil {
-					fmt.Fprintf(c.errw, "warning: unable to reset allow_unsafe_internals to false: %v\n", resetErr)
-				}
-			}()
-		}
-	}
+	defer c.AllowExecuteInternal(ctx)()
 
 	_, newServerVersion, newClusterID, err := c.GetServerMetadata(ctx)
 	if c.conn.IsClosed() {
@@ -519,6 +506,18 @@ func (c *sqlConn) checkServerMetadata(ctx context.Context) error {
 	return c.tryEnableServerExecutionTimings(ctx)
 }
 
+func (c *sqlConn) AllowExecuteInternal(ctx context.Context) func() {
+	oldApplicationName, _ := c.getSessionVariable(ctx, "application_name")
+	if strings.HasPrefix(oldApplicationName, InternalSqlAppName) {
+		// If already an internal app name, we don't need to set a new one
+		return func() {}
+	}
+	_ = c.Exec(ctx, fmt.Sprintf("SET application_name = '%s'", InternalSqlAppName))
+	return func() {
+		_ = c.Exec(ctx, fmt.Sprintf("SET application_name = '%s'", oldApplicationName))
+	}
+}
+
 // GetServerInfo returns a copy of the remote server details.
 func (c *sqlConn) GetServerInfo() ServerInfo {
 	return ServerInfo{

pkg/cli/clisqlshell/statement_diag.go

@@ -25,7 +25,7 @@ func (c *cliState) handleStatementDiag(
 		cmd = args[0]
 		args = args[1:]
 	}
-
+	defer c.conn.AllowExecuteInternal(context.Background())()
 	var cmdErr error
 	switch cmd {
 	case "list":

pkg/cli/debug_job_cleanup.go

@@ -13,11 +13,14 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cli/clierrorplus"
 	"github.com/cockroachdb/cockroach/pkg/jobs"
 	"github.com/cockroachdb/cockroach/pkg/server/telemetry"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	"github.com/spf13/cobra"
 )
 
+const debugJobInfoAppName = catconstants.InternalAppNamePrefix + " cockroach debug job-cleanup-job-info"
+
 var debugJobCleanupInfoRows = &cobra.Command{
 	Use:   "job-cleanup-job-info",
 	Short: "cleans up system.job_info rows with no related system.jobs entry",
@@ -34,7 +37,7 @@ var jobCleanupInfoRowOpts = struct {
 
 func runDebugJobInfoCleanup(_ *cobra.Command, args []string) (resErr error) {
 	ctx := context.Background()
-	sqlConn, err := makeSQLClient(ctx, "cockroach debug job-cleanup-job-info", useSystemDb)
+	sqlConn, err := makeSQLClient(ctx, debugJobInfoAppName, useSystemDb)
 	if err != nil {
 		return errors.Wrap(err, "could not establish connection to cluster")
 	}

pkg/cli/debug_job_trace.go

@@ -15,11 +15,14 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/cli/clierrorplus"
 	"github.com/cockroachdb/cockroach/pkg/cli/clisqlclient"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants"
 	tracezipper "github.com/cockroachdb/cockroach/pkg/util/tracing/zipper"
 	"github.com/cockroachdb/errors"
 	"github.com/spf13/cobra"
 )
 
+const debugJobTraceAppName = catconstants.InternalAppNamePrefix + "  cockroach debug job-trace"
+
 var debugJobTraceFromClusterCmd = &cobra.Command{
 	Use:   "job-trace <job_id> --url=<cluster connection string>",
 	Short: "get the trace payloads for the executing job",
@@ -35,7 +38,7 @@ func runDebugJobTrace(_ *cobra.Command, args []string) (resErr error) {
 		return err
 	}
 	ctx := context.Background()
-	sqlConn, err := makeSQLClient(ctx, "cockroach debug job-trace", useSystemDb)
+	sqlConn, err := makeSQLClient(ctx, debugJobTraceAppName, useSystemDb)
 	if err != nil {
 		return errors.Wrap(err, "could not establish connection to cluster")
 	}

pkg/cli/doctor.go

@@ -32,6 +32,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/doctor"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/builtins"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants"
 	"github.com/cockroachdb/cockroach/pkg/util/buildutil"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
 	"github.com/cockroachdb/cockroach/pkg/util/protoutil"
@@ -72,6 +73,8 @@ tables are queried either from a live cluster or from an unzipped debug.zip.
 `,
 }
 
+const doctorAppName = catconstants.InternalAppNamePrefix + " cockroach doctor"
+
 type doctorFn = func(
 	version *clusterversion.ClusterVersion,
 	descTable doctor.DescriptorTable,
@@ -117,7 +120,7 @@ Run the doctor tool system data from a live cluster specified by --url.
 		Args: cobra.NoArgs,
 		RunE: clierrorplus.MaybeDecorateError(
 			func(cmd *cobra.Command, args []string) (resErr error) {
-				sqlConn, err := makeSQLClient(context.Background(), "cockroach doctor", useSystemDb)
+				sqlConn, err := makeSQLClient(context.Background(), doctorAppName, useSystemDb)
 				if err != nil {
 					return errors.Wrap(err, "could not establish connection to cluster")
 				}

pkg/cli/node.go

@@ -31,6 +31,8 @@ import (
 	"google.golang.org/grpc/status"
 )
 
+const nodeAppName = catconstants.InternalAppNamePrefix + " cockroach node"
+
 var lsNodesColumnHeaders = []string{
 	"id",
 }
@@ -50,7 +52,7 @@ func runLsNodes(cmd *cobra.Command, args []string) (resErr error) {
 	ctx := context.Background()
 	// TODO(ssd): We can potentially make this work against
 	// secondary tenants using sql_instances.
-	conn, err := makeTenantSQLClient(ctx, "cockroach node ls", useSystemDb, catconstants.SystemTenantName)
+	conn, err := makeTenantSQLClient(ctx, nodeAppName+" ls", useSystemDb, catconstants.SystemTenantName)
 	if err != nil {
 		return err
 	}
@@ -210,7 +212,7 @@ SELECT node_id AS id,
 FROM crdb_internal.gossip_liveness LEFT JOIN crdb_internal.gossip_nodes USING (node_id)`
 
 	ctx := context.Background()
-	conn, err := makeTenantSQLClient(ctx, "cockroach node status", useSystemDb, catconstants.SystemTenantName)
+	conn, err := makeTenantSQLClient(ctx, nodeAppName+" status", useSystemDb, catconstants.SystemTenantName)
 	if err != nil {
 		return nil, nil, err
 	}

pkg/cli/nodelocal.go

@@ -17,6 +17,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cli/clisqlclient"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/sql"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants"
 	"github.com/cockroachdb/errors"
 	"github.com/spf13/cobra"
 )
@@ -33,9 +34,11 @@ Uploads a file to a gateway node's local file system using a SQL connection.
 	RunE: clierrorplus.MaybeShoutError(runUpload),
 }
 
+const nodeLocalAppName = catconstants.InternalAppNamePrefix + " cockroach nodelocal"
+
 func runUpload(cmd *cobra.Command, args []string) (resErr error) {
 	ctx := context.Background()
-	conn, err := makeSQLClient(ctx, "cockroach nodelocal", useSystemDb)
+	conn, err := makeSQLClient(ctx, nodeLocalAppName, useSystemDb)
 	if err != nil {
 		return err
 	}

pkg/cli/statement_diag.go

@@ -14,10 +14,13 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/cli/clierrorplus"
 	"github.com/cockroachdb/cockroach/pkg/cli/clisqlclient"
+	"github.com/cockroachdb/cockroach/pkg/sql/sem/catconstants"
 	"github.com/cockroachdb/errors"
 	"github.com/spf13/cobra"
 )
 
+const stmtDiagAppName = catconstants.InternalAppNamePrefix + " cockroach statement-diag"
+
 var stmtDiagCmd = &cobra.Command{
 	Use:   "statement-diag [command] [options]",
 	Short: "commands for managing statement diagnostics bundles",
@@ -39,7 +42,7 @@ diagnostics activation requests.`,
 func runStmtDiagList(cmd *cobra.Command, args []string) (resErr error) {
 	const timeFmt = "2006-01-02 15:04:05 MST"
 	ctx := context.Background()
-	conn, err := makeSQLClient(ctx, "cockroach statement-diag", useSystemDb)
+	conn, err := makeSQLClient(ctx, stmtDiagAppName, useSystemDb)
 	if err != nil {
 		return err
 	}
@@ -128,7 +131,7 @@ func runStmtDiagDownload(cmd *cobra.Command, args []string) (resErr error) {
 		filename = fmt.Sprintf("stmt-bundle-%d.zip", id)
 	}
 	ctx := context.Background()
-	conn, err := makeSQLClient(ctx, "cockroach statement-diag", useSystemDb)
+	conn, err := makeSQLClient(ctx, stmtDiagAppName, useSystemDb)
 	if err != nil {
 		return err
 	}
@@ -153,7 +156,7 @@ command, or delete all bundles.`,
 
 func runStmtDiagDelete(cmd *cobra.Command, args []string) (resErr error) {
 	ctx := context.Background()
-	conn, err := makeSQLClient(ctx, "cockroach statement-diag", useSystemDb)
+	conn, err := makeSQLClient(ctx, stmtDiagAppName, useSystemDb)
 	if err != nil {
 		return err
 	}
@@ -188,7 +191,7 @@ list command, or cancel all outstanding requests.`,
 
 func runStmtDiagCancel(cmd *cobra.Command, args []string) (resErr error) {
 	ctx := context.Background()
-	conn, err := makeSQLClient(ctx, "cockroach statement-diag", useSystemDb)
+	conn, err := makeSQLClient(ctx, stmtDiagAppName, useSystemDb)
 	if err != nil {
 		return err
 	}