Skip to content

Commit

Permalink
sql: amortize schema descriptor lookups in vtable generation
Browse files Browse the repository at this point in the history 
This gets rid of some deprecated descs.Collection method calls.

Informs cockroachdb#64089.

Release note: None
  • Loading branch information
Marius Posta committed Dec 14, 2022
1 parent 42e6370 commit 192b26e
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 115 deletions.
18 changes: 9 additions & 9 deletions pkg/bench/rttanalysis/testdata/benchmark_expectations
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,9 @@ exp,benchmark
11,GrantRole/grant_1_role
15,GrantRole/grant_2_roles
4,ORMQueries/activerecord_type_introspection_query
8,ORMQueries/django_table_introspection_1_table
8,ORMQueries/django_table_introspection_4_tables
8,ORMQueries/django_table_introspection_8_tables
6,ORMQueries/django_table_introspection_1_table
6,ORMQueries/django_table_introspection_4_tables
6,ORMQueries/django_table_introspection_8_tables
2,ORMQueries/has_column_privilege_using_attnum
2,ORMQueries/has_column_privilege_using_column_name
1,ORMQueries/has_schema_privilege_1
Expand All @@ -69,15 +69,15 @@ exp,benchmark
6,ORMQueries/has_table_privilege_5
85,ORMQueries/hasura_column_descriptions
85,ORMQueries/hasura_column_descriptions_8_tables
8,ORMQueries/hasura_column_descriptions_modified
6,ORMQueries/hasura_column_descriptions_modified
6,ORMQueries/information_schema._pg_index_position
5,ORMQueries/pg_attribute
5,ORMQueries/pg_class
9,ORMQueries/pg_is_other_temp_schema
9,ORMQueries/pg_is_other_temp_schema_multiple_times
6,ORMQueries/pg_my_temp_schema
6,ORMQueries/pg_my_temp_schema_multiple_times
6,ORMQueries/pg_namespace
7,ORMQueries/pg_is_other_temp_schema
7,ORMQueries/pg_is_other_temp_schema_multiple_times
4,ORMQueries/pg_my_temp_schema
4,ORMQueries/pg_my_temp_schema_multiple_times
4,ORMQueries/pg_namespace
5,ORMQueries/pg_type
10,Revoke/revoke_all_on_1_table
10,Revoke/revoke_all_on_2_tables
Expand Down
6 changes: 3 additions & 3 deletions pkg/sql/crdb_internal.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2738,7 +2738,7 @@ CREATE TABLE crdb_internal.create_schema_statements (
dbDescs = append(dbDescs, db)
}
for _, db := range dbDescs {
return forEachSchema(ctx, p, db, func(schemaDesc catalog.SchemaDescriptor) error {
return forEachSchema(ctx, p, db, true /* requiresPrivileges */, func(schemaDesc catalog.SchemaDescriptor) error {
switch schemaDesc.SchemaKind() {
case catalog.SchemaUserDefined:
node := &tree.CreateSchema{
Expand Down Expand Up @@ -2794,7 +2794,7 @@ CREATE TABLE crdb_internal.create_function_statements (
fnIDToDBName := make(map[descpb.ID]string)
fnIDToDBID := make(map[descpb.ID]descpb.ID)
for _, curDB := range dbDescs {
err := forEachSchema(ctx, p, curDB, func(sc catalog.SchemaDescriptor) error {
err := forEachSchema(ctx, p, curDB, true /* requiresPrivileges */, func(sc catalog.SchemaDescriptor) error {
return sc.ForEachFunctionOverload(func(overload descpb.SchemaDescriptor_FunctionOverload) error {
fnIDs = append(fnIDs, overload.ID)
fnIDToScName[overload.ID] = sc.GetName()
Expand Down Expand Up @@ -5661,7 +5661,7 @@ CREATE TABLE crdb_internal.default_privileges (
return err
}

return forEachSchema(ctx, p, descriptor, func(schema catalog.SchemaDescriptor) error {
return forEachSchema(ctx, p, descriptor, true /* requiresPrivileges */, func(schema catalog.SchemaDescriptor) error {
return addRowsForSchema(schema.GetDefaultPrivilegeDescriptor(), tree.NewDString(schema.GetName()))
})
})
Expand Down
150 changes: 53 additions & 97 deletions pkg/sql/information_schema.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ import (
"unicode/utf8"

"github.com/cockroachdb/cockroach/pkg/docs"
"github.com/cockroachdb/cockroach/pkg/keys"
"github.com/cockroachdb/cockroach/pkg/kv"
"github.com/cockroachdb/cockroach/pkg/security/username"
"github.com/cockroachdb/cockroach/pkg/sql/catalog"
Expand All @@ -39,6 +38,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/sql/sqlutil"
"github.com/cockroachdb/cockroach/pkg/sql/types"
"github.com/cockroachdb/cockroach/pkg/sql/vtable"
"github.com/cockroachdb/cockroach/pkg/util/iterutil"
"github.com/cockroachdb/cockroach/pkg/util/timeutil/pgdate"
"github.com/cockroachdb/errors"
"github.com/lib/pq/oid"
Expand Down Expand Up @@ -980,7 +980,7 @@ https://www.postgresql.org/docs/9.5/infoschema-schemata.html`,
populate: func(ctx context.Context, p *planner, dbContext catalog.DatabaseDescriptor, addRow func(...tree.Datum) error) error {
return forEachDatabaseDesc(ctx, p, dbContext, true, /* requiresPrivileges */
func(db catalog.DatabaseDescriptor) error {
return forEachSchema(ctx, p, db, func(sc catalog.SchemaDescriptor) error {
return forEachSchema(ctx, p, db, true /* requiresPrivileges */, func(sc catalog.SchemaDescriptor) error {
return addRow(
tree.NewDString(db.GetName()), // catalog_name
tree.NewDString(sc.GetName()), // schema_name
Expand Down Expand Up @@ -1075,7 +1075,7 @@ var informationSchemaSchemataTablePrivileges = virtualSchemaTable{
populate: func(ctx context.Context, p *planner, dbContext catalog.DatabaseDescriptor, addRow func(...tree.Datum) error) error {
return forEachDatabaseDesc(ctx, p, dbContext, true, /* requiresPrivileges */
func(db catalog.DatabaseDescriptor) error {
return forEachSchema(ctx, p, db, func(sc catalog.SchemaDescriptor) error {
return forEachSchema(ctx, p, db, true /* requiresPrivileges */, func(sc catalog.SchemaDescriptor) error {
privs := sc.GetPrivileges().Show(privilege.Schema, true /* showImplicitOwnerPrivs */)
dbNameStr := tree.NewDString(db.GetName())
scNameStr := tree.NewDString(sc.GetName())
Expand Down Expand Up @@ -2247,69 +2247,57 @@ var informationSchemaViewTableUsageTable = virtualSchemaTable{
func forEachSchema(
ctx context.Context,
p *planner,
db catalog.DatabaseDescriptor,
dbContext catalog.DatabaseDescriptor,
requiresPrivileges bool,
fn func(sc catalog.SchemaDescriptor) error,
) error {
schemaNames, err := getSchemaNames(ctx, p, db)
if err != nil {
return err
}

vtableEntries := p.getVirtualTabler().getSchemas()
schemas := make([]catalog.SchemaDescriptor, 0, len(schemaNames)+len(vtableEntries))
var userDefinedSchemaIDs []descpb.ID
for id, name := range schemaNames {
switch {
case strings.HasPrefix(name, catconstants.PgTempSchemaName):
schemas = append(schemas, schemadesc.NewTemporarySchema(name, id, db.GetID()))
case name == tree.PublicSchema:
// TODO(richardjcai): Remove this in 22.2. In 22.2, only the system
// public schema will continue to use keys.PublicSchemaID (29).
if id == keys.PublicSchemaID {
schemas = append(schemas, schemadesc.GetPublicSchema())
} else {
// The default case is a user defined schema. Collect the ID to get the
// descriptor later.
userDefinedSchemaIDs = append(userDefinedSchemaIDs, id)
forEachDatabase := func(db catalog.DatabaseDescriptor) error {
c, err := p.Descriptors().GetAllSchemasInDatabase(ctx, p.txn, db)
if err != nil {
return err
}
var schemas []catalog.SchemaDescriptor
if err := c.ForEachDescriptor(func(desc catalog.Descriptor) error {
if requiresPrivileges {
canSeeDescriptor, err := userCanSeeDescriptor(ctx, p, desc, db, false /* allowAdding */)
if err != nil {
return err
}
if !canSeeDescriptor {
return nil
}
}
sc, err := catalog.AsSchemaDescriptor(desc)
schemas = append(schemas, sc)
return err
}); err != nil {
return err
}
sort.Slice(schemas, func(i int, j int) bool {
return schemas[i].GetName() < schemas[j].GetName()
})
for _, sc := range schemas {
if err := fn(sc); err != nil {
return err
}
default:
// The default case is a user defined schema. Collect the ID to get the
// descriptor later.
userDefinedSchemaIDs = append(userDefinedSchemaIDs, id)
}
return nil
}

userDefinedSchemas, err := p.Descriptors().Direct().GetSchemaDescriptorsFromIDs(ctx, p.txn, userDefinedSchemaIDs)
if dbContext != nil {
return iterutil.Map(forEachDatabase(dbContext))
}
c, err := p.Descriptors().GetAllDatabases(ctx, p.txn)
if err != nil {
return err
}
for i := range userDefinedSchemas {
desc := userDefinedSchemas[i]
canSeeDescriptor, err := userCanSeeDescriptor(ctx, p, desc, db, false /* allowAdding */)
return c.ForEachDescriptor(func(desc catalog.Descriptor) error {
db, err := catalog.AsDatabaseDescriptor(desc)
if err != nil {
return err
}
if !canSeeDescriptor {
continue
}
schemas = append(schemas, desc)
}

for _, schema := range vtableEntries {
schemas = append(schemas, schema.Desc())
}

sort.Slice(schemas, func(i int, j int) bool {
return schemas[i].GetName() < schemas[j].GetName()
return forEachDatabase(db)
})

for _, sc := range schemas {
if err := fn(sc); err != nil {
return err
}
}

return nil
}

// forEachDatabaseDesc calls a function for the given DatabaseDescriptor, or if
Expand Down Expand Up @@ -2499,32 +2487,6 @@ func forEachTableDescWithTableLookup(
)
}

func getSchemaNames(
ctx context.Context, p *planner, dbContext catalog.DatabaseDescriptor,
) (map[descpb.ID]string, error) {
if dbContext != nil {
return p.Descriptors().GetSchemasForDatabase(ctx, p.txn, dbContext)
}
ret := make(map[descpb.ID]string)
allDbDescs, err := p.Descriptors().GetAllDatabaseDescriptors(ctx, p.txn)
if err != nil {
return nil, err
}
for _, db := range allDbDescs {
if db == nil {
return nil, catalog.ErrDescriptorNotFound
}
schemas, err := p.Descriptors().GetSchemasForDatabase(ctx, p.txn, db)
if err != nil {
return nil, err
}
for id, name := range schemas {
ret[id] = name
}
}
return ret, nil
}

// forEachTableDescWithTableLookupInternal is the logic that supports
// forEachTableDescWithTableLookup.
//
Expand Down Expand Up @@ -2655,27 +2617,21 @@ func forEachTableDescWithTableLookupInternalFromDescriptors(
// missing temporary schema name. Temporary schemas have namespace
// entries. The below code will go and lookup schema names from the
// namespace table if needed to qualify the name of a temporary table.
namesForSchema, err := getSchemaNames(ctx, p, dbDesc)
if err != nil {
return errors.Wrapf(err, "failed to look up schema id %d", table.GetParentSchemaID())
}
for id, n := range namesForSchema {
if id != table.GetParentSchemaID() {
continue
if err := forEachSchema(ctx, p, dbDesc, false /* requiresPrivileges*/, func(schema catalog.SchemaDescriptor) error {
if schema.GetID() != table.GetParentSchemaID() {
return nil
}
_, exists, err := lCtx.GetSchemaName(ctx, id, dbDesc.GetID(), p.ExecCfg().Settings.Version)
if err != nil {
_, exists, err := lCtx.GetSchemaName(ctx, schema.GetID(), dbDesc.GetID(), p.ExecCfg().Settings.Version)
if err != nil || exists {
return err
}
if exists {
continue
}
if strings.HasPrefix(n, catconstants.PgTempSchemaName) {
sc = schemadesc.NewTemporarySchema(n, id, dbDesc.GetID())
lCtx.schemaNames[id] = n
lCtx.schemaDescs[id] = sc
lCtx.schemaIDs = append(lCtx.schemaIDs, id)
}
sc = schema
lCtx.schemaNames[sc.GetID()] = sc.GetName()
lCtx.schemaDescs[sc.GetID()] = sc
lCtx.schemaIDs = append(lCtx.schemaIDs, sc.GetID())
return nil
}); err != nil {
return errors.Wrapf(err, "failed to look up schema id %d", table.GetParentSchemaID())
}
if sc == nil {
sc = schemadesc.NewTemporarySchema(catconstants.PgTempSchemaName, table.GetParentSchemaID(), dbDesc.GetID())
Expand Down
19 changes: 13 additions & 6 deletions pkg/sql/pg_catalog.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/sql/sqlerrors"
"github.com/cockroachdb/cockroach/pkg/sql/types"
"github.com/cockroachdb/cockroach/pkg/sql/vtable"
"github.com/cockroachdb/cockroach/pkg/util/iterutil"
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/cockroachdb/errors"
"github.com/lib/pq/oid"
Expand Down Expand Up @@ -2058,7 +2059,7 @@ https://www.postgresql.org/docs/9.5/catalog-pg-namespace.html`,
populate: func(ctx context.Context, p *planner, dbContext catalog.DatabaseDescriptor, addRow func(...tree.Datum) error) error {
return forEachDatabaseDesc(ctx, p, dbContext, true, /* requiresPrivileges */
func(db catalog.DatabaseDescriptor) error {
return forEachSchema(ctx, p, db, func(sc catalog.SchemaDescriptor) error {
return forEachSchema(ctx, p, db, true /* requiresPrivileges */, func(sc catalog.SchemaDescriptor) error {
ownerOID := tree.DNull
if sc.SchemaKind() == catalog.SchemaUserDefined {
var err error
Expand Down Expand Up @@ -2102,12 +2103,18 @@ https://www.postgresql.org/docs/9.5/catalog-pg-namespace.html`,
return nil, false, err
}
// Fallback to looking for temporary schemas.
schemaNames, err := getSchemaNames(ctx, p, db)
if err != nil {
var tempSchema catalog.SchemaDescriptor
if err := forEachSchema(ctx, p, db, false /* requiresPrivileges */, func(schema catalog.SchemaDescriptor) error {
if schema.GetID() != descpb.ID(ooid) {
return nil
}
tempSchema = schema
return iterutil.StopIteration()
}); err != nil {
return nil, false, err
}
if scName, ok := schemaNames[descpb.ID(ooid)]; ok {
return schemadesc.NewTemporarySchema(scName, descpb.ID(ooid), db.GetID()), true, nil
if tempSchema != nil {
return tempSchema, true, nil
}
return nil, false, nil
}()
Expand Down Expand Up @@ -2543,7 +2550,7 @@ https://www.postgresql.org/docs/9.5/catalog-pg-proc.html`,
}
return forEachDatabaseDesc(ctx, p, dbContext, false, /* requiresPrivileges */
func(dbDesc catalog.DatabaseDescriptor) error {
return forEachSchema(ctx, p, dbDesc, func(scDesc catalog.SchemaDescriptor) error {
return forEachSchema(ctx, p, dbDesc, true /* requiresPrivileges */, func(scDesc catalog.SchemaDescriptor) error {
return scDesc.ForEachFunctionOverload(func(overload descpb.SchemaDescriptor_FunctionOverload) error {
fnDesc, err := p.Descriptors().GetImmutableFunctionByID(
ctx, p.Txn(), overload.ID,
Expand Down

0 comments on commit 192b26e

Please sign in to comment.