Skip to content

Commit c6b8621

Browse files
TCeasonTCeason
authored
fix(rbac): create or replace ownership_object should delete the old ownership key (#18667)
1 parent c1e8b01 commit c6b8621

File tree

15 files changed

+275
-44
lines changed

15 files changed

+275
-44
lines changed

‎src/meta/api/src/database_api.rs‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,8 @@ where
108108
}
109109

110110
let mut trials = txn_backoff(None, func_name!());
111+
let mut old_db_id = None;
112+
111113
loop {
112114
trials.next().unwrap()?.await;
113115

@@ -118,6 +120,7 @@ where
118120
let mut txn = TxnRequest::default();
119121

120122
if let Some(ref curr_seq_db_id) = curr_seq_db_id {
123+
old_db_id = Some(curr_seq_db_id.data.into_inner());
121124
match req.create_option {
122125
CreateOption::Create => {
123126
return Err(KVAppError::AppError(AppError::DatabaseAlreadyExists(
@@ -130,6 +133,7 @@ where
130133
CreateOption::CreateIfNotExists => {
131134
return Ok(CreateDatabaseReply {
132135
db_id: curr_seq_db_id.data.into_inner(),
136+
old_db_id: None,
133137
});
134138
}
135139
CreateOption::CreateOrReplace => {
@@ -187,7 +191,10 @@ where
187191
);
188192

189193
if succ {
190-
return Ok(CreateDatabaseReply { db_id: id_key });
194+
return Ok(CreateDatabaseReply {
195+
db_id: id_key,
196+
old_db_id,
197+
});
191198
}
192199
}
193200
}

‎src/meta/api/src/schema_api_test_suite.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4841,6 +4841,7 @@ impl SchemaApiTestSuite {
48414841
new_table: true,
48424842
orphan_table_name: None,
48434843
spec_vec: None,
4844+
old_table_id: None,
48444845
};
48454846
let cur_db = util.get_database().await?;
48464847
assert!(old_db.meta.seq < cur_db.meta.seq);

‎src/meta/api/src/table_api.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -272,6 +272,8 @@ where
272272
}
273273

274274
let mut trials = txn_backoff(None, func_name!());
275+
let mut old_table_id = None;
276+
275277
loop {
276278
trials.next().unwrap()?.await;
277279

@@ -322,6 +324,7 @@ where
322324
spec_vec: None,
323325
prev_table_id: None,
324326
orphan_table_name: None,
327+
old_table_id: None,
325328
});
326329
}
327330
CreateOption::CreateOrReplace => {
@@ -333,6 +336,7 @@ where
333336

334337
SeqV::new(id.seq, *id.data)
335338
} else {
339+
old_table_id = Some(*id.data);
336340
let (seq, id) = construct_drop_table_txn_operations(
337341
self,
338342
req.name_ident.table_name.clone(),
@@ -470,6 +474,7 @@ where
470474
spec_vec: None,
471475
prev_table_id,
472476
orphan_table_name,
477+
old_table_id,
473478
});
474479
} else {
475480
// re-run txn with re-fetched data

‎src/meta/app/src/schema/database.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ impl Display for CreateDatabaseReq {
217217
#[derive(Clone, Debug, Eq, PartialEq)]
218218
pub struct CreateDatabaseReply {
219219
pub db_id: DatabaseId,
220+
pub old_db_id: Option<DatabaseId>,
220221
}
221222

222223
#[derive(Clone, Debug, PartialEq, Eq)]

‎src/meta/app/src/schema/table.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -659,6 +659,7 @@ pub struct CreateTableReply {
659659
pub spec_vec: Option<(u64, u64)>,
660660
pub prev_table_id: Option<u64>,
661661
pub orphan_table_name: Option<String>,
662+
pub old_table_id: Option<u64>,
662663
}
663664

664665
/// Drop table by id.

‎src/query/service/src/catalogs/default/mutable_catalog.rs‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -333,7 +333,10 @@ impl Catalog for MutableCatalog {
333333
});
334334
let database = self.build_db_instance(&db_info)?;
335335
database.init_database(req.name_ident.tenant_name()).await?;
336-
Ok(CreateDatabaseReply { db_id: res.db_id })
336+
Ok(CreateDatabaseReply {
337+
db_id: res.db_id,
338+
old_db_id: res.old_db_id,
339+
})
337340
}
338341

339342
#[async_backtrace::framed]

‎src/query/service/src/interpreters/interpreter_database_create.rs‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,16 @@ impl Interpreter for CreateDatabaseInterpreter {
8888
&current_role.name,
8989
)
9090
.await?;
91+
92+
// if old_db_id is some means create or replace is success, we should delete the old db id's ownership key
93+
if let Some(old_db_id) = reply.old_db_id {
94+
role_api
95+
.revoke_ownership(&OwnershipObject::Database {
96+
catalog_name: self.plan.catalog.clone(),
97+
db_id: *old_db_id,
98+
})
99+
.await?;
100+
}
91101
RoleCacheManager::instance().invalidate_cache(&tenant);
92102
}
93103
}

‎src/query/service/src/interpreters/interpreter_table_create.rs‎

