builtins: improve performance of has_table_privilege and has_column_privilege

pkg/bench/ddl_analysis/orm_queries_bench_test.go

@@ -159,6 +159,26 @@ WHERE
 			setup: `CREATE TABLE t1(a int primary key, b int);`,
 			stmt:  `SELECT * FROM pg_attribute`,
 		},
+		{
+			name:  "has_table_privilege real table",
+			setup: `CREATE TABLE t(a int primary key, b int)`,
+			stmt:  `SELECT has_table_privilege('t', 'SELECT')`,
+		},
+		{
+			name:  "has_table_privilege virtual table",
+			setup: `CREATE TABLE t(a int primary key, b int)`,
+			stmt:  `SELECT has_table_privilege('t', 'SELECT')`,
+		},
+		{
+			name:  "has_column_privilege using attnum",
+			setup: `CREATE TABLE t(a int primary key, b int)`,
+			stmt:  `SELECT has_column_privilege('t', 1, 'INSERT')`,
+		},
+		{
+			name:  "has_column_privilege using column name",
+			setup: `CREATE TABLE t(a int primary key, b int)`,
+			stmt:  `SELECT has_column_privilege('t', 'a', 'INSERT')`,
+		},
 	}
 
 	RunRoundTripBenchmark(b, tests)

pkg/bench/ddl_analysis/testdata/benchmark_expectations

@@ -56,6 +56,10 @@ exp,benchmark
 2,ORMQueries/django_table_introspection_1_table
 2,ORMQueries/django_table_introspection_4_tables
 2,ORMQueries/django_table_introspection_8_tables
+2,ORMQueries/has_column_privilege_using_attnum
+2,ORMQueries/has_column_privilege_using_column_name
+2,ORMQueries/has_table_privilege_real_table
+2,ORMQueries/has_table_privilege_virtual_table
 2,ORMQueries/pg_attribute
 2,ORMQueries/pg_class
 2,ORMQueries/pg_namespace

pkg/ccl/importccl/import_table_creation.go

@@ -25,6 +25,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/resolver"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
+	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
@@ -268,6 +269,17 @@ func (so *importSequenceOperators) IsTypeVisible(
 	return false, false, errors.WithStack(errSequenceOperators)
 }
 
