diff --git a/pkg/server/admin.go b/pkg/server/admin.go index a910bb7dba53..33ff29b1b9e7 100644 --- a/pkg/server/admin.go +++ b/pkg/server/admin.go @@ -1066,7 +1066,7 @@ func (s *adminServer) getUIData( if i != 0 { query.Append(",") } - query.Append("$", tree.NewDString(key)) + query.Append("$", tree.NewDString(makeUIKey(userName, key))) } query.Append(");") if err := query.Errors(); err != nil { @@ -1074,7 +1074,7 @@ func (s *adminServer) getUIData( } rows, err := s.server.internalExecutor.QueryEx( ctx, "admin-getUIData", nil, /* txn */ - sqlbase.InternalExecutorSessionDataOverride{User: userName}, + sqlbase.InternalExecutorSessionDataOverride{User: security.RootUser}, query.String(), query.QueryArguments()..., ) if err != nil { @@ -1088,6 +1088,9 @@ func (s *adminServer) getUIData( if !ok { return nil, s.serverErrorf("unexpected type for UI key: %T", row[0]) } + _, key := splitUIKey(string(dKey)) + dKey = tree.DString(key) + dValue, ok := row[1].(*tree.DBytes) if !ok { return nil, s.serverErrorf("unexpected type for UI value: %T", row[1]) @@ -1105,6 +1108,21 @@ func (s *adminServer) getUIData( return &resp, nil } +// makeUIKey combines username and key to form a lookup key in +// system.ui. +// The username is combined to ensure that different users +// can use different customizations. +func makeUIKey(username, key string) string { + return username + "$" + key +} + +// splitUIKey is the inverse of makeUIKey. +// The caller must ensure that the value was produced by makeUIKey. +func splitUIKey(combined string) (string, string) { + pair := strings.SplitN(combined, "$", 2) + return pair[0], pair[1] +} + // 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( @@ -1128,9 +1146,9 @@ func (s *adminServer) SetUIData( rowsAffected, err := s.server.internalExecutor.ExecEx( ctx, "admin-set-ui-data", nil, /* txn */ sqlbase.InternalExecutorSessionDataOverride{ - User: userName, + User: security.RootUser, }, - query, key, val) + query, makeUIKey(userName, key), val) if err != nil { return nil, s.serverError(err) } diff --git a/pkg/server/admin_test.go b/pkg/server/admin_test.go index 60fe6a389a21..78159b166337 100644 --- a/pkg/server/admin_test.go +++ b/pkg/server/admin_test.go @@ -71,7 +71,16 @@ func getAdminJSONProtoWithAdminOption( func postAdminJSONProto( ts serverutils.TestServerInterface, path string, request, response protoutil.Message, ) error { - return serverutils.PostJSONProto(ts, adminPrefix+path, request, response) + return postAdminJSONProtoWithAdminOption(ts, path, request, response, true) +} + +func postAdminJSONProtoWithAdminOption( + ts serverutils.TestServerInterface, + path string, + request, response protoutil.Message, + isAdmin bool, +) error { + return serverutils.PostJSONProtoWithAdminOption(ts, adminPrefix+path, request, response, isAdmin) } // getText fetches the HTTP response body as text in the form of a @@ -1027,107 +1036,150 @@ func TestAdminAPISettings(t *testing.T) { }) } +// TestAdminAPIUIData checks that UI customizations are properly +// persisted for both admin and non-admin users. func TestAdminAPIUIData(t *testing.T) { defer leaktest.AfterTest(t)() s, _, _ := serverutils.StartServer(t, base.TestServerArgs{}) defer s.Stopper().Stop(context.TODO()) - start := timeutil.Now() + testutils.RunTrueAndFalse(t, "isAdmin", func(t *testing.T, isAdmin bool) { + start := timeutil.Now() - mustSetUIData := func(keyValues map[string][]byte) { - if err := postAdminJSONProto(s, "uidata", &serverpb.SetUIDataRequest{ - KeyValues: keyValues, - }, &serverpb.SetUIDataResponse{}); err != nil { - t.Fatal(err) + mustSetUIData := func(keyValues map[string][]byte) { + if err := postAdminJSONProtoWithAdminOption(s, "uidata", &serverpb.SetUIDataRequest{ + KeyValues: keyValues, + }, &serverpb.SetUIDataResponse{}, isAdmin); err != nil { + t.Fatal(err) + } } - } - expectKeyValues := func(expKeyValues map[string][]byte) { - var resp serverpb.GetUIDataResponse - queryValues := make(url.Values) - for key := range expKeyValues { - queryValues.Add("keys", key) - } - url := "uidata?" + queryValues.Encode() - if err := getAdminJSONProto(s, url, &resp); err != nil { - t.Fatal(err) - } - // Do a two-way comparison. We can't use reflect.DeepEqual(), because - // resp.KeyValues has timestamps and expKeyValues doesn't. - for key, actualVal := range resp.KeyValues { - if a, e := actualVal.Value, expKeyValues[key]; !bytes.Equal(a, e) { - t.Fatalf("key %s: value = %v, expected = %v", key, a, e) + expectKeyValues := func(expKeyValues map[string][]byte) { + var resp serverpb.GetUIDataResponse + queryValues := make(url.Values) + for key := range expKeyValues { + queryValues.Add("keys", key) } - } - for key, expVal := range expKeyValues { - if a, e := resp.KeyValues[key].Value, expVal; !bytes.Equal(a, e) { - t.Fatalf("key %s: value = %v, expected = %v", key, a, e) + url := "uidata?" + queryValues.Encode() + if err := getAdminJSONProtoWithAdminOption(s, url, &resp, isAdmin); err != nil { + t.Fatal(err) } - } - - // Sanity check LastUpdated. - for _, val := range resp.KeyValues { - now := timeutil.Now() - if val.LastUpdated.Before(start) { - t.Fatalf("val.LastUpdated %s < start %s", val.LastUpdated, start) + // Do a two-way comparison. We can't use reflect.DeepEqual(), because + // resp.KeyValues has timestamps and expKeyValues doesn't. + for key, actualVal := range resp.KeyValues { + if a, e := actualVal.Value, expKeyValues[key]; !bytes.Equal(a, e) { + t.Fatalf("key %s: value = %v, expected = %v", key, a, e) + } + } + for key, expVal := range expKeyValues { + if a, e := resp.KeyValues[key].Value, expVal; !bytes.Equal(a, e) { + t.Fatalf("key %s: value = %v, expected = %v", key, a, e) + } } - if val.LastUpdated.After(now) { - t.Fatalf("val.LastUpdated %s > now %s", val.LastUpdated, now) + + // Sanity check LastUpdated. + for _, val := range resp.KeyValues { + now := timeutil.Now() + if val.LastUpdated.Before(start) { + t.Fatalf("val.LastUpdated %s < start %s", val.LastUpdated, start) + } + if val.LastUpdated.After(now) { + t.Fatalf("val.LastUpdated %s > now %s", val.LastUpdated, now) + } } } - } - expectValueEquals := func(key string, expVal []byte) { - expectKeyValues(map[string][]byte{key: expVal}) - } + expectValueEquals := func(key string, expVal []byte) { + expectKeyValues(map[string][]byte{key: expVal}) + } - expectKeyNotFound := func(key string) { - var resp serverpb.GetUIDataResponse - url := "uidata?keys=" + key - if err := getAdminJSONProto(s, url, &resp); err != nil { - t.Fatal(err) + expectKeyNotFound := func(key string) { + var resp serverpb.GetUIDataResponse + url := "uidata?keys=" + key + if err := getAdminJSONProtoWithAdminOption(s, url, &resp, isAdmin); err != nil { + t.Fatal(err) + } + if len(resp.KeyValues) != 0 { + t.Fatal("key unexpectedly found") + } } - if len(resp.KeyValues) != 0 { - t.Fatal("key unexpectedly found") + + // Basic tests. + var badResp serverpb.GetUIDataResponse + const errPattern = "400 Bad Request" + if err := getAdminJSONProtoWithAdminOption(s, "uidata", &badResp, isAdmin); !testutils.IsError(err, errPattern) { + t.Fatalf("unexpected error: %v\nexpected: %s", err, errPattern) } - } - // Basic tests. - var badResp serverpb.GetUIDataResponse - const errPattern = "400 Bad Request" - if err := getAdminJSONProto(s, "uidata", &badResp); !testutils.IsError(err, errPattern) { - t.Fatalf("unexpected error: %v\nexpected: %s", err, errPattern) - } + mustSetUIData(map[string][]byte{"k1": []byte("v1")}) + expectValueEquals("k1", []byte("v1")) - mustSetUIData(map[string][]byte{"k1": []byte("v1")}) - expectValueEquals("k1", []byte("v1")) + expectKeyNotFound("NON_EXISTENT_KEY") - expectKeyNotFound("NON_EXISTENT_KEY") + mustSetUIData(map[string][]byte{ + "k2": []byte("v2"), + "k3": []byte("v3"), + }) + expectValueEquals("k2", []byte("v2")) + expectValueEquals("k3", []byte("v3")) + expectKeyValues(map[string][]byte{ + "k2": []byte("v2"), + "k3": []byte("v3"), + }) - mustSetUIData(map[string][]byte{ - "k2": []byte("v2"), - "k3": []byte("v3"), - }) - expectValueEquals("k2", []byte("v2")) - expectValueEquals("k3", []byte("v3")) - expectKeyValues(map[string][]byte{ - "k2": []byte("v2"), - "k3": []byte("v3"), - }) + mustSetUIData(map[string][]byte{"k2": []byte("v2-updated")}) + expectKeyValues(map[string][]byte{ + "k2": []byte("v2-updated"), + "k3": []byte("v3"), + }) - mustSetUIData(map[string][]byte{"k2": []byte("v2-updated")}) - expectKeyValues(map[string][]byte{ - "k2": []byte("v2-updated"), - "k3": []byte("v3"), + // Write a binary blob with all possible byte values, then verify it. + var buf bytes.Buffer + for i := 0; i < 997; i++ { + buf.WriteByte(byte(i % 256)) + } + mustSetUIData(map[string][]byte{"bin": buf.Bytes()}) + expectValueEquals("bin", buf.Bytes()) }) +} - // Write a binary blob with all possible byte values, then verify it. - var buf bytes.Buffer - for i := 0; i < 997; i++ { - buf.WriteByte(byte(i % 256)) +// TestAdminAPIUISeparateData check that separate users have separate customizations. +func TestAdminAPIUISeparateData(t *testing.T) { + defer leaktest.AfterTest(t)() + s, _, _ := serverutils.StartServer(t, base.TestServerArgs{}) + defer s.Stopper().Stop(context.TODO()) + + // Make a setting for an admin user. + if err := postAdminJSONProtoWithAdminOption(s, "uidata", + &serverpb.SetUIDataRequest{KeyValues: map[string][]byte{"k": []byte("v1")}}, + &serverpb.SetUIDataResponse{}, + true /*isAdmin*/); err != nil { + t.Fatal(err) + } + + // Make a setting for a non-admin user. + if err := postAdminJSONProtoWithAdminOption(s, "uidata", + &serverpb.SetUIDataRequest{KeyValues: map[string][]byte{"k": []byte("v2")}}, + &serverpb.SetUIDataResponse{}, + false /*isAdmin*/); err != nil { + t.Fatal(err) + } + + var resp serverpb.GetUIDataResponse + url := "uidata?keys=k" + + if err := getAdminJSONProtoWithAdminOption(s, url, &resp, true /* isAdmin */); err != nil { + t.Fatal(err) + } + if len(resp.KeyValues) != 1 || !bytes.Equal(resp.KeyValues["k"].Value, []byte("v1")) { + t.Fatalf("unexpected admin values: %+v", resp.KeyValues) + } + if err := getAdminJSONProtoWithAdminOption(s, url, &resp, false /* isAdmin */); err != nil { + t.Fatal(err) + } + if len(resp.KeyValues) != 1 || !bytes.Equal(resp.KeyValues["k"].Value, []byte("v2")) { + t.Fatalf("unexpected non-admin values: %+v", resp.KeyValues) } - mustSetUIData(map[string][]byte{"bin": buf.Bytes()}) - expectValueEquals("bin", buf.Bytes()) } func TestClusterAPI(t *testing.T) { diff --git a/pkg/testutils/serverutils/test_server_shim.go b/pkg/testutils/serverutils/test_server_shim.go index 87a7f307e8c4..be71e9cabe84 100644 --- a/pkg/testutils/serverutils/test_server_shim.go +++ b/pkg/testutils/serverutils/test_server_shim.go @@ -249,10 +249,13 @@ func GetJSONProtoWithAdminOption( return httputil.GetJSON(httpClient, ts.AdminURL()+path, response) } -// PostJSONProto uses the supplied client to POST request to the URL specified by -// the parameters and unmarshals the result into response. -func PostJSONProto(ts TestServerInterface, path string, request, response protoutil.Message) error { - httpClient, err := ts.GetAdminAuthenticatedHTTPClient() +// PostJSONProtoWithAdminOption is like PostJSONProto but the caller +// can customize whether the request is performed with admin +// privilege. +func PostJSONProtoWithAdminOption( + ts TestServerInterface, path string, request, response protoutil.Message, isAdmin bool, +) error { + httpClient, err := ts.GetAuthenticatedHTTPClient(isAdmin) if err != nil { return err }