Lines changed: 42 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ use databend_common_management::RoleApi;
3333
use databend_common_meta_app::principal::OwnershipObject;
3434
use databend_common_meta_app::schema::CommitTableMetaReq;
3535
use databend_common_meta_app::schema::CreateOption;
36+
use databend_common_meta_app::schema::CreateTableReply;
3637
use databend_common_meta_app::schema::CreateTableReq;
3738
use databend_common_meta_app::schema::TableIdent;
3839
use databend_common_meta_app::schema::TableIndexType;
@@ -41,6 +42,7 @@ use databend_common_meta_app::schema::TableMeta;
4142
use databend_common_meta_app::schema::TableNameIdent;
4243
use databend_common_meta_app::schema::TablePartition;
4344
use databend_common_meta_app::schema::TableStatistics;
45+
use databend_common_meta_app::tenant::Tenant;
4446
use databend_common_meta_types::MatchSeq;
4547
use databend_common_pipeline_core::always_callback;
4648
use databend_common_pipeline_core::ExecutionInfo;
@@ -228,22 +230,7 @@ impl CreateTableInterpreter {
228230
let db_id = reply.db_id;
229231

230232
if !req.table_meta.options.contains_key(OPT_KEY_TEMP_PREFIX) {
231-
// grant the ownership of the table to the current role.
232-
let current_role = self.ctx.get_current_role();
233-
if let Some(current_role) = current_role {
234-
let role_api = UserApiProvider::instance().role_api(&tenant);
235-
role_api
236-
.grant_ownership(
237-
&OwnershipObject::Table {
238-
catalog_name: self.plan.catalog.clone(),
239-
db_id,
240-
table_id,
241-
},
242-
&current_role.name,
243-
)
244-
.await?;
245-
RoleCacheManager::instance().invalidate_cache(&tenant);
246-
}
233+
self.process_ownership(&tenant, reply).await?;
247234
}
248235

249236
// If the table creation query contains column definitions, like 'CREATE TABLE t1(a int) AS SELECT * from t2',
@@ -336,6 +323,42 @@ impl CreateTableInterpreter {
336323
Ok(pipeline)
337324
}
338325

326+
async fn process_ownership(&self, tenant: &Tenant, reply: CreateTableReply) -> Result<()> {
327+
// grant the ownership of the table to the current role.
328+
let mut invalid_cache = false;
329+
let current_role = self.ctx.get_current_role();
330+
let role_api = UserApiProvider::instance().role_api(tenant);
331+
if let Some(current_role) = current_role {
332+
role_api
333+
.grant_ownership(
334+
&OwnershipObject::Table {
335+
catalog_name: self.plan.catalog.clone(),
336+
db_id: reply.db_id,
337+
table_id: reply.table_id,
338+
},
339+
&current_role.name,
340+
)
341+
.await?;
342+
invalid_cache = true;
343+
}
344+
345+
// if old_table_id is some means create or replace is success, we should delete the old table id's ownership key
346+
if let Some(old_table_id) = reply.old_table_id {
347+
role_api
348+
.revoke_ownership(&OwnershipObject::Table {
349+
catalog_name: self.plan.catalog.clone(),
350+
db_id: reply.db_id,
351+
table_id: old_table_id,
352+
})
353+
.await?;
354+
invalid_cache = true;
355+
}
356+
if invalid_cache {
357+
RoleCacheManager::instance().invalidate_cache(tenant);
358+
}
359+
Ok(())
360+
}
361+
339362
#[async_backtrace::framed]
340363
async fn create_table(&self) -> Result<PipelineBuildResult> {
341364
let catalog = self.ctx.get_catalog(self.plan.catalog.as_str()).await?;
@@ -390,27 +413,10 @@ impl CreateTableInterpreter {
390413
self.register_temp_table(prefix).await?;
391414
}
392415

416+
// iceberg table do not need to generate ownership.
393417
if !req.table_meta.options.contains_key(OPT_KEY_TEMP_PREFIX) && !catalog.is_external() {
394-
// iceberg table do not need to generate ownership.
395-
// grant the ownership of the table to the current role, the above req.table_meta.owner could be removed in the future.
396-
if let Some(current_role) = self.ctx.get_current_role() {
397-
let tenant = self.ctx.get_tenant();
398-
let db = catalog.get_database(&tenant, &self.plan.database).await?;
399-
let db_id = db.get_db_info().database_id.db_id;
400-
401-
let role_api = UserApiProvider::instance().role_api(&tenant);
402-
role_api
403-
.grant_ownership(
404-
&OwnershipObject::Table {
405-
catalog_name: self.plan.catalog.clone(),
406-
db_id,
407-
table_id: reply.table_id,
408-
},
409-
&current_role.name,
410-
)
411-
.await?;
412-
RoleCacheManager::instance().invalidate_cache(&tenant);
413-
}
418+
let tenant = self.ctx.get_tenant();
419+
self.process_ownership(&tenant, reply).await?;
414420
}
415421

416422
Ok(PipelineBuildResult::create())

‎src/query/service/src/test_kits/fixture.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -128,8 +128,8 @@ pub trait Setup {
128128
async fn setup(&self) -> Result<InnerConfig>;
129129
}
130130

131-
struct OSSSetup {
132-
config: InnerConfig,
131+
pub struct OSSSetup {
132+
pub config: InnerConfig,
133133
}
134134

135135
#[async_trait::async_trait]

0 commit comments

Comments
 (0)