Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: fix missing rewrite of table fragments during migration to sql backend #19813

Open
wants to merge 4 commits into
base: release-2.0
Choose a base branch
from
+94 −18
Open

fix: fix missing rewrite of table fragments during migration to sql backend #19813

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
59e8019
fix: fix missing rewrite of table fragments during migration to sql b…
yezizp2012 Dec 16, 2024
6b3785a
clippy
yezizp2012 Dec 17, 2024
6f29822
fix schema id not found
yezizp2012 Dec 17, 2024
daba688
fmt
yezizp2012 Dec 17, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
112 changes: 94 additions & 18 deletions src/ctl/src/cmd_impl/meta/migration.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
// See the License for the specific language governing permissions and
// limitations under the License.

use std::collections::{BTreeSet, HashMap, HashSet};
use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
use std::time::Duration;

use anyhow::Context;
Expand Down Expand Up @@ -59,14 +59,17 @@ use risingwave_meta_model_v2::{
use risingwave_pb::catalog::table::PbTableType;
use risingwave_pb::catalog::{
PbConnection, PbDatabase, PbFunction, PbIndex, PbSchema, PbSecret, PbSink, PbSource,
PbSubscription, PbTable, PbView,
PbStreamSourceInfo, PbSubscription, PbTable, PbView,
};
use risingwave_pb::common::WorkerType;
use risingwave_pb::hummock::{
CompactTaskAssignment, HummockPinnedSnapshot, HummockPinnedVersion, HummockVersionStats,
};
use risingwave_pb::meta::table_fragments::State;
use risingwave_pb::meta::PbSystemParams;
use risingwave_pb::secret::PbSecretRef;
use risingwave_pb::stream_plan::stream_node::PbNodeBody;
use risingwave_pb::stream_plan::PbStreamNode;
use risingwave_pb::user::grant_privilege::PbObject as GrantObject;
use risingwave_pb::user::PbUserInfo;
use sea_orm::ActiveValue::Set;
Expand Down Expand Up @@ -198,15 +201,15 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
.chain(subscriptions.iter().map(|s| s.id))
.collect::<BTreeSet<_>>();

// Helper function to get next available id.
let mut next_available_id = || -> u32 {
let id = inuse_obj_ids
// Helper function to get the next available id.
let next_available_id = |objs: &mut BTreeSet<u32>| -> u32 {
let id = objs
.iter()
.enumerate()
.find(|(i, id)| i + 1 != **id as usize)
.map(|(i, _)| i + 1)
.unwrap_or(inuse_obj_ids.len() + 1) as u32;
inuse_obj_ids.insert(id);
.unwrap_or(objs.len() + 1) as u32;
objs.insert(id);
id
};

Expand All @@ -233,7 +236,7 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
// database
let mut db_rewrite = HashMap::new();
for mut db in databases {
let id = next_available_id();
let id = next_available_id(&mut inuse_obj_ids);
db_rewrite.insert(db.id, id);
db.id = id as _;

Expand All @@ -253,7 +256,7 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
// schema
let mut schema_rewrite = HashMap::new();
for mut schema in schemas {
let id = next_available_id();
let id = next_available_id(&mut inuse_obj_ids);
schema_rewrite.insert(schema.id, id);
schema.id = id as _;

Expand All @@ -274,7 +277,7 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
// function
let mut function_rewrite = HashMap::new();
for mut function in functions {
let id = next_available_id();
let id = next_available_id(&mut inuse_obj_ids);
function_rewrite.insert(function.id, id);
function.id = id as _;

Expand All @@ -296,7 +299,7 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
// connection mapping
let mut connection_rewrite = HashMap::new();
for mut connection in connections {
let id = next_available_id();
let id = next_available_id(&mut inuse_obj_ids);
connection_rewrite.insert(connection.id, id);
connection.id = id as _;

Expand All @@ -320,7 +323,7 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
// secret mapping
let mut secret_rewrite = HashMap::new();
for mut secret in secrets {
let id = next_available_id();
let id = next_available_id(&mut inuse_obj_ids);
secret_rewrite.insert(secret.id, id);
secret.id = id as _;

Expand Down Expand Up @@ -466,8 +469,13 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
let mut stream_node = fragment.stream_node.to_protobuf();
visit_stream_node_tables(&mut stream_node, |table, _| {
table.database_id = *db_rewrite.get(&table.database_id).unwrap();
table.schema_id = *schema_rewrite.get(&table.schema_id).unwrap();
// Note: The schema could be altered and not updated here.
// It's safe to leave it as 0 if not found.
table.schema_id = *schema_rewrite.get(&table.schema_id).unwrap_or(&0);
});
// rewrite secret ids.
visit_secret_ref_mut(&mut stream_node, &secret_rewrite);

let mut fragment = fragment.into_active_model();
fragment.stream_node = Set((&stream_node).into());
Fragment::insert(fragment)
Expand Down Expand Up @@ -561,7 +569,12 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
if let Some(info) = &mut src.info {
for secret_ref in info.format_encode_secret_refs.values_mut() {
secret_ref.secret_id = *secret_rewrite.get(&secret_ref.secret_id).unwrap();
dependent_secret_ids.insert(secret_ref.secret_id);
}
if let Some(table) = &mut info.external_table {
for secret_ref in table.secret_refs.values_mut() {
secret_ref.secret_id =
*secret_rewrite.get(&secret_ref.secret_id).unwrap();
}
}
}
object_dependencies.extend(dependent_secret_ids.into_iter().map(|secret_id| {
Expand Down Expand Up @@ -707,16 +720,24 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
for user in users {
for gp in user.grant_privileges {
let id = match gp.get_object()? {
GrantObject::DatabaseId(id) => *db_rewrite.get(id).unwrap(),
GrantObject::SchemaId(id) => *schema_rewrite.get(id).unwrap(),
GrantObject::FunctionId(id) => *function_rewrite.get(id).unwrap(),
GrantObject::DatabaseId(id) => *db_rewrite.get(id).unwrap_or(&0),
GrantObject::SchemaId(id) => *schema_rewrite.get(id).unwrap_or(&0),
GrantObject::FunctionId(id) => *function_rewrite.get(id).unwrap_or(&0),
GrantObject::TableId(id)
| GrantObject::SourceId(id)
| GrantObject::SinkId(id)
| GrantObject::ViewId(id)
| GrantObject::SubscriptionId(id) => *id,
ty => unreachable!("invalid object type: {:?}", ty),
};
if id == 0 {
tracing::warn!("grant object not found");
continue;
} else if !inuse_obj_ids.contains(&id) {
tracing::warn!("grant object id {} not found", id);
continue;
}

for action_with_opt in &gp.action_with_opts {
privileges.push(user_privilege::ActiveModel {
user_id: Set(user.id as _),
Expand Down Expand Up @@ -923,7 +944,7 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
format!("ALTER TABLE worker AUTO_INCREMENT = {next_worker_id};"),
))
.await?;
let next_object_id = next_available_id();
let next_object_id = next_available_id(&mut inuse_obj_ids);
meta_store_sql
.conn
.execute(Statement::from_string(
Expand Down Expand Up @@ -968,6 +989,61 @@ pub async fn migrate(from: EtcdBackend, target: String, force_clean: bool) -> an
Ok(())
}

fn visit_secret_ref_mut(stream_node: &mut PbStreamNode, secret_rewrite: &HashMap<u32, u32>) {
let visit_map_secret_refs = |refs: &mut BTreeMap<String, PbSecretRef>| {
for secret_ref in refs.values_mut() {
secret_ref.secret_id = *secret_rewrite.get(&secret_ref.secret_id).unwrap();
}
};
let visit_info = |info: &mut PbStreamSourceInfo| {
for secret_ref in info.format_encode_secret_refs.values_mut() {
secret_ref.secret_id = *secret_rewrite.get(&secret_ref.secret_id).unwrap();
}
if let Some(table) = &mut info.external_table {
for secret_ref in table.secret_refs.values_mut() {
secret_ref.secret_id = *secret_rewrite.get(&secret_ref.secret_id).unwrap();
}
}
};

let visit_body = |body: &mut PbNodeBody| match body {
PbNodeBody::Source(node) => {
if let Some(inner) = &mut node.source_inner {
visit_map_secret_refs(&mut inner.secret_refs);
inner.info.as_mut().map(visit_info);
}
}
PbNodeBody::Sink(node) => {
if let Some(desc) = &mut node.sink_desc {
visit_map_secret_refs(&mut desc.secret_refs);
if let Some(desc) = &mut desc.format_desc {
visit_map_secret_refs(&mut desc.secret_refs)
}
}
}
PbNodeBody::StreamFsFetch(node) => {
if let Some(inner) = &mut node.node_inner {
visit_map_secret_refs(&mut inner.secret_refs);
inner.info.as_mut().map(visit_info);
}
}
PbNodeBody::StreamCdcScan(node) => {
if let Some(desc) = &mut node.cdc_table_desc {
visit_map_secret_refs(&mut desc.secret_refs)
}
}
PbNodeBody::SourceBackfill(node) => {
visit_map_secret_refs(&mut node.secret_refs);
node.info.as_mut().map(visit_info);
}
_ => {}
};
visit_body(stream_node.node_body.as_mut().unwrap());
for input in &mut stream_node.input {
visit_secret_ref_mut(input, secret_rewrite);
}
}

async fn load_current_id(meta_store: &MetaStoreRef, category: &str, start: Option<u64>) -> u64 {
let category_gen_key = format!("{}_id_next_generator", category);
let res = meta_store
Expand Down
Loading