Skip to main content

common_meta/ddl/
alter_database.rs

1// Copyright 2023 Greptime Team
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 async_trait::async_trait;
16use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu};
17use common_procedure::{
18    Context as ProcedureContext, EventRuntimeContext, EventTrigger, LockKey, Procedure, Status,
19};
20use common_telemetry::tracing::info;
21use serde::{Deserialize, Serialize};
22use snafu::{ResultExt, ensure};
23use strum::AsRefStr;
24
25use crate::cache_invalidator::Context;
26use crate::ddl::DdlContext;
27use crate::ddl::event::database::{ALTER_DATABASE_EVENT_TYPE, DatabaseDdlEvent};
28use crate::ddl::utils::map_to_procedure_error;
29use crate::error::{Result, SchemaNotFoundSnafu};
30use crate::instruction::CacheIdent;
31use crate::key::DeserializedValueWithBytes;
32use crate::key::schema_name::{SchemaName, SchemaNameKey, SchemaNameValue};
33use crate::lock_key::{CatalogLock, SchemaLock};
34use crate::rpc::ddl::UnsetDatabaseOption::{self};
35use crate::rpc::ddl::{AlterDatabaseKind, AlterDatabaseTask, EventContext, SetDatabaseOption};
36
37pub struct AlterDatabaseProcedure {
38    pub context: DdlContext,
39    pub data: AlterDatabaseData,
40}
41
42fn build_new_schema_value(
43    mut value: SchemaNameValue,
44    alter_kind: &AlterDatabaseKind,
45) -> Result<SchemaNameValue> {
46    match alter_kind {
47        AlterDatabaseKind::SetDatabaseOptions(options) => {
48            for option in options.0.iter() {
49                match option {
50                    SetDatabaseOption::Ttl(ttl) => {
51                        value.ttl = Some(*ttl);
52                    }
53                    SetDatabaseOption::Other(key, val) => {
54                        value.extra_options.insert(key.clone(), val.clone());
55                    }
56                }
57            }
58        }
59        AlterDatabaseKind::UnsetDatabaseOptions(keys) => {
60            for key in keys.0.iter() {
61                match key {
62                    UnsetDatabaseOption::Ttl => value.ttl = None,
63                    UnsetDatabaseOption::Other(key) => {
64                        value.extra_options.remove(key);
65                    }
66                }
67            }
68        }
69    }
70    Ok(value)
71}
72
73impl AlterDatabaseProcedure {
74    pub const TYPE_NAME: &'static str = "metasrv-procedure::AlterDatabase";
75
76    pub fn new(
77        task: AlterDatabaseTask,
78        event_context: EventContext,
79        context: DdlContext,
80    ) -> Result<Self> {
81        Ok(Self {
82            context,
83            data: AlterDatabaseData::new(task, event_context)?,
84        })
85    }
86
87    pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult<Self> {
88        let data = serde_json::from_str(json).context(FromJsonSnafu)?;
89
90        Ok(Self { context, data })
91    }
92
93    pub async fn on_prepare(&mut self) -> Result<Status> {
94        let value = self
95            .context
96            .table_metadata_manager
97            .schema_manager()
98            .get(SchemaNameKey::new(self.data.catalog(), self.data.schema()))
99            .await?;
100
101        ensure!(
102            value.is_some(),
103            SchemaNotFoundSnafu {
104                table_schema: self.data.schema(),
105            }
106        );
107
108        self.data.schema_value = value;
109        self.data.state = AlterDatabaseState::UpdateMetadata;
110
111        Ok(Status::executing(true))
112    }
113
114    pub async fn on_update_metadata(&mut self) -> Result<Status> {
115        let schema_name = SchemaNameKey::new(self.data.catalog(), self.data.schema());
116
117        // Safety: schema_value is not None.
118        let current_schema_value = self.data.schema_value.as_ref().unwrap();
119
120        let new_schema_value = build_new_schema_value(
121            current_schema_value.get_inner_ref().clone(),
122            &self.data.kind,
123        )?;
124
125        self.context
126            .table_metadata_manager
127            .schema_manager()
128            .update(schema_name, current_schema_value, &new_schema_value)
129            .await?;
130
131        info!("Updated database metadata for schema {schema_name}");
132        self.data.state = AlterDatabaseState::InvalidateSchemaCache;
133        Ok(Status::executing(true))
134    }
135
136    pub async fn on_invalidate_schema_cache(&mut self) -> Result<Status> {
137        let cache_invalidator = &self.context.cache_invalidator;
138        cache_invalidator
139            .invalidate(
140                &Context::default(),
141                &[CacheIdent::SchemaName(SchemaName {
142                    catalog_name: self.data.catalog().to_string(),
143                    schema_name: self.data.schema().to_string(),
144                })],
145            )
146            .await?;
147
148        Ok(Status::done())
149    }
150}
151
152#[async_trait]
153impl Procedure for AlterDatabaseProcedure {
154    fn type_name(&self) -> &str {
155        Self::TYPE_NAME
156    }
157
158    async fn execute(&mut self, _ctx: &ProcedureContext) -> ProcedureResult<Status> {
159        match self.data.state {
160            AlterDatabaseState::Prepare => self.on_prepare().await,
161            AlterDatabaseState::UpdateMetadata => self.on_update_metadata().await,
162            AlterDatabaseState::InvalidateSchemaCache => self.on_invalidate_schema_cache().await,
163        }
164        .map_err(map_to_procedure_error)
165    }
166
167    fn dump(&self) -> ProcedureResult<String> {
168        serde_json::to_string(&self.data).context(ToJsonSnafu)
169    }
170
171    fn lock_key(&self) -> LockKey {
172        let catalog = self.data.catalog();
173        let schema = self.data.schema();
174
175        let lock_key = vec![
176            CatalogLock::Read(catalog).into(),
177            SchemaLock::write(catalog, schema).into(),
178        ];
179
180        LockKey::new(lock_key)
181    }
182
183    fn event(
184        &self,
185        ctx: &EventRuntimeContext<'_>,
186    ) -> Option<Box<dyn common_event_recorder::Event>> {
187        if !ctx.event_type_filter.allows(ALTER_DATABASE_EVENT_TYPE) {
188            return None;
189        }
190
191        let event = if matches!(&ctx.trigger, EventTrigger::Submitted) {
192            DatabaseDdlEvent::alter_submitted(
193                self.data.catalog(),
194                self.data.schema(),
195                &self.data.kind,
196                self.data.event_context.clone(),
197            )
198        } else {
199            DatabaseDdlEvent::alter_lifecycle()
200        };
201        Some(Box::new(event))
202    }
203}
204
205#[derive(Debug, Serialize, Deserialize, AsRefStr)]
206enum AlterDatabaseState {
207    Prepare,
208    UpdateMetadata,
209    InvalidateSchemaCache,
210}
211
212/// The data of alter database procedure.
213#[derive(Debug, Serialize, Deserialize)]
214pub struct AlterDatabaseData {
215    state: AlterDatabaseState,
216    kind: AlterDatabaseKind,
217    catalog_name: String,
218    schema_name: String,
219    schema_value: Option<DeserializedValueWithBytes<SchemaNameValue>>,
220    #[serde(default)]
221    event_context: EventContext,
222}
223
224impl AlterDatabaseData {
225    pub fn new(task: AlterDatabaseTask, event_context: EventContext) -> Result<Self> {
226        Ok(Self {
227            state: AlterDatabaseState::Prepare,
228            kind: AlterDatabaseKind::try_from(task.alter_expr.kind.unwrap())?,
229            catalog_name: task.alter_expr.catalog_name,
230            schema_name: task.alter_expr.schema_name,
231            schema_value: None,
232            event_context,
233        })
234    }
235
236    pub fn catalog(&self) -> &str {
237        &self.catalog_name
238    }
239
240    pub fn schema(&self) -> &str {
241        &self.schema_name
242    }
243}
244
245#[cfg(test)]
246mod tests {
247    use std::time::Duration;
248
249    use crate::ddl::alter_database::build_new_schema_value;
250    use crate::key::schema_name::SchemaNameValue;
251    use crate::rpc::ddl::{
252        AlterDatabaseKind, SetDatabaseOption, SetDatabaseOptions, UnsetDatabaseOption,
253        UnsetDatabaseOptions,
254    };
255
256    #[test]
257    fn test_build_new_schema_value() {
258        let set_ttl = AlterDatabaseKind::SetDatabaseOptions(SetDatabaseOptions(vec![
259            SetDatabaseOption::Ttl(Duration::from_secs(10).into()),
260        ]));
261        let current_schema_value = SchemaNameValue::default();
262        let new_schema_value =
263            build_new_schema_value(current_schema_value.clone(), &set_ttl).unwrap();
264        assert_eq!(new_schema_value.ttl, Some(Duration::from_secs(10).into()));
265
266        let unset_ttl_alter_kind =
267            AlterDatabaseKind::UnsetDatabaseOptions(UnsetDatabaseOptions(vec![
268                UnsetDatabaseOption::Ttl,
269            ]));
270        let new_schema_value =
271            build_new_schema_value(current_schema_value, &unset_ttl_alter_kind).unwrap();
272        assert_eq!(new_schema_value.ttl, None);
273    }
274
275    #[test]
276    fn test_build_new_schema_value_with_compaction_options() {
277        let set_compaction = AlterDatabaseKind::SetDatabaseOptions(SetDatabaseOptions(vec![
278            SetDatabaseOption::Other("compaction.type".to_string(), "twcs".to_string()),
279            SetDatabaseOption::Other("compaction.twcs.time_window".to_string(), "1d".to_string()),
280        ]));
281
282        let current_schema_value = SchemaNameValue::default();
283        let new_schema_value =
284            build_new_schema_value(current_schema_value.clone(), &set_compaction).unwrap();
285
286        assert_eq!(
287            new_schema_value.extra_options.get("compaction.type"),
288            Some(&"twcs".to_string())
289        );
290        assert_eq!(
291            new_schema_value
292                .extra_options
293                .get("compaction.twcs.time_window"),
294            Some(&"1d".to_string())
295        );
296
297        let unset_compaction = AlterDatabaseKind::UnsetDatabaseOptions(UnsetDatabaseOptions(vec![
298            UnsetDatabaseOption::Other("compaction.type".to_string()),
299        ]));
300
301        let new_schema_value = build_new_schema_value(new_schema_value, &unset_compaction).unwrap();
302
303        assert_eq!(new_schema_value.extra_options.get("compaction.type"), None);
304        assert_eq!(
305            new_schema_value
306                .extra_options
307                .get("compaction.twcs.time_window"),
308            Some(&"1d".to_string())
309        );
310    }
311}