Skip to content

Commit

Permalink
test
Browse files Browse the repository at this point in the history
  • Loading branch information
@ZhouXing19
ZhouXing19 committed Aug 18, 2022
1 parent 2aeed01 commit b8caec5
Show file tree
Hide file tree
Showing 6 changed files with 42 additions and 40 deletions.
18 changes: 9 additions & 9 deletions pkg/ccl/backupccl/restore_job.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1497,7 +1497,7 @@ func (r *restoreResumer) doResume(ctx context.Context, execCtx interface{}) erro
// Reload the details as we may have updated the job.
details = r.job.Details().(jobspb.RestoreDetails)

if err := r.cleanupTempSystemTables(ctx, nil /* txn */); err != nil {
if err := r.cleanupTempSystemTables(ctx); err != nil {
return err
}
} else if isSystemUserRestore(details) {
Expand All @@ -1506,7 +1506,7 @@ func (r *restoreResumer) doResume(ctx context.Context, execCtx interface{}) erro
}
details = r.job.Details().(jobspb.RestoreDetails)

if err := r.cleanupTempSystemTables(ctx, nil /* txn */); err != nil {
if err := r.cleanupTempSystemTables(ctx); err != nil {
return err
}
}
Expand Down Expand Up @@ -2012,12 +2012,12 @@ func (r *restoreResumer) OnFailOrCancel(
// We've dropped defaultdb and postgres in the planning phase, we must
// recreate them now if the full cluster restore failed.
ie := p.ExecCfg().InternalExecutor
_, err := ie.Exec(ctx, "recreate-defaultdb", txn, "CREATE DATABASE IF NOT EXISTS defaultdb")
_, err := ie.Exec(ctx, "recreate-defaultdb", nil /* txn */, "CREATE DATABASE IF NOT EXISTS defaultdb")
if err != nil {
return err
}

_, err = ie.Exec(ctx, "recreate-postgres", txn, "CREATE DATABASE IF NOT EXISTS postgres")
_, err = ie.Exec(ctx, "recreate-postgres", nil /* txn */, "CREATE DATABASE IF NOT EXISTS postgres")
if err != nil {
return err
}
Expand All @@ -2030,7 +2030,7 @@ func (r *restoreResumer) OnFailOrCancel(
if details.DescriptorCoverage == tree.AllDescriptors {
// The temporary system table descriptors should already have been dropped
// in `dropDescriptors` but we still need to drop the temporary system db.
if err := execCfg.DB.Txn(ctx, r.cleanupTempSystemTables); err != nil {
if err := r.cleanupTempSystemTables(ctx); err != nil {
return err
}
}
Expand Down Expand Up @@ -2551,12 +2551,12 @@ func (r *restoreResumer) restoreSystemTables(
return nil
}

func (r *restoreResumer) cleanupTempSystemTables(ctx context.Context, txn *kv.Txn) error {
func (r *restoreResumer) cleanupTempSystemTables(ctx context.Context) error {
executor := r.execCfg.InternalExecutor
// Check if the temp system database has already been dropped. This can happen
// if the restore job fails after the system database has cleaned up.
checkIfDatabaseExists := "SELECT database_name FROM [SHOW DATABASES] WHERE database_name=$1"
if row, err := executor.QueryRow(ctx, "checking-for-temp-system-db" /* opName */, txn, checkIfDatabaseExists, restoreTempSystemDB); err != nil {
if row, err := executor.QueryRow(ctx, "checking-for-temp-system-db" /* opName */, nil /* txn */, checkIfDatabaseExists, restoreTempSystemDB); err != nil {
return errors.Wrap(err, "checking for temporary system db")
} else if row == nil {
// Temporary system DB might already have been dropped by the restore job.
Expand All @@ -2566,11 +2566,11 @@ func (r *restoreResumer) cleanupTempSystemTables(ctx context.Context, txn *kv.Tx
// After restoring the system tables, drop the temporary database holding the
// system tables.
gcTTLQuery := fmt.Sprintf("ALTER DATABASE %s CONFIGURE ZONE USING gc.ttlseconds=1", restoreTempSystemDB)
if _, err := executor.Exec(ctx, "altering-gc-ttl-temp-system" /* opName */, txn, gcTTLQuery); err != nil {
if _, err := executor.Exec(ctx, "altering-gc-ttl-temp-system" /* opName */, nil /* txn */, gcTTLQuery); err != nil {
log.Errorf(ctx, "failed to update the GC TTL of %q: %+v", restoreTempSystemDB, err)
}
dropTableQuery := fmt.Sprintf("DROP DATABASE %s CASCADE", restoreTempSystemDB)
if _, err := executor.Exec(ctx, "drop-temp-system-db" /* opName */, txn, dropTableQuery); err != nil {
if _, err := executor.Exec(ctx, "drop-temp-system-db" /* opName */, nil /* txn */, dropTableQuery); err != nil {
return errors.Wrap(err, "dropping temporary system db")
}
return nil
Expand Down
36 changes: 18 additions & 18 deletions pkg/cloud/userfile/filetable/file_table_read_writer.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -146,14 +146,14 @@ type FileToTableSystem struct {
}

// FileTable which contains records for every uploaded file.
const fileTableSchema = `CREATE TABLE %s (filename STRING PRIMARY KEY,
const fileTableSchema = `CREATE TABLE IF NOT EXISTS %s (filename STRING PRIMARY KEY,
file_id UUID UNIQUE NOT NULL DEFAULT gen_random_uuid(),
file_size INT NOT NULL,
username STRING NOT NULL,
upload_time TIMESTAMP DEFAULT now())`

// PayloadTable contains the chunked payloads of each file.
const payloadTableSchema = `CREATE TABLE %s (file_id UUID,
const payloadTableSchema = `CREATE TABLE IF NOT EXISTS %s (file_id UUID,
byte_offset INT,
payload BYTES,
PRIMARY KEY(file_id, byte_offset))`
Expand Down Expand Up @@ -247,21 +247,21 @@ func NewFileToTableSystem(
// TODO(adityamaru): Handle scenario where the user has already created
// tables with the same names not via the FileToTableSystem
// object. Not sure if we want to error out or work around it.
tablesExist, err := f.checkIfFileAndPayloadTableExist(ctx, txn, e.ie)
tablesExist, err := f.checkIfFileAndPayloadTableExist(ctx, nil /* txn */, e.ie)
if err != nil {
return err
}

if !tablesExist {
if err := f.createFileAndPayloadTables(ctx, txn, e.ie); err != nil {
if err := f.createFileAndPayloadTables(ctx, e.ie); err != nil {
return err
}

if err := f.grantCurrentUserTablePrivileges(ctx, txn, e.ie); err != nil {
if err := f.grantCurrentUserTablePrivileges(ctx, e.ie); err != nil {
return err
}

if err := f.revokeOtherUserTablePrivileges(ctx, txn, e.ie); err != nil {
if err := f.revokeOtherUserTablePrivileges(ctx, e.ie); err != nil {
return err
}
}
Expand Down Expand Up @@ -365,15 +365,15 @@ func DestroyUserFileSystem(ctx context.Context, f *FileToTableSystem) error {
if err := e.db.Txn(ctx,
func(ctx context.Context, txn *kv.Txn) error {
dropPayloadTableQuery := fmt.Sprintf(`DROP TABLE %s`, f.GetFQPayloadTableName())
_, err := e.ie.ExecEx(ctx, "drop-payload-table", txn,
_, err := e.ie.ExecEx(ctx, "drop-payload-table", nil, /* txn */
sessiondata.InternalExecutorOverride{User: f.username},
dropPayloadTableQuery)
if err != nil {
return errors.Wrap(err, "failed to drop payload table")
}

dropFileTableQuery := fmt.Sprintf(`DROP TABLE %s CASCADE`, f.GetFQFileTableName())
_, err = e.ie.ExecEx(ctx, "drop-file-table", txn,
_, err = e.ie.ExecEx(ctx, "drop-file-table", nil, /* txn */
sessiondata.InternalExecutorOverride{User: f.username},
dropFileTableQuery)
if err != nil {
Expand Down Expand Up @@ -821,28 +821,28 @@ func (f *FileToTableSystem) checkIfFileAndPayloadTableExist(
}

func (f *FileToTableSystem) createFileAndPayloadTables(
ctx context.Context, txn *kv.Txn, ie sqlutil.InternalExecutor,
ctx context.Context, ie sqlutil.InternalExecutor,
) error {
// Create the File and Payload tables to hold the file chunks.
fileTableCreateQuery := fmt.Sprintf(fileTableSchema, f.GetFQFileTableName())
_, err := ie.ExecEx(ctx, "create-file-table", txn,
_, err := ie.ExecEx(ctx, "create-file-table", nil, /* txn */
sessiondata.InternalExecutorOverride{User: f.username},
fileTableCreateQuery)
if err != nil {
return errors.Wrap(err, "failed to create file table to store uploaded file names")
}

payloadTableCreateQuery := fmt.Sprintf(payloadTableSchema, f.GetFQPayloadTableName())
_, err = ie.ExecEx(ctx, "create-payload-table", txn,
_, err = ie.ExecEx(ctx, "create-payload-table", nil, /* txn */
sessiondata.InternalExecutorOverride{User: f.username},
payloadTableCreateQuery)
if err != nil {
return errors.Wrap(err, "failed to create table to store chunks of uploaded files")
}

addFKQuery := fmt.Sprintf(`ALTER TABLE %s ADD CONSTRAINT file_id_fk FOREIGN KEY (
addFKQuery := fmt.Sprintf(`ALTER TABLE %s ADD CONSTRAINT IF NOT EXISTS file_id_fk FOREIGN KEY (
file_id) REFERENCES %s (file_id)`, f.GetFQPayloadTableName(), f.GetFQFileTableName())
_, err = ie.ExecEx(ctx, "create-payload-table", txn,
_, err = ie.ExecEx(ctx, "create-payload-table", nil, /* txn */
sessiondata.InternalExecutorOverride{User: f.username},
addFKQuery)
if err != nil {
Expand All @@ -855,11 +855,11 @@ file_id) REFERENCES %s (file_id)`, f.GetFQPayloadTableName(), f.GetFQFileTableNa
// Grant the current user all read/edit privileges for the file and payload
// tables.
func (f *FileToTableSystem) grantCurrentUserTablePrivileges(
ctx context.Context, txn *kv.Txn, ie sqlutil.InternalExecutor,
ctx context.Context, ie sqlutil.InternalExecutor,
) error {
grantQuery := fmt.Sprintf(`GRANT SELECT, INSERT, DROP, DELETE ON TABLE %s, %s TO %s`,
f.GetFQFileTableName(), f.GetFQPayloadTableName(), f.username.SQLIdentifier())
_, err := ie.ExecEx(ctx, "grant-user-file-payload-table-access", txn,
_, err := ie.ExecEx(ctx, "grant-user-file-payload-table-access", nil, /* txn */
sessiondata.InternalExecutorOverride{User: username.RootUserName()},
grantQuery)
if err != nil {
Expand All @@ -872,12 +872,12 @@ func (f *FileToTableSystem) grantCurrentUserTablePrivileges(
// Revoke all privileges from every user and role except root/admin and the
// current user.
func (f *FileToTableSystem) revokeOtherUserTablePrivileges(
ctx context.Context, txn *kv.Txn, ie sqlutil.InternalExecutor,
ctx context.Context, ie sqlutil.InternalExecutor,
) error {
getUsersQuery := `SELECT username FROM system.
users WHERE NOT "username" = 'root' AND NOT "username" = 'admin' AND NOT "username" = $1`
it, err := ie.QueryIteratorEx(
ctx, "get-users", txn,
ctx, "get-users", nil, /* txn */
sessiondata.InternalExecutorOverride{User: username.RootUserName()},
getUsersQuery, f.username,
)
Expand All @@ -899,7 +899,7 @@ users WHERE NOT "username" = 'root' AND NOT "username" = 'admin' AND NOT "userna
for _, user := range users {
revokeQuery := fmt.Sprintf(`REVOKE ALL ON TABLE %s, %s FROM %s`,
f.GetFQFileTableName(), f.GetFQPayloadTableName(), user.SQLIdentifier())
_, err = ie.ExecEx(ctx, "revoke-user-privileges", txn,
_, err = ie.ExecEx(ctx, "revoke-user-privileges", nil, /* txn */
sessiondata.InternalExecutorOverride{User: username.RootUserName()},
revokeQuery)
if err != nil {
Expand Down
2 changes: 1 addition & 1 deletion pkg/server/api_v2_sql.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -413,7 +413,7 @@ func (a *apiV2Server) execSQL(w http.ResponseWriter, r *http.Request) {
}
}()

it, err := a.admin.ie.QueryIteratorEx(ctx, "run-query-via-api", txn,
it, err := a.admin.ie.QueryIteratorEx(ctx, "run-query-via-api", nil, /* txn */
sessiondata.InternalExecutorOverride{
User: username,
Database: requestPayload.Database,
Expand Down
19 changes: 9 additions & 10 deletions pkg/sql/authorization_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ func TestCheckAnyPrivilegeForNodeUser(t *testing.T) {
defer log.Scope(t).Close(t)

ctx := context.Background()
s, _, kv := serverutils.StartServer(t, base.TestServerArgs{})
s, _, _ := serverutils.StartServer(t, base.TestServerArgs{})

defer s.Stopper().Stop(ctx)

Expand All @@ -40,40 +40,39 @@ func TestCheckAnyPrivilegeForNodeUser(t *testing.T) {

ie := ts.InternalExecutor().(sqlutil.InternalExecutor)

txn := kv.NewTxn(ctx, "get-all-databases")
row, err := ie.QueryRowEx(
ctx, "get-all-databases", txn, sessiondata.NodeUserSessionDataOverride,
ctx, "get-all-databases", nil /* txn */, sessiondata.NodeUserSessionDataOverride,
"SELECT count(1) FROM crdb_internal.databases",
)
require.NoError(t, err)
// 3 databases (system, defaultdb, postgres).
require.Equal(t, row.String(), "(3)")

_, err = ie.ExecEx(ctx, "create-database1", txn, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
_, err = ie.ExecEx(ctx, "create-database1", nil /* txn */, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
"CREATE DATABASE test1")
require.NoError(t, err)

_, err = ie.ExecEx(ctx, "create-database2", txn, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
_, err = ie.ExecEx(ctx, "create-database2", nil /* txn */, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
"CREATE DATABASE test2")
require.NoError(t, err)

// Revoke CONNECT on all non-system databases and ensure that when querying
// with node, we can still see all the databases.
_, err = ie.ExecEx(ctx, "revoke-privileges", txn, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
_, err = ie.ExecEx(ctx, "revoke-privileges", nil /* txn */, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
"REVOKE CONNECT ON DATABASE test1 FROM public")
require.NoError(t, err)
_, err = ie.ExecEx(ctx, "revoke-privileges", txn, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
_, err = ie.ExecEx(ctx, "revoke-privileges", nil /* txn */, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
"REVOKE CONNECT ON DATABASE test2 FROM public")
require.NoError(t, err)
_, err = ie.ExecEx(ctx, "revoke-privileges", txn, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
_, err = ie.ExecEx(ctx, "revoke-privileges", nil /* txn */, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
"REVOKE CONNECT ON DATABASE defaultdb FROM public")
require.NoError(t, err)
_, err = ie.ExecEx(ctx, "revoke-privileges", txn, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
_, err = ie.ExecEx(ctx, "revoke-privileges", nil /* txn */, sessiondata.InternalExecutorOverride{User: username.RootUserName()},
"REVOKE CONNECT ON DATABASE postgres FROM public")
require.NoError(t, err)

row, err = ie.QueryRowEx(
ctx, "get-all-databases", txn, sessiondata.NodeUserSessionDataOverride,
ctx, "get-all-databases", nil /* txn */, sessiondata.NodeUserSessionDataOverride,
"SELECT count(1) FROM crdb_internal.databases",
)
require.NoError(t, err)
Expand Down
3 changes: 3 additions & 0 deletions pkg/sql/internal.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -871,6 +871,9 @@ func (ie *InternalExecutor) execInternal(
timeReceived := timeutil.Now()
parseStart := timeReceived
parsed, err := parser.ParseOne(stmt)
if tree.CanModifySchema(parsed.AST) && txn != nil {
return nil, errors.New("DDL statement is disallowed with internal executor")
}
if err != nil {
return nil, err
}
Expand Down
4 changes: 2 additions & 2 deletions pkg/sql/temporary_schema.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -345,7 +345,7 @@ func cleanupSchemaObjects(
_, err = ie.ExecEx(
ctx,
"delete-temp-dependent-col",
txn,
nil, /* txn */
override,
fmt.Sprintf(
"ALTER TABLE %s ALTER COLUMN %s DROP DEFAULT",
Expand Down Expand Up @@ -385,7 +385,7 @@ func cleanupSchemaObjects(
query.WriteString(tbName.FQString())
}
query.WriteString(" CASCADE")
_, err = ie.ExecEx(ctx, "delete-temp-"+toDelete.typeName, txn, override, query.String())
_, err = ie.ExecEx(ctx, "delete-temp-"+toDelete.typeName, nil /* txn */, override, query.String())
if err != nil {
return err
}
Expand Down

0 comments on commit b8caec5

Please sign in to comment.