1use 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 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#[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}