risingwave_meta/backup_restore/restore_impl/
v2.rs

1// Copyright 2023 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 std::cmp;
16
17use itertools::Itertools;
18use risingwave_backup::MetaSnapshotId;
19use risingwave_backup::error::{BackupError, BackupResult};
20use risingwave_backup::meta_snapshot::MetaSnapshot;
21use risingwave_backup::meta_snapshot_v2::{
22    MetaSnapshotV2, MetadataV2, decode_hummock_sequences_from_stream,
23};
24use risingwave_backup::storage::{MetaSnapshotStorage, MetaSnapshotStorageRef};
25use sea_orm::{DbErr, EntityTrait};
26
27use crate::backup_restore::restore_impl::{Loader, Writer};
28use crate::controller::SqlMetaStore;
29
30pub struct LoaderV2 {
31    backup_store: MetaSnapshotStorageRef,
32}
33
34impl LoaderV2 {
35    pub fn new(backup_store: MetaSnapshotStorageRef) -> Self {
36        Self { backup_store }
37    }
38}
39
40#[async_trait::async_trait]
41impl Loader<MetadataV2> for LoaderV2 {
42    async fn load(&self, target_id: MetaSnapshotId) -> BackupResult<MetaSnapshot<MetadataV2>> {
43        let snapshot_list = &self.backup_store.manifest().await.snapshot_metadata;
44        let mut target_snapshot: MetaSnapshotV2 = self.backup_store.get(target_id).await?;
45        tracing::debug!(
46            "snapshot {} before rewrite:\n{}",
47            target_id,
48            target_snapshot
49        );
50        let newest_id = snapshot_list
51            .iter()
52            .map(|m| m.id)
53            .max()
54            .expect("should exist");
55        assert!(
56            newest_id >= target_id,
57            "newest_id={}, target_id={}",
58            newest_id,
59            target_id
60        );
61
62        // validate and rewrite seq
63        if newest_id > target_id {
64            let newest_hummock_sequences = decode_hummock_sequences_from_stream(
65                self.backup_store.get_bytes_stream(newest_id).await?,
66            )
67            .await?;
68            for seq in &target_snapshot.metadata.hummock_sequences {
69                let newest = newest_hummock_sequences
70                    .iter()
71                    .find(|s| s.name == seq.name)
72                    .unwrap_or_else(|| {
73                        panic!(
74                            "violate superset requirement. Hummock sequence name {}",
75                            seq.name
76                        )
77                    });
78                assert!(newest.seq >= seq.seq, "violate monotonicity requirement");
79            }
80            target_snapshot.metadata.hummock_sequences = newest_hummock_sequences;
81            tracing::info!(
82                "snapshot {} is rewritten by snapshot {}:\n",
83                target_id,
84                newest_id,
85            );
86            tracing::debug!("{target_snapshot}");
87        }
88        Ok(target_snapshot)
89    }
90}
91
92pub struct WriterModelV2ToMetaStoreV2 {
93    meta_store: SqlMetaStore,
94}
95
96impl WriterModelV2ToMetaStoreV2 {
97    pub fn new(meta_store: SqlMetaStore) -> Self {
98        Self { meta_store }
99    }
100}
101
102#[async_trait::async_trait]
103impl Writer<MetadataV2> for WriterModelV2ToMetaStoreV2 {
104    async fn write(&self, target_snapshot: MetaSnapshot<MetadataV2>) -> BackupResult<()> {
105        let metadata = target_snapshot.metadata;
106        let db = &self.meta_store.conn;
107        ensure_all_meta_store_tables_are_empty(db).await?;
108        insert_models(metadata.seaql_migrations.clone(), db).await?;
109        insert_models(metadata.clusters.clone(), db).await?;
110        insert_models(metadata.version_stats.clone(), db).await?;
111        insert_models(metadata.compaction_configs.clone(), db).await?;
112        insert_models(metadata.hummock_sequences.clone(), db).await?;
113        insert_models(metadata.workers.clone(), db).await?;
114        insert_models(metadata.worker_properties.clone(), db).await?;
115        insert_models(metadata.users.clone(), db).await?;
116        // The sort is required to pass table's foreign key check.
117        use risingwave_meta_model::object::ObjectType;
118        insert_models(
119            metadata
120                .objects
121                .iter()
122                .sorted_by(|a, b| match (a.obj_type, b.obj_type) {
123                    (ObjectType::Database, ObjectType::Database) => a.oid.cmp(&b.oid),
124                    (ObjectType::Database, _) => cmp::Ordering::Less,
125                    (_, ObjectType::Database) => cmp::Ordering::Greater,
126                    (ObjectType::Schema, ObjectType::Schema) => a.oid.cmp(&b.oid),
127                    (ObjectType::Schema, _) => cmp::Ordering::Less,
128                    (_, ObjectType::Schema) => cmp::Ordering::Greater,
129                    (_, _) => a.oid.cmp(&b.oid),
130                })
131                .cloned(),
132            db,
133        )
134        .await?;
135        insert_models(
136            metadata
137                .user_privileges
138                .iter()
139                .sorted_by_key(|u| u.id)
140                .cloned(),
141            db,
142        )
143        .await?;
144        insert_models(metadata.object_dependencies.clone(), db).await?;
145        insert_models(metadata.databases.clone(), db).await?;
146        insert_models(metadata.schemas.clone(), db).await?;
147        insert_models(metadata.streaming_jobs.clone(), db).await?;
148        insert_models(metadata.fragments.clone(), db).await?;
149        insert_models(metadata.fragment_relation.clone(), db).await?;
150        insert_models(metadata.connections.clone(), db).await?;
151        insert_models(metadata.sources.clone(), db).await?;
152        insert_models(metadata.tables.clone(), db).await?;
153        insert_models(metadata.sinks.clone(), db).await?;
154        insert_models(metadata.views.clone(), db).await?;
155        insert_models(metadata.indexes.clone(), db).await?;
156        insert_models(metadata.functions.clone(), db).await?;
157        insert_models(metadata.system_parameters.clone(), db).await?;
158        insert_models(metadata.catalog_versions.clone(), db).await?;
159        insert_models(metadata.subscriptions.clone(), db).await?;
160        insert_models(metadata.session_parameters.clone(), db).await?;
161        insert_models(metadata.secrets.clone(), db).await?;
162        insert_models(metadata.exactly_once_iceberg_sinks.clone(), db).await?;
163        insert_models(metadata.iceberg_tables.clone(), db).await?;
164        insert_models(metadata.iceberg_namespace_properties.clone(), db).await?;
165        insert_models(metadata.user_default_privilege.clone(), db).await?;
166        insert_models(metadata.fragment_splits.clone(), db).await?;
167        insert_models(metadata.pending_sink_state.clone(), db).await?;
168        insert_models(metadata.refresh_jobs.clone(), db).await?;
169        insert_models(metadata.cdc_table_snapshot_splits.clone(), db).await?;
170        insert_models(metadata.hummock_table_change_logs.clone(), db).await?;
171        // update_auto_inc must be called last.
172        update_auto_inc(&metadata, db).await?;
173        Ok(())
174    }
175
176    async fn overwrite(
177        &self,
178        new_storage_url: &str,
179        new_storage_dir: &str,
180        new_backup_url: &str,
181        new_backup_dir: &str,
182    ) -> BackupResult<()> {
183        let kvs = [
184            ("state_store", new_storage_url),
185            ("data_directory", new_storage_dir),
186            ("backup_storage_url", new_backup_url),
187            ("backup_storage_directory", new_backup_dir),
188        ];
189        for (k, v) in kvs {
190            let Some(model) = risingwave_meta_model::system_parameter::Entity::find_by_id(k)
191                .one(&self.meta_store.conn)
192                .await
193                .map_err(map_db_err)?
194            else {
195                return Err(BackupError::MetaStorage(
196                    anyhow::anyhow!("{k} not found in system_parameter table").into(),
197                ));
198            };
199            let mut kv: risingwave_meta_model::system_parameter::ActiveModel = model.into();
200            kv.value = sea_orm::ActiveValue::Set(v.to_owned());
201            risingwave_meta_model::system_parameter::Entity::update(kv)
202                .exec(&self.meta_store.conn)
203                .await
204                .map_err(map_db_err)?;
205        }
206        Ok(())
207    }
208}
209
210async fn ensure_all_meta_store_tables_are_empty(
211    db: &impl sea_orm::ConnectionTrait,
212) -> BackupResult<()> {
213    macro_rules! ensure_entity_empty {
214        ($($entity_mod:ident),* $(,)?) => {
215            $(
216                if risingwave_meta_model::$entity_mod::Entity::find()
217                    .one(db)
218                    .await
219                    .map_err(map_db_err)?
220                    .is_some()
221                {
222                    return Err(BackupError::NonemptyMetaStorage);
223                }
224            )*
225        };
226    }
227
228    risingwave_meta_model::for_all_meta_model_entities!(ensure_entity_empty);
229    Ok(())
230}
231
232fn map_db_err(e: DbErr) -> BackupError {
233    BackupError::MetaStorage(e.into())
234}
235
236#[macro_export]
237macro_rules! for_all_auto_increment {
238    ($metadata:ident, $db:ident, $macro:ident) => {
239        $macro! ($metadata, $db,
240            {"worker", workers, worker_id},
241            {"object", objects, oid},
242            {"user", users, user_id},
243            {"user_privilege", user_privileges, id},
244            {"fragment", fragments, fragment_id},
245            {"object_dependency", object_dependencies, id}
246        )
247    };
248}
249
250macro_rules! reset_sql_sequence {
251    ($metadata:ident, $db:ident, $( {$table:expr, $model:ident, $id_field:ident} ),*) => {
252        $(
253        match $db.get_database_backend() {
254            sea_orm::DbBackend::MySql => {
255                if let Some(v) = $metadata.$model.iter().map(|w| w.$id_field + 1).max() {
256                    $db.execute(sea_orm::Statement::from_string(
257                        sea_orm::DatabaseBackend::MySql,
258                        format!("ALTER TABLE {} AUTO_INCREMENT = {};", $table, v),
259                    ))
260                    .await
261                    .map_err(map_db_err)?;
262                }
263            }
264            sea_orm::DbBackend::Postgres => {
265                $db.execute(sea_orm::Statement::from_string(
266                    sea_orm::DatabaseBackend::Postgres,
267                    format!("SELECT setval('{}_{}_seq', (SELECT MAX({}) FROM \"{}\"));", $table, stringify!($id_field), stringify!($id_field), $table),
268                ))
269                .await
270                .map_err(map_db_err)?;
271            }
272            sea_orm::DbBackend::Sqlite => {}
273            }
274        )*
275    };
276}
277
278/// Fixes `auto_increment` fields.
279async fn update_auto_inc(
280    metadata: &MetadataV2,
281    db: &impl sea_orm::ConnectionTrait,
282) -> BackupResult<()> {
283    for_all_auto_increment!(metadata, db, reset_sql_sequence);
284    Ok(())
285}
286
287async fn insert_models<S, A>(
288    models: impl IntoIterator<Item = S>,
289    db: &impl sea_orm::ConnectionTrait,
290) -> BackupResult<()>
291where
292    S: sea_orm::ModelTrait + Sync + Send + Sized + sea_orm::IntoActiveModel<A>,
293    A: sea_orm::ActiveModelTrait + sea_orm::ActiveModelBehavior + Send + Sync + From<S>,
294    <<A as sea_orm::ActiveModelTrait>::Entity as sea_orm::EntityTrait>::Model:
295        sea_orm::IntoActiveModel<A>,
296{
297    use sea_orm::EntityTrait;
298    if <S as sea_orm::ModelTrait>::Entity::find()
299        .one(db)
300        .await
301        .map_err(map_db_err)?
302        .is_some()
303    {
304        return Err(BackupError::NonemptyMetaStorage);
305    }
306    for m in models {
307        m.into_active_model().insert(db).await.map_err(map_db_err)?;
308    }
309    Ok(())
310}
risingwave_meta/backup_restore/restore_impl/v2.rs

risingwave_meta/backup_restore/restore_impl/
v2.rs
