Skip to content

Commit

Permalink
sql: add session variable to not hold up role membership changes
Browse files Browse the repository at this point in the history
Epic: CRDB-22371

Release note (sql change): A new session variable has been introduced which
can be used to make the granting and revoking of role memberships faster at
the cost of some isolation claims. When granting or revoking a role from
another role, the system waits until all transactions which are consulting the
current set of role memberships to complete. This means, by default, that by
the time the transaction which performed the grant or revoke operation returns
successfully, the user has a proof that no ongoing transaction is relying on
the state that existed prior to the change. The downside of this waiting is
that it means that grant and revoke will take longer than the longest currently
executing transaction. In some cases, users do not care about whether concurrent
transactions will immediately see the side-effects of the operation, but would
instead prefer that the grant or revoke finish rapidly. In order to aid in
those cases, a new session variable
`allow_role_memberships_to_change_during_transaction` has been added. Now, the
grant or revoke will only need to wait for the completion of statements in
sessions which do not have that option set. One can set the option as enabled
by default in all sessions in order to accelerate and grant and revoke role
operations.
  • Loading branch information
ajwerner committed Mar 14, 2023
1 parent 25c00d4 commit 8db91e5
Show file tree
Hide file tree
Showing 9 changed files with 185 additions and 1 deletion.
17 changes: 16 additions & 1 deletion pkg/sql/authorization.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import (
"github.com/cockroachdb/cockroach/pkg/sql/catalog/descpb"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/descs"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/funcdesc"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/lease"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/schemadesc"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/tabledesc"
"github.com/cockroachdb/cockroach/pkg/sql/catalog/typedesc"
Expand Down Expand Up @@ -492,7 +493,7 @@ func (p *planner) MemberOfWithAdminOption(
// Requires a valid transaction to be open.
func MemberOfWithAdminOption(
ctx context.Context, execCfg *ExecutorConfig, txn descs.Txn, member username.SQLUsername,
) (map[username.SQLUsername]bool, error) {
) (_ map[username.SQLUsername]bool, retErr error) {
if txn == nil {
return nil, errors.AssertionFailedf("cannot use MemberOfWithAdminoption without a txn")
}
Expand All @@ -511,6 +512,20 @@ func MemberOfWithAdminOption(
if tableDesc.IsUncommittedVersion() {
return resolveMemberOfWithAdminOption(ctx, member, txn, useSingleQueryForRoleMembershipCache.Get(execCfg.SV()))
}
if txn.SessionData().AllowRoleMembershipsToChangeDuringTransaction {
defer func() {
if retErr != nil {
return
}
txn.Descriptors().ReleaseSpecifiedLeases(ctx, []lease.IDVersion{
{
Name: tableDesc.GetName(),
ID: tableDesc.GetID(),
Version: tableVersion,
},
})
}()
}

// Check version and maybe clear cache while holding the mutex.
// We use a closure here so that we release the lock here, then keep
Expand Down
4 changes: 4 additions & 0 deletions pkg/sql/exec_util.go
Original file line number Diff line number Diff line change
Expand Up @@ -3474,6 +3474,10 @@ func (m *sessionDataMutator) SetEnableCreateStatsUsingExtremes(val bool) {
m.data.EnableCreateStatsUsingExtremes = val
}

func (m *sessionDataMutator) SetAllowRoleMembershipsToChangeDuringTransaction(val bool) {
m.data.AllowRoleMembershipsToChangeDuringTransaction = val
}

// Utility functions related to scrubbing sensitive information on SQL Stats.

// quantizeCounts ensures that the Count field in the
Expand Down
1 change: 1 addition & 0 deletions pkg/sql/logictest/testdata/logic_test/information_schema
Original file line number Diff line number Diff line change
Expand Up @@ -4946,6 +4946,7 @@ WHERE
variable value
allow_ordinal_column_references off
allow_prepare_as_opt_plan off
allow_role_memberships_to_change_during_transaction off
alter_primary_region_super_region_override off
application_name ·
avoid_buffering off
Expand Down
3 changes: 3 additions & 0 deletions pkg/sql/logictest/testdata/logic_test/pg_catalog
Original file line number Diff line number Diff line change
Expand Up @@ -2566,6 +2566,7 @@ WHERE
----
name setting category short_desc extra_desc vartype
allow_ordinal_column_references off NULL NULL NULL string
allow_role_memberships_to_change_during_transaction off NULL NULL NULL string
alter_primary_region_super_region_override off NULL NULL NULL string
application_name · NULL NULL NULL string
avoid_buffering off NULL NULL NULL string
Expand Down Expand Up @@ -2715,6 +2716,7 @@ WHERE
----
name setting unit context enumvals boot_val reset_val
allow_ordinal_column_references off NULL user NULL off off
allow_role_memberships_to_change_during_transaction off NULL user NULL off off
alter_primary_region_super_region_override off NULL user NULL off off
application_name · NULL user NULL · ·
avoid_buffering off NULL user NULL false false
Expand Down Expand Up @@ -2858,6 +2860,7 @@ SELECT name, source, min_val, max_val, sourcefile, sourceline FROM pg_catalog.pg
----
name source min_val max_val sourcefile sourceline
allow_ordinal_column_references NULL NULL NULL NULL NULL
allow_role_memberships_to_change_during_transaction NULL NULL NULL NULL NULL
alter_primary_region_super_region_override NULL NULL NULL NULL NULL
application_name NULL NULL NULL NULL NULL
avoid_buffering NULL NULL NULL NULL NULL
Expand Down
1 change: 1 addition & 0 deletions pkg/sql/logictest/testdata/logic_test/show_source
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ WHERE variable NOT IN ('optimizer', 'crdb_version', 'session_id', 'distsql_workm
----
variable value
allow_ordinal_column_references off
allow_role_memberships_to_change_during_transaction off
alter_primary_region_super_region_override off
application_name ·
avoid_buffering off
Expand Down
8 changes: 8 additions & 0 deletions pkg/sql/sessiondatapb/local_only_session_data.proto
Original file line number Diff line number Diff line change
Expand Up @@ -349,6 +349,14 @@ message LocalOnlySessionData {
// EnableCreateStatsUsingExtremes, when true, allows the use of CREATE
// STATISTICS .. USING EXTREMES.
bool enable_create_stats_using_extremes = 95;
// AllowRoleMembershipsToChangeDuringTransaction, when true, means that
// operations which consult the role membership cache do not retain their
// lease on that version of the cache throughout the transaction. The
// consequence of this is that the transaction may not experience a singular
// view of role membership, and it may be able to commit after the revocation
// of a role membership which the transaction relied on has successfully been
// committed and acknowledged to the user.
bool allow_role_memberships_to_change_during_transaction = 96;

///////////////////////////////////////////////////////////////////////////
// WARNING: consider whether a session parameter you're adding needs to //
Expand Down
1 change: 1 addition & 0 deletions pkg/sql/tests/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ go_test(
name = "tests_test",
size = "large",
srcs = [
"allow_role_memberships_to_change_during_transaction_test.go",
"autocommit_extended_protocol_test.go",
"bank_test.go",
"copy_file_upload_test.go",
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
// Copyright 2023 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.

package tests

import (
"context"
gosql "database/sql"
"net/url"
"testing"
"time"

"github.com/cockroachdb/cockroach/pkg/base"
"github.com/cockroachdb/cockroach/pkg/testutils"
"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
"github.com/cockroachdb/cockroach/pkg/util/leaktest"
"github.com/cockroachdb/cockroach/pkg/util/log"
"github.com/stretchr/testify/require"
)

// TestAllowRoleMembershipsToChangeDuringTransaction ensures that when the
// corresponding session variable is set for all sessions using cached
// role memberships, performing new grants do not need to wait for open
// transactions to complete.
func TestAllowRoleMembershipsToChangeDuringTransaction(t *testing.T) {
defer leaktest.AfterTest(t)()
defer log.Scope(t).Close(t)

ctx := context.Background()
s, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{})
defer s.Stopper().Stop(ctx)

openUser := func(username, dbName string) (_ *gosql.DB, cleanup func()) {
pgURL, testuserCleanupFunc := sqlutils.PGUrlWithOptionalClientCerts(
t, s.ServingSQLAddr(), username,
url.UserPassword(username, username),
false, /* withClientCerts */
)
pgURL.Path = dbName
db, err := gosql.Open("postgres", pgURL.String())
if err != nil {
t.Fatal(err)
}
return db, func() {
require.NoError(t, db.Close())
testuserCleanupFunc()
}
}

// Create three users: foo, bar, and baz.
// Use one of these users to hold open a transaction which uses a lease
// on the role_memberships table. Ensure that initially granting does
// wait on that transaction. Then set the session variable and ensure
// that it does not.
tdb := sqlutils.MakeSQLRunner(sqlDB)
tdb.Exec(t, "CREATE USER foo PASSWORD 'foo'")
tdb.Exec(t, "CREATE USER bar")
tdb.Exec(t, "CREATE USER baz")
tdb.Exec(t, "GRANT admin TO foo")
tdb.Exec(t, "CREATE DATABASE db2")
tdb.Exec(t, "CREATE TABLE db2.public.t (i int primary key)")

fooDB, cleanupFoo := openUser("foo", "db2")
defer cleanupFoo()

// In this first test, we want to make sure that the query cannot proceed
// until the outer transaction commits. We launch the grant while the
// transaction is open, wait a tad, make sure it hasn't made progress,
// then we commit the relevant transaction and ensure that it does finish.
t.Run("normal path waits", func(t *testing.T) {
fooTx, err := fooDB.BeginTx(ctx, nil)
require.NoError(t, err)
var count int
require.NoError(t,
fooTx.QueryRow("SELECT count(*) FROM t").Scan(&count))
require.Equal(t, 0, count)
errCh := make(chan error, 1)
go func() {
_, err := sqlDB.Exec("GRANT bar TO baz;")
errCh <- err
}()
select {
case <-time.After(10 * time.Millisecond):
// Wait a tad and make sure nothing happens. This test will be
// flaky if we mess up the leasing.
case err := <-errCh:
t.Fatalf("expected transaction to block, got err: %v", err)
}
require.NoError(t, fooTx.Commit())
require.NoError(t, <-errCh)
})
// In this test we ensure that we can perform role grant and revoke
// operations while the transaction which uses the relevant roles
// remains open. We ensure that the transaction still succeeds and
// that the operations occur in a timely manner.
t.Run("session variable prevents waiting", func(t *testing.T) {
fooConn, err := fooDB.Conn(ctx)
require.NoError(t, err)
defer func() { _ = fooConn.Close() }()
_, err = fooConn.ExecContext(ctx, "SET allow_role_memberships_to_change_during_transaction = true;")
require.NoError(t, err)
fooTx, err := fooConn.BeginTx(ctx, nil)
require.NoError(t, err)
var count int
require.NoError(t,
fooTx.QueryRow("SELECT count(*) FROM t").Scan(&count))
require.Equal(t, 0, count)
conn, err := sqlDB.Conn(ctx)
require.NoError(t, err)
defer func() { _ = conn.Close() }()
// Set a timeout on the SQL operations to ensure that they both
// happen in a timely manner.
grantRevokeTimeout, cancel := context.WithTimeout(
ctx, testutils.DefaultSucceedsSoonDuration,
)
defer cancel()

_, err = conn.ExecContext(grantRevokeTimeout, "GRANT foo TO baz;")
require.NoError(t, err)
_, err = conn.ExecContext(grantRevokeTimeout, "REVOKE bar FROM baz;")
require.NoError(t, err)

// Ensure the transaction we held open commits without issue.
require.NoError(t, fooTx.Commit())
})
}
17 changes: 17 additions & 0 deletions pkg/sql/vars.go
Original file line number Diff line number Diff line change
Expand Up @@ -2549,6 +2549,23 @@ var varGen = map[string]sessionVar{
},
GlobalDefault: globalFalse,
},

// CockroachDB extension.
`allow_role_memberships_to_change_during_transaction`: {
GetStringVal: makePostgresBoolGetStringValFn(`allow_role_memberships_to_change_during_transaction`),
Set: func(_ context.Context, m sessionDataMutator, s string) error {
b, err := paramparse.ParseBoolVar("allow_role_memberships_to_change_during_transaction", s)
if err != nil {
return err
}
m.SetAllowRoleMembershipsToChangeDuringTransaction(b)
return nil
},
Get: func(evalCtx *extendedEvalContext, _ *kv.Txn) (string, error) {
return formatBoolAsPostgresSetting(evalCtx.SessionData().AllowRoleMembershipsToChangeDuringTransaction), nil
},
GlobalDefault: globalFalse,
},
}

// We want test coverage for this on and off so make it metamorphic.
Expand Down

0 comments on commit 8db91e5

Please sign in to comment.