+// HasTablePrivilege is part of the tree.EvalDatabase interface.
+func (so *importSequenceOperators) HasPrivilege(
+	ctx context.Context,
+	specifier tree.HasPrivilegeSpecifier,
+	user security.SQLUsername,
+	kind privilege.Kind,
+	withGrantOpt bool,
+) (bool, error) {
+	return false, errors.WithStack(errSequenceOperators)
+}
+
 // Implements the tree.SequenceOperators interface.
 func (so *importSequenceOperators) IncrementSequence(
 	ctx context.Context, seqName *tree.TableName,

pkg/sql/faketreeeval/BUILD.bazel

@@ -11,6 +11,7 @@ go_library(
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/pgwire/pgnotice",
+        "//pkg/sql/privilege",
         "//pkg/sql/roleoption",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sessiondata",

pkg/sql/faketreeeval/evalctx.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgnotice"
+	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
 	"github.com/cockroachdb/cockroach/pkg/sql/roleoption"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
 	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
@@ -91,6 +92,17 @@ func (so *DummySequenceOperators) IsTypeVisible(
 	return false, false, errors.WithStack(errEvalPlanner)
 }
 
+// HasPrivilege is part of the tree.EvalDatabase interface.
+func (so *DummySequenceOperators) HasPrivilege(
+	ctx context.Context,
+	specifier tree.HasPrivilegeSpecifier,
+	user security.SQLUsername,
+	kind privilege.Kind,
+	withGrantOpt bool,
+) (bool, error) {
+	return false, errors.WithStack(errEvalPlanner)
+}
+
 // IncrementSequence is part of the tree.SequenceOperators interface.
 func (so *DummySequenceOperators) IncrementSequence(
 	ctx context.Context, seqName *tree.TableName,
@@ -231,6 +243,17 @@ func (ep *DummyEvalPlanner) IsTypeVisible(
 	return false, false, errors.WithStack(errEvalPlanner)
 }
 
+// HasPrivilege is part of the tree.EvalDatabase interface.
+func (ep *DummyEvalPlanner) HasPrivilege(
+	ctx context.Context,
+	specifier tree.HasPrivilegeSpecifier,
+	user security.SQLUsername,
+	kind privilege.Kind,
+	withGrantOpt bool,
+) (bool, error) {
+	return false, errors.WithStack(errEvalPlanner)
+}
+
 // ResolveTableName is part of the tree.EvalDatabase interface.
 func (ep *DummyEvalPlanner) ResolveTableName(
 	ctx context.Context, tn *tree.TableName,

pkg/sql/logictest/testdata/logic_test/pg_builtins

@@ -15,7 +15,7 @@ SELECT pg_my_temp_schema()
 ----
 0
 
-# Regression test for #49072.
+# Regression test for #43166.
 statement ok
 SELECT has_table_privilege('root'::NAME, 0, 'select')
 

pkg/sql/logictest/testdata/logic_test/privilege_builtins

@@ -182,9 +182,12 @@ true  true  true  true
 query error pgcode 42P01 relation "does_not_exist" does not exist
 SELECT has_column_privilege('does_not_exist', 1, 'SELECT')
 
-query error pgcode 42703 column 100 of relation 'pg_type' does not exist
+query error pgcode 42703 column 100 of relation pg_type does not exist
 SELECT has_column_privilege('pg_type', 100, 'SELECT')
 
+query error pgcode 42703 column 100 of relation pg_type does not exist
+SELECT has_column_privilege('pg_type'::regclass, 100, 'SELECT')
+
 query BBBBB
 SELECT has_column_privilege('pg_type', 1, 'SELECT'),
        has_column_privilege('pg_type', 1, 'INSERT'),
@@ -263,9 +266,12 @@ true  true  true  true
 query error pgcode 42P01 relation "does_not_exist" does not exist
 SELECT has_column_privilege('does_not_exist', 'col', 'SELECT')
 
-query error pgcode 42703 column 'does not exist' of relation 'pg_type' does not exist
+query error pgcode 42703 column "does not exist" does not exist
 SELECT has_column_privilege('pg_type', 'does not exist', 'SELECT')
 
+query error pgcode 42703 column "does not exist" does not exist
+SELECT has_column_privilege('pg_type'::regclass, 'does not exist', 'SELECT')
+
 query BBBBB
 SELECT has_column_privilege('pg_type', 'typname', 'SELECT'),
        has_column_privilege('pg_type', 'typname', 'INSERT'),
@@ -1135,7 +1141,7 @@ SELECT has_database_privilege('testuser', 'my_db', 'create'),
        has_schema_privilege('testuser', 'public', 'usage'),
        has_table_privilege('testuser', 's.t', 'select')
 ----
-true false false
+true  false  true
 
 user root
 
@@ -1164,4 +1170,4 @@ SELECT has_database_privilege('testuser2', 'my_db', 'create'),
        has_schema_privilege('testuser2', 'public', 'usage'),
        has_table_privilege('testuser2', 's.t', 'select')
 ----
-true false false
+true  false  true

pkg/sql/resolver.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv"
+	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/catalogkv"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/dbdesc"
@@ -25,6 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
 	"github.com/cockroachdb/cockroach/pkg/sql/opt/cat"
+	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
 	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
@@ -301,6 +303,160 @@ func (p *planner) IsTypeVisible(
 	return false, true, nil
 }
 
+// HasPrivilege is part of the tree.EvalDatabase interface.
+func (p *planner) HasPrivilege(
+	ctx context.Context,
+	specifier tree.HasPrivilegeSpecifier,
+	user security.SQLUsername,
+	kind privilege.Kind,
+	withGrantOpt bool,
+) (bool, error) {
+	desc, err := p.ResolveDescriptorForPrivilegeSpecifier(
+		ctx,
+		specifier,
+	)
+	if err != nil {
+		return false, err
+	}
+
+	// hasPrivilegeFunc checks whether any role has the given privilege.
+	hasPrivilegeFunc := func(priv privilege.Kind) (bool, error) {
+		err := p.CheckPrivilegeForUser(ctx, desc, priv, user)
+		if err != nil {
+			if pgerror.GetPGCode(err) == pgcode.InsufficientPrivilege {
+				return false, nil
+			}
+			return false, err
+		}
+		return true, nil
+	}
+
+	hasPrivilege, err := hasPrivilegeFunc(privilege.ALL)
+	if err != nil {
+		return false, err
+	}
+	if hasPrivilege {
+		return true, nil
+	}
+	// For WITH GRANT OPTION, check the roles also has the GRANT privilege.
+	if withGrantOpt {
+		hasPrivilege, err := hasPrivilegeFunc(privilege.GRANT)
+		if err != nil {
+			return false, err
+		}
+		if !hasPrivilege {
+			return false, nil
+		}
+	}
+	return hasPrivilegeFunc(kind)
+}
+
+// ResolveDescriptorForPrivilegeSpecifier resolves a tree.HasPrivilegeSpecifier
+// and returns the descriptor for the given object.
+func (p *planner) ResolveDescriptorForPrivilegeSpecifier(
+	ctx context.Context, specifier tree.HasPrivilegeSpecifier,
+) (catalog.Descriptor, error) {
+	if specifier.TableName != nil {
+		tn, err := parser.ParseQualifiedTableName(*specifier.TableName)
+		if err != nil {
+			return nil, err
+		}
+		if _, err := p.ResolveTableName(ctx, tn); err != nil {
+			return nil, err
+		}
+
+		if p.SessionData().Database != "" && p.SessionData().Database != string(tn.CatalogName) {
+			// Postgres does not allow cross-database references in these
+			// functions, so we don't either.
+			return nil, pgerror.Newf(pgcode.FeatureNotSupported,
+				"cross-database references are not implemented: %s", tn)
+		}
+
+		table, err := p.getVirtualTabler().getVirtualTableDesc(tn)
+		if err != nil {
+			return nil, err
+		}
+		if table == nil {
+			_, table, err = p.Descriptors().GetImmutableTableByName(
+				ctx,
+				p.txn,
+				tn,
+				tree.ObjectLookupFlags{
+					CommonLookupFlags: tree.CommonLookupFlags{
+						Required: true,
+					},
+				},
+			)
+			if err != nil {
+				return nil, err
+			}
+		}
+		if err := validateColumnForHasPrivilegeSpecifier(
+			table,
+			specifier,
+		); err != nil {
+			return nil, err
+		}
+		return table, nil
+	}
+	if specifier.TableOID == nil {
+		return nil, errors.AssertionFailedf("no table name or oid found")
+	}
+	virtualDesc, err := p.getVirtualTabler().getVirtualTableEntryByID(descpb.ID(*specifier.TableOID))
+	if err != nil && !errors.Is(err, catalog.ErrDescriptorNotFound) {
+		return nil, err
+	}
+	var table catalog.TableDescriptor
+	if virtualDesc != nil {
+		table = virtualDesc.desc
+	} else {
+		table, err = p.Descriptors().GetImmutableTableByID(
+			ctx,
+			p.txn,
+			descpb.ID(*specifier.TableOID),
+			tree.ObjectLookupFlags{
+				CommonLookupFlags: tree.CommonLookupFlags{
+					Required: true,
+				},
+			},
+		)
+		if err != nil {
+			return nil, err
+		}
+	}
+	if err := validateColumnForHasPrivilegeSpecifier(
+		table,
+		specifier,
+	); err != nil {
+		return nil, err
+	}
+	return table, nil
+}
+
+func validateColumnForHasPrivilegeSpecifier(
+	table catalog.TableDescriptor, specifier tree.HasPrivilegeSpecifier,
+) error {
+	if specifier.ColumnName != nil {
+		_, err := table.FindColumnWithName(*specifier.ColumnName)
+		return err
+	}
+	if specifier.ColumnAttNum != nil {
+		for _, col := range table.PublicColumns() {
+			if col.GetPGAttributeNum() == *specifier.ColumnAttNum {
+				return nil
+			}
+		}
+		return pgerror.Newf(
+			pgcode.UndefinedColumn,
+			"column %d of relation %s does not exist",
+			*specifier.ColumnAttNum,
+			tree.Name(table.GetName()),
+		)
+
+	}
+	return nil
+}
+
 // GetTypeDescriptor implements the descpb.TypeDescriptorResolver interface.
 func (p *planner) GetTypeDescriptor(
 	ctx context.Context, id descpb.ID,