diff --git a/pkg/cli/BUILD.bazel b/pkg/cli/BUILD.bazel index 72d289d2d57f..195d4c1fe224 100644 --- a/pkg/cli/BUILD.bazel +++ b/pkg/cli/BUILD.bazel @@ -315,6 +315,7 @@ go_test( "//pkg/util", "//pkg/util/bitarray", "//pkg/util/duration", + "//pkg/util/envutil", "//pkg/util/ipaddr", "//pkg/util/json", "//pkg/util/leaktest", diff --git a/pkg/cli/start.go b/pkg/cli/start.go index f298ae766947..ae78b03028a1 100644 --- a/pkg/cli/start.go +++ b/pkg/cli/start.go @@ -22,6 +22,7 @@ import ( "os/signal" "path/filepath" "runtime" + "strconv" "strings" "text/tabwriter" "time" @@ -35,6 +36,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/rpc" "github.com/cockroachdb/cockroach/pkg/security" "github.com/cockroachdb/cockroach/pkg/server" + "github.com/cockroachdb/cockroach/pkg/server/serverpb" "github.com/cockroachdb/cockroach/pkg/server/status" "github.com/cockroachdb/cockroach/pkg/server/telemetry" "github.com/cockroachdb/cockroach/pkg/settings/cluster" @@ -600,6 +602,10 @@ If problems persist, please see %s.` if err := s.AcceptClients(ctx); err != nil { return err } + // Configure UI settings. + if err := setUIDataFromEnv(ctx, s); err != nil { + return err + } // Now inform the user that the server is running and tell the // user about its run-time derived parameters. @@ -877,6 +883,19 @@ If problems persist, please see %s.` return returnErr } +// setUIDataFromEnv toggles presence of release notes signup banner +// based on an environment variable. +func setUIDataFromEnv(ctx context.Context, s *server.Server) error { + b := envutil.EnvOrDefaultBool("COCKROACH_UI_RELEASE_NOTES_SIGNUP_DISMISSED", false) + _, err := s.SetUIData(ctx, &serverpb.SetUIDataRequest{ + KeyValues: map[string][]byte{"release_notes_signup_dismissed": []byte(strconv.FormatBool(b))}}, + security.RootUserName()) + if err != nil { + return err + } + return nil +} + // expandTabsInRedactableBytes expands tabs in the redactable byte // slice, so that columns are aligned. The correctness of this // function depends on the assumption that the `tabwriter` does not diff --git a/pkg/cli/start_test.go b/pkg/cli/start_test.go index 11f451afc613..088c2090f5dc 100644 --- a/pkg/cli/start_test.go +++ b/pkg/cli/start_test.go @@ -11,12 +11,21 @@ package cli import ( + "context" + "fmt" + "os" "testing" + "github.com/cockroachdb/cockroach/pkg/base" + "github.com/cockroachdb/cockroach/pkg/rpc" "github.com/cockroachdb/cockroach/pkg/server" + "github.com/cockroachdb/cockroach/pkg/server/serverpb" "github.com/cockroachdb/cockroach/pkg/testutils" + "github.com/cockroachdb/cockroach/pkg/testutils/serverutils" + "github.com/cockroachdb/cockroach/pkg/util/envutil" "github.com/cockroachdb/cockroach/pkg/util/leaktest" "github.com/cockroachdb/cockroach/pkg/util/log" + "github.com/stretchr/testify/require" ) func TestInitInsecure(t *testing.T) { @@ -143,3 +152,52 @@ func TestAddrWithDefaultHost(t *testing.T) { } } } + +func TestSetUIFromEnv(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + ctx := context.Background() + s, _, _ := serverutils.StartServer(t, base.TestServerArgs{}) + defer s.Stopper().Stop(ctx) + ts := s.(*server.TestServer) + + const name = "COCKROACH_UI_RELEASE_NOTES_SIGNUP_DISMISSED" + const key = "release_notes_signup_dismissed" + for v, expected := range map[string]string{"true": "true", "false": "false", "": "false"} { + t.Run(fmt.Sprintf("%s=%s", name, v), func(t *testing.T) { + var err error + if v == "" { + err = os.Unsetenv(name) + } else { + err = os.Setenv(name, v) + } + if err != nil { + t.Fatal(err) + } + defer envutil.ClearEnvCache() + + err = setUIDataFromEnv(ctx, ts.Server) + if err != nil { + t.Fatal(err) + } + + cc, err := ts.RPCContext().GRPCDialNode( + ts.RPCAddr(), + 1, + rpc.DefaultClass, + ).Connect(ctx) + if err != nil { + t.Fatal(err) + } + adminClient := serverpb.NewAdminClient(cc) + + resp, err := adminClient.GetUIData(ctx, &serverpb.GetUIDataRequest{Keys: []string{key}}) + if err != nil { + t.Fatal(err) + } + actual := string(resp.KeyValues[key].Value) + require.Equal(t, expected, actual) + }) + } +} diff --git a/pkg/server/admin.go b/pkg/server/admin.go index cafc0838736a..9a26882ba2a2 100644 --- a/pkg/server/admin.go +++ b/pkg/server/admin.go @@ -1194,17 +1194,12 @@ func splitUIKey(combined string) (string, string) { return pair[0], pair[1] } -// SetUIData is an endpoint that stores the given key/value pairs in the +// SetUIData stores the given key/value pairs in the // system.ui table. See GetUIData for more details on semantics. -func (s *adminServer) SetUIData( - ctx context.Context, req *serverpb.SetUIDataRequest, +// This function writes to system.ui using the internal SQL executor of the provided server. +func (s *Server) SetUIData( + ctx context.Context, req *serverpb.SetUIDataRequest, userName security.SQLUsername, ) (*serverpb.SetUIDataResponse, error) { - ctx = s.server.AnnotateCtx(ctx) - - userName, err := userFromContext(ctx) - if err != nil { - return nil, err - } if len(req.KeyValues) == 0 { return nil, status.Errorf(codes.InvalidArgument, "KeyValues cannot be empty") @@ -1214,23 +1209,40 @@ func (s *adminServer) SetUIData( // Do an upsert of the key. We update each key in a separate transaction to // avoid long-running transactions and possible deadlocks. query := `UPSERT INTO system.ui (key, value, "lastUpdated") VALUES ($1, $2, now())` - rowsAffected, err := s.server.sqlServer.internalExecutor.ExecEx( + rowsAffected, err := s.sqlServer.internalExecutor.ExecEx( ctx, "admin-set-ui-data", nil, /* txn */ sessiondata.InternalExecutorOverride{ User: security.RootUserName(), }, query, makeUIKey(userName, key), val) if err != nil { - return nil, s.serverError(err) + return nil, err } if rowsAffected != 1 { - return nil, s.serverErrorf("rows affected %d != expected %d", rowsAffected, 1) + return nil, errors.Newf("rows affected %d != expected %d", rowsAffected, 1) } } return &serverpb.SetUIDataResponse{}, nil } +// SetUIData is an endpoint that stores the given key/value pairs in the +// system.ui table. See GetUIData for more details on semantics. +func (s *adminServer) SetUIData( + ctx context.Context, req *serverpb.SetUIDataRequest, +) (*serverpb.SetUIDataResponse, error) { + ctx = s.server.AnnotateCtx(ctx) + userName, err := userFromContext(ctx) + if err != nil { + return nil, err + } + _, err = s.server.SetUIData(ctx, req, userName) + if err != nil { + return nil, s.serverError(err) + } + return &serverpb.SetUIDataResponse{}, nil +} + // GetUIData returns data associated with the given keys, which was stored // earlier through SetUIData. //