risingwave_meta_model_migration/
m20261013_000000_backfill_iceberg_writer_fragment_flag.rs

1// Copyright 2026 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use prost::Message;
16use risingwave_pb::id::FragmentId;
17use risingwave_pb::stream_plan::StreamNode;
18use risingwave_pb::stream_plan::stream_node::NodeBody;
19use sea_orm::{FromQueryResult, Statement};
20use sea_orm_migration::prelude::*;
21use thiserror_ext::AsReport;
22
23// Keep this synchronized with `FragmentTypeFlag::IcebergWithPkIndexWriter`.
24const ICEBERG_WITH_PK_INDEX_WRITER_FLAG: i32 = 1 << 18;
25
26#[derive(DeriveMigrationName)]
27pub struct Migration;
28
29#[async_trait::async_trait]
30impl MigrationTrait for Migration {
31    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
32        let connection = manager.get_connection();
33        let database_backend = connection.get_database_backend();
34        let (sql, values) = Query::select()
35            .columns([
36                Fragment::FragmentId,
37                Fragment::FragmentTypeMask,
38                Fragment::StreamNode,
39            ])
40            .from(Fragment::Table)
41            .build_any(&*database_backend.get_query_builder());
42
43        for row in connection
44            .query_all(Statement::from_sql_and_values(
45                database_backend,
46                sql,
47                values,
48            ))
49            .await?
50        {
51            let fragment = FragmentEntity::from_query_result(&row, "")?;
52            if fragment.fragment_type_mask & ICEBERG_WITH_PK_INDEX_WRITER_FLAG != 0 {
53                continue;
54            }
55            let stream_node =
56                StreamNode::decode(fragment.stream_node.as_slice()).map_err(|err| {
57                    DbErr::Custom(format!(
58                        "failed to decode stream node for fragment {}: {}",
59                        fragment.fragment_id,
60                        err.as_report()
61                    ))
62                })?;
63            if contains_iceberg_writer(&stream_node) {
64                manager
65                    .exec_stmt(
66                        Query::update()
67                            .table(Fragment::Table)
68                            .value(
69                                Fragment::FragmentTypeMask,
70                                fragment.fragment_type_mask | ICEBERG_WITH_PK_INDEX_WRITER_FLAG,
71                            )
72                            .and_where(Expr::col(Fragment::FragmentId).eq(fragment.fragment_id))
73                            .to_owned(),
74                    )
75                    .await?;
76            }
77        }
78        Ok(())
79    }
80
81    async fn down(&self, _manager: &SchemaManager) -> Result<(), DbErr> {
82        Err(DbErr::Custom(
83            "cannot rollback Iceberg writer fragment flag migration".to_owned(),
84        ))
85    }
86}
87
88fn contains_iceberg_writer(node: &StreamNode) -> bool {
89    matches!(node.node_body, Some(NodeBody::IcebergWithPkIndexWriter(_)))
90        || node.input.iter().any(contains_iceberg_writer)
91}
92
93#[derive(Debug, FromQueryResult)]
94#[sea_orm(entity = "Fragment")]
95struct FragmentEntity {
96    fragment_id: FragmentId,
97    fragment_type_mask: i32,
98    stream_node: Vec<u8>,
99}
100
101#[derive(DeriveIden)]
102enum Fragment {
103    Table,
104    FragmentId,
105    FragmentTypeMask,
106    StreamNode,
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn finds_nested_iceberg_writer() {
115        let node = StreamNode {
116            input: vec![StreamNode {
117                node_body: Some(NodeBody::IcebergWithPkIndexWriter(Default::default())),
118                ..Default::default()
119            }],
120            ..Default::default()
121        };
122        assert!(contains_iceberg_writer(&node));
123        assert!(!contains_iceberg_writer(&StreamNode::default()));
124    }
125}
risingwave_meta_model_migration/m20261013_000000_backfill_iceberg_writer_fragment_flag.rs

risingwave_meta_model_migration/
m20261013_000000_backfill_iceberg_writer_fragment_flag.rs
