Skip to content

Commit 35ea6d5

Browse files
committed
sql,multitenant: new package 'mtinfo' to share SQL decoding logic
Release note: None
1 parent f8f8349 commit 35ea6d5

File tree

7 files changed

+162
-116
lines changed

7 files changed

+162
-116
lines changed

pkg/BUILD.bazel

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1424,6 +1424,7 @@ GO_TARGETS = [
14241424
"//pkg/kv/kvserver:kvserver_test",
14251425
"//pkg/kv:kv",
14261426
"//pkg/kv:kv_test",
1427+
"//pkg/multitenant/mtinfo:mtinfo",
14271428
"//pkg/multitenant/mtinfopb:mtinfopb",
14281429
"//pkg/multitenant/multitenantcpu:multitenantcpu",
14291430
"//pkg/multitenant/multitenantio:multitenantio",

pkg/multitenant/mtinfo/BUILD.bazel

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
load("@io_bazel_rules_go//go:def.bzl", "go_library")
2+
3+
go_library(
4+
name = "mtinfo",
5+
srcs = ["info.go"],
6+
importpath = "github.com/cockroachdb/cockroach/pkg/multitenant/mtinfo",
7+
visibility = ["//visibility:public"],
8+
deps = [
9+
"//pkg/multitenant/mtinfopb",
10+
"//pkg/multitenant/tenantcapabilities",
11+
"//pkg/roachpb",
12+
"//pkg/sql/sem/tree",
13+
"//pkg/util/protoutil",
14+
"@com_github_cockroachdb_errors//:errors",
15+
],
16+
)

pkg/multitenant/mtinfo/info.go

Lines changed: 125 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,125 @@
1+
// Copyright 2023 The Cockroach Authors.
2+
//
3+
// Use of this software is governed by the Business Source License
4+
// included in the file licenses/BSL.txt.
5+
//
6+
// As of the Change Date specified in that file, in accordance with
7+
// the Business Source License, use of this software will be governed
8+
// by the Apache License, Version 2.0, included in the file
9+
// licenses/APL.txt.
10+
11+
package mtinfo
12+
13+
import (
14+
"github.com/cockroachdb/cockroach/pkg/multitenant/mtinfopb"
15+
"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
16+
"github.com/cockroachdb/cockroach/pkg/roachpb"
17+
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
18+
"github.com/cockroachdb/cockroach/pkg/util/protoutil"
19+
"github.com/cockroachdb/errors"
20+
)
21+
22+
// GetTenantInfoFromSQLRow synthetizes a TenantInfo from a SQL row
23+
// extracted from system.tenants. The caller is responsible for
24+
// passing a []tree.Datum of length 2 or more (ID, info, name*, data
25+
// state*, service mode*). The values besides ID and info can be
26+
// absent or NULL; the logic knows how to handle data collected across
27+
// the v23.1 version boundary that introduced the name, data state and
28+
// service mode columns.
29+
func GetTenantInfoFromSQLRow(
30+
row tree.Datums,
31+
) (tid roachpb.TenantID, info *mtinfopb.TenantInfo, err error) {
32+
if len(row) < 2 {
33+
return tid, nil, errors.AssertionFailedf("expected at least 2 columns, got %d", len(row))
34+
}
35+
info = &mtinfopb.TenantInfo{}
36+
37+
idval, ok := tree.AsDInt(row[0])
38+
if !ok {
39+
return tid, nil, errors.AssertionFailedf("tenant ID: expected int, got %T", row[0])
40+
}
41+
tid, err = roachpb.MakeTenantID(uint64(idval))
42+
if err != nil {
43+
return tid, nil, errors.NewAssertionErrorWithWrappedErrf(err, "%v", idval)
44+
}
45+
46+
info.ID = tid.ToUint64()
47+
48+
// For the benefit of pre-23.1 BACKUP/RESTORE.
49+
info.DeprecatedID = info.ID
50+
51+
if len(row) < 2 || row[1] == tree.DNull {
52+
return tid, nil, errors.AssertionFailedf("%v: missing data in info column", tid)
53+
}
54+
ival, ok := tree.AsDBytes(row[1])
55+
if !ok {
56+
return tid, nil, errors.AssertionFailedf("%v: info: expected bytes, got %T", tid, row[1])
57+
}
58+
infoBytes := []byte(ival)
59+
if err := protoutil.Unmarshal(infoBytes, &info.ProtoInfo); err != nil {
60+
return tid, nil, errors.NewAssertionErrorWithWrappedErrf(err, "%v: decoding info column", tid)
61+
}
62+
63+
// If we are loading the entry for the system tenant, inject all the
64+
// capabilities that the system tenant should have. This will allow
65+
// us to reduce our dependency on a tenant ID check to retain the
66+
// system tenant's access to all services.
67+
if tid.IsSystem() {
68+
tenantcapabilities.EnableAll(&info.ProtoInfo.Capabilities)
69+
}
70+
71+
// Load the name if defined.
72+
if len(row) > 2 && row[2] != tree.DNull {
73+
name, ok := tree.AsDString(row[2])
74+
if !ok {
75+
return tid, nil, errors.AssertionFailedf("%v: name: expected string, got %T", tid, row[2])
76+
}
77+
info.Name = roachpb.TenantName(name)
78+
}
79+
80+
// Load the data state column if defined.
81+
// Compute a suitable default value, from the pre-v23.1 info struct.
82+
switch info.ProtoInfo.DeprecatedDataState {
83+
case mtinfopb.ProtoInfo_READY:
84+
info.DataState = mtinfopb.DataStateReady
85+
case mtinfopb.ProtoInfo_ADD:
86+
info.DataState = mtinfopb.DataStateAdd
87+
case mtinfopb.ProtoInfo_DROP:
88+
info.DataState = mtinfopb.DataStateDrop
89+
default:
90+
return tid, nil, errors.AssertionFailedf("%v: unhandled: %d", tid, info.ProtoInfo.DeprecatedDataState)
91+
}
92+
if len(row) > 3 && row[3] != tree.DNull {
93+
val, ok := tree.AsDInt(row[3])
94+
if !ok {
95+
return tid, nil, errors.AssertionFailedf("%v: data state: expected int, got %T", tid, row[3])
96+
}
97+
if val < 0 || mtinfopb.TenantDataState(val) > mtinfopb.MaxDataState {
98+
return tid, nil, errors.AssertionFailedf("%v: invalid data state: %d", tid, val)
99+
} else {
100+
info.DataState = mtinfopb.TenantDataState(val)
101+
}
102+
}
103+
104+
// Load the service mode if defined.
105+
if info.DataState == mtinfopb.DataStateReady {
106+
// Suitable default for records created for CC Serverless pre-v23.1.
107+
info.ServiceMode = mtinfopb.ServiceModeExternal
108+
} else {
109+
// Suitable default for in-progress records created for CC Serverless pre-v23.1.
110+
info.ServiceMode = mtinfopb.ServiceModeNone
111+
}
112+
if len(row) > 4 && row[4] != tree.DNull {
113+
val, ok := tree.AsDInt(row[4])
114+
if !ok {
115+
return tid, nil, errors.AssertionFailedf("%v: service mode: expected int, got %T", tid, row[4])
116+
}
117+
if val < 0 || mtinfopb.TenantServiceMode(val) > mtinfopb.MaxServiceMode {
118+
return tid, nil, errors.AssertionFailedf("%v: invalid service mode: %d", tid, val)
119+
} else {
120+
info.ServiceMode = mtinfopb.TenantServiceMode(val)
121+
}
122+
}
123+
124+
return tid, info, nil
125+
}

pkg/multitenant/tenantcapabilities/tenantcapabilitieswatcher/BUILD.bazel

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ go_library(
1616
"//pkg/kv/kvclient/rangefeed/rangefeedbuffer",
1717
"//pkg/kv/kvclient/rangefeed/rangefeedcache",
1818
"//pkg/kv/kvpb",
19-
"//pkg/multitenant/mtinfopb",
19+
"//pkg/multitenant/mtinfo",
2020
"//pkg/multitenant/tenantcapabilities",
2121
"//pkg/multitenant/tenantcapabilities/tenantcapabilitiespb",
2222
"//pkg/roachpb",
@@ -30,7 +30,6 @@ go_library(
3030
"//pkg/util/hlc",
3131
"//pkg/util/log",
3232
"//pkg/util/log/logcrash",
33-
"//pkg/util/protoutil",
3433
"//pkg/util/stop",
3534
"//pkg/util/syncutil",
3635
"@com_github_cockroachdb_errors//:errors",

pkg/multitenant/tenantcapabilities/tenantcapabilitieswatcher/decoder.go

Lines changed: 13 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ import (
1616
"github.com/cockroachdb/cockroach/pkg/keys"
1717
"github.com/cockroachdb/cockroach/pkg/kv/kvclient/rangefeed/rangefeedbuffer"
1818
"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
19-
"github.com/cockroachdb/cockroach/pkg/multitenant/mtinfopb"
19+
"github.com/cockroachdb/cockroach/pkg/multitenant/mtinfo"
2020
"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
2121
"github.com/cockroachdb/cockroach/pkg/roachpb"
2222
"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -26,9 +26,7 @@ import (
2626
"github.com/cockroachdb/cockroach/pkg/sql/rowenc/valueside"
2727
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
2828
"github.com/cockroachdb/cockroach/pkg/sql/types"
29-
"github.com/cockroachdb/cockroach/pkg/util/log"
3029
"github.com/cockroachdb/cockroach/pkg/util/log/logcrash"
31-
"github.com/cockroachdb/cockroach/pkg/util/protoutil"
3230
"github.com/cockroachdb/errors"
3331
)
3432

@@ -43,7 +41,7 @@ type decoder struct {
4341

4442
// newDecoder constructs and returns a decoder.
4543
func newDecoder(st *cluster.Settings) *decoder {
46-
columns := systemschema.TenantsTable.PublicColumns()
44+
columns := systemschema.TenantsTable.VisibleColumns()
4745
return &decoder{
4846
columns: columns,
4947
decoder: valueside.MakeDecoder(columns),
@@ -55,7 +53,6 @@ func (d *decoder) decode(
5553
ctx context.Context, kv roachpb.KeyValue,
5654
) (tenantcapabilities.Entry, error) {
5755
// First we decode the tenantID from the key.
58-
var tenID roachpb.TenantID
5956
types := []*types.T{d.columns[0].GetType()}
6057
tenantIDRow := make([]rowenc.EncDatum, 1)
6158
if _, err := rowenc.DecodeIndexKey(keys.SystemSQLCodec, tenantIDRow, nil /* colDirs */, kv.Key); err != nil {
@@ -65,16 +62,10 @@ func (d *decoder) decode(
6562
return tenantcapabilities.Entry{},
6663
errors.NewAssertionErrorWithWrappedErrf(err, "failed to decode key in system.tenants %v", kv.Key)
6764
}
68-
tenID, err := roachpb.MakeTenantID(uint64(tree.MustBeDInt(tenantIDRow[0].Datum)))
69-
if err != nil {
70-
return tenantcapabilities.Entry{}, err
71-
}
7265

73-
// The remaining columns are stored in the value; we're just interested in the
74-
// info column.
7566
if !kv.Value.IsPresent() {
7667
return tenantcapabilities.Entry{},
77-
errors.AssertionFailedf("missing value for tenant: %v", tenID)
68+
errors.AssertionFailedf("missing value for tenant: %v", tenantIDRow[0].Datum)
7869
}
7970

8071
bytes, err := kv.Value.GetTuple()
@@ -86,55 +77,20 @@ func (d *decoder) decode(
8677
return tenantcapabilities.Entry{}, err
8778
}
8879

89-
var tenantInfo mtinfopb.ProtoInfo
90-
if i := datums[2]; i != tree.DNull {
91-
infoBytes := tree.MustBeDBytes(i)
92-
if err := protoutil.Unmarshal([]byte(infoBytes), &tenantInfo); err != nil {
93-
return tenantcapabilities.Entry{}, errors.Wrapf(err, "failed to unmarshall tenant info")
94-
}
95-
}
80+
// The tenant ID is the first column and comes from the PK decoder above.
81+
datums[0] = tenantIDRow[0].Datum
9682

97-
// The name, data state and service mode columns only exist after the
98-
// V23_1TenantNamesStateAndServiceMode migration has run. We need to
99-
// keep it optional here until we're not supporting running against
100-
// v23.1 versions any more.
101-
var name roachpb.TenantName
102-
// Compatibility with rows prior to the
103-
// V23_1TenantNamesStateAndServiceMode migration.
104-
dataState := mtinfopb.DataStateReady
105-
serviceMode := mtinfopb.ServiceModeExternal
106-
if len(datums) >= 6 {
107-
if i := datums[3]; i != tree.DNull {
108-
name = roachpb.TenantName(tree.MustBeDString(i))
109-
}
110-
if i := datums[4]; i != tree.DNull {
111-
rawDataState := tree.MustBeDInt(i)
112-
if rawDataState >= 0 && rawDataState <= tree.DInt(mtinfopb.MaxDataState) {
113-
dataState = mtinfopb.TenantDataState(rawDataState)
114-
} else {
115-
// This can happen if e.g. an invalid value was added into the
116-
// table manually.
117-
log.Warningf(ctx, "invalid data state %d for tenant %d", rawDataState, tenID)
118-
}
119-
}
120-
if i := datums[5]; i != tree.DNull {
121-
rawServiceMode := tree.MustBeDInt(i)
122-
if rawServiceMode >= 0 && rawServiceMode <= tree.DInt(mtinfopb.MaxServiceMode) {
123-
serviceMode = mtinfopb.TenantServiceMode(rawServiceMode)
124-
} else {
125-
// This can happen if e.g. an invalid value was added into the
126-
// table manually.
127-
log.Warningf(ctx, "invalid service mode %d for tenant %d", rawServiceMode, tenID)
128-
}
129-
}
83+
tid, info, err := mtinfo.GetTenantInfoFromSQLRow(datums)
84+
if err != nil {
85+
return tenantcapabilities.Entry{}, err
13086
}
13187

13288
return tenantcapabilities.Entry{
133-
TenantID: tenID,
134-
TenantCapabilities: &tenantInfo.Capabilities,
135-
Name: name,
136-
DataState: dataState,
137-
ServiceMode: serviceMode,
89+
TenantID: tid,
90+
TenantCapabilities: &info.Capabilities,
91+
Name: info.Name,
92+
DataState: info.DataState,
93+
ServiceMode: info.ServiceMode,
13894
}, nil
13995
}
14096

pkg/sql/BUILD.bazel

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -327,6 +327,7 @@ go_library(
327327
"//pkg/kv/kvserver/liveness/livenesspb",
328328
"//pkg/kv/kvserver/protectedts",
329329
"//pkg/multitenant",
330+
"//pkg/multitenant/mtinfo",
330331
"//pkg/multitenant/mtinfopb",
331332
"//pkg/multitenant/multitenantcpu",
332333
"//pkg/multitenant/tenantcapabilities",

pkg/sql/tenant_accessors.go

Lines changed: 5 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -16,8 +16,8 @@ import (
1616
"github.com/cockroachdb/cockroach/pkg/clusterversion"
1717
"github.com/cockroachdb/cockroach/pkg/keys"
1818
"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
19+
"github.com/cockroachdb/cockroach/pkg/multitenant/mtinfo"
1920
"github.com/cockroachdb/cockroach/pkg/multitenant/mtinfopb"
20-
"github.com/cockroachdb/cockroach/pkg/multitenant/tenantcapabilities"
2121
"github.com/cockroachdb/cockroach/pkg/roachpb"
2222
"github.com/cockroachdb/cockroach/pkg/settings"
2323
"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -26,7 +26,6 @@ import (
2626
"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
2727
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
2828
"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
29-
"github.com/cockroachdb/cockroach/pkg/util/protoutil"
3029
"github.com/cockroachdb/errors"
3130
)
3231

@@ -102,60 +101,8 @@ func GetTenantRecordByName(
102101
} else if row == nil {
103102
return nil, pgerror.Newf(pgcode.UndefinedObject, "tenant %q does not exist", tenantName)
104103
}
105-
return getTenantInfoFromRow(row)
106-
}
107-
108-
func getTenantInfoFromRow(row tree.Datums) (*mtinfopb.TenantInfo, error) {
109-
info := &mtinfopb.TenantInfo{}
110-
info.ID = uint64(tree.MustBeDInt(row[0]))
111-
112-
// For the benefit of pre-23.1 BACKUP/RESTORE.
113-
info.DeprecatedID = info.ID
114-
115-
infoBytes := []byte(tree.MustBeDBytes(row[1]))
116-
if err := protoutil.Unmarshal(infoBytes, &info.ProtoInfo); err != nil {
117-
return nil, err
118-
}
119-
120-
// If we are loading the entry for the system tenant, inject all the
121-
// capabilities that the system tenant should have. This will allow
122-
// us to reduce our dependency on a tenant ID check to retain the
123-
// system tenant's access to all services.
124-
if roachpb.IsSystemTenantID(info.ID) {
125-
tenantcapabilities.EnableAll(&info.ProtoInfo.Capabilities)
126-
}
127-
128-
// Load the name if defined.
129-
if row[2] != tree.DNull {
130-
info.Name = roachpb.TenantName(tree.MustBeDString(row[2]))
131-
}
132-
133-
// Load the data state column if defined.
134-
if row[3] != tree.DNull {
135-
info.DataState = mtinfopb.TenantDataState(tree.MustBeDInt(row[3]))
136-
} else {
137-
// Pre-v23.1 info struct.
138-
switch info.ProtoInfo.DeprecatedDataState {
139-
case mtinfopb.ProtoInfo_READY:
140-
info.DataState = mtinfopb.DataStateReady
141-
case mtinfopb.ProtoInfo_ADD:
142-
info.DataState = mtinfopb.DataStateAdd
143-
case mtinfopb.ProtoInfo_DROP:
144-
info.DataState = mtinfopb.DataStateDrop
145-
default:
146-
return nil, errors.AssertionFailedf("unhandled: %d", info.ProtoInfo.DeprecatedDataState)
147-
}
148-
}
149-
150-
// Load the service mode if defined.
151-
info.ServiceMode = mtinfopb.ServiceModeNone
152-
if row[4] != tree.DNull {
153-
info.ServiceMode = mtinfopb.TenantServiceMode(tree.MustBeDInt(row[4]))
154-
} else if info.DataState == mtinfopb.DataStateReady {
155-
// Records created for CC Serverless pre-v23.1.
156-
info.ServiceMode = mtinfopb.ServiceModeExternal
157-
}
158-
return info, nil
104+
_, info, err := mtinfo.GetTenantInfoFromSQLRow(row)
105+
return info, err
159106
}
160107

161108
// GetTenantRecordByID retrieves a tenant in system.tenants.
@@ -176,7 +123,8 @@ func GetTenantRecordByID(
176123
return nil, pgerror.Newf(pgcode.UndefinedObject, "tenant \"%d\" does not exist", tenID.ToUint64())
177124
}
178125

179-
return getTenantInfoFromRow(row)
126+
_, info, err := mtinfo.GetTenantInfoFromSQLRow(row)
127+
return info, err
180128
}
181129

182130
// LookupTenantID implements the tree.TenantOperator interface.

0 commit comments

Comments
 (0)