risingwave_meta_model_migration/
m20241025_062548_singleton_vnode_count.rs

1use sea_orm_migration::prelude::*;
2
3#[derive(DeriveMigrationName)]
4pub struct Migration;
5
6#[async_trait::async_trait]
7impl MigrationTrait for Migration {
8    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
9        // Hack for unifying the test for different backends.
10        let json_is_empty_array = |col| {
11            Expr::col(col)
12                .cast_as(Alias::new("char(2)"))
13                .eq(Expr::value("[]"))
14        };
15
16        // Fill vnode count with 1 for singleton tables.
17        manager
18            .exec_stmt(
19                UpdateStatement::new()
20                    .table(Table::Table)
21                    .values([(Table::VnodeCount, Expr::value(1))])
22                    .and_where(json_is_empty_array(Table::DistributionKey))
23                    .and_where(json_is_empty_array(Table::DistKeyInPk))
24                    .and_where(Expr::col(Table::VnodeColIndex).is_null())
25                    .to_owned(),
26            )
27            .await?;
28
29        // Fill vnode count with 1 for singleton fragments.
30        manager
31            .exec_stmt(
32                UpdateStatement::new()
33                    .table(Fragment::Table)
34                    .values([(Fragment::VnodeCount, Expr::value(1))])
35                    .and_where(Expr::col(Fragment::DistributionType).eq(Expr::value("SINGLE")))
36                    .to_owned(),
37            )
38            .await?;
39
40        Ok(())
41    }
42
43    async fn down(&self, _manager: &SchemaManager) -> Result<(), DbErr> {
44        Err(DbErr::Migration(
45            "cannot rollback singleton vnode count migration".to_owned(),
46        ))?
47    }
48}
49
50#[derive(DeriveIden)]
51enum Fragment {
52    Table,
53    VnodeCount,
54    DistributionType,
55}
56
57#[derive(DeriveIden)]
58enum Table {
59    Table,
60    VnodeCount,
61    DistributionKey,
62    DistKeyInPk,
63    VnodeColIndex,
64}