// 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) {
if len(req.KeyValues) == 0 {
return nil, grpc.Errorf(codes.InvalidArgument, "KeyValues cannot be empty")
}
args := sql.SessionArgs{User: s.getUser(req)}
session := s.NewSessionForRPC(ctx, args)
defer session.Finish(s.server.sqlExecutor)
for key, val := range req.KeyValues {
// Do an upsert of the key. We update each key in a separate transaction to
// avoid long-running transactions and possible deadlocks.
br := s.server.sqlExecutor.ExecuteStatements(session, "BEGIN;", nil)
defer br.Close()
if err := s.checkQueryResults(br.ResultList, 1); err != nil {
return nil, s.serverError(err)
}
// See if the key already exists.
resp, err := s.getUIData(session, s.getUser(req), []string{key})
if err != nil {
return nil, s.serverError(err)
}
_, alreadyExists := resp.KeyValues[key]
// INSERT or UPDATE as appropriate.
if alreadyExists {
query := "UPDATE system.ui SET value = $1, lastUpdated = NOW() WHERE key = $2; COMMIT;"
qargs := parser.NewPlaceholderInfo()
qargs.SetValue(`1`, parser.NewDString(string(val)))
qargs.SetValue(`2`, parser.NewDString(key))
r := s.server.sqlExecutor.ExecuteStatements(session, query, qargs)
defer r.Close()
if err := s.checkQueryResults(r.ResultList, 2); err != nil {
return nil, s.serverError(err)
}
if a, e := r.ResultList[0].RowsAffected, 1; a != e {
return nil, s.serverErrorf("rows affected %d != expected %d", a, e)
}
} else {
query := "INSERT INTO system.ui (key, value, lastUpdated) VALUES ($1, $2, NOW()); COMMIT;"
qargs := parser.NewPlaceholderInfo()
qargs.SetValue(`1`, parser.NewDString(key))
qargs.SetValue(`2`, parser.NewDBytes(parser.DBytes(val)))
r := s.server.sqlExecutor.ExecuteStatements(session, query, qargs)
defer r.Close()
if err := s.checkQueryResults(r.ResultList, 2); err != nil {
return nil, s.serverError(err)
}
if a, e := r.ResultList[0].RowsAffected, 1; a != e {
return nil, s.serverErrorf("rows affected %d != expected %d", a, e)
}
}
}
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) {
if len(req.KeyValues) == 0 {
return nil, grpc.Errorf(codes.InvalidArgument, "KeyValues cannot be empty")
}
args := sql.SessionArgs{User: s.getUser(req)}
session := s.NewSessionForRPC(ctx, args)
defer session.Finish(s.server.sqlExecutor)
for key, val := range req.KeyValues {
// 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())"
qargs := parser.NewPlaceholderInfo()
qargs.SetValue(`1`, parser.NewDString(key))
qargs.SetValue(`2`, parser.NewDBytes(parser.DBytes(val)))
r := s.server.sqlExecutor.ExecuteStatements(session, query, qargs)
defer r.Close()
if err := s.checkQueryResults(r.ResultList, 1); err != nil {
return nil, s.serverError(err)
}
if a, e := r.ResultList[0].RowsAffected, 1; a != e {
return nil, s.serverErrorf("rows affected %d != expected %d", a, e)
}
}
return &serverpb.SetUIDataResponse{}, nil
}
// queryNamespaceID queries for the ID of the namespace with the given name and
// parent ID.
func (s *adminServer) queryNamespaceID(
session *sql.Session, parentID sqlbase.ID, name string,
) (sqlbase.ID, error) {
const query = `SELECT id FROM system.namespace WHERE parentID = $1 AND name = $2`
params := parser.NewPlaceholderInfo()
params.SetValue(`1`, parser.NewDInt(parser.DInt(parentID)))
params.SetValue(`2`, parser.NewDString(name))
r := s.server.sqlExecutor.ExecuteStatements(session, query, params)
defer r.Close()
if err := s.checkQueryResults(r.ResultList, 1); err != nil {
return 0, err
}
result := r.ResultList[0]
if result.Rows.Len() == 0 {
return 0, errors.Errorf("namespace %s with ParentID %d not found", name, parentID)
}
var id int64
scanner := resultScanner{}
err := scanner.ScanIndex(result.Rows.At(0), 0, &id)
if err != nil {
return 0, err
}
return sqlbase.ID(id), nil
}
// queryZone retrieves the specific ZoneConfig associated with the supplied ID,
// if it exists.
func (s *adminServer) queryZone(
session *sql.Session, id sqlbase.ID,
) (config.ZoneConfig, bool, error) {
const query = `SELECT config FROM system.zones WHERE id = $1`
params := parser.NewPlaceholderInfo()
params.SetValue(`1`, parser.NewDInt(parser.DInt(id)))
r := s.server.sqlExecutor.ExecuteStatements(session, query, params)
defer r.Close()
if err := s.checkQueryResults(r.ResultList, 1); err != nil {
return config.ZoneConfig{}, false, err
}
result := r.ResultList[0]
if result.Rows.Len() == 0 {
return config.ZoneConfig{}, false, nil
}
var zoneBytes []byte
scanner := resultScanner{}
err := scanner.ScanIndex(result.Rows.At(0), 0, &zoneBytes)
if err != nil {
return config.ZoneConfig{}, false, err
}
var zone config.ZoneConfig
if err := zone.Unmarshal(zoneBytes); err != nil {
return config.ZoneConfig{}, false, err
}
return zone, true, nil
}
// TestAdminAPIZoneDetails verifies the zone configuration information returned
// for both DatabaseDetailsResponse AND TableDetailsResponse.
func TestAdminAPIZoneDetails(t *testing.T) {
defer leaktest.AfterTest(t)()
s, _, _ := serverutils.StartServer(t, base.TestServerArgs{})
defer s.Stopper().Stop()
ts := s.(*TestServer)
// Create database and table.
ac := log.AmbientContext{Tracer: tracing.NewTracer()}
ctx, span := ac.AnnotateCtxWithSpan(context.Background(), "test")
defer span.Finish()
session := sql.NewSession(
ctx, sql.SessionArgs{User: security.RootUser}, ts.sqlExecutor, nil, &sql.MemoryMetrics{})
session.StartUnlimitedMonitor()
setupQueries := []string{
"CREATE DATABASE test",
"CREATE TABLE test.tbl (val STRING)",
}
for _, q := range setupQueries {
res := ts.sqlExecutor.ExecuteStatements(session, q, nil)
defer res.Close()
if res.ResultList[0].Err != nil {
t.Fatalf("error executing '%s': %s", q, res.ResultList[0].Err)
}
}
// Function to verify the zone for table "test.tbl" as returned by the Admin
// API.
verifyTblZone := func(
expectedZone config.ZoneConfig, expectedLevel serverpb.ZoneConfigurationLevel,
) {
var resp serverpb.TableDetailsResponse
if err := getAdminJSONProto(s, "databases/test/tables/tbl", &resp); err != nil {
t.Fatal(err)
}
if a, e := &resp.ZoneConfig, &expectedZone; !proto.Equal(a, e) {
t.Errorf("actual table zone config %v did not match expected value %v", a, e)
}
if a, e := resp.ZoneConfigLevel, expectedLevel; a != e {
t.Errorf("actual table ZoneConfigurationLevel %s did not match expected value %s", a, e)
}
if t.Failed() {
t.FailNow()
}
}
// Function to verify the zone for database "test" as returned by the Admin
// API.
verifyDbZone := func(
expectedZone config.ZoneConfig, expectedLevel serverpb.ZoneConfigurationLevel,
) {
var resp serverpb.DatabaseDetailsResponse
if err := getAdminJSONProto(s, "databases/test", &resp); err != nil {
t.Fatal(err)
}
if a, e := &resp.ZoneConfig, &expectedZone; !proto.Equal(a, e) {
t.Errorf("actual db zone config %v did not match expected value %v", a, e)
}
if a, e := resp.ZoneConfigLevel, expectedLevel; a != e {
t.Errorf("actual db ZoneConfigurationLevel %s did not match expected value %s", a, e)
}
if t.Failed() {
t.FailNow()
}
}
// Function to store a zone config for a given object ID.
setZone := func(zoneCfg config.ZoneConfig, id sqlbase.ID) {
zoneBytes, err := zoneCfg.Marshal()
if err != nil {
t.Fatal(err)
}
const query = `INSERT INTO system.zones VALUES($1, $2)`
params := parser.NewPlaceholderInfo()
params.SetValue(`1`, parser.NewDInt(parser.DInt(id)))
params.SetValue(`2`, parser.NewDBytes(parser.DBytes(zoneBytes)))
res := ts.sqlExecutor.ExecuteStatements(session, query, params)
defer res.Close()
if res.ResultList[0].Err != nil {
t.Fatalf("error executing '%s': %s", query, res.ResultList[0].Err)
}
}
// Verify zone matches cluster default.
verifyDbZone(config.DefaultZoneConfig(), serverpb.ZoneConfigurationLevel_CLUSTER)
verifyTblZone(config.DefaultZoneConfig(), serverpb.ZoneConfigurationLevel_CLUSTER)
// Get ID path for table. This will be an array of three IDs, containing the ID of the root namespace,
// the database, and the table (in that order).
idPath, err := ts.admin.queryDescriptorIDPath(session, []string{"test", "tbl"})
if err != nil {
t.Fatal(err)
}
// Apply zone configuration to database and check again.
dbZone := config.ZoneConfig{
RangeMinBytes: 456,
}
setZone(dbZone, idPath[1])
verifyDbZone(dbZone, serverpb.ZoneConfigurationLevel_DATABASE)
verifyTblZone(dbZone, serverpb.ZoneConfigurationLevel_DATABASE)
// Apply zone configuration to table and check again.
tblZone := config.ZoneConfig{
RangeMinBytes: 789,
}
setZone(tblZone, idPath[2])
verifyDbZone(dbZone, serverpb.ZoneConfigurationLevel_DATABASE)
verifyTblZone(tblZone, serverpb.ZoneConfigurationLevel_TABLE)
}