1use async_trait::async_trait;
16use common_procedure::error::{FromJsonSnafu, Result as ProcedureResult, ToJsonSnafu};
17use common_procedure::{
18 Context as ProcedureContext, EventContext, 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, 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(task: AlterDatabaseTask, context: DdlContext) -> Result<Self> {
77 Ok(Self {
78 context,
79 data: AlterDatabaseData::new(task)?,
80 })
81 }
82
83 pub fn from_json(json: &str, context: DdlContext) -> ProcedureResult<Self> {
84 let data = serde_json::from_str(json).context(FromJsonSnafu)?;
85
86 Ok(Self { context, data })
87 }
88
89 pub async fn on_prepare(&mut self) -> Result<Status> {
90 let value = self
91 .context
92 .table_metadata_manager
93 .schema_manager()
94 .get(SchemaNameKey::new(self.data.catalog(), self.data.schema()))
95 .await?;
96
97 ensure!(
98 value.is_some(),
99 SchemaNotFoundSnafu {
100 table_schema: self.data.schema(),
101 }
102 );
103
104 self.data.schema_value = value;
105 self.data.state = AlterDatabaseState::UpdateMetadata;
106
107 Ok(Status::executing(true))
108 }
109
110 pub async fn on_update_metadata(&mut self) -> Result<Status> {
111 let schema_name = SchemaNameKey::new(self.data.catalog(), self.data.schema());
112
113 let current_schema_value = self.data.schema_value.as_ref().unwrap();
115
116 let new_schema_value = build_new_schema_value(
117 current_schema_value.get_inner_ref().clone(),
118 &self.data.kind,
119 )?;
120
121 self.context
122 .table_metadata_manager
123 .schema_manager()
124 .update(schema_name, current_schema_value, &new_schema_value)
125 .await?;
126
127 info!("Updated database metadata for schema {schema_name}");
128 self.data.state = AlterDatabaseState::InvalidateSchemaCache;
129 Ok(Status::executing(true))
130 }
131
132 pub async fn on_invalidate_schema_cache(&mut self) -> Result<Status> {
133 let cache_invalidator = &self.context.cache_invalidator;
134 cache_invalidator
135 .invalidate(
136 &Context::default(),
137 &[CacheIdent::SchemaName(SchemaName {
138 catalog_name: self.data.catalog().to_string(),
139 schema_name: self.data.schema().to_string(),
140 })],
141 )
142 .await?;
143
144 Ok(Status::done())
145 }
146}
147
148#[async_trait]
149impl Procedure for AlterDatabaseProcedure {
150 fn type_name(&self) -> &str {
151 Self::TYPE_NAME
152 }
153
154 async fn execute(&mut self, _ctx: &ProcedureContext) -> ProcedureResult<Status> {
155 match self.data.state {
156 AlterDatabaseState::Prepare => self.on_prepare().await,
157 AlterDatabaseState::UpdateMetadata => self.on_update_metadata().await,
158 AlterDatabaseState::InvalidateSchemaCache => self.on_invalidate_schema_cache().await,
159 }
160 .map_err(map_to_procedure_error)
161 }
162
163 fn dump(&self) -> ProcedureResult<String> {
164 serde_json::to_string(&self.data).context(ToJsonSnafu)
165 }
166
167 fn lock_key(&self) -> LockKey {
168 let catalog = self.data.catalog();
169 let schema = self.data.schema();
170
171 let lock_key = vec![
172 CatalogLock::Read(catalog).into(),
173 SchemaLock::write(catalog, schema).into(),
174 ];
175
176 LockKey::new(lock_key)
177 }
178
179 fn event(&self, ctx: &EventContext<'_>) -> Option<Box<dyn common_event_recorder::Event>> {
180 if !ctx.event_type_filter.allows(ALTER_DATABASE_EVENT_TYPE) {
181 return None;
182 }
183
184 let event = if matches!(&ctx.trigger, EventTrigger::Submitted) {
185 DatabaseDdlEvent::alter_submitted(
186 self.data.catalog(),
187 self.data.schema(),
188 &self.data.kind,
189 )
190 } else {
191 DatabaseDdlEvent::alter_lifecycle(self.data.catalog(), self.data.schema())
192 };
193 Some(Box::new(event))
194 }
195}
196
197#[derive(Debug, Serialize, Deserialize, AsRefStr)]
198enum AlterDatabaseState {
199 Prepare,
200 UpdateMetadata,
201 InvalidateSchemaCache,
202}
203
204#[derive(Debug, Serialize, Deserialize)]
206pub struct AlterDatabaseData {
207 state: AlterDatabaseState,
208 kind: AlterDatabaseKind,
209 catalog_name: String,
210 schema_name: String,
211 schema_value: Option<DeserializedValueWithBytes<SchemaNameValue>>,
212}
213
214impl AlterDatabaseData {
215 pub fn new(task: AlterDatabaseTask) -> Result<Self> {
216 Ok(Self {
217 state: AlterDatabaseState::Prepare,
218 kind: AlterDatabaseKind::try_from(task.alter_expr.kind.unwrap())?,
219 catalog_name: task.alter_expr.catalog_name,
220 schema_name: task.alter_expr.schema_name,
221 schema_value: None,
222 })
223 }
224
225 pub fn catalog(&self) -> &str {
226 &self.catalog_name
227 }
228
229 pub fn schema(&self) -> &str {
230 &self.schema_name
231 }
232}
233
234#[cfg(test)]
235mod tests {
236 use std::time::Duration;
237
238 use crate::ddl::alter_database::build_new_schema_value;
239 use crate::key::schema_name::SchemaNameValue;
240 use crate::rpc::ddl::{
241 AlterDatabaseKind, SetDatabaseOption, SetDatabaseOptions, UnsetDatabaseOption,
242 UnsetDatabaseOptions,
243 };
244
245 #[test]
246 fn test_build_new_schema_value() {
247 let set_ttl = AlterDatabaseKind::SetDatabaseOptions(SetDatabaseOptions(vec![
248 SetDatabaseOption::Ttl(Duration::from_secs(10).into()),
249 ]));
250 let current_schema_value = SchemaNameValue::default();
251 let new_schema_value =
252 build_new_schema_value(current_schema_value.clone(), &set_ttl).unwrap();
253 assert_eq!(new_schema_value.ttl, Some(Duration::from_secs(10).into()));
254
255 let unset_ttl_alter_kind =
256 AlterDatabaseKind::UnsetDatabaseOptions(UnsetDatabaseOptions(vec![
257 UnsetDatabaseOption::Ttl,
258 ]));
259 let new_schema_value =
260 build_new_schema_value(current_schema_value, &unset_ttl_alter_kind).unwrap();
261 assert_eq!(new_schema_value.ttl, None);
262 }
263
264 #[test]
265 fn test_build_new_schema_value_with_compaction_options() {
266 let set_compaction = AlterDatabaseKind::SetDatabaseOptions(SetDatabaseOptions(vec![
267 SetDatabaseOption::Other("compaction.type".to_string(), "twcs".to_string()),
268 SetDatabaseOption::Other("compaction.twcs.time_window".to_string(), "1d".to_string()),
269 ]));
270
271 let current_schema_value = SchemaNameValue::default();
272 let new_schema_value =
273 build_new_schema_value(current_schema_value.clone(), &set_compaction).unwrap();
274
275 assert_eq!(
276 new_schema_value.extra_options.get("compaction.type"),
277 Some(&"twcs".to_string())
278 );
279 assert_eq!(
280 new_schema_value
281 .extra_options
282 .get("compaction.twcs.time_window"),
283 Some(&"1d".to_string())
284 );
285
286 let unset_compaction = AlterDatabaseKind::UnsetDatabaseOptions(UnsetDatabaseOptions(vec![
287 UnsetDatabaseOption::Other("compaction.type".to_string()),
288 ]));
289
290 let new_schema_value = build_new_schema_value(new_schema_value, &unset_compaction).unwrap();
291
292 assert_eq!(new_schema_value.extra_options.get("compaction.type"), None);
293 assert_eq!(
294 new_schema_value
295 .extra_options
296 .get("compaction.twcs.time_window"),
297 Some(&"1d".to_string())
298 );
299 }
300}