risingwave_meta_model_migration/
m20260624_000000_pending_sink_state_object_fk.rs

1use sea_orm_migration::prelude::*;
2
3use crate::m20230908_072257_init::Object;
4use crate::sea_orm::{ConnectionTrait, DatabaseBackend, Statement, TransactionTrait};
5use crate::utils::ColumnDefExt;
6
7#[derive(DeriveMigrationName)]
8pub struct Migration;
9
10const FK_NAME: &str = "FK_pending_sink_state_object_id";
11const NEW_TABLE: &str = "pending_sink_state_new";
12
13#[async_trait::async_trait]
14impl MigrationTrait for Migration {
15    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
16        // The model declares the `sink_id -> object(oid)` cascade but the foreign key was never
17        // created, leaking rows on DROP SINK. Add the key and clean up the leaked rows.
18        // https://github.com/risingwavelabs/risingwave/issues/26044
19        let backend = manager.get_database_backend();
20        match backend {
21            DatabaseBackend::MySql | DatabaseBackend::Postgres => {
22                manager
23                    .get_connection()
24                    .execute(Statement::from_string(
25                        backend,
26                        "DELETE FROM pending_sink_state WHERE sink_id NOT IN (SELECT oid FROM object)",
27                    ))
28                    .await?;
29                manager
30                    .alter_table(
31                        Table::alter()
32                            .table(PendingSinkState::Table)
33                            .add_foreign_key(&object_foreign_key())
34                            .to_owned(),
35                    )
36                    .await?;
37            }
38            DatabaseBackend::Sqlite => {
39                // SQLite can't `ALTER TABLE ADD FOREIGN KEY`; recreate the table with it instead.
40                recreate_table(manager, true).await?;
41            }
42        }
43        Ok(())
44    }
45
46    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
47        match manager.get_database_backend() {
48            DatabaseBackend::MySql | DatabaseBackend::Postgres => {
49                manager
50                    .alter_table(
51                        Table::alter()
52                            .table(PendingSinkState::Table)
53                            .drop_foreign_key(Alias::new(FK_NAME))
54                            .to_owned(),
55                    )
56                    .await?;
57            }
58            DatabaseBackend::Sqlite => {
59                recreate_table(manager, false).await?;
60            }
61        }
62        Ok(())
63    }
64}
65
66fn object_foreign_key() -> TableForeignKey {
67    TableForeignKey::new()
68        .name(FK_NAME)
69        .from_tbl(PendingSinkState::Table)
70        .from_col(PendingSinkState::SinkId)
71        .to_tbl(Object::Table)
72        .to_col(Object::Oid)
73        .on_delete(ForeignKeyAction::Cascade)
74        .to_owned()
75}
76
77/// Recreate `pending_sink_state` with (or without) the object foreign key, dropping leaked rows
78/// during the copy. The swap runs in one transaction so it stays atomic and re-runnable even
79/// though the migrator opens none for SQLite. SQLite-only: it can't alter foreign keys in place.
80async fn recreate_table(manager: &SchemaManager<'_>, with_foreign_key: bool) -> Result<(), DbErr> {
81    let backend = manager.get_database_backend();
82    // TODO(#26046): on sea-orm 2.0, drop this manual transaction and override
83    // `MigrationTrait::use_transaction()` to return `Some(true)` instead.
84    let txn = manager.get_connection().begin().await?;
85    {
86        let txn_manager = SchemaManager::new(&txn);
87
88        let mut create = Table::create();
89        create
90            .table(Alias::new(NEW_TABLE))
91            .col(
92                ColumnDef::new(PendingSinkState::SinkId)
93                    .integer()
94                    .not_null(),
95            )
96            .col(
97                ColumnDef::new(PendingSinkState::Epoch)
98                    .big_integer()
99                    .not_null(),
100            )
101            .col(
102                ColumnDef::new(PendingSinkState::SinkState)
103                    .string()
104                    .not_null(),
105            )
106            .col(
107                ColumnDef::new(PendingSinkState::Metadata)
108                    .rw_binary(&txn_manager)
109                    .null(),
110            )
111            .col(
112                ColumnDef::new(PendingSinkState::SchemaChange)
113                    .rw_binary(&txn_manager)
114                    .null(),
115            )
116            .primary_key(
117                Index::create()
118                    .col(PendingSinkState::SinkId)
119                    .col(PendingSinkState::Epoch),
120            );
121        if with_foreign_key {
122            create.foreign_key(
123                ForeignKey::create()
124                    .name(FK_NAME)
125                    .from_tbl(Alias::new(NEW_TABLE))
126                    .from_col(PendingSinkState::SinkId)
127                    .to_tbl(Object::Table)
128                    .to_col(Object::Oid)
129                    .on_delete(ForeignKeyAction::Cascade),
130            );
131        }
132        txn_manager.create_table(create).await?;
133
134        txn.execute(Statement::from_string(
135            backend,
136            "INSERT INTO pending_sink_state_new \
137             (sink_id, epoch, sink_state, metadata, schema_change) \
138             SELECT sink_id, epoch, sink_state, metadata, schema_change \
139             FROM pending_sink_state WHERE sink_id IN (SELECT oid FROM object)",
140        ))
141        .await?;
142
143        txn_manager
144            .drop_table(Table::drop().table(PendingSinkState::Table).to_owned())
145            .await?;
146        txn_manager
147            .rename_table(
148                Table::rename()
149                    .table(Alias::new(NEW_TABLE), PendingSinkState::Table)
150                    .to_owned(),
151            )
152            .await?;
153    }
154    txn.commit().await?;
155    Ok(())
156}
157
158#[derive(DeriveIden)]
159enum PendingSinkState {
160    Table,
161    SinkId,
162    Epoch,
163    SinkState,
164    Metadata,
165    SchemaChange,
166}
risingwave_meta_model_migration/m20260624_000000_pending_sink_state_object_fk.rs

risingwave_meta_model_migration/
m20260624_000000_pending_sink_state_object_fk.rs
